Compile grounded subpatterns into hash-cons lookups

examples/src/main/scala/foresight/eqsat/examples/incremental/IncrementalSaturation.scala

@@ -2,7 +2,7 @@ package foresight.eqsat.examples.incremental
 
 import foresight.eqsat.parallel.ParallelMap
 import foresight.eqsat.rewriting.{EClassSearcher, EClassesToSearch, Rule, Searcher}
-import foresight.eqsat.rewriting.patterns.{CompiledPattern, Instruction, MachineEClassSearcher, PatternMatch}
+import foresight.eqsat.rewriting.patterns.{CompiledPattern, Instruction, PatternCompiler, MachineEClassSearcher, PatternMatch}
 import foresight.eqsat.{EClassCall, EClassRef, ENode}
 import foresight.eqsat.readonly.{EGraph, EGraphWithMetadata}
 import foresight.util.collections.UnsafeSeqFromArray
@@ -91,7 +91,14 @@ object IncrementalSaturation {
   ): Searcher[NodeT, PatternMatch[NodeT], EGraphWithMetadata[NodeT, EGraphT]] = {
     val newInstructions = ArrayBuffer[Instruction[NodeT, EGraphWithMetadata[NodeT, EGraphT]]]()
     var bindings = 0
-    for (instr <- pattern.instructions) {
+    // Lookups hide node bindings that require individual top-k checks. Expand only
+    // compiler-generated lookup tapes; preserve custom instructions otherwise.
+    val instructions = if (pattern.instructions.exists(_.isInstanceOf[Instruction.Lookup[_, _]])) {
+      require(pattern.instructions == PatternCompiler.compile[NodeT, EGraphWithMetadata[NodeT, EGraphT]](pattern.pattern),
+        "Compile with groundedLookups = false before adding custom instructions to an incremental pattern")
+      PatternCompiler.compile[NodeT, EGraphWithMetadata[NodeT, EGraphT]](pattern.pattern, groundedLookups = false)
+    } else pattern.instructions
+    for (instr <- instructions) {
       newInstructions.append(instr)
       instr match {
         case Instruction.BindNode(register, nodeType, definitions, uses, arity) =>

foresight/src/main/scala/foresight/eqsat/rewriting/patterns/Instruction.scala

@@ -381,6 +381,56 @@ object Instruction {
     }
   }
 
+  /**
+   * Checks a fully bound subtree through the hash-cons.
+   * @param register The machine register containing the e-class call that the subtree must match.
+   * @param term The subtree to look up, with atoms identifying machine registers containing its bound e-class calls.
+   */
+  final case class Lookup[NodeT, EGraphT <: EGraph[NodeT]](
+    register: Int,
+    term: MixedTree[NodeT, Int]
+  ) extends Instruction[NodeT, EGraphT] {
+    override val effects: Effects = Effects.none
+
+    override def execute(ctx: Execution[NodeT, EGraphT]): Boolean = {
+      def hasSlottedArgument(tree: MixedTree[NodeT, Int]): Boolean = tree match {
+        case MixedTree.Atom(r) => !ctx.machine.registerAt(r).args.isEmpty
+        case MixedTree.Node(_, _, _, args) => args.exists(hasSlottedArgument)
+      }
+
+      def lookup(tree: MixedTree[NodeT, Int]): EClassCall = tree match {
+        case MixedTree.Atom(r) => ctx.machine.registerAt(r)
+        case MixedTree.Node(t, definitions, uses, args) =>
+          val children = new Array[EClassCall](args.length)
+          var i = 0
+          while (i < children.length) {
+            val child = lookup(args(i))
+            if (child eq null) return null
+            children(i) = child
+            i += 1
+          }
+          ctx.graph.findOrNull(ENode(t, definitions, uses, ArraySeq.unsafeWrapArray(children)))
+      }
+
+      def matches(tree: MixedTree[NodeT, Int], call: EClassCall): Boolean = tree match {
+        case MixedTree.Atom(r) => ctx.graph.areSame(call, ctx.machine.registerAt(r))
+        case MixedTree.Node(t, definitions, uses, args) =>
+          ctx.graph.nodes(call, t).exists { node =>
+            node.definitions == definitions && node.uses == uses &&
+              node.args.length == args.length && args.indices.forall(i => matches(args(i), node.args(i)))
+          }
+      }
+
+      val expected = ctx.machine.registerAt(register)
+      val found = if (hasSlottedArgument(term)) matches(term, expected) else {
+        val actual = lookup(term)
+        (actual ne null) && ctx.graph.areSame(expected, actual)
+      }
+      if (found) ctx.continue()
+      else ctx.error(MachineError.LookupFailed(this, expected))
+    }
+  }
+
   /**
    * An instruction that binds a pattern variable to an e-class call in register.
    * @param register The index of the register to bind the variable to.

foresight/src/main/scala/foresight/eqsat/rewriting/patterns/MachineError.scala

@@ -14,6 +14,18 @@ trait MachineError[NodeT]
  * A companion object for the `MachineError` trait.
  */
 object MachineError {
+  /**
+   * An error that occurs when the fully bound subtree of a `Lookup` instruction is absent from the e-graph
+   * or does not match the expected e-class call.
+   * @param instruction The `Lookup` instruction that failed to find a matching subtree.
+   * @param call The e-class call that the subtree was expected to match.
+   * @tparam NodeT The type of the nodes in the e-graph.
+   * @tparam EGraphT The type of the e-graph.
+   */
+  final case class LookupFailed[NodeT, EGraphT <: EGraph[NodeT]](
+    instruction: Instruction.Lookup[NodeT, EGraphT], call: EClassCall
+  ) extends MachineError[NodeT]
+
   /**
    * An error that occurs when no e-node in an e-class matches the node type, definitions, uses, and argument count
    * of a `BindNode` instruction.

foresight/src/main/scala/foresight/eqsat/rewriting/patterns/PatternCompiler.scala

@@ -12,10 +12,12 @@ object PatternCompiler {
   /**
    * Compiles a pattern into a list of instructions for the pattern-matching virtual machine.
    * @param pattern The pattern to compile.
+   * @param groundedLookups Replace fully bound, slot-free subtrees with lookups.
    * @return The list of instructions.
    */
-  def compile[NodeT, EGraphT <: EGraph[NodeT]](pattern: MixedTree[NodeT, Pattern.Var]): List[Instruction[NodeT, EGraphT]] = {
-    val compiler = new PatternCompiler[NodeT, EGraphT]()
+  def compile[NodeT, EGraphT <: EGraph[NodeT]](pattern: MixedTree[NodeT, Pattern.Var],
+                                                        groundedLookups: Boolean = true): List[Instruction[NodeT, EGraphT]] = {
+    val compiler = new PatternCompiler[NodeT, EGraphT](groundedLookups)
     compiler.compile(pattern, 0)
   }
 }
@@ -24,7 +26,7 @@ object PatternCompiler {
 /**
  * A compiler that compiles patterns into a list of instructions for the pattern-matching virtual machine.
  */
-private final class PatternCompiler[NodeT, EGraphT <: EGraph[NodeT]] {
+private final class PatternCompiler[NodeT, EGraphT <: EGraph[NodeT]](groundedLookups: Boolean) {
   /**
    * A mapping from expression variable IDs to registers.
    */
@@ -35,6 +37,18 @@ private final class PatternCompiler[NodeT, EGraphT <: EGraph[NodeT]] {
    */
   private var tapeLength: Int = 0
 
+  private def isGrounded(pattern: MixedTree[NodeT, Pattern.Var]): Boolean = pattern match {
+    case MixedTree.Atom(v) => varToReg.contains(v)
+    case MixedTree.Node(_, definitions, uses, args) =>
+      definitions.isEmpty && uses.isEmpty && args.forall(isGrounded)
+  }
+
+  private def lookupTerm(pattern: MixedTree[NodeT, Pattern.Var]): MixedTree[NodeT, Int] = pattern match {
+    case MixedTree.Atom(v) => MixedTree.Atom(varToReg(v))
+    case MixedTree.Node(t, definitions, uses, args) =>
+      MixedTree.Node(t, definitions, uses, args.map(lookupTerm))
+  }
+
   def compile(pattern: MixedTree[NodeT, Pattern.Var], out: Int): List[Instruction[NodeT, EGraphT]] = {
     pattern match {
       case MixedTree.Atom(p) => p match {
@@ -51,6 +65,9 @@ private final class PatternCompiler[NodeT, EGraphT <: EGraph[NodeT]] {
 
       // If we encountered any other type of expression, we want to unpack it and compile its children.
       case MixedTree.Node(nodeType, definitions, uses, args) =>
+        if (groundedLookups && isGrounded(pattern)) {
+          return List(Instruction.Lookup(out, lookupTerm(pattern)))
+        }
         val intro = Instruction.BindNode[NodeT, EGraphT](out, nodeType, definitions, uses, args.length)
         val tapeLengthAtIntro = tapeLength + 1
         tapeLength += args.length

foresight/src/test/scala-3/foresight/eqsat/rewriting/patterns/GroundedLookupTest.scala

@@ -0,0 +1,130 @@
+package foresight.eqsat.rewriting.patterns
+
+import foresight.eqsat.{EClassCall, EClassRef, ENode, ShapeCall, Slot}
+import foresight.eqsat.{mutable, readonly}
+import foresight.eqsat.lang.*
+import scala.language.implicitConversions
+import org.junit.Test
+import org.junit.Assert._
+
+class GroundedLookupTest {
+  sealed trait Expr derives Language
+  final case class Literal(value: String) extends Expr
+  final case class F(lhs: Expr, rhs: Expr) extends Expr
+  final case class Four(a: Expr, b: Expr, c: Expr, d: Expr) extends Expr
+  final case class GNode(arg: Expr) extends Expr
+  final case class H(arg: Expr) extends Expr
+  final case class Missing(arg: Expr) extends Expr
+  final case class Variable(slot: Use[Slot]) extends Expr
+  final case class Ref(call: EClassCall) extends Expr derives Box
+  final case class PatternVar(variable: Pattern.Var) extends Expr derives Box
+
+  private val lang = summon[Language[Expr]]
+  type Op = LanguageOp[Expr]
+  type G = readonly.EGraph[Op]
+  private def compiled(p: Expr, optimized: Boolean): CompiledPattern[Op, G] = {
+    val tree = lang.toTree[Pattern.Var](p)
+    CompiledPattern(tree, PatternCompiler.compile[Op, G](tree, groundedLookups = optimized))
+  }
+
+  extension (g: mutable.EGraph[Op])
+    private def addExpr(expr: Expr): EClassCall = g.add(lang.toTree[EClassCall](expr))
+
+  private class CountingGraph(g: G) extends readonly.EGraph[Op] {
+    var candidates = 0
+    var lookups = 0
+    def canonicalizeOrNull(r: EClassRef): EClassCall = g.canonicalizeOrNull(r)
+    def canonicalize(n: ENode[Op]): ShapeCall[Op] = g.canonicalize(n)
+    def classes: Iterable[EClassRef] = g.classes
+    def nodes(c: EClassCall): Iterable[ENode[Op]] = g.nodes(c)
+    override def nodes(c: EClassCall, t: Op): Iterable[ENode[Op]] = new Iterable[ENode[Op]] {
+      def iterator: Iterator[ENode[Op]] = g.nodes(c, t).iterator.map { n => candidates += 1; n }
+    }
+    def users(r: EClassRef): Iterable[ENode[Op]] = g.users(r)
+    def findOrNull(n: ENode[Op]): EClassCall = { lookups += 1; g.findOrNull(n) }
+    def areSame(a: EClassCall, b: EClassCall): Boolean = g.areSame(a, b)
+  }
+
+  @Test def avoidsEnumeratingGroundedAlternatives(): Unit = {
+    val g = mutable.EGraph.empty[Op]
+    val leaves = (0 until 100).map(i => g.addExpr(Literal("a" + i)))
+    val hs = leaves.map(a => g.addExpr(H(Ref(a))))
+    g.unionMany(hs.tail.map(h => (hs.head, h)))
+    val gx = g.addExpr(GNode(Ref(leaves.head)))
+    val root = g.addExpr(F(Ref(gx), Ref(hs.head)))
+    val x = PatternVar(Pattern.Var.fresh())
+    val pattern = F(GNode(x), H(x))
+    val oldGraph = new CountingGraph(g)
+    val newGraph = new CountingGraph(g)
+    val expected = compiled(pattern, false).search(root, oldGraph).toSet
+    val actual = compiled(pattern, true).search(root, newGraph).toSet
+    assertEquals(expected, actual)
+    assertEquals(1, actual.size)
+    assertEquals(102, oldGraph.candidates)
+    assertEquals(2, newGraph.candidates)
+    assertEquals(1, newGraph.lookups)
+    var streamed = 0
+    compiled(pattern, true).search(root, g, (_, _) => { streamed += 1; false })
+    assertEquals(1, streamed)
+  }
+
+  @Test def nestedConstantsMissingNodesAndWrongClasses(): Unit = {
+    val g = mutable.EGraph.empty[Op]
+    val a = g.addExpr(Literal("a"))
+    val b = g.addExpr(Literal("b"))
+    val ha = g.addExpr(H(Ref(a)))
+    val hb = g.addExpr(H(Ref(b)))
+    val root = g.addExpr(F(Ref(a), Ref(hb)))
+    val x = PatternVar(Pattern.Var.fresh())
+    val p = F(x, H(x))
+    assertFalse(compiled(p, true).matches(root, g)) // h(a) exists, but in the wrong class
+    assertFalse(compiled(F(x, Missing(x)), true).matches(root, g))
+    g.unionMany(Seq((ha, hb)))
+    assertEquals(compiled(p, false).search(root, g).toSet, compiled(p, true).search(root, g).toSet)
+    assertTrue(compiled(p, true).matches(root, g))
+    val constant = H(Literal("a"))
+    assertTrue(compiled(constant, true).matches(ha, g))
+    assertFalse(compiled(H(Literal("absent")), true).matches(ha, g))
+    val (ir, ig) = lang.toEGraph(Literal("a"))
+    assertTrue(compiled(Literal("a"), true).matches(ir, ig))
+  }
+
+  @Test def preservesRegistersAfterLookupAndReportsFailures(): Unit = {
+    val g = mutable.EGraph.empty[Op]
+    val a = g.addExpr(Literal("a"))
+    val b = g.addExpr(Literal("b"))
+    val ha = g.addExpr(H(Ref(a)))
+    val gb = g.addExpr(GNode(Ref(b)))
+    val hb = g.addExpr(H(Ref(b)))
+    val root = g.addExpr(Four(Ref(a), Ref(ha), Ref(gb), Ref(hb)))
+    val x = PatternVar(Pattern.Var.fresh())
+    val y = PatternVar(Pattern.Var.fresh())
+    val p = Four(x, H(x), GNode(y), H(y))
+    val optimized = compiled(p, true)
+    assertEquals(compiled(p, false).search(root, g).toSet, optimized.search(root, g).toSet)
+    assertEquals(1, optimized.search(root, g).size)
+    val failure = compiled(Literal("absent"), true)
+    val machine = MutableMachineState[Op](a, Instruction.Effects.from(failure.instructions))
+    val results = Machine.tryRun(g, machine, failure.instructions)
+    assertEquals(1, results.size)
+    results.head match {
+      case MachineResult.Failure(_, _: MachineError.LookupFailed[?, ?], _) =>
+      case other => fail("Expected a lookup failure, got " + other)
+    }
+  }
+
+  @Test def slottedBindingsRetainStructuralMatching(): Unit = {
+    val g = mutable.EGraph.empty[Op]
+    val a = g.addExpr(Variable(Slot.fresh()))
+    val h = g.addExpr(H(Ref(a)))
+    val root = g.addExpr(F(Ref(a), Ref(h)))
+    val x = PatternVar(Pattern.Var.fresh())
+    val p = F(x, H(x))
+    val counted = new CountingGraph(g)
+    assertEquals(compiled(p, false).search(root, g).toSet, compiled(p, true).search(root, counted).toSet)
+    assertTrue(compiled(p, true).matches(root, g))
+    assertEquals(0, counted.lookups)
+    val slotPattern = Variable(Slot.fresh())
+    assertFalse(compiled(slotPattern, true).instructions.exists(_.isInstanceOf[Instruction.Lookup[?, ?]]))
+  }
+}