From 8bc8709d0e55ea6ef219da0c3ab47d99bea76785 Mon Sep 17 00:00:00 2001
From: Calvin Rose <calsrose@gmail.com>
Date: Thu, 29 Aug 2019 19:09:43 -0500
Subject: [PATCH] Try to address memoization problem in pegs.

---
 src/core/peg.c    | 24 ++++++++++++++++--------
 test/suite3.janet | 10 ++++++++++
 2 files changed, 26 insertions(+), 8 deletions(-)

diff --git a/src/core/peg.c b/src/core/peg.c
index faf66f2b..fcc66a02 100644
--- a/src/core/peg.c
+++ b/src/core/peg.c
@@ -426,7 +426,7 @@ tail:
                     Janet capture = s->captures->data[i];
                     if (!janet_checktype(capture, JANET_STRING))
                         return NULL;
-                    const uint8_t *bytes = janet_unwrap_string(capture); 
+                    const uint8_t *bytes = janet_unwrap_string(capture);
                     int32_t len = janet_string_length(bytes);
                     if (text + len > s->text_end)
                         return NULL;
@@ -879,10 +879,14 @@ static const SpecialPair peg_specials[] = {
 static uint32_t peg_compile1(Builder *b, Janet peg) {
 
     /* Check for already compiled rules */
-    Janet check = janet_table_get(b->memoized, peg);
-    if (!janet_checktype(check, JANET_NIL)) {
-        uint32_t rule = (uint32_t) janet_unwrap_number(check);
-        return rule;
+    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 */
@@ -896,10 +900,10 @@ static uint32_t peg_compile1(Builder *b, Janet peg) {
 
     /* The final rule to return */
     uint32_t rule = janet_v_count(b->bytecode);
-    if (!janet_checktype(peg, JANET_KEYWORD) &&
-            !janet_checktype(peg, JANET_STRUCT)) {
+
+    /* Cache keywords for recursion points (loops) */
+    if (is_keyword)
         janet_table_put(b->memoized, peg, janet_wrap_number(rule));
-    }
 
     switch (janet_type(peg)) {
         default:
@@ -960,6 +964,10 @@ 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;
diff --git a/test/suite3.janet b/test/suite3.janet
index accdb833..3d27852d 100644
--- a/test/suite3.janet
+++ b/test/suite3.janet
@@ -421,4 +421,14 @@
 (assert (= (tuple/type (-> '(1 2 3) marshal unmarshal)) :parens) "normal tuple marshalled/unmarshalled")
 (assert (= (tuple/type (-> '[1 2 3] marshal unmarshal)) :brackets) "normal tuple marshalled/unmarshalled")
 
+# Check for bad memoization (+ :a) should mean different things in different contexts.
+(def redef-a
+  ~{:a "abc"
+    :c (+ :a)
+    :main (* :c {:a "def" :main (+ :a)} -1)})
+
+(check-match redef-a "abcdef" true)
+(check-match redef-a "abcabc" false)
+(check-match redef-a "defdef" false)
+
 (end-suite)