Use ErrorPrinter in PatternRewrite facilities (#445)

When pattern rewriters catch an exception, they use the new
`ErrorPrinter` to print the current state of IR, indicate which exact op
was worked on by the failing pattern, and the full exception that
happened.

Author: PapyChacal

core/src/ir/IRUtils.scala

@@ -1,5 +1,6 @@
 package scair.ir
 
+import scala.annotation.tailrec
 import scala.collection.mutable.LinkedHashMap
 import scala.collection.mutable.ListBuffer
 import scala.collection.mutable.Map
@@ -25,6 +26,12 @@ import scala.collection.mutable.Map
 trait IRNode:
   def parent: Option[IRNode]
 
+  @tailrec
+  final def topLevel: IRNode =
+    parent match
+      case Some(p) => p.topLevel
+      case None    => this
+
   final def isAncestor(other: IRNode): Boolean =
     other.parent match
       case Some(parent) if parent == this => true

core/src/print/ErrorPrinter.scala

@@ -61,7 +61,33 @@ private final class ErrorPrinterFilter(writer: Writer)
   override def write(str: String): Unit =
     writeRec(str, 0)
 
-final class ErrorPrinter(
+object ErrorPrinter:
+
+  def apply(
+      error: Err,
+      w: Writer = PrintWriter(System.out),
+      _indent: String = "  ",
+      _valueNextID: Int = 0,
+      _blockNextID: Int = 0,
+      _valueNameMap: mutable.Map[Value[? <: Attribute], String] = mutable.Map
+        .empty,
+      _blockNameMap: mutable.Map[Block, String] = mutable.Map.empty,
+      _aliasesMap: Map[Attribute, String] = Map.empty,
+      _indentLevel: Int = 0,
+  ): ErrorPrinter =
+    new ErrorPrinter(
+      error,
+      new ErrorPrinterFilter(w),
+      _indent,
+      _valueNextID,
+      _blockNextID,
+      _valueNameMap,
+      _blockNameMap,
+      _aliasesMap,
+      _indentLevel,
+    )
+
+final class ErrorPrinter private (
     error: Err,
     w: ErrorPrinterFilter = ErrorPrinterFilter(new PrintWriter(System.out)),
     _indent: String = "  ",

core/src/transformations/PatternRewriter.scala

@@ -2,7 +2,10 @@ package scair.transformations
 
 import scair.helpers.*
 import scair.ir.*
+import scair.print.ErrorPrinter
+import scair.utils.Err
 
+import java.io.StringWriter
 import scala.annotation.tailrec
 import scala.collection.mutable.LinkedHashSet
 
@@ -24,6 +27,21 @@ import scala.collection.mutable.LinkedHashSet
 ||   Utils realm   ||
 \*≡==---==≡==---==≡*/
 
+private def augmentException(
+    e: Exception,
+    op: Operation,
+    pattern: RewritePattern,
+): Nothing =
+  val errorStringWriter = StringWriter()
+  val printer = ErrorPrinter(
+    Err(s"Caught exception applying $pattern", Some(op)),
+    errorStringWriter,
+  )
+  printer.printTopLevel(op.topLevel.asInstanceOf[Operation])
+  val errorString = errorStringWriter.toString()
+  errorStringWriter.close()
+  throw Exception(errorString, e)
+
 object InsertPoint:
 
   def before(op: Operation) =
@@ -199,7 +217,8 @@ case class GreedyRewritePatternApplier(patterns: Seq[RewritePattern])
     patterns match
       case Nil    => ()
       case h :: t =>
-        h.matchAndRewrite(op, rewriter)
+        try h.matchAndRewrite(op, rewriter)
+        catch case e: Exception => augmentException(e, op, h)
         if !rewriter.hasDoneAction then matchAndRewriteRec(op, rewriter, t)
 
   override def matchAndRewrite(
@@ -345,8 +364,7 @@ class PatternRewriteWalker(
       rewriter.currentOp = op
 
       try pattern.matchAndRewrite(op, rewriter)
-      catch
-        case e: Exception => throw e // throw new Exception("Caught exception!")
+      catch case e: Exception => augmentException(e, op, pattern)
 
       rewriter_done_action |= rewriter.hasDoneAction
 

tools/opt/src/ScairOpt.scala

@@ -116,8 +116,8 @@ trait ScairOptBase extends ScairToolBase[ScairOptArgs]:
   ): OK[Operation] =
     error match
       case Err(msg, Some(_)) =>
-        val p = new ErrorPrinter(error)
-        p.print(operation)
+        val p = ErrorPrinter(error)
+        p.printTopLevel(operation)
         if verifyDiagnostics then error else sys.exit(42)
 
   def main(args: Array[String]): Unit =