closures

1 files changed, 164 insertions, 22 deletions

diff --git a/pez.c b/pez.c
index c9d8f5e..40d097b 100644
--- a/pez.c
+++ b/pez.c
@@ -72,12 +72,13 @@ typedef struct Proto {
     uint8_t variadic : 1,
             named : 1;
     short nvars, nparams;
-    short ncon;
+    short ncon, nupval;
     uint ncode;
     int linebegin, lineend;
     const char *file;
     const uint8_t *code;
     const Val *con;
+    struct cmupval *upval;
     char name[];
 } Proto;
 
@@ -101,6 +102,11 @@ typedef struct Local Local;
 
 enum { NAMEMAX = 80 };
 
+struct cmupval {
+    uint8_t idx;
+    bool arg : 1, local : 1;
+};
+
 typedef struct Comp {
     PezContext *cx;
     int (*readcb)(void *);
@@ -130,10 +136,13 @@ typedef struct Comp {
     vec_of(char) spool; // for var names
     vec_of(Val) kpool; // for consts
     struct fenv {
+        const char *name;
         struct fenv *prev;
         vec_of(Local) locals; // also params
         uint16_t scope;
         uint16_t nvars; // num of locals in scope not counting konsts
+        uint16_t nupval;
+        struct cmupval upvals[256];
     } fenv;
 } Comp;
 
@@ -144,7 +153,8 @@ struct Local {
     uint16_t index : 10;
     uint16_t isparam : 1,
              has_k : 1,
-             mutable : 1;
+             mutable : 1,
+             captured : 1;
 };
 
 static struct Str strpool_deleted;
@@ -172,6 +182,7 @@ struct PezContext {
     char errstr[140];
     Val *stack, *stktop, *stkend;
     Obj *heap;
+    Upval *openup;
     uint nalloc, gcthresh, gccanrun : 1;
     StrPool strpool;
     Globals globals;
@@ -454,22 +465,62 @@ delproto(PezContext *cx, Proto *pr)
 {
     cxfree(cx, (void *)pr->code, pr->ncode);
     cxfree(cx, (void *)pr->con, pr->ncon * sizeof(Val));
+    cxfree(cx, (void *)pr->upval, pr->nupval * sizeof(struct cmupval));
 }
 
-static Fn *
-newfn(PezContext *cx, Proto *pr)
+static void
+closeups(PezContext *cx, Val *vals)
 {
-    Fn *fn = newobj(cx, PEZ_TFn, sizeof *fn);
-    if (fn) {
-        fn->proto = pr;
+    for (Upval *up = cx->openup; up; up = up->nextup) {
+        if (up->ptr >= vals) {
+            up->slot = *up->ptr;
+            up->ptr = &up->slot;
+            cx->openup = up->nextup;
+        }
     }
-    return fn;
 }
 
 static inline size_t
 sizeoffn(Fn *fn)
 {
-    return sizeof(*fn);
+    return sizeof(*fn) + fn->nupval * sizeof(struct Upval *);
+}
+
+static Fn *
+newfn(PezContext *cx, Proto *pr, Val *args, Val *locals, Fn *parent)
+{
+    Fn *fn = newobj(cx, PEZ_TFn, sizeof *fn + pr->nupval * sizeof(struct Upval *));
+    if (fn) {
+        fn->proto = pr;
+        fn->nupval = pr->nupval;
+        for (int i = 0; i < pr->nupval; ++i) {
+            struct cmupval *upinfo = &pr->upval[i];
+            if (upinfo->local || upinfo->arg) {
+                Upval *up;
+                Val *vals = upinfo->local ? locals : args;
+                for (up = cx->openup; up; up = up->nextup) {
+                    if (up->ptr == &vals[upinfo->idx]) {
+                        break;
+                    }
+                }
+                if (!up) {
+                    up = newobj(cx, PEZ_TUpval, sizeof *up);
+                    if (!up) {
+                        cxfree(cx, fn, sizeoffn(fn));
+                        return NULL;
+                    }
+                    up->ptr = &vals[upinfo->idx];
+                    up->nextup = cx->openup;
+                    cx->openup = up;
+                }
+                fn->upval[i] = up;
+            } else {
+                assert(fn && upinfo->idx < parent->nupval);
+                fn->upval[i] = parent->upval[upinfo->idx];
+            }
+        }
+    }
+    return fn;
 }
 
 static void
@@ -693,6 +744,9 @@ delobj(PezContext *cx, Obj *o)
         sz = sizeofstr((Str *)o);
         delstring(cx, (Str *)o);
         break;
+    case PEZ_TUpval:
+        sz = sizeof(Upval);
+        break;
     }
     assert(sz);
     cxfree(cx, o, sz);
@@ -714,6 +768,9 @@ static void
 markfn(PezContext *cx, Fn *fn)
 {
     gcmark(cx, (Obj *)fn->proto);
+    for (int i = 0; i < fn->nupval; ++i) {
+        gcmark(cx, (Obj *)fn->upval[i]);
+    }
 }
 
 static void
@@ -911,6 +968,9 @@ fixtrunc(fixnum f)
     _(setarg, 0)              \
     _(local,  1)              \
     _(setloc,-1)              \
+    _(upval,  1)              \
+    _(setupv,-1)              \
+    _(close,  0)              \
     _(global, 0)              \
     _(setglo,-2)              \
     _(putglo,-2)              \
@@ -1044,7 +1104,9 @@ inspectproto(Proto *pr)
             ip += 4;
             fprintf(stderr, "%f", fixtof(num));
             break;
-        case Olocal: case Osetloc: case Oarg: case Osetarg:
+        case Olocal: case Osetloc: case Oarg:
+        case Osetarg: case Oupval: case Osetupv:
+        case Oclose:
             ++ip;
             fprintf(stderr, "#%d", *argp);
             break;
@@ -1289,6 +1351,7 @@ exefn(PezContext *cx, Fn *fn, uint nargs)
             cx->stktop -= pr->nvars;
             assert(cx->stktop >= cx->stack);
             push(cx, ret);
+            closeups(cx, args);
             return 1;
         }
         CASE(Opop) {
@@ -1367,12 +1430,12 @@ exefn(PezContext *cx, Fn *fn, uint nargs)
         CASE(Olambda) {
             uint8_t idx = code[ip++];
             Val k;
-            Fn *fn;
+            Fn *ofn;
             assert(idx < pr->ncon);
             k = pr->con[idx];
             assert(isobj_of(k, PEZ_TFnProto));
-            TRY(fn = newfn(cx, unbox_obj(k)));
-            TRY(push(cx, box_obj(fn)));
+            TRY(ofn = newfn(cx, unbox_obj(k), args, locals, fn));
+            TRY(push(cx, box_obj(ofn)));
         }
         CASE(Onot) {
             Val *p = peek(cx);
@@ -1460,7 +1523,7 @@ exefn(PezContext *cx, Fn *fn, uint nargs)
         }
         CASE(Osetarg) {
             uint8_t idx = code[ip++];
-            args[idx] = *--cx->stktop;
+            args[idx] = pop(cx);
         }
         CASE(Olocal) {
             uint8_t idx = code[ip++];
@@ -1470,6 +1533,20 @@ exefn(PezContext *cx, Fn *fn, uint nargs)
             uint8_t idx = code[ip++];
             locals[idx] = pop(cx);
         }
+        CASE(Oupval) {
+            uint8_t idx = code[ip++];
+            assert(idx < fn->nupval);
+            TRY(push(cx, *fn->upval[idx]->ptr));
+        }
+        CASE(Osetupv) {
+            uint8_t idx = code[ip++];
+            assert(idx < fn->nupval);
+            *fn->upval[idx]->ptr = pop(cx);
+        }
+        CASE(Oclose) {
+            uint8_t idx = code[ip++];
+            closeups(cx, &locals[idx]);
+        }
         CASE(Oglobal) {
             Val k = pop(cx), *v;
             if ((v = getglobal(cx, k))) {
@@ -1856,7 +1933,7 @@ deinitcomp(Comp *cm)
     delvec(cm->cx, &cm->kpool);
 }
 
-static void
+static bool
 fincompfn(Comp *cm)
 {
     Proto *pr = cm->proto;
@@ -1869,6 +1946,13 @@ fincompfn(Comp *cm)
     pr->ncode = cm->code.len;
     pr->con = con;
     pr->ncon = cm->con.len;
+    pr->upval = cxalloc(cm->cx, cm->fenv.nupval * sizeof *pr->upval);
+    pr->nupval = cm->fenv.nupval;
+    // fprintf(stderr, "fini %s %d\n", pr->name, pr->nupval);
+    if (pr->upval) {
+        memcpy(pr->upval, cm->fenv.upvals, pr->nupval * sizeof(struct cmupval));
+    }
+    return 1;
 }
 
 static void
@@ -2044,9 +2128,57 @@ findlocal(Comp *cm, const char *name)
             return l;
         }
     }
+    // search for foldablable consts in enclosing functions
+    for (struct fenv *fenv = cm->fenv.prev; fenv; fenv = fenv->prev) {
+        for (int i = fenv->locals.len - 1; i >= 0; --i) {
+            Local *l = &fenv->locals.at[i];
+            if (l->has_k && !strcmp(&cm->spool.at[l->sref], name)) {
+                return l;
+            }
+        }
+    }
     return NULL;
 }
 
