Extend power of the each form

This changes the implementation of the `next` function which
is now used to implement each. This let's us iterate over
more types, not just tables and structs.

src/boot/boot.janet

@@ -333,17 +333,16 @@
          (set ,i (,delta ,i ,step))))))
 
 (defn- each-template
-  [binding in body]
+  [binding inx body]
-  (with-syms [i len]
+  (with-syms [k]
-    (def ds (if (idempotent? in) in (gensym)))
+    (def ds (if (idempotent? inx) inx (gensym)))
     ~(do
-       (var ,i 0)
+       ,(unless (= ds inx) ~(def ,ds ,inx))
-       ,(unless (= ds in) ~(def ,ds ,in))
+       (var ,k (,next ,ds nil))
-       (def ,len (,length ,ds))
+       (while (,not= nil ,k)
-       (while (,< ,i ,len)
+         (def ,binding (,in ,ds ,k))
-         (def ,binding (in ,ds ,i))
          ,;body
-         (++ ,i)))))
+         (set ,k (,next ,ds ,k))))))
 
 (defn- keys-template
   [binding in pair? body]
@@ -352,7 +351,7 @@
     ~(do
        ,(unless (= ds in) ~(def ,ds ,in))
        (var ,k (,next ,ds nil))
-       (while ,k
+       (while (,not= nil ,k)
          (def ,binding ,(if pair? ~(tuple ,k (in ,ds ,k)) k))
          ,;body
          (set ,k (,next ,ds ,k))))))

src/core/asm.c

