From 486752d19efbfefc5eca04b14af2d209548d8eeb Mon Sep 17 00:00:00 2001
From: tiferet <tiferet@github.com>
Date: Wed, 30 Nov 2022 12:56:59 -0800
Subject: [PATCH] Code improvements:

- Replace the containment logic with built in `Location` functionality.
- Generalize `tokenize` to output the tokens that fall within any location.
---
 .../EndpointFeatures.qll                      | 44 +++++--------------
 .../ql/lib/semmle/javascript/Locations.qll    | 33 ++++++++++++++
 2 files changed, 43 insertions(+), 34 deletions(-)

diff --git a/javascript/ql/experimental/adaptivethreatmodeling/lib/experimental/adaptivethreatmodeling/EndpointFeatures.qll b/javascript/ql/experimental/adaptivethreatmodeling/lib/experimental/adaptivethreatmodeling/EndpointFeatures.qll
index 5b3eed44497..23cd82612d5 100644
--- a/javascript/ql/experimental/adaptivethreatmodeling/lib/experimental/adaptivethreatmodeling/EndpointFeatures.qll
+++ b/javascript/ql/experimental/adaptivethreatmodeling/lib/experimental/adaptivethreatmodeling/EndpointFeatures.qll
@@ -226,43 +226,12 @@ class CodexPrompt extends EndpointFeature, TCodexPrompt {
   }
 
   /**
-   * Holds if the location of `node` contains the location of `token`
+   * Gets the reconstructed source code text for a range of locations.
    */
-  cached
-  predicate containsToken(AstNode node, Token token) {
-    exists(
-      string file, int node_start_line, int node_start_column, int node_end_line,
-      int node_end_column, int token_start_line, int token_start_column, int token_end_line,
-      int token_end_column
-    |
-      node.getLocation()
-          .hasLocationInfo(file, node_start_line, node_start_column, node_end_line, node_end_column) and
-      token
-          .getLocation()
-          .hasLocationInfo(file, token_start_line, token_start_column, token_end_line,
-            token_end_column) and
-      (
-        node_start_line < token_start_line
-        or
-        node_start_line = token_start_line and
-        node_start_column <= token_start_column
-      ) and
-      (
-        node_end_line > token_end_line
-        or
-        node_end_line = token_end_line and
-        node_end_column >= token_end_column
-      )
-    )
-  }
-
-  /**
-   * Gets the reconstructed source code text for `node`.
-   */
-  string tokenise(DataFlow::Node node) {
+  string tokenize(Location location) {
     result =
       strictconcat(Token token |
-        containsToken(node.getAstNode(), token)
+        location.containsLoosely(token.getLocation())
       |
         token.getValue(),
           // Use space as the separator, since that is most likely.
@@ -274,6 +243,13 @@ class CodexPrompt extends EndpointFeature, TCodexPrompt {
           token.getLocation().getStartLine(), token.getLocation().getStartColumn()
       )
   }
+
+  /**
+   * Gets the reconstructed source code text for `node`.
+   */
+  string tokenizeEndpoint(DataFlow::Node node) {
+    result = tokenize(node.getAstNode().getLocation())
+  }
 }
 
 /**
diff --git a/javascript/ql/lib/semmle/javascript/Locations.qll b/javascript/ql/lib/semmle/javascript/Locations.qll
index c0748f7b3e7..b3cb275695d 100644
--- a/javascript/ql/lib/semmle/javascript/Locations.qll
+++ b/javascript/ql/lib/semmle/javascript/Locations.qll
@@ -40,6 +40,18 @@ class Location extends @location {
     )
   }
 
+  /** Holds if this location starts before or at the same place as location `that`. */
+  predicate startsBeforeOrWith(Location that) {
+    exists(File f, int sl1, int sc1, int sl2, int sc2 |
+      locations_default(this, f, sl1, sc1, _, _) and
+      locations_default(that, f, sl2, sc2, _, _)
+    |
+      sl1 < sl2
+      or
+      sl1 = sl2 and sc1 <= sc2
+    )
+  }
+
   /** Holds if this location ends after location `that`. */
   pragma[inline]
   predicate endsAfter(Location that) {
@@ -53,12 +65,33 @@ class Location extends @location {
     )
   }
 
+  /** Holds if this location ends after or at the same place as location `that`. */
+  pragma[inline]
+  predicate endsAfterOrWith(Location that) {
+    exists(File f, int el1, int ec1, int el2, int ec2 |
+      locations_default(this, f, _, _, el1, ec1) and
+      locations_default(that, f, _, _, el2, ec2)
+    |
+      el1 > el2
+      or
+      el1 = el2 and ec1 >= ec2
+    )
+  }
+
   /**
    * Holds if this location contains location `that`, meaning that it starts
    * before and ends after it.
    */
   predicate contains(Location that) { this.startsBefore(that) and this.endsAfter(that) }
 
+  /**
+   * Holds if this location contains location `that`, meaning that it starts
+   * before or at the same place and ends after or at the same place.
+   */
+  predicate containsLoosely(Location that) {
+    this.startsBeforeOrWith(that) and this.endsAfterOrWith(that)
+  }
+
   /** Holds if this location is empty. */
   predicate isEmpty() { exists(int l, int c | locations_default(this, _, l, c, l, c - 1)) }