Scope memo expansions to sub-branches rather than globally

src/production/concat.js

@@ -5,15 +5,18 @@ import memo from "./memo.js";
 import unique from "./unique.js";
 
 class Concat {
-  constructor(expansion) {
+  constructor(expansion, registry) {
     this.expansion = expansion;
+    this.registry = registry;
   }
 
   evaluate(options) {
+    this.registry.pushBranch();
     const concat = this.expansion.reduce((accumulator, production) => {
       accumulator.push(production.evaluate(options));
       return accumulator;
     }, []);
+    this.registry.popBranch();
 
     return [Symbol.for("concat"), concat];
   }
@@ -36,7 +39,11 @@ function concat(production, registry) {
       let rule;
 
       if (expr[0][0] == MEMO_SIGIL) {
-        rule = memo(expr[0].slice(1, fragment.length - 1), registry);
+        if (expr[0][1] == MEMO_SIGIL) {
+          rule = memo(expr[0].slice(2, fragment.length - 1), true, registry);
+        } else {
+          rule = memo(expr[0].slice(1, fragment.length - 1), false, registry);
+        }
       } else if (expr[0][0] == UNIQUE_SIGIL) {
         rule = unique(expr[0].slice(1, fragment.length - 1), registry);
       } else {
@@ -53,7 +60,7 @@ function concat(production, registry) {
     }
   });
 
-  return new Concat(expansion);
+  return new Concat(expansion, registry);
 }
 
 export default concat;

src/production/memo.js

@@ -1,16 +1,20 @@
 class Memo {
-  constructor(symbol, registry) {
+  constructor(symbol, isBranchScoped, registry) {
     this.symbol = symbol;
     this.registry = registry;
+    this.isBranchScoped = isBranchScoped;
   }
 
   evaluate(options) {
+    if (this.isBranchScoped) {
+      this.registry.setBranchScopedMemo();
+    }
     return [Symbol.for(this.symbol), this.registry.evaluateMemo(this.symbol)];
   }
 }
 
-function memo(symbol, registry) {
-  return new Memo(symbol, registry);
+function memo(symbol, isBranchScoped, registry) {
+  return new Memo(symbol, isBranchScoped, registry);
 }
 
 export default memo;

src/registry.js

@@ -8,6 +8,7 @@ class Registry {
     this.memos = {};
     this.uniques = {};
     this.context = {};
+    this.level = 0;
 
     for (const key in rules) {
       this.rules[key] = rule(key, rules[key], this);
@@ -36,6 +37,26 @@ class Registry {
     return expansion;
   }
 
+  hasMemoExpanded(symbol) {
+    return this.memos[symbol];
+  }
+
+  pushBranch() {
+    this.level++;
+  }
+
+  popBranch() {
+    this.level--;
+    if (this.level < this.memoLevel) {
+      this.memoLevel = 0;
+      this.memos = {};
+    }
+  }
+
+  setBranchScopedMemo() {
+    this.memoLevel = this.level;
+  }
+
   evaluateMemo(symbol) {
     if (!this.memos[symbol]) {
       this.memos[symbol] = this.expand(symbol).evaluate(this.options);

test/grammar/memo-test.js

@@ -0,0 +1,37 @@
+import test from "ava";
+import grammar from "../../src/grammar.js";
+
+const _ = 0.5;
+const rng = () => 0.4;
+const fakeRandom = rngResults => () => rngResults.shift();
+
+test("global scoped memoized expansions", (t) => {
+  const g = grammar({
+    "start": "{report};{report}",
+    "report": "{@name}:{@name}",
+    "name": ["fluorite", "pyrite", "sodalite"]
+  }, { rng });
+
+  t.is(g.generate().text, "pyrite:pyrite;pyrite:pyrite");
+});
+
+test("branch scoped memoized expansions", (t) => {
+  const g = grammar({
+    "start": "{report};{report}",
+    "report": "{@@name}:{@@name}",
+    "name": ["fluorite", "pyrite", "sodalite"]
+  }, { rng: fakeRandom([_, _, 0.1, _, 0.7]) });
+
+  t.is(g.generate().text, "fluorite:fluorite;sodalite:sodalite");
+});
+
+test("branch scoped memoized expansions with nesting", (t) => {
+  const g = grammar({
+    "start": "{report};{report}",
+    "report": "{@@name}:{description}",
+    "description": "{@@name}",
+    "name": ["fluorite", "pyrite", "sodalite"]
+  }, { rng: fakeRandom([_, _, 0.7, _, _, 0.1, _]) });
+
+  t.is(g.generate().text, "sodalite:sodalite;fluorite:fluorite");
+});

test/production/choices-test.js

@@ -1,5 +1,6 @@
 import test from "ava";
 import Options from "../../src/options.js";
+import Registry from "../../src/registry.js";
 import choices from "../../src/production/choices.js";
 
 // TODO: result tree for this use case is undefined in the Ruby implementation
@@ -10,7 +11,7 @@ import choices from "../../src/production/choices.js";
 // })
 
 test("construct choices from list of atoms", (t) => {
-  const production = choices(["atom", "atom", "atom"]);
+  const production = choices(["atom", "atom", "atom"], new Registry({}, {}));
 
   t.deepEqual(production.evaluate(Options), [
     Symbol.for("choice"),

test/production/memo-test.js

@@ -2,20 +2,16 @@ import test from "ava";
 import Registry from "../../src/registry.js";
 import memo from "../../src/production/memo.js";
 
-// TODO: this can be replaced once the registry supports injectable rng options
-const stubRand = () => 0.1;
-// const nativeRand = Math.random
-//
-// test.before(t => Math.random = stubRand)
-// test.after(t => Math.random = nativeRand)
+const rng = () => 0.1;
 
 test("uses the registry to memoize expansions", (t) => {
   const production = memo(
     "num",
-    new Registry({ rng: stubRand }, { num: ["1", "2", "3"] })
+    false,
+    new Registry({ rng }, { num: ["1", "2", "3"] })
   );
 
-  t.deepEqual(production.evaluate({ rng: stubRand }), [
+  t.deepEqual(production.evaluate({ rng }), [
     Symbol.for("num"),
     [Symbol.for("choice"), [Symbol.for("concat"), [[Symbol.for("atom"), "1"]]]],
   ]);

test/production/weighted-choices-test.js

@@ -1,11 +1,12 @@
 import test from "ava";
 import Options from "../../src/options.js";
+import Registry from "../../src/registry.js";
 import weightedChoices from "../../src/production/weighted-choices.js";
 
 const rng = () => 0.1;
 
 test("selects rules with a weighted choice", (t) => {
-  const production = weightedChoices({ "20%": 0.2, "80%": 0.8 }, rng);
+  const production = weightedChoices({ "20%": 0.2, "80%": 0.8 }, new Registry({}, {}));
 
   t.deepEqual(production.evaluate({ rng }), [
     Symbol.for("weighted_choice"),
@@ -14,13 +15,13 @@ test("selects rules with a weighted choice", (t) => {
 });
 
 test("raises error if weighted choices do not sum to 1", (t) => {
-  const production = () => weightedChoices({ "10%": 0.1, "70%": 0.7 }, rng);
+  const production = () => weightedChoices({ "10%": 0.1, "70%": 0.7 }, new Registry({}, {}));
 
   t.throws(production, { instanceOf: Error }, "Weights must sum to 1");
 });
 
 test("selects rules with integer weights", (t) => {
-  const production = weightedChoices({ "20%": 20, "80%": 80 }, rng);
+  const production = weightedChoices({ "20%": 20, "80%": 80 }, new Registry({}, {}));
 
   t.deepEqual(production.evaluate({ rng }), [
     Symbol.for("weighted_choice"),