Better register const prop through speculative de-optimization

* Fixes #1240
* Add failing reg const prop test case from #1240

2 files changed, 31 insertions, 0 deletions

diff --git a/src/main/scala/firrtl/transforms/ConstantPropagation.scala b/src/main/scala/firrtl/transforms/ConstantPropagation.scala
index 66cf24b8..ad784559 100644
--- a/src/main/scala/firrtl/transforms/ConstantPropagation.scala
+++ b/src/main/scala/firrtl/transforms/ConstantPropagation.scala
@@ -490,6 +490,12 @@ class ConstantPropagation extends Transform with ResolvedAnnotationPaths {
           * NonConstant values.  When encountering a node reference, it expands the node by to its
           * RHS assignment and recurses.
           *
+          * @note Some optimization of Mux trees turn 1-bit mux operators into boolean operators. This
+          * can stifle register constant propagations, which looks at drivers through value-preserving
+          * Muxes and Connects only. By speculatively expanding some 1-bit Or and And operations into
+          * muxes, we can obtain the best possible insight on the value of the mux with a simple peephole
+          * de-optimization that does not actually appear in the output code.
+          *
           * @return a RegCPEntry describing the constant prop-compatible sources driving this expression
           */
           def regConstant(e: Expression): RegCPEntry = e match {
@@ -498,6 +504,8 @@ class ConstantPropagation extends Transform with ResolvedAnnotationPaths {
             case WRef(nodeName, _, NodeKind, _) if nodeMap.contains(nodeName) =>
               nodeRegCPEntries.getOrElseUpdate(nodeName, { regConstant(nodeMap(nodeName)) })
             case Mux(_, tval, fval, _) => regConstant(tval).resolve(regConstant(fval))
+            case DoPrim(Or, Seq(a, b), Nil, BoolType) => regConstant(Mux(a, one, b, BoolType))
+            case DoPrim(And, Seq(a, b), Nil, BoolType) => regConstant(Mux(a, b, zero, BoolType))
             case _ => RegCPEntry(NonConstant, NonConstant)
           }
 
diff --git a/src/test/scala/firrtlTests/ConstantPropagationTests.scala b/src/test/scala/firrtlTests/ConstantPropagationTests.scala
index af186cda..99b8c2f0 100644
--- a/src/test/scala/firrtlTests/ConstantPropagationTests.scala
+++ b/src/test/scala/firrtlTests/ConstantPropagationTests.scala
@@ -1148,6 +1148,29 @@ class ConstantPropagationIntegrationSpec extends LowTransformSpec {
     execute(input, check, Seq.empty)
   }
 
+  "Const prop of registers" should "do limited speculative expansion of optimized muxes to absorb bigger cones" in {
+      val input =
+        """circuit Top :
+          |  module Top :
+          |    input clock : Clock
+          |    input reset : UInt<1>
+          |    input en : UInt<1>
+          |    output out : UInt<1>
+          |    reg r : UInt<1>, clock
+          |    when en :
+          |      r <= UInt<1>(1)
+          |    out <= r""".stripMargin
+      val check =
+        """circuit Top :
+          |  module Top :
+          |    input clock : Clock
+          |    input reset : UInt<1>
+          |    input en : UInt<1>
+          |    output out : UInt<1>
+          |    out <= UInt<1>("h1")""".stripMargin
+    execute(input, check, Seq.empty)
+  }
+
   "A register with constant reset and all connection to either itself or the same constant" should "be replaced with that constant" in {
       val input =
         """circuit Top :