From b209a630539208a097c638d3a50b65eb28d6aed3 Mon Sep 17 00:00:00 2001
From: dante <c.dante.federici@gmail.com>
Date: Thu, 22 Dec 2022 15:44:06 -0500
Subject: [PATCH] Only preserve directives already at the start of the program

---
 src/utils/extract-ast-nodes.ts                | 17 ++++++---
 .../__snapshots__/ppsi.spec.js.snap           | 34 ++++++++++++++++++
 .../imports-with-directives.ts                | 17 +++++++++
 .../__snapshots__/ppsi.spec.js.snap           | 36 +++++++++++++++++++
 .../imports-with-directives.ts                | 18 ++++++++++
 5 files changed, 118 insertions(+), 4 deletions(-)

diff --git a/src/utils/extract-ast-nodes.ts b/src/utils/extract-ast-nodes.ts
index 14e19b1..55ba311 100644
--- a/src/utils/extract-ast-nodes.ts
+++ b/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>) {
diff --git a/tests/ImportsNotSeparated/__snapshots__/ppsi.spec.js.snap b/tests/ImportsNotSeparated/__snapshots__/ppsi.spec.js.snap
index 46ab77d..a428268 100644
--- a/tests/ImportsNotSeparated/__snapshots__/ppsi.spec.js.snap
+++ b/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;
+})();
 
 `;
 
diff --git a/tests/ImportsNotSeparated/imports-with-directives.ts b/tests/ImportsNotSeparated/imports-with-directives.ts
index 4782ede..a1b09b6 100644
--- a/tests/ImportsNotSeparated/imports-with-directives.ts
+++ b/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;
+})();
diff --git a/tests/ImportsSeparated/__snapshots__/ppsi.spec.js.snap b/tests/ImportsSeparated/__snapshots__/ppsi.spec.js.snap
index 0e3a6ef..0b31f66 100644
--- a/tests/ImportsSeparated/__snapshots__/ppsi.spec.js.snap
+++ b/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`] = `
diff --git a/tests/ImportsSeparated/imports-with-directives.ts b/tests/ImportsSeparated/imports-with-directives.ts
index 5188471..221ec60 100644
--- a/tests/ImportsSeparated/imports-with-directives.ts
+++ b/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;
+})();