Fix some more recursion issues with pegs.

A keyword reference only counts as visited if we have
it as cached in the memoized->table, and we know it was
originally referenced from the same grammar table. If these
two conditions are true, then compilation must work correctly.

Also add janet_table_get_ex.

src/core/peg.c

@@ -446,6 +446,7 @@ tail:
 typedef struct {
     JanetTable *grammar;
     JanetTable *memoized;
+    JanetTable *memoized_scopes;
     JanetTable *tags;
     Janet *constants;
     uint32_t *bytecode;
@@ -878,17 +879,6 @@ static const SpecialPair peg_specials[] = {
 /* Compile a janet value into a rule and return the rule index. */
 static uint32_t peg_compile1(Builder *b, Janet peg) {
 
-    /* Check for already compiled rules */
-    int is_keyword = janet_checktype(peg, JANET_KEYWORD);
-    Janet old_memo = janet_wrap_nil(); /* for compiler warnings */
-    if (is_keyword) {
-        old_memo = janet_table_get(b->memoized, peg);
-        if (!janet_checktype(old_memo, JANET_NIL)) {
-            uint32_t rule = (uint32_t) janet_unwrap_number(old_memo);
-            return rule;
-        }
-    }
-
     /* Keep track of the form being compiled for error purposes */
     Janet old_form = b->form;
     b->form = peg;
@@ -901,10 +891,6 @@ static uint32_t peg_compile1(Builder *b, Janet peg) {
     /* The final rule to return */
     uint32_t rule = janet_v_count(b->bytecode);
 
-    /* Cache keywords for recursion points (loops) */
-    if (is_keyword)
-        janet_table_put(b->memoized, peg, janet_wrap_number(rule));
-
     switch (janet_type(peg)) {
         default:
             peg_panic(b, "unexpected peg source");
@@ -926,20 +912,35 @@ static uint32_t peg_compile1(Builder *b, Janet peg) {
             break;
         }
         case JANET_KEYWORD: {
-            Janet check = janet_table_get(b->grammar, peg);
-            if (janet_checktype(check, JANET_NIL))
+            /* Find rule in grammar */
+            JanetTable *scope = NULL;
+            Janet check = janet_table_get_ex(b->grammar, peg, &scope);
+            if (scope == NULL)
                 peg_panic(b, "unknown rule");
+
+            /* Check if we should compile as a recursion */
+            Janet memo_rule = janet_table_get(b->memoized, peg);
+            Janet memo_scope = janet_table_get(b->memoized_scopes, peg);
+            if (!janet_checktype(memo_rule, JANET_NIL) &&
+                    scope == janet_unwrap_table(memo_scope)) {
+                rule = (uint32_t) janet_unwrap_number(memo_rule);
+                break;
+            }
+
+            /* Compile with rule and current scope memoized. This will
+             * let child rules refer to this rule for recursion */
+            janet_table_put(b->memoized, peg, janet_wrap_number(rule));
+            janet_table_put(b->memoized_scopes, peg, janet_wrap_table(scope));
             rule = peg_compile1(b, check);
+            janet_table_put(b->memoized, peg, memo_rule);
+            janet_table_put(b->memoized_scopes, peg, memo_scope);
             break;
         }
         case JANET_STRUCT: {
             JanetTable *grammar = janet_struct_to_table(janet_unwrap_struct(peg));
             grammar->proto = b->grammar;
             b->grammar = grammar;
-            Janet main_rule = janet_table_get(grammar, janet_ckeywordv("main"));
-            if (janet_checktype(main_rule, JANET_NIL))
-                peg_panic(b, "grammar requires :main rule");
-            rule = peg_compile1(b, main_rule);
+            rule = peg_compile1(b, janet_ckeywordv("main"));
             b->grammar = grammar->proto;
             break;
         }
@@ -964,10 +965,6 @@ static uint32_t peg_compile1(Builder *b, Janet peg) {
         }
     }
 
-    /* Reset old cached rule */
-    if (is_keyword)
-        janet_table_put(b->memoized, peg, old_memo);
-
     /* Increase depth again */
     b->depth++;
     b->form = old_form;
@@ -1198,6 +1195,7 @@ static Peg *compile_peg(Janet x) {
     Builder builder;
     builder.grammar = janet_table(0);
     builder.memoized = janet_table(0);
+    builder.memoized_scopes = janet_table(0);
     builder.tags = janet_table(0);
     builder.constants = NULL;
     builder.bytecode = NULL;

src/core/table.c

@@ -137,6 +137,27 @@ Janet janet_table_get(JanetTable *t, Janet key) {
     return janet_wrap_nil();
 }
 
+/* Get a value out of the table, and record which prototype it was from. */
+Janet janet_table_get_ex(JanetTable *t, Janet key, JanetTable **which) {
+    JanetKV *bucket = janet_table_find(t, key);
+    if (NULL != bucket && !janet_checktype(bucket->key, JANET_NIL)) {
+        *which = t;
+        return bucket->value;
+    }
+    /* Check prototypes */
+    {
+        int i;
+        for (i = JANET_MAX_PROTO_DEPTH, t = t->proto; t && i; t = t->proto, --i) {
+            bucket = janet_table_find(t, key);
+            if (NULL != bucket && !janet_checktype(bucket->key, JANET_NIL)) {
+                *which = t;
+                return bucket->value;
+            }
+        }
+    }
+    return janet_wrap_nil();
+}
+
 /* Get a value out of the table. Don't check prototype tables. */
 Janet janet_table_rawget(JanetTable *t, Janet key) {
     JanetKV *bucket = janet_table_find(t, key);

src/include/janet.h

@@ -1194,6 +1194,7 @@ JANET_API JanetTable *janet_table(int32_t capacity);
 JANET_API JanetTable *janet_table_init(JanetTable *table, int32_t capacity);
 JANET_API void janet_table_deinit(JanetTable *table);
 JANET_API Janet janet_table_get(JanetTable *t, Janet key);
+JANET_API Janet janet_table_get_ex(JanetTable *t, Janet key, JanetTable **which);
 JANET_API Janet janet_table_rawget(JanetTable *t, Janet key);
 JANET_API Janet janet_table_remove(JanetTable *t, Janet key);
 JANET_API void janet_table_put(JanetTable *t, Janet key, Janet value);

test/suite3.janet

@@ -431,4 +431,12 @@
 (check-match redef-a "abcabc" false)
 (check-match redef-a "defdef" false)
 
+(def redef-b
+  ~{:pork {:pork "beef" :main (+ -1 (* 1 :pork))}
+    :main :pork})
+
+(check-match redef-b "abeef" true)
+(check-match redef-b "aabeef" false)
+(check-match redef-b "aaaaaa" false)
+
 (end-suite)