perf: use bytes for pos comparison

Yash-Singh1 · Apr 11, 2024 · 2b2ffd3 · 2b2ffd3
1 parent 8894f83
commit 2b2ffd3
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 66 deletions.
diff --git a/server/src/capabilities/completions.ts b/server/src/capabilities/completions.ts
@@ -58,6 +58,11 @@ async function onCompletion({
   const stateManager = new StateManager();
   let moduleStart = 0;
 
+  // Precalculate the byte offset of the parameter
+  const paramByte = byteRepresentation.charIndexToByteOffset(
+    textDocument.offsetAt(params.position),
+  );
+
   await walk<{
     propertyName: string | undefined;
     callInside: string | null | undefined;
@@ -80,24 +85,13 @@ async function onCompletion({
       },
 
       "*"(node) {
-        if ("span" in node && node.type !== "VariableDeclaration") {
-          const startSpanRelative = textDocument.positionAt(
-            byteRepresentation.byteOffsetToCharIndex(
-              node.span.start - moduleStart,
-            ),
-          );
-          const endSpanRelative = textDocument.positionAt(
-            byteRepresentation.byteOffsetToCharIndex(
-              node.span.end - moduleStart,
-            ),
-          );
-
-          if (
-            params.position.line > endSpanRelative.line ||
-            params.position.line < startSpanRelative.line
-          ) {
-            return false;
-          }
+        if (
+          "span" in node &&
+          node.type !== "VariableDeclaration" &&
+          paramByte < node.span.start - moduleStart &&
+          paramByte > node.span.end - moduleStart
+        ) {
+          return false;
         }
       },
 
@@ -199,19 +193,10 @@ async function onCompletion({
               node.span.start - moduleStart,
             ),
           );
-          const endSpanRelative = textDocument.positionAt(
-            byteRepresentation.byteOffsetToCharIndex(
-              node.span.end - moduleStart,
-            ),
-          );
 
           if (
-            params.position.line > endSpanRelative.line ||
-            params.position.line < startSpanRelative.line ||
-            (params.position.line === endSpanRelative.line &&
-              params.position.character > endSpanRelative.character) ||
-            (params.position.line === startSpanRelative.line &&
-              params.position.character < startSpanRelative.character)
+            paramByte < node.span.start - moduleStart ||
+            paramByte > node.span.end - moduleStart
           ) {
             return false;
           }

diff --git a/server/src/capabilities/hover.ts b/server/src/capabilities/hover.ts
@@ -65,6 +65,10 @@ async function onHover({
   // Resulting hover
   let hover: Hover | null = null;
 
+  const paramByte = byteRepresentation.charIndexToByteOffset(
+    textDocument.offsetAt(params.position),
+  );
+
   await walk<{
     parentClass: string[];
     callInside: string | null | undefined;
@@ -78,24 +82,13 @@ async function onHover({
       },
 
       "*"(node) {
-        if ("span" in node && node.type !== "VariableDeclaration") {
-          const startSpanRelative = textDocument.positionAt(
-            byteRepresentation.byteOffsetToCharIndex(
-              node.span.start - moduleStart,
-            ),
-          );
-          const endSpanRelative = textDocument.positionAt(
-            byteRepresentation.byteOffsetToCharIndex(
-              node.span.end - moduleStart,
-            ),
-          );
-
-          if (
-            params.position.line > endSpanRelative.line ||
-            params.position.line < startSpanRelative.line
-          ) {
-            return false;
-          }
+        if (
+          "span" in node &&
+          node.type !== "VariableDeclaration" &&
+          paramByte < node.span.start &&
+          paramByte > node.span.end
+        ) {
+          return false;
         }
       },
 
@@ -247,27 +240,10 @@ async function onHover({
             };
           }
 
-          const startSpanRelative = textDocument.positionAt(
-            byteRepresentation.byteOffsetToCharIndex(
-              node.key.span.start - moduleStart,
-            ),
-          );
-          const endSpanRelative = textDocument.positionAt(
-            byteRepresentation.byteOffsetToCharIndex(
-              node.key.span.end - moduleStart,
-            ),
-          );
-
           // Don't use out of range nodes
           if (
-            !(
-              params.position.line >= startSpanRelative.line &&
-              params.position.line <= endSpanRelative.line &&
-              (params.position.line !== startSpanRelative.line ||
-                params.position.character >= startSpanRelative.character) &&
-              (params.position.line !== endSpanRelative.line ||
-                params.position.character <= endSpanRelative.character)
-            )
+            paramByte > node.key.span.end - moduleStart ||
+            paramByte < node.key.span.start - moduleStart
           ) {
             return state;
           }

diff --git a/server/src/lib/string-bytes.ts b/server/src/lib/string-bytes.ts
@@ -25,6 +25,7 @@ export class StringAsBytes {
   private prefixArray: Uint32Array;
   private stringSegments: string[];
   private preStringLength: Uint32Array;
+  private stringVal: string;
 
   constructor(string: string) {
     this.encoder = new TextEncoder();
@@ -34,6 +35,7 @@ export class StringAsBytes {
     this.stringSegments = [
       ...new Intl.Segmenter("en", { granularity: "grapheme" }).segment(string),
     ].map((segment) => segment.segment);
+    this.stringVal = string;
 
     this.calculatePrefixArray();
   }
@@ -74,6 +76,11 @@ export class StringAsBytes {
     this.stringLength = prefixArray[prefixArray.length - 1];
   }
 
+  // We can make this O(n) time because it is called only once per request
+  public charIndexToByteOffset(charIndex: number) {
+    return this.encoder.encode(this.stringVal.slice(0, charIndex)).length;
+  }
+
   /**
    * Binary searches the prefix sum array to find the character index for a given byte offset.
    * @param byteOffset Byte offset to convert to char index