Add tests for exhaustive MuxLookup optimization

Signed-off-by: Schuyler Eldridge <schuyler.eldridge@ibm.com>
Co-authored-by: Schuyler Eldridge <schuyler.eldridge@ibm.com>
Co-authored-by: Albert Magyar <albert.magyar@gmail.com>

1 files changed, 43 insertions, 0 deletions

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)
+  }
+
+}