Add function parsing (#3)

* added source code function parsing functionality - still needs to be tested

* added parameter to slowest function parser, can now test multiple prompt formatter functionality

* cleaned up function parsing formatter

* now importing function grabber function instead of starting subprocess

---------

Co-authored-by: Mazin Karjikar <mazin@pssg-mordor.umiacs.umd.edu>

requirements.txt

@@ -3,4 +3,5 @@ hatchet==1.4.0
 google-generativeai==0.7.2
 openai==1.37.0
 tiktoken==0.7.0
-llama-cpp-python==0.2.90
+llama-cpp-python==0.2.90
+clang==17.0.6

src/function_grabber.py

@@ -0,0 +1,100 @@
+""" This script extracts the function that contains a specified line number in a C++ file.
+"""
+from argparse import ArgumentParser
+import os
+from typing import List, Tuple
+import clang.cindex
+
+
+def get_functions_at_lines(fpath: os.PathLike, lines: List[int], clang_path: os.PathLike = None) -> List[Tuple[str, Tuple[int, int]]]:
+    """ Find all the functions that contain the specified lines in a file.
+    """
+    functions = []
+    for line in lines:
+
+        # check if we've already found this line
+        if any(start <= line <= end for _, (start, end) in functions):
+            continue
+
+        function_body, range = get_function_at_line(fpath, line, clang_path=clang_path)
+        if function_body:
+            functions.append((function_body, range))
+    return functions
+
+
+def remove_macros(filename: str, line_numbers: List[int]) -> List[int]:
+    """ Remove all macros from a file. Lines is a list of numbers that you would like to have mapped
+        to their new line numbers after the macros are removed.
+    """
+    with open(filename, 'r') as f:
+        lines = f.readlines()
+
+    new_line_numbers = []
+    new_lines = []
+    num_removed = 0
+    for i, line in enumerate(lines):
+        if line.startswith('#'):
+            num_removed += 1
+        else:
+            new_lines.append(line)
+
+        if i in line_numbers:
+            new_line_numbers.append(i - num_removed)
+    
+    with open(filename, 'w') as f:
+        f.write(''.join(new_lines))
+
+    return new_line_numbers
+
+
+def get_function_at_line(filename, line_number, clang_path=None):
+    if clang_path and not clang.cindex.Config.loaded:
+        clang.cindex.Config.set_library_file(clang_path)
+    index = clang.cindex.Index.create()
+
+    try:
+        translation_unit = index.parse(filename)
+    except clang.cindex.TranslationUnitLoadError:
+        return None, None
+
+    def find_function(node, line_number):
+        # Check if node is function-like and contains the line number
+        if node.kind == clang.cindex.CursorKind.FUNCTION_DECL or node.kind == clang.cindex.CursorKind.CXX_METHOD:
+            start_line = node.extent.start.line
+            end_line = node.extent.end.line
+
+            #print(f"Checking function {node.spelling} at lines {start_line} - {end_line}")
+
+            if start_line <= line_number <= end_line:
+                return node
+
+        for child in node.get_children():
+            result = find_function(child, line_number)
+            if result:
+                return result
+        return None
+
+    # Start from the root node (translation unit) and find the function
+    function_node = find_function(translation_unit.cursor, line_number)
+
+    if function_node:
+        start_line = function_node.extent.start.line
+        end_line = function_node.extent.end.line
+        with open(filename, 'r') as f:
+            lines = f.readlines()
+        return ''.join(lines[start_line - 1:end_line]), (start_line, end_line)
+    else:
+        return None, None
+
+
+if __name__ == "__main__":
+    parser = ArgumentParser(description="Extract the function that contains a specified line number in a C++ file.")
+    parser.add_argument("filename", help="The C++ file to analyze")
+    parser.add_argument("line_number", type=int, help="The line number to search for")
+    parser.add_argument("--clang_path", help="Path to libclang.so if necessary")
+    args = parser.parse_args()
+
+    result, rnge = get_function_at_line(args.filename, args.line_number, clang_path=args.clang_path)
+    if result is None:
+        result = f"No function found at line {args.line_number}"
+    print(result, rnge)

src/rag.py

@@ -8,6 +8,7 @@ from typing import Optional, List, Mapping
 from os import PathLike
 from os.path import basename
 import random
+from function_grabber import get_function_at_line
 
 from profiles import Profile
 
@@ -63,8 +64,9 @@ class BasicPromptFormatter(PerfGuruPromptFormatter):
     
 class SlowestFunctionPromptFormatter(PerfGuruPromptFormatter):
 
-    def __init__(self):
+    def __init__(self, k):
         super().__init__("slowest_function")
+        self.k = k
 
     def format_prompt(self, prompt: str, code_paths: List[PathLike], profile_path: Optional[PathLike] = None, profile_type: Optional[str] = None, error_fn: Optional[callable] = None) -> str:
         if not code_paths:
@@ -83,11 +85,10 @@ class SlowestFunctionPromptFormatter(PerfGuruPromptFormatter):
                 if error_fn:
                     error_fn("Profile type must be provided if a profile file is provided.")
                 return None
-            k = 1
             profile = self._read_profile(profile_path, profile_type)
-            slowest = profile.gf.dataframe.nlargest(k, 'time')
-            function_names = [slowest['name'].values[i] for i in range(k) if i < len(slowest['name'].values)]
-            execution_times = [slowest['time'].values[i] for i in range(k) if i < len(slowest['name'].values)]
+            slowest = profile.gf.dataframe.nlargest(self.k, 'time')
+            function_names = [slowest['name'].values[i] for i in range(self.k) if i < len(slowest['name'].values)]
+            execution_times = [slowest['time'].values[i] for i in range(self.k) if i < len(slowest['name'].values)]
             # print(profile_content)
             hot_path = profile.gf.hot_path()
             hot_path_functions = []
@@ -95,7 +96,7 @@ class SlowestFunctionPromptFormatter(PerfGuruPromptFormatter):
             for node in hot_path:
                 if "name" in node.frame.attrs: 
                     hot_path_functions.append(node.frame["name"])
-            hot_path_functions = hot_path_functions[:k]
+            hot_path_functions = hot_path_functions[:self.k]
 
             profile_content = (f"The slowest functions are {function_names} and they took {execution_times} seconds, respectively." + 
             f" Also, these functions were in the hot path: {hot_path_functions}.")
@@ -106,9 +107,63 @@ class SlowestFunctionPromptFormatter(PerfGuruPromptFormatter):
 
         return f"Code:\n{concatenated_code}\n\n{profile_type} Profile:\n{profile_content}\n\n{prompt}"
 
-AVAILABLE_FORMATTERS = [SlowestFunctionPromptFormatter()]
-# AVAILABLE_FORMATTERS.append(BasicPromptFormatter())
+class SlowestFunctionParsedPromptFormatter(PerfGuruPromptFormatter):
 
+    def __init__(self):
+        super().__init__("slowest_function_parsed")
+
+    def format_prompt(self, prompt: str, code_paths: List[PathLike], profile_path: Optional[PathLike] = None, profile_type: Optional[str] = None, error_fn: Optional[callable] = None) -> str:
+        if not code_paths:
+            if error_fn:
+                error_fn("No code files provided. At least one code file must be provided.")
+            return None
+
+        concatenated_code = ""
+        profile_content = ""
+
+        if profile_path:
+            if not profile_type:
+                if error_fn:
+                    error_fn("Profile type must be provided if a profile file is provided.")
+                return None
+            
+            k = 1
+            profile = self._read_profile(profile_path, profile_type)
+            slowest = profile.gf.dataframe.nlargest(k, 'time')
+            function_name = slowest['name'].values[0] if len(slowest['name'].values) > 0 else None
+            line_number = slowest['line'].values[0] if len(slowest['line'].values) > 0 else None
+            code = None
+            
+            if line_number:
+                filename = ""
+                code_file_contents = self._read_code_files(code_paths)
+                for code_path, content in code_file_contents.items():
+                    filename = basename(code_path)
+                    code, _ = get_function_at_line(filename, str(line_number))
+                    if code:
+                        break
+                
+                if code:
+                    concatenated_code = f"{fname}:\n{code}\n\n"
+                    print("Only function code:", concatenated_code)
+
+                profile_content = (f"The slowest function is {function_name}.")
+                print(profile_content)
+        
+        if concatenated_code == "":
+            code_file_contents = self._read_code_files(code_paths)
+            for code_path, content in code_file_contents.items():
+                fname = basename(code_path)
+                concatenated_code += f"{fname}:\n{content}\n\n"
+
+        return f"Code:\n{concatenated_code}\n\n{profile_type} Profile:\n{profile_content}\n\n{prompt}"
+
+AVAILABLE_FORMATTERS = []
+AVAILABLE_FORMATTERS.append(SlowestFunctionPromptFormatter(k=1))
+AVAILABLE_FORMATTERS.append(SlowestFunctionPromptFormatter(k=5))
+AVAILABLE_FORMATTERS.append(SlowestFunctionPromptFormatter(k=10))
+# AVAILABLE_FORMATTERS.append(BasicPromptFormatter())
+AVAILABLE_FORMATTERS.append(SlowestFunctionParsedPromptFormatter())
 
 def select_random_formatter() -> PerfGuruPromptFormatter:
     return random.choice(AVAILABLE_FORMATTERS)