+static int
+addupval(Comp *cm, struct fenv *fenv, uint8_t idx, bool arg, bool local)
+{
+    struct cmupval up = { idx, arg, local };
+    for (int i = 0; i < fenv->nupval; ++i) {
+        struct cmupval *up = &fenv->upvals[i];
+        if (up->idx == idx && up->local == local) {
+            return up->idx;
+        }
+    }
+    assert(fenv->nupval < 256 && "upval limit");
+    fenv->upvals[fenv->nupval] = up;
+    return fenv->nupval++;
+}
+
+static int
+findupval(Comp *cm, Local **pl, struct fenv *fenv, const char *name)
+{
+    int idx;
+    if (!fenv->prev) {
+        return -1;
+    }
+
+    for (int i = fenv->prev->locals.len - 1; i >= 0; --i) {
+        Local *l = &fenv->prev->locals.at[i];
+        if (!strcmp(&cm->spool.at[l->sref], name)) {
+            assert(!l->has_k); // would've been found by findlocal earlier
+            *pl = l;
+            return addupval(cm, fenv, l->index, l->isparam, !l->isparam);
+        }
+    }
+    if ((idx = findupval(cm, pl, fenv->prev, name)) != -1) {
+        return addupval(cm, fenv, idx, 0, 0);
+    }
+
+    return -1;
+
+}
+
 static bool
 endscope(Comp *cm)
 {
@@ -2057,6 +2189,9 @@ endscope(Comp *cm)
         } else {
             if (!l->has_k) {
                 --cm->fenv.nvars;
+                if (l->mutable) {
+                    TRY(compop(cm, Oclose) && compbyte(cm, l->index));
+                }
             }
             --cm->fenv.locals.len;
         }
