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

Resolves #841

2 files changed, 9 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)) }