Install codealign

Dockerfile

@@ -40,6 +40,8 @@ RUN --mount=type=secret,id=GH_TOKEN,mode=0444,required=true \
 
 RUN conda env create -f idioms/env.yml 
 
+RUN git clone https://github.com/squaresLab/codealign.git && pip install codealign/.
+
 ENV PATH="/home/user/miniconda3/envs/idioms/bin:${PATH}"
 ENV PYTHONPATH="${PYTHONPATH}:${HOME}/app/idioms"
 

main.py

@@ -30,8 +30,14 @@ def canonicalize_output(decompiled: str):
 from prepare import (
     parser, 
     FileTypeMapping, 
+    PreprocessedFunction, 
     has_error, 
+    TypeNotDefinedError, 
+    TypeNotFoundError, 
+    UnsupportedFeatureError,
+    get_all_user_defined_types
 )
+from evaluator import add_placeholders_for_nonvariable_types
 from tree_sitter import Node
 
 class InvalidPredictionError(Exception):
@@ -70,6 +76,26 @@ def clean_output(raw_prediction: str) -> str:
     fn_node = fn_nodes[0]
     assert fn_node.text is not None # To make mypy happy
 
+    # We don't track or measure the accuracy of types not associated with variables (though this is taken into account when computing an alignment).
+    # However PreprocessedFunction will throw an exception if it encounters an unrecognized type in a typecast.
+    # To prevent this, we add generic placholder types.
+    predicted_body = fn_node.child_by_field_name("body")
+    assert predicted_body is not None
+    try: # Add placholders in a try-except in case a non-variable type is unparsable.
+        add_placeholders_for_nonvariable_types(predicted_body, types)
+        predicted_fn = PreprocessedFunction(fn_node, types)
+    except (TypeNotFoundError, TypeNotDefinedError, UnsupportedFeatureError, AssertionError):
+        # We need the variable-type mapping that PreprocessedFunction provides to standardize the field names
+        # for the alignment. We don't return here, however, because PreprocessedFunction is very strict about
+        # types; a function may be syntatically, or arguably, semantically valid but fail here. Thus, we
+        # assign predicted_fn to None here and then do the early return later.
+        raise InvalidPredictionError("Invalid type used in predicted function.")
+
+    predicted_code = fn_node.text.decode()
+    predicted_udts = get_all_user_defined_types(predicted_fn)
+
+    return predicted_code + "\n" + "\n".join(udt.declaration("") for udt in predicted_udts)
+
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline #, BitsAndBytesConfig
 from peft import PeftModel