From c25f826741a747b216d0cacb8f6a7c3207c574d8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Daniel=20Kro=CC=88ni?= <daniel.kroeni@fhnw.ch>
Date: Mon, 27 Sep 2021 14:06:24 +0200
Subject: [PATCH] Allow future function label assignment.

---
 src/starkware/cairo/lang/compiler/ast/expr.py |  2 +-
 .../preprocessor/compound_expressions.py      |  7 ++++
 .../compiler/preprocessor/preprocessor.py     |  2 +-
 .../preprocessor/preprocessor_test.py         | 32 +++++++++++++++++++
 4 files changed, 41 insertions(+), 2 deletions(-)

diff --git a/src/starkware/cairo/lang/compiler/ast/expr.py b/src/starkware/cairo/lang/compiler/ast/expr.py
index 7f75378a..31020eef 100644
--- a/src/starkware/cairo/lang/compiler/ast/expr.py
+++ b/src/starkware/cairo/lang/compiler/ast/expr.py
@@ -406,7 +406,7 @@ def to_expr_str(self):
         return self.identifier.to_expr_str()
 
     @property
-    def locaion(self):
+    def location(self):
         return self.identifier.location
 
     def get_children(self) -> Sequence[Optional[AstNode]]:
diff --git a/src/starkware/cairo/lang/compiler/preprocessor/compound_expressions.py b/src/starkware/cairo/lang/compiler/preprocessor/compound_expressions.py
index af8f217f..182ad330 100644
--- a/src/starkware/cairo/lang/compiler/preprocessor/compound_expressions.py
+++ b/src/starkware/cairo/lang/compiler/preprocessor/compound_expressions.py
@@ -12,6 +12,7 @@
     ExprDeref,
     Expression,
     ExprHint,
+    ExprFutureLabel,
     ExprIdentifier,
     ExprNeg,
     ExprOperator,
@@ -161,6 +162,12 @@ def rewrite_ExprDeref(self, expr: ExprDeref, sim: SimplicityLevel):
         )
         return expr if sim is SimplicityLevel.OPERATION else self.wrap(expr)
 
+    def rewrite_ExprFutureLabel(self, expr: ExprFutureLabel, sim: SimplicityLevel):
+        # Treat this as a constant.
+        if sim in [SimplicityLevel.DEREF_CONST, SimplicityLevel.OPERATION]:
+            return expr
+        return self.wrap(expr)
+
     def rewrite_ExprHint(self, expr: ExprHint, sim: SimplicityLevel):
         return self.wrap(expr)
 
diff --git a/src/starkware/cairo/lang/compiler/preprocessor/preprocessor.py b/src/starkware/cairo/lang/compiler/preprocessor/preprocessor.py
index c02a26e1..10747b46 100644
--- a/src/starkware/cairo/lang/compiler/preprocessor/preprocessor.py
+++ b/src/starkware/cairo/lang/compiler/preprocessor/preprocessor.py
@@ -1899,7 +1899,7 @@ def get_variable(self, var: ExprIdentifier):
         assert identifier_definition is not None
 
         if isinstance(identifier_definition, FutureIdentifierDefinition):
-            if identifier_definition.identifier_type == LabelDefinition:
+            if identifier_definition.identifier_type in [LabelDefinition, FunctionDefinition]:
                 # Allow future label assignment.
                 return ExprFutureLabel(identifier=var)
             raise PreprocessorError(
diff --git a/src/starkware/cairo/lang/compiler/preprocessor/preprocessor_test.py b/src/starkware/cairo/lang/compiler/preprocessor/preprocessor_test.py
index 63275ff1..b06006ea 100644
--- a/src/starkware/cairo/lang/compiler/preprocessor/preprocessor_test.py
+++ b/src/starkware/cairo/lang/compiler/preprocessor/preprocessor_test.py
@@ -5,6 +5,7 @@
 from starkware.cairo.lang.compiler.error_handling import LocationError
 from starkware.cairo.lang.compiler.identifier_definition import (
     ConstDefinition,
+    FunctionDefinition,
     LabelDefinition,
     ReferenceDefinition,
 )
@@ -167,6 +168,37 @@ def test_assign_future_label():
     )
 
 
+def test_assign_future_function_label():
+    code = """\
+g(f)
+g((f + 1) * 2 + 3)
+
+func f() -> ():
+    ret
+end
+func g(x: felt) -> ():
+    ret
+end
+"""
+    program = preprocess_str(code=code, prime=PRIME)
+    f_definition = program.identifiers.get_by_full_name(ScopedName.from_string("test_scope.f"))
+    assert isinstance(f_definition, FunctionDefinition)
+    assert (
+        program.format()
+        == f"""\
+[ap] = {f_definition.pc}; ap++
+call rel 13
+[ap] = {f_definition.pc}; ap++
+[ap] = [ap + (-1)] + 1; ap++
+[ap] = [ap + (-1)] * 2; ap++
+[ap] = [ap + (-1)] + 3; ap++
+call rel 3
+ret
+ret
+"""
+    )
+
+
 def test_temporary_variable():
     code = """\
 struct T: