fixup! Add SparseVec

chipsalliance · Nov 3, 2023 · 3a8d7a9 · 3a8d7a9
1 parent 8c2efa6
commit 3a8d7a9
Showing 1 changed file with 29 additions and 17 deletions.
diff --git a/src/main/scala/chisel3/util/SparseVec.scala b/src/main/scala/chisel3/util/SparseVec.scala
@@ -8,12 +8,20 @@ import chisel3.util.{log2Up, BitPat, Mux1H}
 import chisel3.util.experimental.decode.{decoder, TruthTable}
 import scala.collection.mutable.{LinkedHashMap, LinkedHashSet}
 
+sealed trait SparseVec[A] {
+
+  def apply(addr: UInt)(implicit _sourceInfo: SourceInfo): A
+
+}
+
 object SparseVec {
 
   def binary[A <: Data](
     size:    Int,
     mapping: Seq[(Int, A)]
-  ): (UInt) => A = {
+  )(
+    implicit _sourceInfo: SourceInfo
+  ): SparseVec[A] = {
     val addrWidth = log2Up(size)
 
     val valueMap = LinkedHashMap.from(mapping)
@@ -31,8 +39,8 @@ object SparseVec {
       BitPat.N(addrWidth)
     )
 
-    (addr: UInt) => {
-      values(decoder(addr, ttable))
+    new SparseVec[A] {
+      override def apply(addr: UInt)(implicit _sourceInfo: SourceInfo): A = values(decoder(addr, ttable))
     }
   }
 
@@ -41,7 +49,7 @@ object SparseVec {
     mapping: Seq[(Int, A)]
   )(
     implicit _sourceInfo: SourceInfo
-  ): (UInt) => A = {
+  ): SparseVec[A] = {
     val addrWidth = log2Up(size)
 
     val dataMap = LinkedHashMap
@@ -65,14 +73,18 @@ object SparseVec {
     )
 
     val default = WireInit(chiselTypeOf(mapping.head._2), DontCare)
-    (addr: UInt) => Mux1H(decoder(addr, ttable), default +: dataMap.keys.toSeq)
+
+    new SparseVec[A] {
+      override def apply(addr: UInt)(implicit _sourceInfo: SourceInfo): A =
+        Mux1H(decoder(addr, ttable), default +: dataMap.keys.toSeq)
+    }
   }
 
   def ifelse[A <: Data](
     mapping: Seq[(Int, A)]
   )(
     implicit _sourceInfo: SourceInfo
-  ): (UInt) => A = {
+  ): SparseVec[A] = {
     val tpe = chiselTypeOf(mapping.head._2)
 
     val dataMap = LinkedHashMap
@@ -84,19 +96,19 @@ object SparseVec {
         dataMap(v) = dataMap(v) += k
     }
 
-    def index(addr: UInt): A = {
-      val result = Wire(tpe)
-      result := 0.U
-      dataMap.foreach {
-        case (k, v) =>
-          when(VecInit(v.map(addr === _.U).toSeq).reduce(_ || _)) {
-            result := k
-          }
+    new SparseVec[A] {
+      override def apply(addr: UInt)(implicit _sourceInfo: SourceInfo): A = {
+        val result = Wire(tpe)
+        result := 0.U
+        dataMap.foreach {
+          case (k, v) =>
+            when(VecInit(v.map(addr === _.U).toSeq).reduce(_ || _)) {
+              result := k
+            }
+        }
+        result
       }
-      result
     }
-
-    index
   }
 
 }