From fe5891ac43fa9bcebead950e50159c6c5115137f Mon Sep 17 00:00:00 2001
From: Adam Izraelevitz <adam.izraelevitz@sifive.com>
Date: Fri, 13 Oct 2023 14:37:33 -0700
Subject: [PATCH] Added disalowIOCreation as a public API

---
 core/src/main/scala/chisel3/IO.scala          |  3 ++-
 core/src/main/scala/chisel3/Module.scala      |  2 +-
 .../scala/chiselTests/FixedIOModuleSpec.scala | 23 ++++++++++++++++---
 3 files changed, 23 insertions(+), 5 deletions(-)

diff --git a/core/src/main/scala/chisel3/IO.scala b/core/src/main/scala/chisel3/IO.scala
index bfb0cdc5637..0eae2d06625 100644
--- a/core/src/main/scala/chisel3/IO.scala
+++ b/core/src/main/scala/chisel3/IO.scala
@@ -23,7 +23,8 @@ object IO {
     */
   def apply[T <: Data](iodef: => T)(implicit sourceInfo: SourceInfo): T = {
     val module = Module.currentModule.get // Impossible to fail
-    require(module.isIOCreationAllowed, "This module cannot have user-created IO")
+    if (!module.isIOCreationAllowed)
+      Builder.error("This module cannot have IOs instantiated after calling disallowIOCreation()")
     require(!module.isClosed, "Can't add more ports after module close")
     val prevId = Builder.idGen.value
     val data = iodef // evaluate once (passed by name)
diff --git a/core/src/main/scala/chisel3/Module.scala b/core/src/main/scala/chisel3/Module.scala
index df4db7c2dfa..5457822cf2d 100644
--- a/core/src/main/scala/chisel3/Module.scala
+++ b/core/src/main/scala/chisel3/Module.scala
@@ -403,7 +403,7 @@ package experimental {
     private[chisel3] def isIOCreationAllowed = _isIOCreationAllowed
 
     /** Disallow any more IO creation for this module. */
-    private[chisel3] def disallowIOCreation(): Unit = {
+    def disallowIOCreation(): Unit = {
       _isIOCreationAllowed = false
     }
 
diff --git a/src/test/scala/chiselTests/FixedIOModuleSpec.scala b/src/test/scala/chiselTests/FixedIOModuleSpec.scala
index 616104b253e..1784408f8ab 100644
--- a/src/test/scala/chiselTests/FixedIOModuleSpec.scala
+++ b/src/test/scala/chiselTests/FixedIOModuleSpec.scala
@@ -20,10 +20,10 @@ class FixedIOModuleSpec extends ChiselFlatSpec with Utils with MatchesAndOmits {
     class Foo extends FixedIORawModule[Bool](Bool()) {
       val a = IO(Bool())
     }
-    val exception = intercept[IllegalArgumentException] {
-      ChiselStage.emitCHIRRTL(new Foo)
+    val exception = intercept[ChiselException] {
+      ChiselStage.emitCHIRRTL(new Foo, Array("--throw-on-first-error"))
     }
-    exception.getMessage should include("This module cannot have user-created IO")
+    exception.getMessage should include("This module cannot have IOs instantiated after calling disallowIOCreation()")
   }
 
   "FixedIOBlackBox" should "create a module with flattend IO" in {
@@ -47,4 +47,21 @@ class FixedIOModuleSpec extends ChiselFlatSpec with Utils with MatchesAndOmits {
     )()
   }
 
+  "User defined RawModules" should "be able to lock down their ios" in {
+
+    class Bar extends RawModule {
+      val in = IO(Input(UInt(1.W)))
+      val out = IO(Output(UInt(1.W)))
+
+      disallowIOCreation()
+
+      val other = IO(Input(UInt(1.W)))
+    }
+
+    val exception = intercept[ChiselException] {
+      ChiselStage.emitCHIRRTL(new Bar, Array("--throw-on-first-error"))
+    }
+    exception.getMessage should include("This module cannot have IOs instantiated after calling disallowIOCreation()")
+  }
+
 }