Optimize recursive autogen

HerodotusDev · Sep 17, 2024 · 7fe65a4 · 7fe65a4
1 parent d6ddabc
commit 7fe65a4
Show file tree

Hide file tree

Showing 2 changed files with 462 additions and 945 deletions.
diff --git a/src/air/layouts/_generator/main.py b/src/air/layouts/_generator/main.py
@@ -4,7 +4,9 @@
 from starkware.cairo.lang.compiler.ast.expr_func_call import *
 import requests
 
-OPTIMIZE_VALUE_ARRAY = False
+settings = {
+    'OPTIMIZE_VALUE_ARRAY': False,
+}
 
 global array_read_offset
 global constants
@@ -74,10 +76,12 @@ class LineTypeTotalSum(LineType):
 
 
 def optimize(lines: list[tuple[str, LineType]]) -> str:
+    global settings
     acc = ''
     acc_var_pops = []
     var_pops_varname = None
     latest_value_comment = None
+    total_sum_defined = False
     acc_values = []
     acc_total_sum = []
     for (line, line_type), (_, next_line_type) in zip(lines, lines[1:] + [(None, None)]):
@@ -90,16 +94,16 @@ def optimize(lines: list[tuple[str, LineType]]) -> str:
                 vars_len = str(len(acc_var_pops))
                 acc += F"let [{vars_arr}] = (*{var_pops_varname}.multi_pop_front::<{vars_len}>().unwrap()).unbox();\n\t"
                 acc_var_pops = []
-        elif OPTIMIZE_VALUE_ARRAY and isinstance(line_type, LineTypeValueCalc):
+        elif settings['OPTIMIZE_VALUE_ARRAY'] and isinstance(line_type, LineTypeValueCalc):
             line_type.comment = latest_value_comment
             acc_values.append(line_type)
-        elif OPTIMIZE_VALUE_ARRAY and isinstance(line_type, LineTypeTotalSum):
+        elif settings['OPTIMIZE_VALUE_ARRAY'] and isinstance(line_type, LineTypeTotalSum):
             acc_total_sum.append(line)
-        elif OPTIMIZE_VALUE_ARRAY and isinstance(line_type, LineTypeComment) and isinstance(next_line_type, LineTypeValueCalc):
+        elif settings['OPTIMIZE_VALUE_ARRAY'] and isinstance(line_type, LineTypeComment) and isinstance(next_line_type, LineTypeValueCalc):
             # comments before value calculations are moved to the value calculation line
             pass
         else:
-            if OPTIMIZE_VALUE_ARRAY and not isinstance(line_type, LineTypeEmpty) and acc_values:
+            if settings['OPTIMIZE_VALUE_ARRAY'] and not isinstance(line_type, LineTypeEmpty) and acc_values:
                 total_sum_line = None
                 # assert that all total sum calculations are the same
                 for x,y in zip(acc_total_sum, acc_total_sum[1:]):
@@ -112,7 +116,12 @@ def optimize(lines: list[tuple[str, LineType]]) -> str:
                 acc += 'let values = [\n\t\t'
                 acc += '\n\t\t'.join([f"{v.expr},{' '+v.comment.rstrip('\n\t') if v.comment is not None else ''}" for v in acc_values])
                 acc += '\n\t].span();\n\t\n\t'
-                acc += 'for value in values {\n\t\t' + total_sum_line + '};\n\t'
+                if not total_sum_defined:
+                    acc += 'let mut total_sum = 0;\n\t'
+                    total_sum_defined = True
+                acc += 'for value in values {\n\t\t'
+                acc += total_sum_line.replace('let total_sum = total_sum + ', 'total_sum += ').replace('value', '*value')
+                acc += '};\n\t'
                 acc_values = []
                 acc_total_sum = []
 
@@ -282,7 +291,11 @@ def remove_parenthesis(arg):
     return ''
 
 
-def handle_github_file(url, output_file, layout):
+def handle_github_file(url, output_file, layout, settings_override={}):
+    global settings
+    old_settings = settings.copy()
+    settings = {**settings, **settings_override}
+
     global array_read_offset
     response = requests.get(url)
     if response.status_code != 200:
@@ -309,19 +322,24 @@ def handle_github_file(url, output_file, layout):
             case CodeElementConst(identifier=ExprIdentifier(name=name), expr=expr):
                 constants[name] = eval(expr)
 
+    settings = old_settings
     with open(output_file, 'w') as f:
         f.write(imports(layout) + '\n'.join(functions_result.values()))
 
 
 def main():
     # layouts = ('recursive', 'recursive_with_poseidon', 'small', 'dex', 'starknet', 'starknet_with_keccak')
-    layouts = ('recursive', 'recursive_with_poseidon', 'small', 'dex')
+    layouts = ('recursive', )
+    optimizations = {
+        'OPTIMIZE_VALUE_ARRAY': {'recursive'}
+    }
 
     for layout in layouts:
         handle_github_file(
             f"https://raw.githubusercontent.com/starkware-libs/cairo-lang/master/src/starkware/cairo/stark_verifier/air/layouts/{layout}/autogenerated.cairo",
             f"../{layout}/autogenerated.cairo",
-            layout
+            layout,
+            {'OPTIMIZE_VALUE_ARRAY': layout in optimizations['OPTIMIZE_VALUE_ARRAY']}
         )