different approach to find pragma comments as previous did not work i…

…n all cases

scripts/babel/transform-test-gate-pragma.js

@@ -263,9 +263,54 @@ function transform(babel) {
     }
   }
 
+  function getComments(path) {
+    if (path.node.leadingComments) {
+      // Babel AST includes comments.
+      return path.node.leadingComments;
+    }
+    // In Hermes AST we need to find the comments by range.
+    const comments = [];
+    let line = path.node.loc.start.line;
+    let i = allComments.length - 1;
+    while (i >= 0 && allComments[i].loc.end.line >= line) {
+      i--;
+    }
+    while (i >= 0 && allComments[i].loc.end.line === line - 1) {
+      const comment = allComments[i];
+      line = comment.loc.start.line;
+      comments.unshift(comment);
+      if (usedComments.has(comment)) {
+        throw new Error('expected to find just one place for comment');
+      }
+      usedComments.add(comment);
+      i--;
+    }
+    return comments;
+  }
+
+  let usedComments;
+  let allComments;
   return {
     name: 'test-gate-pragma',
     visitor: {
+      Program: {
+        enter(path) {
+          usedComments = new Set();
+          allComments = path.parent.comments;
+        },
+        exit(path) {
+          let unusedPragmaComment = allComments.find(
+            c => c.value.includes('@gate') && !usedComments.has(c)
+          );
+          if (unusedPragmaComment != null) {
+            throw path.hub.buildError(
+              unusedPragmaComment,
+              '@gate directive needs to be immediately before test/it call',
+              Error
+            );
+          }
+        },
+      },
       ExpressionStatement(path) {
         const statement = path.node;
         const expression = statement.expression;
@@ -331,26 +376,4 @@ function transform(babel) {
   };
 }
 
-function getComments(path) {
-  if (path.node.leadingComments) {
-    // Babel AST includes comments.
-    return path.node.leadingComments;
-  }
-  // In Hermes AST we need to find the comments by range.
-  const comments = path.hub.file.ast.comments;
-  let prevSibling = path.getPrevSibling();
-  let searchStart;
-  if (prevSibling.node) {
-    searchStart = prevSibling.node.end;
-  } else if (path.parentPath.node) {
-    searchStart = path.parentPath.node.start;
-  } else {
-    throw new Error('Unexpected AST structure');
-  }
-  const filteredComments = comments.filter(
-    c => c.start >= searchStart && c.end <= path.node.start
-  );
-  return filteredComments.length > 0 ? filteredComments : undefined;
-}
-
 module.exports = transform;