Updates to scanning and rendering of functions

Author: tmr232

scripts/render-function.ts

@@ -5,8 +5,8 @@ import type Parser from "web-tree-sitter";
 import type { SyntaxNode } from "web-tree-sitter";
 import { type Language, supportedLanguages } from "../src/control-flow/cfg.ts";
 import {
-  deserializeColorList,
-  listToScheme,
+  deserializeColorList, getDarkColorList, getLightColorList,
+  listToScheme
 } from "../src/control-flow/colors.ts";
 import { graphToDot } from "../src/control-flow/render.ts";
 import { getLanguage, iterFunctions } from "../src/file-parsing/bun.ts";
@@ -32,11 +32,24 @@ export function getFuncDef(sourceCode: string, func: SyntaxNode): string {
   return normalizeFuncdef(sourceCode.slice(func.startIndex, body.startIndex));
 }
 
-
 function writeError(message: string): void {
   Bun.write(Bun.stderr, `${message}\n`);
 }
 
+async function getColorScheme(
+  colors?:string
+) {
+  if (!colors || colors === "dark") {
+    return listToScheme(getDarkColorList())
+  }
+  if (colors === "light") {
+    return listToScheme(getLightColorList())
+  }
+  return colors
+    ? listToScheme(deserializeColorList(await Bun.file(colors).text()))
+    : undefined;
+}
+
 async function main() {
   process.on("SIGINT", () => {
     // close watcher when Ctrl-C is pressed
@@ -122,9 +135,7 @@ async function main() {
   const graphviz = await Graphviz.load();
   const cfg = buildCFG(func, language);
 
-  const colorScheme = values.colors
-    ? listToScheme(deserializeColorList(await Bun.file(values.colors).text()))
-    : undefined;
+  const colorScheme = await getColorScheme(values.colors)
 
   const svg = graphviz.dot(graphToDot(cfg, false, undefined, colorScheme));
 

scripts/scan-codebase.ts

@@ -8,54 +8,119 @@ import * as path from "node:path";
  */
 import { parseArgs } from "node:util";
 import { Glob } from "bun";
-import { buildCFG } from "./cfg-helper.ts";
 import {
   fileTypes,
   getLanguage,
   iterFunctions,
 } from "../src/file-parsing/bun.ts";
-import {getFuncDef} from "./render-function.ts";
+import { buildCFG } from "./cfg-helper.ts";
+import { getFuncDef } from "./render-function.ts";
 
 export function iterSourceFiles(root: string): IterableIterator<string> {
   const sourceGlob = new Glob(
     `**/*.{${fileTypes.map(({ ext }) => ext).join(",")}}`,
   );
   return sourceGlob.scanSync(root);
 }
+function* iterFilenames(
+  root: string,
+  dirsToInclude: string[],
+): IterableIterator<string> {
+  if (dirsToInclude.length === 1 && dirsToInclude[0] === "*") {
+    yield* iterSourceFiles(root);
+  } else {
+    for (const dir of dirsToInclude) {
+      for (const filename of iterSourceFiles(path.join(root, dir))) {
+        // We want the path relative to the root
+        yield path.join(dir, filename);
+      }
+    }
+  }
+}
 
-async function* iterInfo(root: string) {
-  for (const filename of iterSourceFiles(root)) {
-    const filepath = path.join(root, filename);
-    const code = await Bun.file(filepath).text();
+async function* iterFunctionInfo(
+  root: string,
+  filenames: IterableIterator<string>,
+): AsyncIterableIterator<{
+  node_count: number;
+  start_position: { row:number, column:number };
+  funcdef: string;
+  filename: string;
+}> {
+  for (const filename of filenames) {
+    const code = await Bun.file(path.join(root, filename)).text();
     const language = getLanguage(filename);
     for (const func of iterFunctions(code, language)) {
       const cfg = buildCFG(func, language);
       yield {
-        file: filename,
-        startIndex: func.startIndex,
-        nodeCount: cfg.graph.order,
-        funcDef: getFuncDef(code, func),
-        startPosition: func.startPosition,
+        node_count: cfg.graph.order,
+        start_position: func.startPosition,
+        funcdef: getFuncDef(code, func),
+        filename: filename.replaceAll("\\", "/"),
       };
     }
   }
 }
 
+async function generateIndex(
+  /** Project name on GitHub */
+  project: string,
+  /** Git ref */
+  ref: string,
+  /** Root on local filesystem */
+  root: string,
+  /** Directories to index, relative to the root */
+  dirsToInclude: string[],
+) {
+  const filenames = iterFilenames(root, dirsToInclude);
+  const functions = await Array.fromAsync(iterFunctionInfo(root, filenames));
+  return {
+    version: 1,
+    content: {
+      index_type: "github",
+      project,
+      ref,
+      functions,
+    },
+  };
+}
+
 async function main() {
-  const { values } = parseArgs({
+  const {
+    values,
+    positionals: [_runtime, _this, ...dirsToInclude],
+  } = parseArgs({
     args: Bun.argv,
     options: {
+      project: {
+        type: "string",
+      },
+      ref: {
+        type: "string",
+      },
       root: {
         type: "string",
       },
+      out: {
+        type: "string",
+      },
     },
     strict: true,
     allowPositionals: true,
   });
 
-  const root = values.root ?? ".";
+  if (!values.project || !values.ref || !values.root) {
+    throw new Error("Missing arguments");
+  }
 
-  process.stdout.write(JSON.stringify(await Array.fromAsync(iterInfo(root))));
+  const output = JSON.stringify(
+    await generateIndex(values.project, values.ref, values.root, dirsToInclude),
+  );
+  if (values.out) {
+    await Bun.write(values.out, output);
+  } else {
+    await Bun.write(Bun.stdout, output);
+  }
 }
 
 if (require.main === module) {