Fix lambda and closures

3 files changed, 131 insertions, 32 deletions

diff --git a/examples/variables.bdl b/examples/variables.bdl
index 6097368..2e2543a 100644
--- a/examples/variables.bdl
+++ b/examples/variables.bdl
@@ -14,3 +14,59 @@
 (print "(error? (set! a 42)) -> ") (error? (set! a 42))
 (print "a -> ") a
 (print "(error? (set! b 99)) -> ") (error? (set! b 99))
+(lambda (a b) (+ 10 a b))
+
+(def a 20)
+((lambda (a b) (+ 10 a b)) 1 2)
+((lambda (a b) (+ 10 a b)) a 3)
+(def myfun (lambda (a b) (+ a b)))
+(myfun 6 9)
+(+ 1 (myfun 10 (myfun a a)) 30)
+(myfun 10 (myfun 5 0))
+
+;; Closures.
+(def make-counter (lambda ()
+    (def value 0)
+    (def counter (lambda ()
+        (set! value (+ value 1))
+        value))
+    counter))
+
+(def counter-a (make-counter))
+(def counter-b (make-counter))
+(counter-a)
+(counter-a)
+(counter-b)
+(counter-b)
+(counter-b)
+(counter-b)
+(counter-a)
+
+;; Fibonacci.
+(def fib (lambda (n)
+  (if (<= n 2)
+      1
+      (+ (fib (- n 1)) (fib (- n 2))))))
+(fib 15)
+
+;; Lambda capture.
+(def a 20)
+(def b (lambda ()
+    (display a)
+    (newline)
+    (def a 42)
+    (display a)
+    (newline)))
+
+(b)
+(b)
+
+
+;; Infinite loop. (For teseting purposes)
+; (def test (lambda (n)
+;     (print "ITER\n")
+;   (if (<= n 2)
+;     'ok
+;     (test (+ n 1)))))
+
+; (test 3)
diff --git a/src/bootstrap/environment.c b/src/bootstrap/environment.c
index 99dd7fd..e111753 100644
--- a/src/bootstrap/environment.c
+++ b/src/bootstrap/environment.c
@@ -60,8 +60,27 @@ env_lookup(Environment *env, Object *symbol) {
     return obj_err;
 }
 
+Object *
+env_update(Environment *env, Object *symbol, Object *value) {
+    while (env != NULL) {
+        for (size_t i = 0; i < env->size; i++) {
+            EnvEntry entry = env->buf[i];
+            if (obj_eq(symbol, entry.symbol)) {
+                env->buf[i].value = obj_duplicate(value);
+                return obj_nil;
+            }
+        }
+        env = env->parent;
+    }
+    error_push((Error){
+        .type = ERR_TYPE_RUNTIME,
+        .value = ERR_SYMBOL_NOT_FOUND,
+    });
+    return obj_err;
+}
+
 ssize_t
