1+ function function_name_from_code ( code : string ) {
2+ let name = code . split ( '(' ) [ 0 ] . trim ( ) ;
3+ name = name
4+ . replace ( / : : / g, '_' )
5+ . replace ( / / g, '_' )
6+ . replace ( / \* / g, '' ) ;
7+ return name ;
8+ }
9+
10+ function tree_sitter_get_functions ( captures : QueryCapture [ ] , code : string ) {
11+ return captures . map ( ( { name, node } ) => ( {
12+ code : node . text ,
13+ start : node . startIndex ,
14+ end : node . endIndex ,
15+ name : function_name_from_code ( node . text )
16+ } ) ) ;
17+ }
18+
19+ function output_name_from_file ( inputFile : WorkspaceFile , name : string ) {
20+ let outputFile = "slice_" + path . basename ( inputFile . filename ) + name + ".cpp" ;
21+ return path . join ( "code_slices" , outputFile ) ;
22+ }
23+
24+ export async function splitFunctions ( inputFile : WorkspaceFile ) {
25+ const { captures : functions } = await parsers . code (
26+ inputFile ,
27+ `(function_definition) @function`
28+ ) ;
29+ return tree_sitter_get_functions ( functions , inputFile . content ) ;
30+ }
31+
32+
33+ export async function saveFunctions ( inputFile : WorkspaceFile , funs : { code : string , name : string } [ ] ) {
34+ for ( const fun of funs ) {
35+ let name = function_name_from_code ( fun . code ) ;
36+ console . log ( name ) ;
37+ let outputFile = output_name_from_file ( inputFile , name ) ;
38+ await workspace . writeText ( outputFile , `//Extracted ${ name } ${ inputFile . filename } ${ fun . code } ) ;
39+ }
40+ }
41+
42+ // given a source file <src>
43+ // list files in code_slice directory based on that name with extension opt
44+ // replace functions in src by the corresponding ones in the opt files.
45+ // Save into <dst>
46+
47+ async function parseOptFunctions ( inputFile : WorkspaceFile ) {
48+ const modifiedFunctions = "slice_" + path . basename ( inputFile . filename ) + "*.opt" ;
49+ console . log ( modifiedFunctions ) ;
50+ const directory_path = path . join ( "code_slices" , modifiedFunctions ) ;
51+ const files = await workspace . findFiles ( directory_path ) ;
52+ let modifiedFunctionsList = [ ] ;
53+ for ( const file of files ) {
54+ console . log ( file . filename ) ;
55+ const code = file . content . match ( / ` ` ` c p p ( [ \s \S ] * ?) ` ` ` / ) ;
56+ if ( ! code ) {
57+ continue ;
58+ }
59+ const modifiedFunction = code [ 1 ] ;
60+ modifiedFunctionsList . push ( modifiedFunction ) ;
61+ }
62+ return modifiedFunctionsList ;
63+ }
64+
65+ import * as fs from 'fs' ;
66+
67+ export async function mergeModifiedFunction ( code :string , funs : { code : string , name : string } [ ] , new_code : string ) {
68+ let name = function_name_from_code ( new_code ) ;
69+ let fun = funs . find ( f => f . name === name ) ;
70+ if ( fun ) {
71+ console . log ( "Updated function: " + name ) ;
72+ code = code . replace ( fun . code , new_code ) ;
73+ }
74+ return code ;
75+ }
76+
77+ async function canCompileCode ( inputFile : WorkspaceFile , code : string ) {
78+
79+ // move input file to a temp file
80+ // move code to the inputFile.filename
81+ // invoke ninja in the build directory: ninja -b build
82+ // move the temp file back to the original file
83+ // return true iff it succeeded
84+
85+ let tempFile = inputFile . filename + ".tmp" ;
86+ let original_content = inputFile . content ;
87+ await workspace . writeText ( tempFile , inputFile . content ) ;
88+ await workspace . writeText ( inputFile . filename , code ) ;
89+ let result = await host . exec ( `cmd /k "C:\Program\ Files/Microsoft\ Visual\ Studio/2022/Enterprise/Common7/Tools/VsDevCmd.bat" -arch=x64 & ninja` , { cwd : "build" } ) ;
90+
91+ // await fs.delete(tempFile);
92+ if ( result . exitCode !== 0 ) {
93+ await workspace . writeText ( inputFile . filename , original_content ) ;
94+ console . log ( result . stderr ) ;
95+ return false ;
96+ }
97+ return true ;
98+ }
99+
100+ export async function mergeCompileFunction ( inputFile : WorkspaceFile , code : string , funs : { code : string , name : string } [ ] , new_code_input : string ) {
101+ let match_new_code = new_code_input . match ( / ` ` ` c p p ( [ \s \S ] * ?) ` ` ` / ) ;
102+ if ( ! match_new_code ) {
103+ console . log ( "Invalid new code" ) ;
104+ return code ;
105+ }
106+ let new_code = match_new_code [ 1 ] ;
107+
108+ let name = function_name_from_code ( new_code ) ;
109+ let fun = funs . find ( f => f . name == name ) ;
110+
111+ if ( ! fun ) {
112+ console . log ( `Function name '${ name } ) ;
113+ for ( const fun of funs )
114+ console . log ( "'" + fun . name + "'" ) ;
115+ return code ;
116+ }
117+ console . log ( "Updated function: " + name ) ;
118+ let modified_code = code . replace ( fun . code , new_code ) ;
119+ if ( code == modified_code ) {
120+ console . log ( "No change in function: " + name ) ;
121+ return code ;
122+ }
123+ let canCompile = await canCompileCode ( inputFile , modified_code ) ;
124+ console . log ( "Can compile: " + canCompile ) ;
125+ if ( canCompile )
126+ return modified_code ;
127+ return code ;
128+ }
129+
130+ export async function mergeFunctionsFromList ( inputCode : string , funs : { code : string , name : string } [ ] , modifiedFunctionList : string [ ] ) {
131+ let code = inputCode ;
132+ for ( const new_code of modifiedFunctionList )
133+ code = await mergeModifiedFunction ( code , funs , new_code ) ;
134+ return code ;
135+ }
136+
137+ export async function mergeFunctions ( inputFile : WorkspaceFile ) {
138+ let funs = await splitFunctions ( inputFile ) ;
139+ let modifiedFunctionList = await parseOptFunctions ( inputFile ) ;
140+ return mergeFunctionsFromList ( inputFile . content , funs , modifiedFunctionList ) ;
141+ }
142+
143+ export async function invokeLLMOpt ( code : string ) {
144+ const answer = await runPrompt (
145+ ( _ ) => {
146+ _ . def ( "CODE" , code ) ;
147+ _ . $ `You are a highly experienced compiler engineer with over 20 years of expertise,
148+ specializing in C and C++ programming. Your deep knowledge of best coding practices
149+ and software engineering principles enables you to produce robust, efficient, and
150+ maintainable code in any scenario.
151+
152+ Please modify the original code in <CODE> to ensure that it uses best practices for optimal code execution.' `
153+ } , {
154+ system : [ ] ,
155+ systemSafety : false
156+ }
157+ ) ;
158+ return answer . text ;
159+ }
0 commit comments