From 58af636293d4c723ab28b024cc100ca1ff25bd5e Mon Sep 17 00:00:00 2001
From: Konstantin Nazarov <mail@knazarov.com>
Date: Thu, 10 Oct 2024 18:57:06 +0100
Subject: [PATCH] Allow defining and redefining local variables with "def"

---
 src/compiler.cpp  | 33 ++++++++++++++++++++++++++-------
 test/function.vli | 10 ++++++++++
 2 files changed, 36 insertions(+), 7 deletions(-)

diff --git a/src/compiler.cpp b/src/compiler.cpp
index 2e76a92..b2564c3 100644
--- a/src/compiler.cpp
+++ b/src/compiler.cpp
@@ -633,13 +633,32 @@ Result<Expression> Compiler::compile_def(Context& context, Symbol& op,
   auto comp = TRY(compile_expr(context, var_expr));
   ex.add_code(comp.code);
 
-  TRY(ex.add_opcode(Oc::GlobalStore, {1, (int64_t)gname},
-                    {0, (int64_t)comp.reg}));
+  if (context.toplevel) {
+    TRY(ex.add_opcode(Oc::GlobalStore, {1, (int64_t)gname},
+                      {0, (int64_t)comp.reg}));
+    ex.reg = comp.reg;
+    return std::move(ex);
+  } else {
+    auto maybe_binding = context.get_var(varname_unwrapped);
 
-  ex.reg = comp.reg;
-  context.maxreg = maxreg + 1;
+    if (maybe_binding.has_value()) {
+      auto reg = maybe_binding.value();
+      TRY(ex.add_opcode(Oc::Mov, {0, (int64_t)reg}, {0, (int64_t)comp.reg}));
+      ex.reg = comp.reg;
 
-  return std::move(ex);
+      context.maxreg = maxreg + 1;
+      return std::move(ex);
+    } else {
+      context.maxreg = maxreg + 1;
+      auto reg = TRY(context.add_var(varname_unwrapped));
+
+      TRY(ex.add_opcode(Oc::Mov, {0, (int64_t)reg}, {0, (int64_t)comp.reg}));
+      ex.reg = reg;
+      context.maxreg = reg + 1;
+
+      return std::move(ex);
+    }
+  }
 }
 
 Result<Expression> Compiler::compile_quote(Context& context, Symbol& op,
@@ -868,7 +887,7 @@ Result<Expression> Compiler::compile_let(Context& context, Symbol& op,
 }
 
 Result<Expression> Compiler::compile_do(Context& context, Symbol& op,
-                                         const Value& expr) {
+                                        const Value& expr) {
   auto rest = TRY(expr.rest());
 
   return compile_body(context, rest);
@@ -905,7 +924,7 @@ Result<Expression> Compiler::compile_body(Context& context, const Value& expr) {
     if (cur.is<Nil>()) {
       ex_res.reg = expr.reg;
     } else {
-      context.maxreg = maxreg;
+      // context.maxreg = maxreg;
     }
   }
 
diff --git a/test/function.vli b/test/function.vli
index 4b414d1..34d70ba 100644
--- a/test/function.vli
+++ b/test/function.vli
@@ -55,3 +55,13 @@
 ;; "do" form returns the value of last expression
 
 (assert (= (do 1 2 3) 3))
+
+;; allow definition and redefinition of variables with "def"
+
+(fn def-redef ()
+    (def x 1)
+    (def x (+ x 1))
+    (def y 4)
+    (+ x y))
+
+(assert (= (def-redef) 6))