From e68c1ea168a32fa3f7677c6b27b3ddd576b73105 Mon Sep 17 00:00:00 2001
From: Qiming Chu <cchuqiming@gmail.com>
Date: Mon, 23 Dec 2024 00:07:18 +0800
Subject: [PATCH] Add input seq length check to Mux1H and PriorityMux

Added a new require to check that the input sequences have the same
length in both Mux1H and PriorityMux implementations.
---
 src/main/scala/chisel3/util/MuxImpl.scala             | 9 +++++++--
 src/test/scala/chiselTests/OneHotMuxSpec.scala        | 6 ++++++
 src/test/scala/chiselTests/util/PriorityMuxSpec.scala | 7 +++++++
 3 files changed, 20 insertions(+), 2 deletions(-)

diff --git a/src/main/scala/chisel3/util/MuxImpl.scala b/src/main/scala/chisel3/util/MuxImpl.scala
index a9c5f4668b..93ddd8264a 100644
--- a/src/main/scala/chisel3/util/MuxImpl.scala
+++ b/src/main/scala/chisel3/util/MuxImpl.scala
@@ -23,8 +23,10 @@ import chisel3.experimental.SourceInfo
   * @note results unspecified unless exactly one select signal is high
   */
 object Mux1H {
-  def apply[T <: Data](sel: Seq[Bool], in: Seq[T]): T =
+  def apply[T <: Data](sel: Seq[Bool], in: Seq[T]): T = {
+    require(sel.size == in.size, "Mux1H: Number of select signals and inputs must match")
     apply(sel.zip(in))
+  }
   def apply[T <: Data](in:  Iterable[(Bool, T)]): T = SeqUtils.oneHotMux(in)
   def apply[T <: Data](sel: UInt, in: Seq[T]): T =
     apply((0 until in.size).map(sel(_)), in)
@@ -46,7 +48,10 @@ object Mux1H {
   */
 object PriorityMux {
   def apply[T <: Data](in:  Seq[(Bool, T)]): T = SeqUtils.priorityMux(in)
-  def apply[T <: Data](sel: Seq[Bool], in: Seq[T]): T = apply(sel.zip(in))
+  def apply[T <: Data](sel: Seq[Bool], in: Seq[T]): T = {
+    require(sel.size == in.size, "PriorityMux: Number of select signals and inputs must match")
+    apply(sel.zip(in))
+  }
   def apply[T <: Data](sel: Bits, in: Seq[T]): T = apply((0 until in.size).map(sel(_)), in)
 }
 
diff --git a/src/test/scala/chiselTests/OneHotMuxSpec.scala b/src/test/scala/chiselTests/OneHotMuxSpec.scala
index da10f34a58..dfc4845fa3 100644
--- a/src/test/scala/chiselTests/OneHotMuxSpec.scala
+++ b/src/test/scala/chiselTests/OneHotMuxSpec.scala
@@ -31,6 +31,12 @@ class OneHotMuxSpec extends AnyFreeSpec with Matchers with ChiselRunners {
     }
     e.getMessage should include("Mux1H must have a non-empty argument")
   }
+  "Mux1H should give a error when given two different size" in {
+    val e = intercept[IllegalArgumentException] {
+      Mux1H(Seq(true.B, true.B), Seq(1.U, 2.U, 3.U))
+    }
+    e.getMessage should include("Mux1H: Number of select signals and inputs must match")
+  }
 
 }
 
diff --git a/src/test/scala/chiselTests/util/PriorityMuxSpec.scala b/src/test/scala/chiselTests/util/PriorityMuxSpec.scala
index 53f0d89658..68889c6952 100644
--- a/src/test/scala/chiselTests/util/PriorityMuxSpec.scala
+++ b/src/test/scala/chiselTests/util/PriorityMuxSpec.scala
@@ -47,6 +47,13 @@ class PriorityMuxSpec extends ChiselFlatSpec {
     assertTesterPasses(new PriorityMuxTester)
   }
 
+  it should "give a error when inputs are two different size" in {
+    val e = intercept[IllegalArgumentException] {
+      PriorityMux(Seq(true.B, true.B), Seq(1.U, 2.U, 3.U))
+    }
+    e.getMessage should include("PriorityMux: Number of select signals and inputs must match")
+  }
+
   it should "be stack safe" in {
     emitCHIRRTL(new RawModule {
       val n = 1 << 15