@@ -87,6 +87,8 @@ static const JanetInstructionDef janet_ops[] = {
     {"in", JOP_IN},
     {"jmp", JOP_JUMP},
     {"jmpif", JOP_JUMP_IF},
+    {"jmpni", JOP_JUMP_IF_NIL},
+    {"jmpnn", JOP_JUMP_IF_NOT_NIL},
     {"jmpno", JOP_JUMP_IF_NOT},
     {"ldc", JOP_LOAD_CONSTANT},
     {"ldf", JOP_LOAD_FALSE},
@@ -110,6 +112,7 @@ static const JanetInstructionDef janet_ops[] = {
     {"movn", JOP_MOVE_NEAR},
     {"mul", JOP_MULTIPLY},
     {"mulim", JOP_MULTIPLY_IMMEDIATE},
+    {"next", JOP_NEXT},
     {"noop", JOP_NOOP},
     {"prop", JOP_PROPAGATE},
     {"push", JOP_PUSH},

src/core/bytecode.c

@@ -56,6 +56,8 @@ enum JanetInstructionType janet_instructions[JOP_INSTRUCTION_COUNT] = {
     JINT_L, /* JOP_JUMP, */
     JINT_SL, /* JOP_JUMP_IF, */
     JINT_SL, /* JOP_JUMP_IF_NOT, */
+    JINT_SL, /* JOP_JUMP_IF_NIL, */
+    JINT_SL, /* JOP_JUMP_IF_NOT_NIL, */
     JINT_SSS, /* JOP_GREATER_THAN, */
     JINT_SSI, /* JOP_GREATER_THAN_IMMEDIATE, */
     JINT_SSS, /* JOP_LESS_THAN, */
@@ -95,7 +97,8 @@ enum JanetInstructionType janet_instructions[JOP_INSTRUCTION_COUNT] = {
     JINT_S, /* JOP_MAKE_TUPLE */
     JINT_S, /* JOP_MAKE_BRACKET_TUPLE */
     JINT_SSS, /* JOP_GREATER_THAN_EQUAL */
-    JINT_SSS /* JOP_LESS_THAN_EQUAL */
+    JINT_SSS, /* JOP_LESS_THAN_EQUAL */
+    JINT_SSS, /* JOP_NEXT */
 };
 
 /* Verify some bytecode */

src/core/cfuns.c

@@ -112,6 +112,9 @@ static JanetSlot do_in(JanetFopts opts, JanetSlot *args) {
 static JanetSlot do_get(JanetFopts opts, JanetSlot *args) {
     return opreduce(opts, args, JOP_GET, janet_wrap_nil());
 }
+static JanetSlot do_next(JanetFopts opts, JanetSlot *args) {
+    return opreduce(opts, args, JOP_NEXT, janet_wrap_nil());
+}
 static JanetSlot do_put(JanetFopts opts, JanetSlot *args) {
     if (opts.flags & JANET_FOPTS_DROP) {
         janetc_emit_sss(opts.compiler, JOP_PUT, args[0], args[1], args[2], 0);
@@ -283,7 +286,8 @@ static const JanetFunOptimizer optimizers[] = {
     {NULL, do_eq},
     {NULL, do_neq},
     {fixarity2, do_propagate},
-    {fixarity2, do_get}
+    {fixarity2, do_get},
+    {fixarity2, do_next}
 };
 
 const JanetFunOptimizer *janetc_funopt(uint32_t flags) {

src/core/compile.h

@@ -57,6 +57,7 @@
 #define JANET_FUN_NEQ 25
 #define JANET_FUN_PROP 26
 #define JANET_FUN_GET 27
+#define JANET_FUN_NEXT 28
 
 /* Compiler typedefs */
 typedef struct JanetCompiler JanetCompiler;

src/core/corelib.c

@@ -427,20 +427,6 @@ static Janet janet_core_type(int32_t argc, Janet *argv) {
     }
 }
 
-static Janet janet_core_next(int32_t argc, Janet *argv) {
-    janet_fixarity(argc, 2);
-    JanetDictView view = janet_getdictionary(argv, 0);
-    const JanetKV *end = view.kvs + view.cap;
-    const JanetKV *kv = janet_checktype(argv[1], JANET_NIL)
-                        ? view.kvs
-                        : janet_dict_find(view.kvs, view.cap, argv[1]) + 1;
-    while (kv < end) {
-        if (!janet_checktype(kv->key, JANET_NIL)) return kv->key;
-        kv++;
-    }
-    return janet_wrap_nil();
-}
-
 static Janet janet_core_hash(int32_t argc, Janet *argv) {
     janet_fixarity(argc, 1);
     return janet_wrap_number(janet_hash(argv[0]));
@@ -630,15 +616,6 @@ static const JanetReg corelib_cfuns[] = {
              "\t:cfunction\n\n"
              "or another symbol for an abstract type.")
     },
-    {
-        "next", janet_core_next,
-        JDOC("(next dict &opt key)\n\n"
-             "Gets the next key in a struct or table. Can be used to iterate through "
-             "the keys of a data structure in an unspecified order. Keys are guaranteed "
-             "to be seen only once per iteration if they data structure is not mutated "
-             "during iteration. If key is nil, next returns the first key. If next "
-             "returns nil, there are no more keys to iterate through. ")
-    },
     {
         "hash", janet_core_hash,
         JDOC("(hash value)\n\n"
@@ -950,6 +927,10 @@ static const uint32_t propagate_asm[] = {
     JOP_PROPAGATE | (1 << 24),
     JOP_RETURN
 };
+static const uint32_t next_asm[] = {
+    JOP_NEXT | (1 << 24),
+    JOP_RETURN
+};
 #endif /* ifdef JANET_BOOTSTRAP */
 
 /*
@@ -992,6 +973,14 @@ static void janet_load_libs(JanetTable *env) {
 
 JanetTable *janet_core_env(JanetTable *replacements) {
     JanetTable *env = (NULL != replacements) ? replacements : janet_table(0);
+    janet_quick_asm(env, JANET_FUN_NEXT,
+                    "next", 2, 2, 2, 2, next_asm, sizeof(next_asm),
+                    JDOC("(next ds &opt key)\n\n"
+                         "Gets the next key in a datastructure. Can be used to iterate through "
+                         "the keys of a data structure in an unspecified order. Keys are guaranteed "
+                         "to be seen only once per iteration if they data structure is not mutated "
+                         "during iteration. If key is nil, next returns the first key. If next "
+                         "returns nil, there are no more keys to iterate through."));
     janet_quick_asm(env, JANET_FUN_PROP,
                     "propagate", 2, 2, 2, 2, propagate_asm, sizeof(propagate_asm),
                     JDOC("(propagate x fiber)\n\n"

src/core/specials.c

@@ -538,6 +538,20 @@ static JanetSlot janetc_break(JanetFopts opts, int32_t argn, const Janet *argv)
     }
 }
 
+/* Check if a form matches the pattern (not= nil _) */
+static int janetc_check_notnil_form(Janet x, Janet *capture) {
+    if (!janet_checktype(x, JANET_TUPLE)) return 0;
+    JanetTuple tup = janet_unwrap_tuple(x);
+    if (!janet_checktype(tup[0], JANET_FUNCTION)) return 0;
+    if (3 != janet_tuple_length(tup)) return 0;
+    JanetFunction *fun = janet_unwrap_function(tup[0]);
+    uint32_t tag = fun->def->flags & JANET_FUNCDEF_FLAG_TAG;
+    if (tag != JANET_FUN_NEQ) return 0;
+    if (!janet_checktype(tup[1], JANET_NIL)) return 0;
+    *capture = tup[2];
+    return 1;
+}
+
 /*
  * :whiletop
  * ...
@@ -554,6 +568,9 @@ static JanetSlot janetc_while(JanetFopts opts, int32_t argn, const Janet *argv)
     JanetScope tempscope;
     int32_t labelwt, labeld, labeljt, labelc, i;
     int infinite = 0;
+    int is_notnil_form = 0;
+    uint8_t ifjmp = JOP_JUMP_IF;
+    uint8_t ifnjmp = JOP_JUMP_IF_NOT;
 
     if (argn < 2) {
         janetc_cerror(c, "expected at least 2 arguments");
@@ -564,13 +581,26 @@ static JanetSlot janetc_while(JanetFopts opts, int32_t argn, const Janet *argv)
 
     janetc_scope(&tempscope, c, JANET_SCOPE_WHILE, "while");
 
+    /* Check for `(not= nil _)` in condition, and if so, use the
+     * jmpnl or jmpnn instructions. This let's us implement `(each ...)`
+     * more efficiently. */
+    Janet condform = argv[0];
+    if (janetc_check_notnil_form(condform, &condform)) {
+        is_notnil_form = 1;
+        ifjmp = JOP_JUMP_IF_NOT_NIL;
+        ifnjmp = JOP_JUMP_IF_NIL;
+    }
+
     /* Compile condition */
-    cond = janetc_value(subopts, argv[0]);
+    cond = janetc_value(subopts, condform);
 
     /* Check for constant condition */
     if (cond.flags & JANET_SLOT_CONSTANT) {
         /* Loop never executes */
-        if (!janet_truthy(cond.constant)) {
+        int never_executes = is_notnil_form
+                             ? janet_checktype(cond.constant, JANET_NIL)
+                             : !janet_truthy(cond.constant);
+        if (never_executes) {
             janetc_popscope(c);
             return janetc_cslot(janet_wrap_nil());
         }
@@ -581,7 +611,7 @@ static JanetSlot janetc_while(JanetFopts opts, int32_t argn, const Janet *argv)
     /* Infinite loop does not need to check condition */
     labelc = infinite
              ? 0
-             : janetc_emit_si(c, JOP_JUMP_IF_NOT, cond, 0, 0);
+             : janetc_emit_si(c, ifnjmp, cond, 0, 0);
 
     /* Compile body */
     for (i = 1; i < argn; i++) {
@@ -603,7 +633,7 @@ static JanetSlot janetc_while(JanetFopts opts, int32_t argn, const Janet *argv)
         cond = janetc_value(subopts, argv[0]);
         if (!(cond.flags & JANET_SLOT_CONSTANT)) {
             /* If not an infinite loop, return nil when condition false */
-            janetc_emit_si(c, JOP_JUMP_IF, cond, 2, 0);
+            janetc_emit_si(c, ifjmp, cond, 2, 0);
             janetc_emit(c, JOP_RETURN_NIL);
         }
         for (i = 1; i < argn; i++) {

src/core/value.c

@@ -22,6 +22,7 @@
 
 #ifndef JANET_AMALG
 #include "features.h"
+#include "util.h"
 #include <janet.h>
 #endif
 
@@ -29,6 +30,67 @@
  * Define a number of functions that can be used internally on ANY Janet.
  */
 
+Janet janet_next(Janet ds, Janet key) {
+    JanetType t = janet_type(ds);
+    switch (t) {
+        default:
+            janet_panicf("expected iterable type, got %v", ds);
+        case JANET_TABLE:
+        case JANET_STRUCT: {
+            const JanetKV *start;
+            int32_t cap;
+            if (t == JANET_TABLE) {
+                JanetTable *tab = janet_unwrap_table(ds);
+                cap = tab->capacity;
+                start = tab->data;
+            } else {
+                JanetStruct st = janet_unwrap_struct(ds);
+                cap = janet_struct_capacity(st);
+                start = st;
+            }
+            const JanetKV *end = start + cap;
+            const JanetKV *kv = janet_checktype(key, JANET_NIL)
+                                ? start
+                                : janet_dict_find(start, cap, key) + 1;
+            while (kv < end) {
+                if (!janet_checktype(kv->key, JANET_NIL)) return kv->key;
+                kv++;
+            }
+            break;
+        }
+        case JANET_STRING:
+        case JANET_KEYWORD:
+        case JANET_SYMBOL:
+        case JANET_BUFFER:
+        case JANET_ARRAY:
+        case JANET_TUPLE: {
+            int32_t i;
+            if (janet_checktype(key, JANET_NIL)) {
+                i = 0;
+            } else if (janet_checkint(key)) {
+                i = janet_unwrap_integer(key) + 1;
+            } else {
+                break;
+            }
+            int32_t len;
+            if (t == JANET_BUFFER) {
+                len = janet_unwrap_buffer(ds)->count;
+            } else if (t == JANET_ARRAY) {
+                len = janet_unwrap_array(ds)->count;
+            } else if (t == JANET_TUPLE) {
+                len = janet_tuple_length(janet_unwrap_tuple(ds));
+            } else {
+                len = janet_string_length(janet_unwrap_string(ds));
+            }
+            if (i < len && i >= 0) {
+                return janet_wrap_integer(i);
+            }
+            break;
+        }
+    }
+    return janet_wrap_nil();
+}
+
 /* Compare two abstract values */
 static int janet_compare_abstract(JanetAbstract xx, JanetAbstract yy) {
     if (xx == yy) return 0;

src/core/vm.c

@@ -264,6 +264,8 @@ static JanetSignal run_vm(JanetFiber *fiber, Janet in) {
         &&label_JOP_JUMP,
         &&label_JOP_JUMP_IF,
         &&label_JOP_JUMP_IF_NOT,
+        &&label_JOP_JUMP_IF_NIL,
+        &&label_JOP_JUMP_IF_NOT_NIL,
         &&label_JOP_GREATER_THAN,
         &&label_JOP_GREATER_THAN_IMMEDIATE,
         &&label_JOP_LESS_THAN,
@@ -304,9 +306,7 @@ static JanetSignal run_vm(JanetFiber *fiber, Janet in) {
         &&label_JOP_MAKE_BRACKET_TUPLE,
         &&label_JOP_GREATER_THAN_EQUAL,
         &&label_JOP_LESS_THAN_EQUAL,
-        &&label_unknown_op,
+        &&label_JOP_NEXT,
-        &&label_unknown_op,
-        &&label_unknown_op,
         &&label_unknown_op,
         &&label_unknown_op,
         &&label_unknown_op,
@@ -648,6 +648,22 @@ static JanetSignal run_vm(JanetFiber *fiber, Janet in) {
     }
     vm_next();
 
+    VM_OP(JOP_JUMP_IF_NIL)
+    if (janet_checktype(stack[A], JANET_NIL)) {
+        pc += ES;
+    } else {
+        pc++;
+    }
+    vm_next();
+
+    VM_OP(JOP_JUMP_IF_NOT_NIL)
+    if (janet_checktype(stack[A], JANET_NIL)) {
+        pc++;
+    } else {
+        pc += ES;
+    }
+    vm_next();
+
     VM_OP(JOP_LESS_THAN)
     vm_compop( <);
 
@@ -680,6 +696,10 @@ static JanetSignal run_vm(JanetFiber *fiber, Janet in) {
     stack[A] = janet_wrap_integer(janet_compare(stack[B], stack[C]));
     vm_pcnext();
 
+    VM_OP(JOP_NEXT)
+    stack[A] = janet_next(stack[B], stack[C]);
+    vm_pcnext();
+
     VM_OP(JOP_LOAD_NIL)
     stack[D] = janet_wrap_nil();
     vm_pcnext();

src/include/janet.h

@@ -1020,6 +1020,8 @@ enum JanetOpCode {
     JOP_JUMP,
     JOP_JUMP_IF,
     JOP_JUMP_IF_NOT,
+    JOP_JUMP_IF_NIL,
+    JOP_JUMP_IF_NOT_NIL,
     JOP_GREATER_THAN,
     JOP_GREATER_THAN_IMMEDIATE,
     JOP_LESS_THAN,
@@ -1060,6 +1062,7 @@ enum JanetOpCode {
     JOP_MAKE_BRACKET_TUPLE,
     JOP_GREATER_THAN_EQUAL,
     JOP_LESS_THAN_EQUAL,
+    JOP_NEXT,
     JOP_INSTRUCTION_COUNT
 };
 
@@ -1312,6 +1315,7 @@ JANET_API int janet_compare(Janet x, Janet y);
 JANET_API int janet_cstrcmp(JanetString str, const char *other);
 JANET_API Janet janet_in(Janet ds, Janet key);
 JANET_API Janet janet_get(Janet ds, Janet key);
+JANET_API Janet janet_next(Janet ds, Janet key);
 JANET_API Janet janet_getindex(Janet ds, int32_t index);
 JANET_API int32_t janet_length(Janet x);
 JANET_API Janet janet_lengthv(Janet x);