Skip to content

Commit

Permalink
add treesitter functionality
Browse files Browse the repository at this point in the history
Signed-off-by: Nikolaj Bjorner <nbjorner@microsoft.com>
  • Loading branch information
NikolajBjorner committed Feb 13, 2025
1 parent bedc95c commit 45f3ea3
Showing 1 changed file with 28 additions and 3 deletions.
31 changes: 28 additions & 3 deletions genaisrc/treesitter.genai.mts
Original file line number Diff line number Diff line change
Expand Up @@ -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:

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`);
}

console.log(JSON.stringify(captures, null,2))

0 comments on commit 45f3ea3

Please sign in to comment.