@@ -2238,7 +2373,7 @@ lambdaexpr(Comp *cm, const char *name)
     vec_of(uint8_t) prevcode;
     vec_of(Val) prevcon;
     struct fenv prevfenv = cm->fenv;;
-    struct fenv fenv = { .prev = &prevfenv };
+    struct fenv fenv = { .prev = &prevfenv, .name = name };
     bool ret = 1;
 
     memcpy(&prevcode, &cm->code, sizeof prevcode);
@@ -2270,7 +2405,7 @@ lambdaexpr(Comp *cm, const char *name)
 
     ETRY(block(cm, '}'));
     ETRY(compop(cm, Oret));
-    fincompfn(cm);
+    ETRY(fincompfn(cm));
 
 Cleanup:
     cm->proto = prevfn;
@@ -2346,6 +2481,7 @@ primaryexpr(Comp *cm)
         return 1;
     }
     if (c == '_' || aisalpha(c)) {
+        int idx;
         // identifier
         *buf = c;
         TRY(readident(cm, buf + 1, sizeof buf - 1));
@@ -2366,6 +2502,12 @@ primaryexpr(Comp *cm)
             cm->lvalue = 1;
             cm->lvalue_const = !local->mutable;
             cm->lvalue_name = &cm->spool.at[local->sref];
+        } else if ((idx = findupval(cm, &local, &cm->fenv, buf)) != -1) {
+            assert(idx >= 0 && idx < 256);
+            TRY(compop(cm, Oupval) && compbyte(cm, idx));
+            cm->lvalue = 1;
+            cm->lvalue_const = !local->mutable;
+            cm->lvalue_name = &cm->spool.at[local->sref];
         } else {
             Val nam;
             TRY(box_str(cm->cx, &nam, buf, strlen(buf)));
@@ -3035,7 +3177,7 @@ setexpr(Comp *cm)
 
         assert(cm->lastop >= 0 && cm->lastop < code->len);
         switch ((opcode = code->at[cm->lastop])) {
-        case Olocal: case Oarg:
+        case Olocal: case Oarg: case Oupval:
             idx = code->at[cm->lastop + 1];
             if (binop == Oset) {
                 code->len -= 2;
@@ -3082,8 +3224,8 @@ setexpr(Comp *cm)
         }
 
         switch (opcode) {
-        case Olocal: case Oarg:
-            TRY(compop(cm, opcode == Olocal ? Osetloc : Osetarg));
+        case Olocal: case Oarg: case Oupval:
+            TRY(compop(cm, opcode == Olocal ? Osetloc : opcode == Oarg ? Osetarg : Osetupv));
             TRY(compbyte(cm, idx));
             break;
         case Oglobal:
@@ -3408,7 +3550,7 @@ pez_eval_cb(PezContext *cx, const char *fname, int (*cb)(void *), void *ud)
     if (!(pr = newproto(cx, fname, "<eval>", /* line */ 1))) {
         return 0;
     }
-    if (!(fn = newfn(cx, pr))) {
+    if (!(fn = newfn(cx, pr, NULL, NULL, NULL))) {
         delproto(cx, pr);
         return 0;
     }
@@ -3416,7 +3558,7 @@ pez_eval_cb(PezContext *cx, const char *fname, int (*cb)(void *), void *ud)
 
     ETRY(block(&cm, EOF));
     ETRY(compop(&cm, Oret));
-    fincompfn(&cm);
+    ETRY(fincompfn(&cm));
 
     if (cx->dbg & DBGbytecode) {
         inspectproto(pr);