diff --git a/examples/src/main/scala/foresight/eqsat/examples/incremental/IncrementalSaturation.scala b/examples/src/main/scala/foresight/eqsat/examples/incremental/IncrementalSaturation.scala index 148a47e9..312bb624 100644 --- a/examples/src/main/scala/foresight/eqsat/examples/incremental/IncrementalSaturation.scala +++ b/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) => diff --git a/foresight/src/main/scala/foresight/eqsat/rewriting/patterns/Instruction.scala b/foresight/src/main/scala/foresight/eqsat/rewriting/patterns/Instruction.scala index 3e2660c7..28d02bfe 100644 --- a/foresight/src/main/scala/foresight/eqsat/rewriting/patterns/Instruction.scala +++ b/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. diff --git a/foresight/src/main/scala/foresight/eqsat/rewriting/patterns/MachineError.scala b/foresight/src/main/scala/foresight/eqsat/rewriting/patterns/MachineError.scala index b30fd956..dd32c73f 100644 --- a/foresight/src/main/scala/foresight/eqsat/rewriting/patterns/MachineError.scala +++ b/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. diff --git a/foresight/src/main/scala/foresight/eqsat/rewriting/patterns/PatternCompiler.scala b/foresight/src/main/scala/foresight/eqsat/rewriting/patterns/PatternCompiler.scala index 6f039830..6cad1d3b 100644 --- a/foresight/src/main/scala/foresight/eqsat/rewriting/patterns/PatternCompiler.scala +++ b/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 diff --git a/foresight/src/test/scala-3/foresight/eqsat/rewriting/patterns/GroundedLookupTest.scala b/foresight/src/test/scala-3/foresight/eqsat/rewriting/patterns/GroundedLookupTest.scala new file mode 100644 index 00000000..36062f4e --- /dev/null +++ b/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[?, ?]])) + } +}