perf(offsets): index line starts once per source, binary-search per offset

ober

25a3d578fbaa59f974ebbc0aa7d064f02dc5d1fe

diff --git a/lib/semgrep/source/offsets.sls b/lib/semgrep/source/offsets.sls
index 0ab8a45..c5db856 100644
--- a/lib/semgrep/source/offsets.sls
+++ b/lib/semgrep/source/offsets.sls
@@ -12,20 +12,6 @@
      with-input-from-string with-output-to-string iota \x31;+
      \x31;- partition make-date make-time meta atom?)
     (except (jerboa prelude) meta atom?))
-  (def (offset->line-col source offset)
-       (let ([len (string-length source)])
-         (let loop ([i 0] [line 1] [col 1])
-           (cond
-             [(or (= i offset) (= i len)) (values line col)]
-             [(char=? (string-ref source i) #\newline)
-              (loop (+ i 1) (+ line 1) 1)]
-             [else
-              (loop
-                (+ i 1)
-                line
-                (+ col
-                   (semgrep-source-byte-length
-                     (string-ref source i))))]))))
   (def (char-utf8-byte-length ch)
        (let ([codepoint (char->integer ch)])
          (cond
@@ -35,20 +21,81 @@
            [else 4])))
   (def (semgrep-source-byte-length ch)
        (if (char=? ch #\�) 1 (char-utf8-byte-length ch)))
-  (def (source-index->semgrep-offset source index)
+  (def (build-source-index source)
        (let ([len (string-length source)])
-         (let loop ([i 0] [offset 0])
+         (let loop ([i 0]
+                    [sb 0]
+                    [ub 0]
+                    [chars '(0)]
+                    [sbytes '(0)]
+                    [ubytes '(0)])
            (cond
-             [(or (>= i index) (>= i len)) offset]
+             [(>= i len)
+              (vector
+                (list->vector (reverse chars))
+                (list->vector (reverse sbytes))
+                (list->vector (reverse ubytes))
+                len)]
+             [(char=? (string-ref source i) #\newline)
+              (loop (+ i 1) (+ sb 1) (+ ub 1) (cons (+ i 1) chars)
+                (cons (+ sb 1) sbytes) (cons (+ ub 1) ubytes))]
              [else
-              (loop
-                (+ i 1)
-                (+ offset
-                   (semgrep-source-byte-length
-                     (string-ref source i))))]))))
+              (let ([ch (string-ref source i)])
+                (loop (+ i 1) (+ sb (semgrep-source-byte-length ch))
+                  (+ ub (char-utf8-byte-length ch)) chars sbytes
+                  ubytes))]))))
+  (def *source-index-cache* (make-weak-eq-hashtable))
+  (def (source-index source)
+       (or (hashtable-ref *source-index-cache* source #f)
+           (let ([idx (build-source-index source)])
+             (hashtable-set! *source-index-cache* source idx)
+             idx)))
+  (def (line-of starts pos)
+       (let loop ([lo 0] [hi (- (vector-length starts) 1)])
+         (if (>= lo hi)
+             lo
+             (let ([mid (quotient (+ lo hi 1) 2)])
+               (if (<= (vector-ref starts mid) pos)
+                   (loop mid hi)
+                   (loop lo (- mid 1)))))))
+  (def (offset->line-col source offset)
+       (let* ([idx (source-index source)]
+              [chars (vector-ref idx 0)]
+              [len (vector-ref idx 3)]
+              [offset (if (> offset len) len offset)]
+              [k (line-of chars offset)])
+         (let loop ([i (vector-ref chars k)] [col 1])
+           (if (>= i offset)
+               (values (+ k 1) col)
+               (loop
+                 (+ i 1)
+                 (+ col
+                    (semgrep-source-byte-length
+                      (string-ref source i))))))))
+  (def (source-index->semgrep-offset source index)
+       (let* ([idx (source-index source)]
+              [chars (vector-ref idx 0)]
+              [sbytes (vector-ref idx 1)]
+              [len (vector-ref idx 3)]
+              [index (if (> index len) len index)]
+              [k (line-of chars index)])
+         (let loop ([i (vector-ref chars k)]
+                    [b (vector-ref sbytes k)])
+           (if (>= i index)
+               b
+               (loop
+                 (+ i 1)
+                 (+ b
+                    (semgrep-source-byte-length
+                      (string-ref source i))))))))
   (def (tree-byte-offset->source-index source offset)
-       (let ([len (string-length source)])
-         (let loop ([i 0] [tree-offset 0])
+       (let* ([idx (source-index source)]
+              [chars (vector-ref idx 0)]
+              [ubytes (vector-ref idx 2)]
+              [len (vector-ref idx 3)]
+              [k (line-of ubytes offset)])
+         (let loop ([i (vector-ref chars k)]
+                    [tree-offset (vector-ref ubytes k)])
            (cond
              [(or (>= i len) (>= tree-offset offset)) i]
              [else
diff --git a/src/.jerbuild-hashes b/src/.jerbuild-hashes
index 045aff8..e244eb4 100644
--- a/src/.jerbuild-hashes
+++ b/src/.jerbuild-hashes
@@ -18,17 +18,17 @@
  ("src/semgrep/targeting/path-filter.ss" . "7EECA69A56137A8")
  ("src/semgrep/engine/js-eval-scan.ss" . "916D40FD4680B9BD")
  ("src/semgrep/result/extras.ss" . "7342000FE621A6A3")
- ("src/semgrep/engine/jsts-positive-dispatch.ss"
-   .
-   "DABB3C66BB249070")
  ("src/semgrep/engine/js-decorator-scan.ss"
    .
    "193758AD2E444FD6")
- ("src/semgrep/schema/lang.ss" . "CAE2CA859C9A9FD0")
- ("src/semgrep/rule.ss" . "E12C108153C181FA")
+ ("src/semgrep/engine/jsts-positive-dispatch.ss"
+   .
+   "DABB3C66BB249070")
  ("src/semgrep/output/text.ss" . "BE476CB84B807FBA")
+ ("src/semgrep/rule.ss" . "E12C108153C181FA")
+ ("src/semgrep/schema/lang.ss" . "CAE2CA859C9A9FD0")
  ("src/semgrep/engine/ts-type-scan.ss" . "B36BAF07D19F4419")
- ("src/semgrep/source/offsets.ss" . "834EFDB706823794")
+ ("src/semgrep/source/offsets.ss" . "87CFA4FE8ADC771E")
  ("src/semgrep/engine/regex-support.ss" . "55D086807DB29774")
  ("src/semgrep/engine/py-constant-scan.ss"
    .
@@ -57,18 +57,18 @@
  ("src/semgrep/parse/parse-target.ss" . "97AA8FFEB12736DA")
  ("src/semgrep/scan.ss" . "E5E3900F4CA79795")
  ("src/semgrep/engine/generic-scan.ss" . "925C9DC8F6C8188F")
- ("src/semgrep/engine/py-constant-prop.ss"
-   .
-   "76462D0EB71F2180")
  ("src/semgrep/engine/js-constructor-scan.ss"
    .
    "6B226B8D0584A7A")
+ ("src/semgrep/engine/py-constant-prop.ss"
+   .
+   "76462D0EB71F2180")
  ("src/semgrep/engine/comparison.ss" . "39F51E8BBC7F761")
  ("src/semgrep/main.ss" . "A4EC9E7F2A09D25E")
  ("src/semgrep/engine/py-import-scan.ss"
    .
    "F1817B0AFD8D9500")
+ ("src/semgrep/engine/text-support.ss" . "644AF29394C53045")
  ("src/semgrep/engine/ts-decorator-scan.ss"
    .
-   "610AA3A90D6A95F8")
- ("src/semgrep/engine/text-support.ss" . "644AF29394C53045"))
+   "610AA3A90D6A95F8"))
diff --git a/src/semgrep/source/offsets.ss b/src/semgrep/source/offsets.ss
index 48e3173..018e5b8 100644
--- a/src/semgrep/source/offsets.ss
+++ b/src/semgrep/source/offsets.ss
@@ -7,20 +7,6 @@
 
 (import (except (jerboa prelude) meta atom?))
 
-(def (offset->line-col source offset)
-  (let ([len (string-length source)])
-    (let loop ([i 0] [line 1] [col 1])
-      (cond
-        [(or (= i offset) (= i len)) (values line col)]
-        [(char=? (string-ref source i) #\newline)
-         (loop (+ i 1) (+ line 1) 1)]
-        [else
-         (loop (+ i 1)
-               line
-               (+ col
-                  (semgrep-source-byte-length
-                    (string-ref source i))))]))))
-
 (def (char-utf8-byte-length ch)
   (let ([codepoint (char->integer ch)])
     (cond
@@ -34,19 +20,84 @@
       1
       (char-utf8-byte-length ch)))
 
-(def (source-index->semgrep-offset source index)
+;; Per-source line index: the char offset, semgrep byte offset, and UTF-8 byte
+;; offset of every line start, plus the source length. Built once per source
+;; object (weakly cached below) so the per-finding offset lookups become a
+;; binary search for the line plus a short within-line scan, instead of an
+;; O(n) full-source scan each (O(n x findings) hung the scanner on a hostile
+;; file with many findings).
+(def (build-source-index source)
   (let ([len (string-length source)])
-    (let loop ([i 0] [offset 0])
+    (let loop ([i 0] [sb 0] [ub 0]
+               [chars '(0)] [sbytes '(0)] [ubytes '(0)])
       (cond
-        [(or (>= i index) (>= i len)) offset]
+        [(>= i len)
+         (vector (list->vector (reverse chars))
+                 (list->vector (reverse sbytes))
+                 (list->vector (reverse ubytes))
+                 len)]
+        [(char=? (string-ref source i) #\newline)
+         (loop (+ i 1) (+ sb 1) (+ ub 1)
+               (cons (+ i 1) chars)
+               (cons (+ sb 1) sbytes)
+               (cons (+ ub 1) ubytes))]
         [else
-         (loop (+ i 1)
-               (+ offset
-                  (semgrep-source-byte-length (string-ref source i))))]))))
+         (let ([ch (string-ref source i)])
+           (loop (+ i 1)
+                 (+ sb (semgrep-source-byte-length ch))
+                 (+ ub (char-utf8-byte-length ch))
+                 chars sbytes ubytes)])))))
+
+(def *source-index-cache* (make-weak-eq-hashtable))
+
+(def (source-index source)
+  (or (hashtable-ref *source-index-cache* source #f)
+      (let ([idx (build-source-index source)])
+        (hashtable-set! *source-index-cache* source idx)
+        idx)))
+
+;; Largest k with (vector-ref starts k) <= pos (starts is ascending, starts[0]=0).
+(def (line-of starts pos)
+  (let loop ([lo 0] [hi (- (vector-length starts) 1)])
+    (if (>= lo hi)
+      lo
+      (let ([mid (quotient (+ lo hi 1) 2)])
+        (if (<= (vector-ref starts mid) pos)
+          (loop mid hi)
+          (loop lo (- mid 1)))))))
+
+(def (offset->line-col source offset)
+  (let* ([idx (source-index source)]
+         [chars (vector-ref idx 0)]
+         [len (vector-ref idx 3)]
+         [offset (if (> offset len) len offset)]
+         [k (line-of chars offset)])
+    (let loop ([i (vector-ref chars k)] [col 1])
+      (if (>= i offset)
+        (values (+ k 1) col)
+        (loop (+ i 1)
+              (+ col (semgrep-source-byte-length (string-ref source i))))))))
+
+(def (source-index->semgrep-offset source index)
+  (let* ([idx (source-index source)]
+         [chars (vector-ref idx 0)]
+         [sbytes (vector-ref idx 1)]
+         [len (vector-ref idx 3)]
+         [index (if (> index len) len index)]
+         [k (line-of chars index)])
+    (let loop ([i (vector-ref chars k)] [b (vector-ref sbytes k)])
+      (if (>= i index)
+        b
+        (loop (+ i 1)
+              (+ b (semgrep-source-byte-length (string-ref source i))))))))
 
 (def (tree-byte-offset->source-index source offset)
-  (let ([len (string-length source)])
-    (let loop ([i 0] [tree-offset 0])
+  (let* ([idx (source-index source)]
+         [chars (vector-ref idx 0)]
+         [ubytes (vector-ref idx 2)]
+         [len (vector-ref idx 3)]
+         [k (line-of ubytes offset)])
+    (let loop ([i (vector-ref chars k)] [tree-offset (vector-ref ubytes k)])
       (cond
         [(or (>= i len) (>= tree-offset offset)) i]
         [else