From e4c7956fde3d9fb53ceecaa9dcd0fb6599858072 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ibrahim=20Numanagic=CC=81?= <ibrahimpasa@gmail.com>
Date: Tue, 25 Feb 2025 19:34:44 -0800
Subject: [PATCH] Fix #337

---
 codon/parser/visitors/typecheck/function.cpp | 10 +++--
 test/parser/typecheck/test_function.codon    | 46 ++++++++++++++++++++
 2 files changed, 53 insertions(+), 3 deletions(-)

diff --git a/codon/parser/visitors/typecheck/function.cpp b/codon/parser/visitors/typecheck/function.cpp
index bcd12186..94a83663 100644
--- a/codon/parser/visitors/typecheck/function.cpp
+++ b/codon/parser/visitors/typecheck/function.cpp
@@ -469,6 +469,7 @@ void TypecheckVisitor::visit(FunctionStmt *stmt) {
 
   // Expression to be used if function binding is modified by captures or decorators
   Expr *finalExpr = nullptr;
+  Expr *selfAssign = nullptr;
   // If there are captures, replace `fn` with `fn(cap1=cap1, cap2=cap2, ...)`
   if (!captures.empty()) {
     if (isClassMember)
@@ -484,9 +485,7 @@ void TypecheckVisitor::visit(FunctionStmt *stmt) {
         a.value = clone(a.getExpr());
     }
     // todo)) right now this adds a capture hook for recursive calls
-    auto assign = N<AssignStmt>(N<IdExpr>(stmt->getName()),
-                                N<CallExpr>(N<IdExpr>(stmt->getName()), pa));
-    f->suite = N<SuiteStmt>(assign, suite);
+    selfAssign = N<CallExpr>(N<IdExpr>(stmt->getName()), pa);
   }
 
   // Parse remaining decorators
@@ -497,9 +496,14 @@ void TypecheckVisitor::visit(FunctionStmt *stmt) {
       // Replace each decorator with `decorator(finalExpr)` in the reverse order
       finalExpr = N<CallExpr>(stmt->decorators[i],
                               finalExpr ? finalExpr : N<IdExpr>(canonicalName));
+      selfAssign = N<CallExpr>(clone(stmt->decorators[i]),
+                               selfAssign ? selfAssign : N<IdExpr>(canonicalName));
     }
   }
 
+  if (selfAssign)
+    f->suite =
+        N<SuiteStmt>(N<AssignStmt>(N<IdExpr>(stmt->getName()), selfAssign), suite);
   if (finalExpr) {
     resultStmt = N<SuiteStmt>(
         f, transform(N<AssignStmt>(N<IdExpr>(stmt->getName()), finalExpr)));
diff --git a/test/parser/typecheck/test_function.codon b/test/parser/typecheck/test_function.codon
index 539896c4..eed9a6e4 100644
--- a/test/parser/typecheck/test_function.codon
+++ b/test/parser/typecheck/test_function.codon
@@ -632,3 +632,49 @@ for fn in l: print(fn(1, 2))
 #: f1:1.2
 #: f2:1+2
 #: f3:<1+2+hey!>
+
+#%% decorator_self_reference
+store = Dict[int, int]()   # need to manually configure cache for now.
+def memoize(func):
+    def inner(val: int) -> int:
+        if val in store:
+            print(f"<- cache[{val}]")
+            return store[val]
+        else:
+            result = func(val)
+            store[val] = result
+            return result
+    return inner
+
+@memoize
+def fib(n: int) -> int:
+    print(f"<- fib[{n}]")
+    if n < 2:
+        return n
+    else:
+        return fib(n - 1) + fib(n - 2)    ## << not accessing decorated function
+
+f4 = fib(4)
+print(f"{f4=} : {store=}")
+#: <- fib[4]
+#: <- fib[3]
+#: <- fib[2]
+#: <- fib[1]
+#: <- fib[0]
+#: <- cache[1]
+#: <- cache[2]
+#: f4=3 : store={0: 0, 1: 1, 2: 1, 3: 2, 4: 3}
+
+f6 = fib(6)
+print(f"{f6=} : {store=}")
+#: <- fib[6]
+#: <- fib[5]
+#: <- cache[4]
+#: <- cache[3]
+#: <- cache[4]
+#: f6=8 : store={0: 0, 1: 1, 2: 1, 3: 2, 4: 3, 5: 5, 6: 8}
+
+f6 = fib(6)
+print(f"{f6=} : {store=}")
+#: <- cache[6]
+#: f6=8 : store={0: 0, 1: 1, 2: 1, 3: 2, 4: 3, 5: 5, 6: 8}