2 files changed, 58 insertions, 4 deletions
diff --git a/src/main/scala/chisel3/util/Mux.scala b/src/main/scala/chisel3/util/Mux.scala
index 342531bd..b9a5f576 100644
--- a/src/main/scala/chisel3/util/Mux.scala
+++ b/src/main/scala/chisel3/util/Mux.scala
@@ -63,10 +63,21 @@ object MuxLookup {
     * @return the value found or the default if not
     */
   def apply[S <: UInt, T <: Data] (key: S, default: T, mapping: Seq[(S, T)]): T = {
-    var res = default
-    for ((k, v) <- mapping.reverse)
-      res = Mux(k === key, v, res)
-    res
+    /* If the mapping is defined for all possible values of the key, then don't use the default value */
+    val (defaultx, mappingx) = key.widthOption match {
+      case Some(width) =>
+        val keySetSize = BigInt(1) << width
+        val keyMask = keySetSize - 1
+        val distinctLitKeys = mapping.flatMap(_._1.litOption).map(_ & keyMask).distinct
+        if (distinctLitKeys.size == keySetSize) {
+          (mapping.head._2, mapping.tail)
+        } else {
+          (default, mapping)
+        }
+      case None => (default, mapping)
+    }
+
+    mappingx.foldLeft(defaultx){ case (d, (k, v)) => Mux(k === key, v, d) }
   }
 }
 
diff --git a/src/test/scala/chiselTests/MuxSpec.scala b/src/test/scala/chiselTests/MuxSpec.scala
index 0ead960a..46710133 100644
--- a/src/test/scala/chiselTests/MuxSpec.scala
+++ b/src/test/scala/chiselTests/MuxSpec.scala
@@ -3,6 +3,7 @@
 package chiselTests
 
 import chisel3._
+import chisel3.util.{MuxLookup, log2Ceil}
 import chisel3.testers.BasicTester
 
 class MuxTester extends BasicTester {
@@ -25,3 +26,45 @@ class MuxSpec extends ChiselFlatSpec {
     assertTesterPasses { new MuxTester }
   }
 }
+
+class MuxLookupWrapper(keyWidth: Int, default: Int, mapping: () => Seq[(UInt, UInt)]) extends RawModule {
+  val outputWidth = log2Ceil(default).max(keyWidth) // make room for default value
+  val key = IO(Input(UInt(keyWidth.W)))
+  val output = IO(Output(UInt(outputWidth.W)))
+  output := MuxLookup(key, default.U, mapping())
+}
+
+class MuxLookupExhaustiveSpec extends ChiselPropSpec {
+  val keyWidth = 2
+  val default = 9 // must be less than 10 to avoid hex/decimal mismatches
+  val firrtlLit = s"""UInt<4>("h0$default")"""
+
+  // Assumes there are no literals with 'UInt<4>("h09")' in the output FIRRTL
+  // Assumes no binary recoding in output
+
+  val incomplete = () => Seq(0.U -> 1.U, 1.U -> 2.U, 2.U -> 3.U)
+  property("The default value should not be optimized away for an incomplete MuxLookup") {
+    Driver.emit { () => new MuxLookupWrapper(keyWidth, default, incomplete) } should include (firrtlLit)
+  }
+
+  val exhaustive = () => (3.U -> 0.U) +: incomplete()
+  property("The default value should be optimized away for an exhaustive MuxLookup") {
+    Driver.emit { () => new MuxLookupWrapper(keyWidth, default, exhaustive) } should not include (firrtlLit)
+  }
+
+  val overlap = () => (4096.U -> 0.U) +: incomplete()
+  property("The default value should not be optimized away for a MuxLookup with 2^{keyWidth} non-distinct mappings") {
+    Driver.emit { () => new MuxLookupWrapper(keyWidth, default, overlap) } should include (firrtlLit)
+  }
+
+  val nonLiteral = () => { val foo = Wire(UInt()); (foo -> 1.U) +: incomplete() }
+  property("The default value should not be optimized away for a MuxLookup with a non-literal") {
+    Driver.emit { () => new MuxLookupWrapper(keyWidth, default, nonLiteral) } should include (firrtlLit)
+  }
+
+  val nonLiteralStillFull = () => { val foo = Wire(UInt()); (foo -> 1.U) +: exhaustive() }
+  property("The default value should be optimized away for a MuxLookup with a non-literal that is still full") {
+    Driver.emit { () => new MuxLookupWrapper(keyWidth, default, nonLiteralStillFull) } should not include (firrtlLit)
+  }
+
+}