From 45f3ea3bf49c12913a1562f6b47737637fa139cb Mon Sep 17 00:00:00 2001 From: Nikolaj Bjorner Date: Thu, 13 Feb 2025 09:56:49 -0800 Subject: [PATCH] add treesitter functionality Signed-off-by: Nikolaj Bjorner --- genaisrc/treesitter.genai.mts | 33 +++++++++++++++++++++++++++++---- 1 file changed, 29 insertions(+), 4 deletions(-) diff --git a/genaisrc/treesitter.genai.mts b/genaisrc/treesitter.genai.mts index b333ede2b..5d8226b66 100644 --- a/genaisrc/treesitter.genai.mts +++ b/genaisrc/treesitter.genai.mts @@ -3,12 +3,37 @@ script({ files: "src/muz/spacer/spacer_qe_project.cpp" }) +// return function names and source code of the functions +function get_functions(captures : QueryCapture[], code : string) { + return captures.map(({ name, node }) => ({ + code : node.text + })) +} const inputFile = env.files[0]; -const { captures } = await parsers.code(inputFile); - -// pretty-print tree sitter parse tree of captures: -console.log(JSON.stringify(captures, null,2)) +const { captures: functions } = await parsers.code( + inputFile, + `(function_definition) @function` +); + + +let funs = get_functions(functions, inputFile.content); + +for (const fun of funs) { + // todo put files in a different directory + let name = fun.code.split('(')[0].trim(); + name = name + .replace(/::/g, '_') + .replace(/ /g, '_'); + let outputFile = path.basename(inputFile.filename) + .replace(/\.cpp$/, `.${name}.cpp`) + .replace(/\.h$/, `.${name}.h`); + outputFile = "slice_" + outputFile; + outputFile = path.join("code_slices", outputFile); + + await workspace.writeText(outputFile, `//Extracted ${name} in ${inputFile.filename}\n${fun.code}\n\n`); +} +