Move stack trimming from Driver to ChiselStage

Signed-off-by: Schuyler Eldridge <schuyler.eldridge@ibm.com>

2 files changed, 30 insertions, 14 deletions

diff --git a/src/main/scala/chisel3/Driver.scala b/src/main/scala/chisel3/Driver.scala
index a78cc92f..158ba65a 100644
--- a/src/main/scala/chisel3/Driver.scala
+++ b/src/main/scala/chisel3/Driver.scala
@@ -6,7 +6,7 @@ import chisel3.internal.ErrorLog
 import chisel3.experimental.RawModule
 import internal.firrtl._
 import firrtl._
-import firrtl.options.{Phase, PhaseManager}
+import firrtl.options.{Phase, PhaseManager, StageError}
 import firrtl.options.phases.DeletedWrapper
 import firrtl.options.Viewer.view
 import firrtl.annotations.JsonProtocol
@@ -226,16 +226,10 @@ object Driver extends BackendCompilationUtilities {
     val annosx = try {
       phases.foldLeft(annos)( (a, p) => p.transform(a) )
     } catch {
-      case ce: ChiselException =>
-        val stackTrace = if (!optionsManager.chiselOptions.printFullStackTrace) {
-          ce.chiselStackTrace
-        } else {
-          val sw = new StringWriter
-          ce.printStackTrace(new PrintWriter(sw))
-          sw.toString
-        }
-        Predef.augmentString(stackTrace).lines.foreach(line => println(s"${ErrorLog.errTag} $line")) // scalastyle:ignore regex line.size.limit
-        annos
+      /* ChiselStage and FirrtlStage can throw StageError. Since Driver is not a StageMain, it cannot catch these. While
+       * Driver is deprecated and removed in 3.2.1+, the Driver catches all errors.
+       */
+      case e: StageError => annos
     }
 
     view[ChiselExecutionResult](annosx)
diff --git a/src/main/scala/chisel3/stage/ChiselStage.scala b/src/main/scala/chisel3/stage/ChiselStage.scala
index 923867d7..aef1abb2 100644
--- a/src/main/scala/chisel3/stage/ChiselStage.scala
+++ b/src/main/scala/chisel3/stage/ChiselStage.scala
@@ -3,9 +3,15 @@
 package chisel3.stage
 
 import firrtl.AnnotationSeq
-import firrtl.options.{Phase, PhaseManager, PreservesAll, Shell, Stage}
+import firrtl.options.{Phase, PhaseManager, PreservesAll, Shell, Stage, StageError, StageMain}
 import firrtl.options.phases.DeletedWrapper
 import firrtl.stage.FirrtlCli
+import firrtl.options.Viewer.view
+
+import chisel3.ChiselException
+import chisel3.internal.ErrorLog
+
+import java.io.{StringWriter, PrintWriter}
 
 class ChiselStage extends Stage with PreservesAll[Phase] {
   val shell: Shell = new Shell("chisel") with ChiselCli with FirrtlCli
@@ -20,11 +26,27 @@ class ChiselStage extends Stage with PreservesAll[Phase] {
          classOf[chisel3.stage.phases.Convert],
          classOf[chisel3.stage.phases.MaybeFirrtlStage] )
 
-  def run(annotations: AnnotationSeq): AnnotationSeq =
-    /* @todo: Should this be wrapped in a try/catch? */
+  def run(annotations: AnnotationSeq): AnnotationSeq = try {
     new PhaseManager(targets) { override val wrappers = Seq( (a: Phase) => DeletedWrapper(a) ) }
       .transformOrder
       .map(firrtl.options.phases.DeletedWrapper(_))
       .foldLeft(annotations)( (a, f) => f.transform(a) )
+  } catch {
+    case ce: ChiselException =>
+      val stackTrace = if (!view[ChiselOptions](annotations).printFullStackTrace) {
+        ce.chiselStackTrace
+      } else {
+        val sw = new StringWriter
+        ce.printStackTrace(new PrintWriter(sw))
+        sw.toString
+      }
+      Predef
+        .augmentString(stackTrace)
+        .lines
+        .foreach(line => println(s"${ErrorLog.errTag} $line")) // scalastyle:ignore regex
+      throw new StageError()
+  }
 
 }
+
+object ChiselMain extends StageMain(new ChiselStage)