add code node extraction from tree-sitter

Author: Calvinn Ng

core/indexing/chunk/code.ts

@@ -1,14 +1,13 @@
 import { SyntaxNode } from "web-tree-sitter";
-import { ChunkWithoutID } from "../..";
-import { countTokens } from "../../llm/countTokens";
-import { getParserForFile } from "../../util/treeSitter";
+import { ChunkWithoutID } from "../../index.js";
+import { countTokens } from "../../llm/countTokens.js";
+import { getParserForFile } from "../../util/treeSitter.js";
 
 function collapsedReplacement(node: SyntaxNode): string {
   if (node.type === "statement_block") {
     return "{ ... }";
-  } else {
-    return "...";
   }
+  return "...";
 }
 
 function firstChild(
@@ -19,9 +18,8 @@ function firstChild(
     return (
       node.children.find((child) => grammarName.includes(child.type)) || null
     );
-  } else {
-    return node.children.find((child) => child.type === grammarName) || null;
   }
+  return node.children.find((child) => child.type === grammarName) || null;
 }
 
 function collapseChildren(
@@ -96,6 +94,14 @@ function collapseChildren(
   return code;
 }
 
+export const FUNCTION_BLOCK_NODE_TYPES = ["block", "statement_block"];
+export const FUNCTION_DECLARATION_NODE_TYPEs = [
+  "method_definition",
+  "function_definition",
+  "function_item",
+  "function_declaration",
+];
+
 function constructClassDefinitionChunk(
   node: SyntaxNode,
   code: string,
@@ -105,8 +111,8 @@ function constructClassDefinitionChunk(
     node,
     code,
     ["block", "class_body", "declaration_list"],
-    ["method_definition", "function_definition", "function_item"],
-    ["block", "statement_block"],
+    FUNCTION_DECLARATION_NODE_TYPEs,
+    FUNCTION_BLOCK_NODE_TYPES,
     maxChunkSize,
   );
 }
@@ -130,12 +136,10 @@ function constructFunctionDefinitionChunk(
     // If inside a class, include the class header
     const classNode = node.parent.parent;
     const classBlock = node.parent;
-    return (
-      code.slice(classNode.startIndex, classBlock.startIndex) +
-      "...\n\n" +
-      " ".repeat(node.startPosition.column) + // ...
-      funcText
-    );
+    return `${code.slice(
+      classNode.startIndex,
+      classBlock.startIndex,
+    )}...\n\n${" ".repeat(node.startPosition.column)}${funcText}`;
   }
   return funcText;
 }
@@ -200,7 +204,7 @@ export async function* codeChunker(
     return;
   }
 
-  let parser = await getParserForFile(filepath);
+  const parser = await getParserForFile(filepath);
   if (parser === undefined) {
     console.warn(`Failed to load parser for file ${filepath}: `);
     return;
@@ -209,4 +213,4 @@ export async function* codeChunker(
   const tree = parser.parse(contents);
 
   yield* getSmartCollapsedChunks(tree.rootNode, contents, maxChunkSize);
-}
+}

core/util/ranges.ts

@@ -1,4 +1,4 @@
-import { Range } from "..";
+import { Position, Range } from "../index.js";
 
 export function getRangeInString(content: string, range: Range): string {
   const lines = content.split("\n");
@@ -23,3 +23,71 @@ export function getRangeInString(content: string, range: Range): string {
 
   return [firstLine, ...middleLines, lastLine].join("\n");
 }
+
+export function intersection(a: Range, b: Range): Range | null {
+  const startLine = Math.max(a.start.line, b.start.line);
+  const endLine = Math.min(a.end.line, b.end.line);
+
+  if (startLine > endLine) {
+    return null;
+  }
+
+  if (startLine === endLine) {
+    const startCharacter = Math.max(a.start.character, b.start.character);
+    const endCharacter = Math.min(a.end.character, b.end.character);
+
+    if (startCharacter > endCharacter) {
+      return null;
+    }
+
+    return {
+      start: { line: startLine, character: startCharacter },
+      end: { line: endLine, character: endCharacter },
+    };
+  }
+
+  const startCharacter =
+    startLine === a.start.line ? a.start.character : b.start.character;
+  const endCharacter =
+    endLine === a.end.line ? a.end.character : b.end.character;
+
+  return {
+    start: { line: startLine, character: startCharacter },
+    end: { line: endLine, character: endCharacter },
+  };
+}
+
+export function union(a: Range, b: Range): Range {
+  const startLine = Math.min(a.start.line, b.start.line);
+  const endLine = Math.max(a.end.line, b.end.line);
+
+  const startCharacter =
+    startLine === a.start.line ? a.start.character : b.start.character;
+  const endCharacter =
+    endLine === a.end.line ? a.end.character : b.end.character;
+
+  return {
+    start: { line: startLine, character: startCharacter },
+    end: { line: endLine, character: endCharacter },
+  };
+}
+
+export function maxPosition(a: Position, b: Position): Position {
+  if (a.line > b.line) {
+    return a;
+  } else if (a.line < b.line) {
+    return b;
+  } else {
+    return a.character > b.character ? a : b;
+  }
+}
+
+export function minPosition(a: Position, b: Position): Position {
+  if (a.line < b.line) {
+    return a;
+  } else if (a.line > b.line) {
+    return b;
+  } else {
+    return a.character < b.character ? a : b;
+  }
+}