-env_symbol_index_in_current_env(Environment *env, Object *symbol) {
+env_index_current(Environment *env, Object *symbol) {
     for (size_t i = 0; i < env->size; i++) {
         EnvEntry entry = env->buf[i];
         if (obj_eq(symbol, entry.symbol)) {
@@ -70,3 +89,34 @@ env_symbol_index_in_current_env(Environment *env, Object *symbol) {
     }
     return -1;
 }
+
+void
+env_add_or_update_current(Environment *env, Object *symbol, Object *value) {
+    ssize_t index = env_index_current(env, symbol);
+    if (index == -1) {
+        env_add_symbol(env, obj_duplicate(symbol), obj_duplicate(value));
+    } else {
+        env->buf[index].value = obj_duplicate(value);
+    }
+}
+
+Environment *
+env_extend(Environment *parent, Environment *extra) {
+    Environment *env = env_create(parent);
+    for (size_t i = 0; i < extra->size; i++) {
+        EnvEntry entry = extra->buf[i];
+        Environment *tmp = env;
+        ssize_t idx = -1;
+        while (tmp != NULL) {
+            idx = env_index_current(tmp, entry.symbol);
+            if (idx != -1) {
+                break;
+            }
+            tmp = tmp->parent;
+        }
+        if (idx == -1) {
+            env_add_symbol(env, obj_duplicate(entry.symbol), obj_duplicate(entry.value));
+        }
+    }
+    return env;
+}
diff --git a/src/bootstrap/primitives.c b/src/bootstrap/primitives.c
index 461799e..4ef56d5 100644
--- a/src/bootstrap/primitives.c
+++ b/src/bootstrap/primitives.c
@@ -1,6 +1,7 @@
+#define DEBUG_OBJ(MSG,OBJ) printf((MSG)); display(OBJ); printf("\n");
+
 Object *
 eval(Environment* env, Object *root) {
-tailcall:
     switch (root->type) {
         case OBJ_TYPE_ERR:
         case OBJ_TYPE_PROCEDURE:
@@ -35,17 +36,26 @@ tailcall:
                 if (val->type == OBJ_TYPE_PROCEDURE) {
                     return val->proc(env, root->cdr);
                 }
+                if (val->type == OBJ_TYPE_LAMBDA) {
+                    goto eval_lambda;
+                }
                 error_push((Error){
                     .type = ERR_TYPE_RUNTIME,
                     .value = ERR_OBJ_NOT_CALLABLE,
                 });
                 return obj_err;
             }
-            Object* lambda = eval(env, root->car);
+            Object* lambda;
+eval_lambda:
+            lambda = eval(env, root->car);
+            if (lambda == obj_err) {
+                return obj_err;
+            }
             if (lambda->type == OBJ_TYPE_LAMBDA) {
                 Object *fun = lambda;
                 Object *args = root->cdr;
                 Object *params = fun->args;
+                env = env_extend(fun->env, env);
                 while (params != obj_nil) {
                     if (args == obj_nil) {
                         error_push((Error){
@@ -66,12 +76,7 @@ tailcall:
                         });
                         return obj_err;
                     }
-                    ssize_t index = env_symbol_index_in_current_env(fun->env, symbol);
-                    if (index == -1) {
-                        env_add_symbol(fun->env, obj_duplicate(symbol), obj_duplicate(value));
-                    } else {
-                        fun->env->buf[index].value = obj_duplicate(value);
-                    }
+                    env_add_or_update_current(env, symbol, value);
                     args = args->cdr;
                     params = params->cdr;
                 }
@@ -82,10 +87,14 @@ tailcall:
                     });
                     return obj_err;
                 }
-
-                env = fun->env;
                 root = fun->body;
-                goto tailcall;
+                while (root->cdr != obj_nil) {
+                    if (eval(env, root->car) == obj_err) {
+                        return obj_err;
+                    };
+                    root = root->cdr;
+                }
+                return eval(env, root->car);
             }
         } break;
     }
@@ -693,13 +702,7 @@ proc_define(Environment *env, Object *obj) {
         return obj_err;
     }
 
-    // Make a copy of the symbol and to make them permanent in the environment.
-    ssize_t index = env_symbol_index_in_current_env(env, symbol);
-    if (index == -1) {
-        env_add_symbol(env, obj_duplicate(symbol), obj_duplicate(value));
-    } else {
-        env->buf[index].value = obj_duplicate(value);
-    }
+    env_add_or_update_current(env, symbol, value);
     return obj_nil;
 }
 
@@ -727,17 +730,7 @@ proc_set(Environment *env, Object *obj) {
         return obj_err;
     }
 
-    ssize_t index = env_symbol_index_in_current_env(env, symbol);
-    if (index == -1) {
-        error_push((Error){
-            .type = ERR_TYPE_RUNTIME,
-            .value = ERR_SYMBOL_NOT_FOUND,
-        });
-        return obj_err;
-    }
-
-    env->buf[index].value = obj_duplicate(value);
-    return obj_nil;
+    return env_update(env, symbol, value);
 }
 
 Object *
@@ -757,11 +750,11 @@ proc_lambda(Environment *env, Object *obj) {
         });
         return obj_err;
     }
-    Object *body = obj->cdr->car;
+    Object *body = obj->cdr;
     Object *fun = alloc_object(OBJ_TYPE_LAMBDA);
     fun->args = obj_duplicate(args);
     fun->body = obj_duplicate(body);
-    fun->env = env_create(env);
+    fun->env = env;
     return fun;
 }