@dallaylaen/ski-interpreter 1.0.0 → 1.1.0

package/lib/expr.js

@@ -1,4 +1,6 @@
-const { missingIndices, isSubset } = require('./util');
+'use strict';
+
+const { skipDup, isSubset } = require('./util');
 
 const globalOptions = {
   terse:   true,
@@ -13,7 +15,6 @@ class Expr {
   constructor () {
     if (new.target === Expr)
       throw new Error('Attempt to instantiate abstract class Expr');
-    this.arity = Infinity;
   }
 
   /**
@@ -79,6 +80,43 @@ class Expr {
     return new Map([[this, 1]]);
   }
 
+  /**
+   *   @desc Given a list of pairs of term, replaces every subtree
+   *         that is equivalent to the first term in pair with the second one.
+   *         If a simgle term is given, it is duplicated into a pair.
+   *
+   *   @example  S(SKK)(SKS).replace('I') = SII // we found 2 subtrees equivalent to I
+   *             and replaced them with I
+   *
+   *   @param {(Expr | [find: Expr, replace: Expr])[]} terms
+   *   @param {Object} [opt] - options
+   *   @return {Expr}
+   */
+  replace (terms, opt = {}) {
+    const pairs = [];
+    if (terms.length === 0)
+      return this; // nothing to replace, return self
+    for (const entry of terms) {
+      const pair = (Array.isArray(entry) ? entry : [entry, entry]);
+      pair[0] = pair[0].guess(opt).expr;
+      if (!pair[0])
+        throw new Error('Failed to canonize term ' + entry);
+      if (pair.length !== 2)
+        throw new Error('Expected a pair of terms to replace, got ' + entry);
+      pairs.push(pair);
+    }
+    return this._replace(pairs, opt) ?? this;
+  }
+
+  _replace (pairs, opt) {
+    const check = this.guess(opt).expr;
+    for (const [canon, term] of pairs) {
+      if (check.equals(canon))
+        return term;
+    }
+    return null;
+  }
+
   /**
    * @desc rought estimate of the complexity of the term
    * @return {number}
@@ -88,57 +126,75 @@ class Expr {
   }
 
   /**
+   * @desc Try to find an equivalent lambda term for the expression,
+   *       returning also the term's arity and some other properties.
+   *
+   *       This is used internally when declaring a Native term,
+   *       unless {canonize: false} is used.
+   *
+   *       As of current it only recognizes terms that have a normal form,
+   *       perhaps after adding some variables. This may change in the future.
    *
-   * @param {{max: number?, maxArgs: number?, bestGuess?: Expr}} options
+   *       Use lambdify() if you want to get a lambda term in any case.
+   *
+   * @param {{max: number?, maxArgs: number?}} options
    * @return {{
-   *    found: boolean,
-   *    proper: boolean,
+   *    normal: boolean,
+   *    steps: number,
+   *    expr: Expr?,
    *    arity: number?,
-   *    linear: boolean?,
-   *    canonical?: Expr,
-   *    steps: number?,
-   *    skip: Set<number>?
+   *    proper: boolean?,
+   *    discard: boolean?,
+   *    duplicate: boolean?,
+   *    skip: Set<number>?,
+   *    dup: Set<number>?
    * }}
    */
-  canonize (options = {}) {
+  guess (options = {}) {
     const max = options.max ?? globalOptions.max;
     const maxArgs = options.maxArgs ?? globalOptions.maxArgs;
+    const out = this._guess({ max, maxArgs, index: 0 });
+    return out;
+  }
 
-    let steps = 0;
-    let expr = this;
-    const jar = [];
-    for (let i = 0; i < maxArgs; i++) {
-      const calc = expr.run({ max });
-      steps += calc.steps;
-      if (!calc.final)
-        break;
-      expr = calc.expr;
-      if (!expr.wantsArgs()) {
-        // found!
-        const symbols = expr.getSymbols();
-        const skip = missingIndices(jar, symbols);
-        const proper = isSubset(symbols.keys(), new Set(jar));
-        const duplicates = [...symbols.entries()].filter(([_, v]) => v > 1);
-        const linear = proper && skip.size === 0 && duplicates.length === 0;
-        return {
-          arity:     i,
-          found:     true,
-          canonical: maybeLambda(jar, expr),
-          proper,
-          linear,
-          steps,
-          ...(skip.size ? { skip } : {}),
-        };
-      }
-      const next = new FreeVar('abcdefgh'[i] ?? 'x' + i);
-      jar.push(next);
-      expr = expr.apply(next);
+  _guess (options, preArgs = [], steps = 0) {
+    if (preArgs.length > options.maxArgs || steps > options.max)
+      return { normal: false, steps };
+
+    // happy case
+    if (this.freeOnly()) {
+      return {
+        normal: true,
+        steps,
+        ...maybeLambda(preArgs, this),
+      };
     }
 
-    const fallback = { found: false, proper: false, steps };
-    if (options.bestGuess)
-      fallback.canonical = options.bestGuess;
-    return fallback;
+    // try reaching the normal form
+    const next = this.run({ max: (options.max - steps) / 3 });
+    steps += next.steps;
+    if (!next.final)
+      return { normal: false, steps };
+
+    // normal form != this, redo exercise
+    if (next.steps !== 0)
+      return next.expr._guess(options, preArgs, steps);
+
+    if (this._firstVar())
+      return { normal: false, steps };
+
+    const push = nthvar(preArgs.length + options.index);
+    return this.apply(push)._guess(options, [...preArgs, push], steps);
+  }
+
+  _aslist () {
+    return [this];
+  }
+
+  _firstVar () {
+    // boolean, whether the expression starts with a free variable
+    // used by guess()
+    return false;
   }
 
   /**
@@ -197,16 +253,6 @@ class Expr {
     return this;
   }
 
-  /**
-   * @desc Whether the term will reduce further if given more arguments.
-   *       In practice, equivalent to "starts with a FreeVar"
-   *       Used by canonize (duh...)
-   * @return {boolean}
-   */
-  wantsArgs () {
-    return true;
-  }
-
   /**
      * Apply self to list of given args.
      * Normally, only native combinators know how to do it.
@@ -249,7 +295,8 @@ class Expr {
     }
     let expr = args ? this.apply(...args) : this;
     let steps = opt.steps ?? 0;
-    const max = (opt.max ?? globalOptions.max) + steps;
+    // make sure we make at least 1 step, to tell whether we've reached the normal form
+    const max = Math.max(opt.max ?? globalOptions.max, 1) + steps;
     let final = false;
     for (; steps < max; ) {
       const next = expr.step();
@@ -293,28 +340,38 @@ class Expr {
   }
 
   /**
-     *
-     * @param {Expr} other
-     * @return {boolean}
-     */
+   *
+   * @param {Expr} other
+   * @return {boolean}
+   */
   equals (other) {
-    return this === other;
+    if (this === other)
+      return true;
+    if (other instanceof Alias)
+      return other.equals(this);
+    return false;
   }
 
   contains (other) {
     return this === other || this.equals(other);
   }
 
-  expect (other) {
-    if (!(other instanceof Expr))
-      throw new Error('Attempt to expect a combinator to equal something else: ' + other);
-    if (this.equals(other))
+  /**
+   * @desc Assert expression equality. Can be used in tests.
+   * @param {Expr} expected
+   * @param {string} comment
+   */
+  expect (expected, comment = '') {
+    comment = comment ? comment + ': ' : '';
+    if (!(expected instanceof Expr))
+      throw new Error(comment + 'attempt to expect a combinator to equal something else: ' + expected);
+    if (this.equals(expected))
       return;
 
     // TODO wanna use AssertionError but webpack doesn't recognize it
     // still the below hack works for mocha-based tests.
-    const poorMans = new Error('Found term ' + this + ' but expected ' + other);
-    poorMans.expected = other.toString();
+    const poorMans = new Error(comment + 'found term ' + this + ' but expected ' + expected);
+    poorMans.expected = expected.toString();
     poorMans.actual = this.toString();
     throw poorMans;
   }
@@ -330,10 +387,11 @@ class Expr {
   }
 
   /**
-   *
+   * @desc Whether the expression needs parentheses when printed.
+   * @param {boolean} [first] - whether this is the first term in a sequence
    * @return {boolean}
    */
-  needsParens () {
+  needsParens (first) {
     return false;
   }
 
@@ -346,22 +404,6 @@ class Expr {
   }
 }
 
-/**
- *   Constants that define when whitespace between terms may be omitted in App.toString()
- */
-const BITS = 4;
-const [T_UNKNOWN, T_PARENS, T_UPPER, T_LOWER]
-    = (function * () { for (let i = 0; ; yield i++); })();
-const canLump = new Set([
-  (T_PARENS  << BITS) + T_PARENS,
-  (T_PARENS  << BITS) + T_UPPER,
-  (T_UPPER   << BITS) + T_PARENS,
-  (T_UPPER   << BITS) + T_UPPER,
-  (T_UPPER   << BITS) + T_LOWER,
-  (T_LOWER   << BITS) + T_PARENS,
-  (T_UNKNOWN << BITS) + T_PARENS,
-]);
-
 class App extends Expr {
   /**
    * @desc Application of fun() to args.
@@ -373,67 +415,104 @@ class App extends Expr {
     if (args.length === 0)
       throw new Error('Attempt to create an application with no arguments (likely interpreter bug)');
     super();
-    this.fun = fun;
-    this.args = args;
+
+    this.arg = args.pop();
+    this.fun = args.length ? new App(fun, ...args) : fun;
     this.final = false;
+    this.arity = this.fun.arity > 0 ? this.fun.arity - 1 : 0;
   }
 
   weight () {
-    return this.args.reduce((acc, x) => acc + x.weight(), this.fun.weight());
+    return this.fun.weight() + this.arg.weight();
   }
 
   getSymbols () {
     const out = this.fun.getSymbols();
-    for (const term of this.args) {
-      for (const [key, value] of term.getSymbols())
-        out.set(key, (out.get(key) ?? 0) + value);
-    }
+    for (const [key, value] of this.arg.getSymbols())
+      out.set(key, (out.get(key) ?? 0) + value);
     return out;
   }
 
-  wantsArgs () {
-    return this.fun.wantsArgs();
+  _guess (options, preArgs = [], steps = 0) {
+    if (preArgs.length > options.maxArgs || steps > options.max)
+      return { normal: false, steps };
+
+    /*
+     * inside and App there are 3 main possibilities:
+     * 1) The parent guess() actually is able to do the job. Then we just proxy the result.
+     * 2) Both `fun` and `arg` form good enough lambda terms. Then lump them together & return.
+     * 3) We literally have no idea, so we just pick the shortest defined term from the above.
+     */
+
+    const proxy = super._guess(options, preArgs, steps);
+    if (proxy.normal)
+      return proxy;
+    steps = proxy.steps; // reimport extra iterations
+
+    const [first, ...list] = this._aslist();
+    if (!(first instanceof FreeVar))
+      return { normal: false, steps }
+    // TODO maybe do it later
+
+    let discard = false;
+    let duplicate = false;
+    const out = [];
+    for (const term of list) {
+      const guess = term._guess({
+        ...options,
+        maxArgs: options.maxArgs - preArgs.length,
+        max:     options.max - steps,
+        index:   preArgs.length + options.index,
+      });
+      steps += guess.steps;
+      if (!guess.normal)
+        return { normal: false, steps };
+      out.push(guess.expr);
+      discard = discard || guess.discard;
+      duplicate = duplicate || guess.duplicate;
+    }
+
+    return {
+      normal: true,
+      steps,
+      ...maybeLambda(preArgs, new App(first, ...out), {
+        discard,
+        duplicate,
+      }),
+    };
+  }
+
+  _firstVar () {
+    return this.fun._firstVar();
   }
 
   apply (...args) {
     if (args.length === 0)
       return this;
-    return this.fun.apply( ...this.args, ...args);
+    return new App(this, ...args);
   }
 
   expand () {
-    return this.fun.expand().apply(...this.args.map(x => x.expand()));
+    return this.fun.expand().apply(this.arg.expand());
   }
 
-  canonize (options = {}) {
-    const [fun, arg] = this.split().map(x => x.canonize(options).canonical);
-    return super.canonize({
-      ...options,
-      ...(fun && arg ? { bestGuess: fun.apply(arg) } : {})
-    });
+  _replace (pairs, opt) {
+    const maybe = super._replace(pairs, opt);
+    if (maybe)
+      return maybe;
+    const [fun, arg] = this.split();
+    return (fun._replace(pairs, opt) ?? fun).apply(arg._replace(pairs, opt) ?? arg);
   }
 
   renameVars (seq) {
-    const fun = this.fun.renameVars(seq);
-    const args = this.args.map(x => x.renameVars(seq));
-    return fun.apply(...args);
+    return this.fun.renameVars(seq).apply(this.arg.renameVars(seq));
   }
 
   subst (plug, value) {
     const fun = this.fun.subst(plug, value);
-    let change = fun === null ? 0 : 1;
-    const args = [];
-    for (const x of this.args) {
-      const next = x.subst(plug, value);
-      if (next === null)
-        args.push(x);
-      else {
-        args.push(next);
-        change++;
-      }
-    }
+    const arg = this.arg.subst(plug, value);
 
-    return change ? (fun ?? this.fun).apply(...args) : null;
+    return (fun || arg) ? (fun ?? this.fun).apply(arg ?? this.arg) : null;
   }
 
   /**
@@ -443,34 +522,39 @@ class App extends Expr {
   step () {
     // normal reduction order: first try root, then at most 1 step
     if (!this.final) {
-      const reduced = this.fun.reduce(this.args);
-      if (reduced)
-        return { expr: reduced, steps: 1, changed: true };
-
+      if (this.arity === 0) {
+        // aha! we have just fulfilled some previous function's argument demands
+        const reduced = this.fun.reduce([this.arg]);
+        // should always be true, but whatever
+        if (reduced)
+          return { expr: reduced, steps: 1, changed: true };
+      }
       // now try recursing
 
       const fun = this.fun.step();
       if (fun.changed)
-        return { expr: fun.expr.apply(...this.args), steps: fun.steps, changed: true };
-
-      for (let i = 0; i < this.args.length; i++) {
-        const next = this.args[i].step();
-        if (!next.changed)
-          continue;
-        const args = this.args.slice();
-        args[i] = next.expr;
-        return { expr: this.fun.apply(...args), steps: next.steps, changed: true };
-      }
+        return { expr: fun.expr.apply(this.arg), steps: fun.steps, changed: true };
+
+      const arg = this.arg.step();
+      if (arg.changed)
+        return { expr: this.fun.apply(arg.expr), steps: arg.steps, changed: true };
+
+      this.final = true;
     }
-    this.final = true;
     return { expr: this, steps: 0, changed: false };
   }
 
+  reduce (args) {
+    return this.fun.reduce([this.arg, ...args]);
+  }
+
   split () {
     // pretend we are an elegant (cons fun arg) and not a sleazy imperative array
-    const args = this.args.slice();
-    const last = args.pop();
-    return [this.fun.apply(...args), last];
+    return [this.fun, this.arg];
+  }
+
+  _aslist () {
+    return [...this.fun._aslist(), this.arg];
   }
 
   /**
@@ -481,63 +565,39 @@ class App extends Expr {
   _rski (options) {
     if (options.steps >= options.max)
       return this;
-    return this.fun._rski(options).apply(...this.args.map(x => x._rski(options)));
+    return this.fun._rski(options).apply(this.arg._rski(options));
   }
 
   equals (other) {
     if (!(other instanceof App))
-      return false;
-    if (other.args.length !== this.args.length)
-      return false;
-    if (!this.fun.equals(other.fun))
-      return false;
-    for (let i = 0; i < this.args.length; i++) {
-      if (!this.args[i].equals(other.args[i]))
-        return false;
-    }
-    return true;
+      return super.equals(other);
+
+    return this.fun.equals(other.fun) && this.arg.equals(other.arg);
   }
 
   contains (other) {
-    if (this.fun.contains(other))
-      return true;
-    for (const subtree of this.args) {
-      if (subtree.contains(other))
-        return true;
-    }
-    return super.contains(other);
+    return this.fun.contains(other) || this.arg.contains(other) || super.contains(other);
+  }
+
+  needsParens (first) {
+    return !first;
   }
 
   toString (opt = {}) {
+    const fun = this.fun.toString(opt);
+    const root = this.fun.needsParens(true) ? '(' + fun + ')' : fun;
     if (opt.terse ?? globalOptions.terse) {
-      const out = [];
-      let oldType = 0;
-      // stupid ad-hoc state machine, see above for constant definitions
-      for (const term of [this.fun, ...this.args]) {
-        let s = term.toString(opt);
-        let newType = T_UNKNOWN;
-        if (s.match(/^[A-Z]$/))
-          newType = T_UPPER;
-        else if (term instanceof FreeVar || s.match(/^[a-z][a-z_0-9]*$/))
-          newType = T_LOWER;
-        else if (s.match(/^[0-9]+$/))
-          // no special treatment for numerals, skip
-          ;
-        else if (out.length !== 0 || term.needsParens()) {
-          s = '(' + s + ')';
-          newType = T_PARENS;
-        }
-        if (!canLump.has((oldType << BITS) | newType) && out.length > 0)
-          out.push(' ');
-        out.push(s);
-        oldType = newType;
-      }
-      return out.join('');
-    } else {
-      const fun = this.fun.toString(opt);
-      const root = this.fun.needsParens() ? '(' + fun + ')' : fun;
-      return root + this.args.map(x => '(' + x.toString(opt) + ')').join('');
-    }
+      // terse mode: omit whitespace and parens if possible
+      let arg = this.arg.toString(opt);
+      if (this.arg.needsParens(false))
+        arg = '(' + arg + ')';
+      const space = (root.match(/\)$/) || arg.match(/^\(/))
+        || (root.match(/[A-Z]$/) && arg.match(/^[a-z]/i))
+        ? ''
+        : ' ';
+      return root + space + arg;
+    } else
+      return root + '(' + this.arg.toString(opt) + ')';
   }
 }
 
@@ -576,8 +636,8 @@ class FreeVar extends Named {
     return 0;
   }
 
-  wantsArgs () {
-    return false;
+  _firstVar () {
+    return true;
   }
 
   toString ( opt = {} ) {
@@ -609,12 +669,12 @@ class Native extends Named {
     this.arity = opt.arity ?? 1;
 
     // try to bootstrap and guess some of our properties
-    const guess = (opt.canonize ?? true) ? this.canonize() : { found: false };
+    const guess = (opt.canonize ?? true) ? this.guess() : { normal: false };
 
     if (!opt.arity)
       this.arity = guess.arity || 1;
 
-    this.note = opt.note ?? guess.canonical?.toString({ terse: true, html: true });
+    this.note = opt.note ?? guess.expr?.toString({ terse: true, html: true });
   }
 
   apply (...args) {
@@ -633,7 +693,7 @@ class Native extends Named {
   _rski (options) {
     if (this === native.I || this === native.K || this === native.S || (options.steps >= options.max))
       return this;
-    const canon = this.canonize().canonical;
+    const canon = this.guess().expr;
     if (!canon)
       return this;
     options.steps++;
@@ -711,6 +771,14 @@ class Lambda extends Expr {
     return this.impl.weight() + 1;
   }
 
+  _guess (options, preArgs = [], steps = 0) {
+    if (preArgs.length > options.maxArgs)
+      return { normal: false, steps };
+
+    const push = nthvar(preArgs.length + options.index);
+    return this.reduce([push])._guess(options, [...preArgs, push], steps + 1);
+  }
+
   reduce (input) {
     if (input.length === 0)
       return null;
@@ -762,9 +830,17 @@ class Lambda extends Expr {
     throw new Error('Don\'t know how to convert to SKI' + this);
   }
 
+  _replace (pairs, opt) {
+    const maybe = super._replace(pairs, opt);
+    if (maybe)
+      return maybe;
+    // TODO filter out terms containing this.arg
+    return new Lambda(this.arg, this.impl._replace(pairs, opt) ?? this.impl);
+  }
+
   equals (other) {
     if (!(other instanceof Lambda))
-      return false;
+      return super.equals(other);
 
     const t = new FreeVar('t');
 
@@ -780,7 +856,7 @@ class Lambda extends Expr {
     return this.arg.toString(opt) + mapsto + this.impl.toString(opt);
   }
 
-  needsParens () {
+  needsParens (first) {
     return true;
   }
 }
@@ -807,7 +883,7 @@ class Church extends Native {
   equals (other) {
     if (other instanceof Church)
       return this.n === other.n;
-    return false;
+    return super.equals(other);
   }
 }
 
@@ -826,12 +902,12 @@ class Alias extends Named {
       this.note = options.note;
 
     const guess = options.canonize
-      ? impl.canonize({ max: options.max, maxArgs: options.maxArgs })
-      : { found: false };
-    this.arity = (guess.found && guess.proper && guess.arity) || 0;
+      ? impl.guess({ max: options.max, maxArgs: options.maxArgs })
+      : { normal: false };
+    this.arity = (guess.proper && guess.arity) || 0;
     this.proper = guess.proper ?? false;
     this.terminal = options.terminal ?? this.proper;
-    this.canonical = guess.canonical;
+    this.canonical = guess.expr;
   }
 
   getSymbols () {
@@ -850,6 +926,10 @@ class Alias extends Named {
     return this.impl.subst(plug, value);
   }
 
+  _guess (options, preArgs = [], steps = 0) {
+    return this.impl._guess(options, preArgs, steps);
+  }
+
   /**
    *
    * @return {{expr: Expr, steps: number}}
@@ -868,8 +948,8 @@ class Alias extends Named {
     return this.impl.apply(...args);
   }
 
-  wantsArgs () {
-    return this.impl.wantsArgs();
+  _firstVar () {
+    return this.impl._firstVar();
   }
 
   equals (other) {
@@ -888,7 +968,7 @@ class Alias extends Named {
     return this.outdated ? this.impl.toString(opt) : super.toString(opt);
   }
 
-  needsParens () {
+  needsParens (first) {
     return this.outdated ? this.impl.needsParens() : false;
   }
 }
@@ -906,15 +986,25 @@ addNative('+', x => y => z => y.apply(x.apply(y, z)), {
   apply: arg => arg instanceof Church ? new Church(arg.n + 1) : null
 });
 
-function maybeLambda (args, expr) {
-  if (args.length === 0)
-    return expr;
-  return new Lambda(args, expr);
+function maybeLambda (args, expr, caps = {}) {
+  const sym = expr.getSymbols();
+
+  const [skip, dup] = skipDup(args, sym);
+
+  return {
+    expr:      args.length ? new Lambda(args, expr) : expr,
+    ...(caps.synth ? {} : { arity: args.length }),
+    ...(skip.size ? { skip } : {}),
+    ...(dup.size ? { dup } : {}),
+    duplicate: !!dup.size  || caps.duplicate || false,
+    discard:   !!skip.size || caps.discard   || false,
+    proper:    isSubset(sym.keys(), new Set(args)),
+  };
 }
 
 function naiveCanonize (expr) {
   if (expr instanceof App)
-    return naiveCanonize(expr.fun).apply(...expr.args.map(naiveCanonize));
+    return naiveCanonize(expr.fun).apply(naiveCanonize(expr.arg));
 
   if (expr instanceof Lambda)
     return new Lambda(expr.arg, naiveCanonize(expr.impl));
@@ -922,9 +1012,9 @@ function naiveCanonize (expr) {
   if (expr instanceof Alias)
     return naiveCanonize(expr.impl);
 
-  const canon = expr.canonize();
-  if (canon.canonical)
-    return canon.canonical;
+  const canon = expr.guess();
+  if (canon.expr)
+    return canon.expr;
 
   throw new Error('Failed to canonize expression: ' + expr);
 }
@@ -934,66 +1024,62 @@ function naiveCanonize (expr) {
  * @param {Expr} expr
  * @param {{max: number?, maxArgs: number?}} options
  * @param {number} maxWeight
- * @return {IterableIterator<{expr: Expr, steps: number?, comment: string?}>}
+ * @yields {{expr: Expr, steps: number?, comment: string?}}
  */
-function * simplifyLambda (expr, options = {}, maxWeight = Infinity) {
+function * simplifyLambda (expr, options = {}, state = { steps: 0 }) {
   // expr is a lambda, free variable, or an application thereof
   // we want to find an equivalent lambda term with less weight
   // which we do sequentially from leaves to the root of the AST
 
+  yield { expr, steps: state.steps, comment: '(self)' };
+
   // short-circuit
-  if (expr.freeOnly()) {
-    if (expr.weight() < maxWeight)
-      yield { expr, steps: 0, comment: 'only free vars' };
+  if (expr.freeOnly())
     return;
-  }
 
-  let steps = 0;
-  let savedSteps = 0;
+  let maxWeight = expr.weight();
+
+  if (expr instanceof Lambda) {
+    for (const term of simplifyLambda(expr.impl, options, state)) {
+      const candidate = new Lambda(expr.arg, term.expr);
+      if (candidate.weight() < maxWeight) {
+        maxWeight = candidate.weight();
+        yield { expr: candidate, steps: state.steps, comment: '(lambda)' + term.comment };
+      }
+    }
+  }
 
   // fun * arg Descartes product
   if (expr instanceof App) {
     // try to split into fun+arg, then try canonization but exposing each step
-    const [fun, arg] = expr.split();
+    let [fun, arg] = expr.split();
 
-    for (const term of simplifyLambda(fun, options, maxWeight - 1)) {
+    for (const term of simplifyLambda(fun, options, state)) {
       const candidate = term.expr.apply(arg);
-      steps = savedSteps + term.steps;
       if (candidate.weight() < maxWeight) {
         maxWeight = candidate.weight();
-        yield { expr: candidate, steps: term.steps, comment: term.comment + '(app)' };
+        fun = term.expr;
+        yield { expr: candidate, steps: state.steps, comment: '(fun)' + term.comment };
       }
     }
-    savedSteps = steps;
 
-    for (const term of simplifyLambda(arg, options, maxWeight - 1)) {
+    for (const term of simplifyLambda(arg, options, state)) {
       const candidate = fun.apply(term.expr);
-      steps = savedSteps + term.steps;
       if (candidate.weight() < maxWeight) {
         maxWeight = candidate.weight();
-        yield { expr: candidate, steps: term.steps, comment: term.comment + '(app)' };
+        yield { expr: candidate, steps: state.steps, comment: '(arg)' + term.comment };
       }
     }
-    savedSteps = steps;
   }
 
-  if (expr instanceof Lambda) {
-    for (const term of simplifyLambda(expr.impl, options, maxWeight - 1)) {
-      const candidate = new Lambda(expr.arg, term.expr);
-      if (candidate.weight() < maxWeight) {
-        maxWeight = candidate.weight();
-        steps = savedSteps + term.steps;
-        yield { expr: candidate, steps: term.steps, comment: term.comment + '(lambda)' };
-      }
-    }
-    savedSteps = steps;
-  }
+  const canon = expr.guess({ max: options.max, maxArgs: options.maxArgs });
+  state.steps += canon.steps;
+  if (canon.expr && canon.expr.weight() < maxWeight)
+    yield { expr: canon.expr, steps: state.steps, comment: '(canonical)' };
+}
 
-  const canon = expr.canonize({ max: options.max, maxArgs: options.maxArgs });
-  if (canon.canonical && canon.canonical.weight() < maxWeight) {
-    maxWeight = canon.canonical.weight();
-    yield { expr: canon.canonical, steps: savedSteps + canon.steps, comment: 'canonical' };
-  }
+function nthvar (n) {
+  return new FreeVar('abcdefgh'[n] ?? 'x' + n);
 }
 
 // A global value meaning "lambda is used somewhere in this expression"