diff --git a/ts/packages/agents/spelunker/src/typescriptChunker.ts b/ts/packages/agents/spelunker/src/typescriptChunker.ts index cdff0cd75..07de0e8d1 100644 --- a/ts/packages/agents/spelunker/src/typescriptChunker.ts +++ b/ts/packages/agents/spelunker/src/typescriptChunker.ts @@ -45,61 +45,51 @@ export async function chunkifyTypeScriptFiles( }; const chunks: Chunk[] = [rootChunk]; const sourceFile: ts.SourceFile = await tsCode.loadSourceFile(fileName); - - // TODO: Also do nested functions, and classes, and interfaces, and modules. - // TODO: For nested things, remove their text from the parent. - function getFunctionsAndClasses(): ( - | ts.FunctionDeclaration - | ts.ClassDeclaration - )[] { - return tsCode.getStatements( - sourceFile, - (s) => ts.isFunctionDeclaration(s) || ts.isClassDeclaration(s), - ); - } - - const things = getFunctionsAndClasses(); - for (const thing of things) { - const treeName = ts.SyntaxKind[thing.kind]; - const codeName = tsCode.getStatementName(thing) ?? ""; - // console.log(` ${treeName}: ${codeName}`); - try { - // console.log( - // "--------------------------------------------------------", - // ); - // console.log(`Name: ${thing.name?.escapedText}`); - // console.log( - // `Parameters: ${thing.parameters.map((p) => p.name?.getFullText(sourceFile))}`, - // ); - // console.log(`Return type: ${thing.type?.getText(sourceFile)}`); - - const chunk: Chunk = { - chunkId: generate_id(), - treeName, - codeName, - blobs: makeBlobs( - sourceFile, - thing.getFullStart(), - thing.getEnd(), - ), - parentId: rootChunk.chunkId, - children: [], - fileName, - }; - chunks.push(chunk); - } catch (e: any) { - results.push({ - error: `${thing.name?.escapedText}: ${e.message}`, - filename: fileName, - }); - } - } - // console.log("========================================================"); + chunks.push(...recursivelyChunkify(sourceFile, rootChunk)); const chunkedFile: ChunkedFile = { fileName, chunks, }; results.push(chunkedFile); + + function recursivelyChunkify( + parentNode: ts.Node, + parentChunk: Chunk, + ): Chunk[] { + const chunks: Chunk[] = []; + for (const childNode of parentNode.getChildren(sourceFile)) { + if ( + ts.isInterfaceDeclaration(childNode) || + ts.isTypeAliasDeclaration(childNode) || + ts.isFunctionDeclaration(childNode) || + ts.isClassDeclaration(childNode) + ) { + // console.log( + // ts.SyntaxKind[childNode.kind], + // tsCode.getStatementName(childNode), + // ); + const chunk: Chunk = { + chunkId: generate_id(), + treeName: ts.SyntaxKind[childNode.kind], + codeName: tsCode.getStatementName(childNode) ?? "", + blobs: makeBlobs( + sourceFile, + childNode.getFullStart(), + childNode.getEnd(), + ), + parentId: parentChunk.chunkId, + children: [], + fileName, + }; + // TODO: Remove chunk.blobs from parentChunk.blobs. + chunks.push(chunk); + recursivelyChunkify(childNode, chunk); + } else { + recursivelyChunkify(childNode, parentChunk); + } + } + return chunks; + } } return results; diff --git a/ts/packages/codeProcessor/src/tsCode.ts b/ts/packages/codeProcessor/src/tsCode.ts index 912a0e322..8b0f0f144 100644 --- a/ts/packages/codeProcessor/src/tsCode.ts +++ b/ts/packages/codeProcessor/src/tsCode.ts @@ -87,6 +87,12 @@ export function getStatementName(statement: ts.Statement): string | undefined { ) { return statement.name?.text; } + if ( + ts.isInterfaceDeclaration(statement) || + ts.isTypeAliasDeclaration(statement) + ) { + return statement.name?.escapedText.toString(); + } return undefined; }