Add check for neighborhood node being top-level enclosing function

javascript/ql/experimental/adaptivethreatmodeling/lib/experimental/adaptivethreatmodeling/CodeToFeatures.qll

@@ -49,6 +49,8 @@ module Raw {
 
     AstNode getParentNode() { result = TAstNode(rawNode.getParent()) }
 
+    raw::ASTNode getNode() { result = rawNode }
+
     raw::StmtContainer getContainer() { result = rawNode.getContainer() }
 
     /**

javascript/ql/experimental/adaptivethreatmodeling/lib/experimental/adaptivethreatmodeling/EndpointFeatures.qll

@@ -184,7 +184,10 @@ module NeighborhoodBodies {
    * in the subtree.
    */
   Raw::AstNode getNeighborhoodAstNode(Raw::AstNode node) {
-    if getNumDescendents(node.getParentNode()) > maxNumDescendants()
+    if
+      // `node` will always have a parent as we start at and endpoint
+      node.getParentNode() = getOutermostEnclosingFunction(node) or
+      getNumDescendents(node.getParentNode()) > maxNumDescendants()
     then result = node
     else result = getNeighborhoodAstNode(node.getParentNode())
   }
@@ -192,6 +195,15 @@ module NeighborhoodBodies {
   /** Count number of descendants of an AST node */
   int getNumDescendents(Raw::AstNode node) { result = count(node.getAChildNode*()) }
 
+  private ASTNode getContainer(ASTNode node) {
+    result = node.getContainer()
+  }
+
+  /** Return the AST node that is outermost enclosing function (as an AST Node) */
+  Raw::AstNode getOutermostEnclosingFunction(Raw::AstNode node) {
+    result = Raw::astNode(getContainer*(node.getNode())) and result.getContainer() instanceof TopLevel
+  }
+
   /**
    * Holds if `childNode` is an AST node under `rootNode` and `token` is a node attribute associated
    * with `childNode`. Note that only AST leaves have node attributes.