diff --git a/javascript/ql/experimental/adaptivethreatmodeling/lib/experimental/adaptivethreatmodeling/CodeToFeatures.qll b/javascript/ql/experimental/adaptivethreatmodeling/lib/experimental/adaptivethreatmodeling/CodeToFeatures.qll
index 3de472147ea..9d8b1632c32 100644
--- a/javascript/ql/experimental/adaptivethreatmodeling/lib/experimental/adaptivethreatmodeling/CodeToFeatures.qll
+++ b/javascript/ql/experimental/adaptivethreatmodeling/lib/experimental/adaptivethreatmodeling/CodeToFeatures.qll
@@ -4,6 +4,8 @@
  * Extracts data about the functions in the database for use in adaptive threat modeling (ATM).
  */
 
+private import EndpointFeatures::NeighborhoodBodies
+
 module Raw {
   private import javascript as raw
 
@@ -30,6 +32,9 @@ module Raw {
     entity.getNumBodyStmt() = 0 and not exists(entity.getAReturnedExpr())
   }
 
+  /**
+   * Wrapper for RawAstNode (could be an alias instead of a newtype)
+   */
   newtype WrappedAstNode = TAstNode(RawAstNode rawNode)
 
   /**
@@ -44,6 +49,8 @@ module Raw {
 
     AstNode getParentNode() { result = TAstNode(rawNode.getParent()) }
 
+    raw::StmtContainer getContainer() { result = rawNode.getContainer() }
+
     /**
      * Holds if the AST node has `result` as its `index`th attribute.
      *
@@ -128,6 +135,11 @@ module Raw {
     }
   }
 
+  /**
+   * Returns the `Raw::AstNode` wrapper of `rawNode`
+   */
+  AstNode astNode(RawAstNode rawNode) { result = TAstNode(rawNode) }
+
   /**
    * Holds if `result` is the `index`'th child of the `parent` entity. Such
    * a node is a root of an AST associated with this entity.
@@ -324,6 +336,11 @@ module Wrapped {
     Raw::Location getLocation() { result = rawNode.getLocation() }
   }
 
+  /**
+   * Returns the `Wrapped::AstNode` for a `rawNode` in the context of `entity`
+   */
+  AstNode astNode(Raw::Entity entity, Raw::AstNode rawNode) { result = TAstNode(entity, rawNode) }
+
   /**
    * A synthetic AST node, created to be a leaf for an otherwise non-leaf attribute.
    */
@@ -383,6 +400,8 @@ module DatabaseFeatures {
     override Location getLocation() { result = entity.getLocation() }
 
     UnderlyingFunction getDefinedFunction() { result = entity.getDefinedFunction() }
+
+    Wrapped::Entity getWrappedEntity() { result = entity }
   }
 
   class AstNode extends EntityOrAstNode, TAstNode {
@@ -392,8 +411,12 @@ module DatabaseFeatures {
 
     AstNode getChild(int index) { result = TAstNode(rawNode.getChild(index)) }
 
+    AstNode getAChild() { result = this.getChild(_) }
+
     string getAttribute(int index) { result = rawNode.getAttribute(index) }
 
+    Wrapped::AstNode getRawNode() { result = rawNode }
+
     override string getType() { result = rawNode.getType() }
 
     override string toString() { result = this.getType() }
@@ -401,6 +424,9 @@ module DatabaseFeatures {
     override Location getLocation() { result = rawNode.getLocation() }
   }
 
+  /** Gets the `DatabaseFeatures::AstNode` that wraps `wrapped` */
+  AstNode astNode(Wrapped::AstNode wrapped) { result.getRawNode() = wrapped }
+
   /** Consistency checks: these predicates should each have no results */
   module Consistency {
     query predicate nonLeafAttribute(AstNode node, int index, string attribute) {