Only preserve directives already at the start of the program

src/utils/extract-ast-nodes.ts

@@ -11,11 +11,20 @@ export function extractASTNodes(ast: ParseResult<File>) {
     const importNodes: ImportDeclaration[] = [];
     const directives: Directive[] = [];
     traverse(ast, {
-        Directive({ node }) {
-            directives.push(node);
+        Directive(path: NodePath<Directive>) {
+            // Only capture directives if they are at the top scope of the source
+            // and their previous siblings are all directives
+            if (
+                path.parent.type === 'Program' &&
+                path.getAllPrevSiblings().every((s) => {
+                    return s.type === 'Directive';
+                })
+            ) {
+                directives.push(path.node);
 
-            // Trailing comments probably shouldn't be attached to the directive
-            node.trailingComments = null;
+                // Trailing comments probably shouldn't be attached to the directive
+                path.node.trailingComments = null;
+            }
         },
 
         ImportDeclaration(path: NodePath<ImportDeclaration>) {

tests/ImportsNotSeparated/__snapshots__/ppsi.spec.js.snap

@@ -160,6 +160,23 @@ import abc from "@core/abc";
 function add(a:number,b:number) {
     return a + b;
 }
+function addStrict(a:number,b:number) {
+    'use strict';
+    return a + b;
+}
+
+'preserve me';
+
+const workletAdd = (a:number,b:number) => {
+    'worklet';
+    return a + b;
+}
+
+(function() {
+    'use strict';
+    // some iffe example
+    return true;
+})();
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 "use strict";
 "use client";
@@ -171,6 +188,23 @@ import otherthing from "@core/otherthing";
 function add(a: number, b: number) {
     return a + b;
 }
+function addStrict(a: number, b: number) {
+    "use strict";
+    return a + b;
+}
+
+("preserve me");
+
+const workletAdd = (a: number, b: number) => {
+    "worklet";
+    return a + b;
+};
+
+(function () {
+    "use strict";
+    // some iffe example
+    return true;
+})();
 
 `;
 

tests/ImportsNotSeparated/imports-with-directives.ts

@@ -6,3 +6,20 @@ import abc from "@core/abc";
 function add(a:number,b:number) {
     return a + b;
 }
+function addStrict(a:number,b:number) {
+    'use strict';
+    return a + b;
+}
+
+'preserve me';
+
+const workletAdd = (a:number,b:number) => {
+    'worklet';
+    return a + b;
+}
+
+(function() {
+    'use strict';
+    // some iffe example
+    return true;
+})();

tests/ImportsSeparated/__snapshots__/ppsi.spec.js.snap

@@ -173,6 +173,24 @@ import abc from "@core/abc";
 function add(a:number,b:number) {
   return a + b;
 }
+
+function addStrict(a:number,b:number) {
+  'use strict';
+  return a + b;
+}
+
+'preserve me';
+
+const workletAdd = (a:number,b:number) => {
+  'worklet';
+  return a + b;
+}
+
+(function() {
+  'use strict';
+  // some iffe example
+  return true;
+})();
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 "use strict";
 "use client";
@@ -187,6 +205,24 @@ function add(a: number, b: number) {
     return a + b;
 }
 
+function addStrict(a: number, b: number) {
+    "use strict";
+    return a + b;
+}
+
+("preserve me");
+
+const workletAdd = (a: number, b: number) => {
+    "worklet";
+    return a + b;
+};
+
+(function () {
+    "use strict";
+    // some iffe example
+    return true;
+})();
+
 `;
 
 exports[`imports-with-file-level-comments.ts - typescript-verify: imports-with-file-level-comments.ts 1`] = `

tests/ImportsSeparated/imports-with-directives.ts

@@ -10,3 +10,21 @@ import abc from "@core/abc";
 function add(a:number,b:number) {
   return a + b;
 }
+
+function addStrict(a:number,b:number) {
+  'use strict';
+  return a + b;
+}
+
+'preserve me';
+
+const workletAdd = (a:number,b:number) => {
+  'worklet';
+  return a + b;
+}
+
+(function() {
+  'use strict';
+  // some iffe example
+  return true;
+})();