perf: use stream for getDeclDeps

graph/decl_deps.ts

@@ -1,3 +1,4 @@
+import { Stream } from "https://deno.land/x/[email protected]/stream/mod.ts"
 import {
 	type ClassDeclaration,
 	type FunctionDeclaration,
@@ -25,7 +26,9 @@ export type DeclDeps = Map<Declaration, Declaration[]>
  * RFC: only track top-level declarations or inside if statements?
  * FIXME: also track `function` statements, are there other types of declarations?
  */
-const getTopDecl = (ref: ReferencedSymbolEntry): Declaration | undefined => {
+export const getTopDecl = (
+	ref: ReferencedSymbolEntry,
+): Declaration | undefined => {
 	const target = ref.getNode().getAncestors()
 
 	// 	.at(-2) // HACK: last is SourceFile, second to last is the variable/class declaration
@@ -42,38 +45,34 @@ const getTopDecl = (ref: ReferencedSymbolEntry): Declaration | undefined => {
 	)
 }
 
-const equals = (a: Node, b: Node) => encodeVSCodeURI(a) === encodeVSCodeURI(b)
+export const equals = (a: Node, b: Node) =>
+	encodeVSCodeURI(a) === encodeVSCodeURI(b)
 
-/**
- * Find all references in other files
- *
- * TODO: use isDefinition()?
- */
-const getActualReferences = (
+export const getActualReferences = (
 	symbol: ReferencedSymbol,
-): ReferencedSymbolEntry[] =>
-	symbol
-		.getReferences()
+): Stream<ReferencedSymbolEntry> =>
+	Stream.from(symbol
+		.getReferences())
 		.filter((ref) => !equals(ref.getNode(), symbol.getDefinition().getNode()))
 		.filter((ref) =>
 			ref.getNode().getParent()?.getKindName() !== "ImportClause"
 		)
 
-const getReferencedDecls = (node: Declaration): Declaration[] =>
-	node.findReferences()
+export const getReferencedDecls = (
+	node: Declaration,
+): Stream<Declaration> =>
+	Stream.from(node.findReferences())
 		.flatMap(getActualReferences)
-		.flatMap((x) => getTopDecl(x) ?? [])
+		.map((x) => getTopDecl(x))
+		.filter((x): x is Declaration => x !== undefined)
 
-/**
- * Recursively query **direct** variable references into key-value map.
- */
 export const getDeclDeps = (links: Declaration[]): DeclDeps => {
 	const graph: DeclDeps = new Map()
 
 	let current = links
 	while (current.length > 0) {
 		current = current.flatMap((node) => {
-			const references = getReferencedDecls(node)
+			const references = getReferencedDecls(node).toArray()
 			graph.set(node, references)
 			return references
 		})

graph/decl_deps_bench.ts

@@ -0,0 +1,69 @@
+import { ReferencedSymbol, ReferencedSymbolEntry } from "../deps/ts_morph.ts"
+import { exampleSrc } from "./_example_project.ts"
+import { inMemoryProject, withSrc } from "./_project.ts"
+import {
+	type Declaration,
+	type DeclDeps,
+	equals,
+	getDeclDeps,
+	getTopDecl,
+} from "./decl_deps.ts"
+import { Stream } from "https://deno.land/x/[email protected]/stream/mod.ts"
+import { getAllDecls } from "./decls.ts"
+
+/**
+ * Find all references in other files
+ *
+ * TODO: use isDefinition()?
+ */
+const getActualReferencesArray = (
+	symbol: ReferencedSymbol,
+): ReferencedSymbolEntry[] =>
+	symbol
+		.getReferences()
+		.filter((ref) => !equals(ref.getNode(), symbol.getDefinition().getNode()))
+		.filter((ref) =>
+			ref.getNode().getParent()?.getKindName() !== "ImportClause"
+		)
+
+const getReferencedDeclsArray = (node: Declaration): Declaration[] =>
+	node.findReferences()
+		.flatMap(getActualReferencesArray)
+		.flatMap((x) => getTopDecl(x) ?? [])
+
+/**
+ * Recursively query **direct** variable references into key-value map.
+ */
+export const getDeclDepsArray = (links: Declaration[]): DeclDeps => {
+	const graph: DeclDeps = new Map()
+
+	let current = links
+	while (current.length > 0) {
+		current = current.flatMap((node) => {
+			const references = getReferencedDeclsArray(node)
+			graph.set(node, references)
+			return references
+		})
+	}
+	return graph
+}
+
+Deno.bench(`array`, (b) => {
+	const project = inMemoryProject()
+	const files = withSrc(project)(exampleSrc)
+	const decls = Stream.fromObjectValues(files).flatMap(getAllDecls).toArray()
+
+	b.start()
+	const _declDeps = getDeclDepsArray(decls)
+	b.end()
+})
+
+Deno.bench(`stream`, { baseline: true }, (b) => {
+	const project = inMemoryProject()
+	const files = withSrc(project)(exampleSrc)
+	const decls = Stream.fromObjectValues(files).flatMap(getAllDecls).toArray()
+
+	b.start()
+	const _declDeps = getDeclDeps(decls)
+	b.end()
+})