Add constant propagation for muxes on UInt

1 files changed, 17 insertions, 0 deletions

diff --git a/src/main/scala/firrtl/passes/Passes.scala b/src/main/scala/firrtl/passes/Passes.scala
index 5c4fbaf6..4f947202 100644
--- a/src/main/scala/firrtl/passes/Passes.scala
+++ b/src/main/scala/firrtl/passes/Passes.scala
@@ -1267,9 +1267,26 @@ object ConstProp extends Pass {
     case _ => e
   }
 
+  private def constPropMuxCond(m: Mux) = (m.cond, tpe(m.tval), tpe(m.fval), m.tpe) match {
+    case (c: UIntValue, ttpe: UIntType, ftpe: UIntType, mtpe: UIntType) =>
+      if (c.value == 1 && ttpe == mtpe) m.tval
+      else if (c.value == 0 && ftpe == mtpe) m.fval
+      else m
+    case _ => m
+  }
+
+  private def constPropMux(m: Mux): Expression = (m.tval, m.fval) match {
+    case (t: UIntValue, f: UIntValue) =>
+      if (t == f) t
+      else if (t.value == 1 && f.value == 0 && long_BANG(m.tpe) == 1) m.cond
+      else constPropMuxCond(m)
+    case _ => constPropMuxCond(m)
+  }
+
   private def constPropExpression(e: Expression): Expression = {
     e map constPropExpression match {
       case p: DoPrim => constPropPrim(p)
+      case m: Mux => constPropMux(m)
       case x => x
     }
   }