Catch returns from within when blocks and provide an error message (#842)

Resolves #841

3 files changed, 31 insertions, 2 deletions

diff --git a/chiselFrontend/src/main/scala/chisel3/core/Module.scala b/chiselFrontend/src/main/scala/chisel3/core/Module.scala
index c4a48fb4..0cf05496 100644
--- a/chiselFrontend/src/main/scala/chisel3/core/Module.scala
+++ b/chiselFrontend/src/main/scala/chisel3/core/Module.scala
@@ -49,7 +49,7 @@ object Module {
     val module: T = bc  // bc is actually evaluated here
 
     if (Builder.whenDepth != 0) {
-      throwException("Internal Error! When depth is != 0, this should not be possible")
+      throwException("Internal Error! when() scope depth is != 0, this should have been caught!")
     }
     if (Builder.readyForModuleConstr) {
       throwException("Error: attempted to instantiate a Module, but nothing happened. " +
diff --git a/chiselFrontend/src/main/scala/chisel3/core/When.scala b/chiselFrontend/src/main/scala/chisel3/core/When.scala
index 55012946..fb246e1b 100644
--- a/chiselFrontend/src/main/scala/chisel3/core/When.scala
+++ b/chiselFrontend/src/main/scala/chisel3/core/When.scala
@@ -71,7 +71,14 @@ final class WhenContext(sourceInfo: SourceInfo, cond: Option[() => Bool], block:
   if (firrtlDepth > 0) { pushCommand(AltBegin(sourceInfo)) }
   cond.foreach( c => pushCommand(WhenBegin(sourceInfo, c().ref)) )
   Builder.whenDepth += 1
-  block
+  try {
+    block
+  } catch {
+    case ret: scala.runtime.NonLocalReturnControl[_] =>
+      throwException("Cannot exit from a when() block with a \"return\"!" +
+        " Perhaps you meant to use Mux or a Wire as a return value?"
+      )
+  }
   Builder.whenDepth -= 1
   cond.foreach( c => pushCommand(WhenEnd(sourceInfo,firrtlDepth)) )
   if (cond.isEmpty) { pushCommand(OtherwiseEnd(sourceInfo,firrtlDepth)) }
diff --git a/src/test/scala/chiselTests/When.scala b/src/test/scala/chiselTests/When.scala
index a1a7afb7..58d5a1bd 100644
--- a/src/test/scala/chiselTests/When.scala
+++ b/src/test/scala/chiselTests/When.scala
@@ -106,4 +106,26 @@ class WhenSpec extends ChiselFlatSpec {
   "Conditional connections to submodule ports" should "be handled properly" in {
     assertTesterPasses(new SubmoduleWhenTester)
   }
+
+  "Returning in a when scope" should "give a reasonable error message" in {
+    val e = the [ChiselException] thrownBy {
+      elaborate(new Module {
+        val io = IO(new Bundle {
+          val foo = Input(UInt(8.W))
+          val bar = Input(UInt(8.W))
+          val cond = Input(Bool())
+          val out = Output(UInt(8.W))
+        })
+        def func(): UInt = {
+          when(io.cond) {
+            // This is bad, do not do this!!!
+            return io.foo
+          }
+          return io.bar
+        }
+        io.out := func()
+      })
+    }
+    e.getMessage should include ("Cannot exit from a when() block with a \"return\"")
+  }
 }