more tree functions

app.py

@@ -6,7 +6,7 @@ import numpy as np
 import torch
 from threading import Thread
 
-from utils.tree_utils import parse_functions, get_docstrings, grab_before_comments, line_chr2char
+from utils.tree_utils import parse_functions, get_docstrings, grab_before_comments, line_chr2char, node_str_idx
 from utils.html_utils import make_iframe, construct_embed
 PIPE = None
 
@@ -235,8 +235,9 @@ def alter_body(old_code, func_id, funcs_list: list, prompt="", temperature=0.2,
     func_start_idx = line_chr2char(old_code, func_node.start_point[0], func_node.start_point[1])
     identifier_str = func_node.child_by_field_name("type").text.decode() + " " + func_node.child_by_field_name("declarator").text.decode() #func_start_idx:body_start_idx?
     body_node = func_node.child_by_field_name("body")
-    body_start_idx = line_chr2char(old_code, body_node.start_point[0], body_node.start_point[1])
-    body_end_idx = line_chr2char(old_code, body_node.end_point[0], body_node.end_point[1])
+    body_start_idx, body_end_idx = node_str_idx(body_node)
+    # body_start_idx = line_chr2char(old_code, body_node.start_point[0], body_node.start_point[1])
+    # body_end_idx = line_chr2char(old_code, body_node.end_point[0], body_node.end_point[1])
     print(f"{old_code[body_start_idx:body_end_idx]=}")
     model_context = identifier_str # base case
     # add any comments at the beginning of the function to the model_context
@@ -267,7 +268,7 @@ def alter_body(old_code, func_id, funcs_list: list, prompt="", temperature=0.2,
         altered_code = old_code[:func_start_idx] + model_context + generation + "//the generation didn't complete the function!\n" + old_code[body_end_idx:] #needs a newline to break out of the comment.
         return altered_code, pipeline
         # raise gr.Error(f"didn't generate a full function: {generation!r}]")
-    print(f"{first_gened_func=}")
+    # print(f"{first_gened_func=}")
     generated_body = first_gened_func.child_by_field_name("body").text.decode()
     print(f"{generated_body=}")
     altered_code = old_code[:func_start_idx] + identifier_str + generated_body + old_code[body_end_idx:]

utils/tree_utils.py

@@ -6,6 +6,38 @@ GLSL_LANGUAGE = Language('./build/my-languages.so', 'glsl')
 parser = Parser()
 parser.set_language(GLSL_LANGUAGE)
 
+def replace_function(old_func_node, new_func_node):
+    """
+    replaces the old function node with the new function node
+    """
+    tree = give_tree(old_func_node)
+    old_func_start, old_func_end = node_str_idx(old_func_node)
+    # new_func_start, new_func_end = node_str_idx(new_func_node)
+    new_code = tree.text[:old_func_start] + new_func_node.text + tree.text[old_func_end:]
+    return new_code
+
+def get_root(node):
+    """
+    returns the root node of a node
+    """
+    while node.parent != None:
+        node = node.parent
+    return node.parent #still need to call parent here
+
+def node_str_idx(node):
+    """
+    returns the character index of start and end of a node
+    """
+    whole_text = get_root(node).text.decode()
+    start_idx = line_chr2char(whole_text, node.start_point[0], node.start_point[1])
+    end_idx = line_chr2char(whole_text, node.end_point[0], node.end_point[1])
+    return start_idx, end_idx
+
+def give_tree(func_node):
+    """
+    return the tree where this function node is in
+    """
+    return parser.parse(func_node.parent.text) #really no better way?
 
 def parse_functions(in_code):
     """