Preserve literals across .asTypeOf

Casting a literal (of any type) to another type with .asTypeOf will result in a literal of the new type. For non-literals, the FIRRTL representation will now be a little bit more efficient.
chipsalliance · Jun 12, 2024 · 78a7e9a · 78a7e9a
1 parent 2997fa8
commit 78a7e9a
Show file tree

Hide file tree

Showing 15 changed files with 232 additions and 91 deletions.
diff --git a/core/src/main/scala/chisel3/Aggregate.scala b/core/src/main/scala/chisel3/Aggregate.scala
@@ -115,24 +115,51 @@ sealed abstract class Aggregate extends Data {
     }
   }
 
-  private[chisel3] override def connectFromBits(
-    that: Bits
-  )(
-    implicit sourceInfo: SourceInfo
-  ): Unit = {
-    var i = 0
-    val bits = if (that.isLit) that else WireDefault(UInt(this.width), that) // handles width padding
-    for (x <- flatten) {
-      val fieldWidth = x.getWidth
-      if (fieldWidth > 0) {
-        x.connectFromBits(bits(i + fieldWidth - 1, i))
-        i += fieldWidth
-      } else {
-        // There's a zero-width field in this bundle.
-        // Zero-width fields can't really be assigned to, but the frontend complains if there are uninitialized fields,
-        // so we assign it to DontCare. We can't use connectFromBits() on DontCare, so use := instead.
-        x := DontCare
+  // Return a literal of the same type from a Seq of literals for each leaf
+  private[chisel3] final def _makeLitFromLeaves(elems: Seq[Element])(implicit sourceInfo: SourceInfo): Data = {
+    // We could use virtual methods instead of matching on concrete subtypes,
+    // but the difference for Record vs Vec is so small, it doesn't seem worth it
+    val clone: Aggregate = this.cloneTypeFull
+    val mapping = clone.flatten.view
+      .zip(elems)
+      .map {
+        case (thisElt, litElt) =>
+          val litArg = litElt.topBindingOpt match {
+            case Some(ElementLitBinding(value)) => value
+            case _                              => throwException(s"Internal Error! For field $thisElt, given non-literal $litElt!")
+          }
+          thisElt -> litArg
+      }
+    val binding = clone match {
+      case r: Record => BundleLitBinding(mapping.to(Map))
+      case v: Vec[_] => VecLitBinding(mapping.to(VectorMap))
+    }
+    clone.bind(binding)
+    clone
+  }
+
+  override private[chisel3] def _fromUInt(that: UInt)(implicit sourceInfo: SourceInfo): Data = {
+    val _asUInt = _resizeToWidth(that, this.widthOption)(identity)
+    // If that is a literal and all constituent Elements can be represented as literals, return a literal
+    val ((_, allLit), rvalues) = {
+      this.flatten.toList.mapAccumulate((0, _asUInt.isLit)) {
+        case ((lo, literal), elt) =>
+          val hi = lo + elt.getWidth
+          // Chisel only supports zero width extraction if hi = -1 and lo = 0, so do it manually
+          val _extracted = if (elt.getWidth == 0) 0.U(0.W) else _asUInt(hi - 1, lo)
+          // _fromUInt returns Data but we know that it is an Element
+          val rhs = elt._fromUInt(_extracted).asInstanceOf[Element]
+          ((hi, literal && rhs.isLit), rhs)
+      }
+    }
+    if (allLit) {
+      this._makeLitFromLeaves(rvalues)
+    } else {
+      val _wire = Wire(this.cloneTypeFull)
+      for ((l, r) <- _wire.flatten.zip(rvalues)) {
+        l := r
       }
+      _wire
     }
   }
 }

diff --git a/core/src/main/scala/chisel3/Bits.scala b/core/src/main/scala/chisel3/Bits.scala
@@ -4,7 +4,7 @@ package chisel3
 
 import scala.language.experimental.macros
 import chisel3.experimental.{requireIsHardware, SourceInfo}
-import chisel3.internal.{throwException, BaseModule}
+import chisel3.internal.{_resizeToWidth, throwException, BaseModule}
 import chisel3.internal.Builder.pushOp
 import chisel3.internal.firrtl.ir._
 import chisel3.internal.sourceinfo.{
@@ -811,12 +811,8 @@ sealed class UInt private[chisel3] (width: Width) extends Bits(width) with Num[U
 
   override private[chisel3] def _asUIntImpl(first: Boolean)(implicit sourceInfo: SourceInfo): UInt = this
 
-  private[chisel3] override def connectFromBits(
-    that: Bits
-  )(
-    implicit sourceInfo: SourceInfo
-  ): Unit = {
-    this := that.asUInt
+  override private[chisel3] def _fromUInt(that: UInt)(implicit sourceInfo: SourceInfo): this.type = {
+    _resizeToWidth(that, this.widthOption)(identity).asInstanceOf[this.type]
   }
 
   private def subtractAsSInt(that: UInt)(implicit sourceInfo: SourceInfo): SInt =
@@ -1071,13 +1067,8 @@ sealed class SInt private[chisel3] (width: Width) extends Bits(width) with Num[S
 
   override def do_asSInt(implicit sourceInfo: SourceInfo): SInt = this
 
-  private[chisel3] override def connectFromBits(
-    that: Bits
-  )(
-    implicit sourceInfo: SourceInfo
-  ): Unit = {
-    this := that.asSInt
-  }
+  override private[chisel3] def _fromUInt(that: UInt)(implicit sourceInfo: SourceInfo): this.type =
+    _resizeToWidth(that.asSInt, this.widthOption)(_.asSInt).asInstanceOf[this.type]
 }
 
 sealed trait Reset extends Element with ToBoolable {
@@ -1118,12 +1109,10 @@ final class ResetType(private[chisel3] val width: Width = Width(1)) extends Elem
     DefPrim(sourceInfo, UInt(this.width), AsUIntOp, ref)
   )
 
-  private[chisel3] override def connectFromBits(
-    that: Bits
-  )(
-    implicit sourceInfo: SourceInfo
-  ): Unit = {
-    this := that
+  override private[chisel3] def _fromUInt(that: UInt)(implicit sourceInfo: SourceInfo): Data = {
+    val _wire = Wire(this.cloneTypeFull)
+    _wire := that
+    _wire
   }
 
   /** @group SourceInfoTransformMacro */
@@ -1162,14 +1151,7 @@ sealed class AsyncReset(private[chisel3] val width: Width = Width(1)) extends El
     DefPrim(sourceInfo, UInt(this.width), AsUIntOp, ref)
   )
 
-  // TODO Is this right?
-  private[chisel3] override def connectFromBits(
-    that: Bits
-  )(
-    implicit sourceInfo: SourceInfo
-  ): Unit = {
-    this := that.asBool.asAsyncReset
-  }
+  override private[chisel3] def _fromUInt(that: UInt)(implicit sourceInfo: SourceInfo): Data = that.asBool.asAsyncReset
 
   /** @group SourceInfoTransformMacro */
   def do_asAsyncReset(implicit sourceInfo: SourceInfo): AsyncReset = this
@@ -1299,4 +1281,8 @@ sealed class Bool() extends UInt(1.W) with Reset {
   /** @group SourceInfoTransformMacro */
   def do_asAsyncReset(implicit sourceInfo: SourceInfo): AsyncReset =
     pushOp(DefPrim(sourceInfo, AsyncReset(), AsAsyncResetOp, ref))
+
+  override private[chisel3] def _fromUInt(that: UInt)(implicit sourceInfo: SourceInfo): this.type = {
+    _resizeToWidth(that, this.widthOption)(identity).asBool.asInstanceOf[this.type]
+  }
 }
diff --git a/core/src/main/scala/chisel3/ChiselEnum.scala b/core/src/main/scala/chisel3/ChiselEnum.scala
@@ -54,13 +54,8 @@ abstract class EnumType(private[chisel3] val factory: ChiselEnum, selfAnnotating
     pushOp(DefPrim(sourceInfo, Bool(), op, this.ref, other.ref))
   }
 
-  private[chisel3] override def connectFromBits(
-    that: Bits
-  )(
-    implicit sourceInfo: SourceInfo
-  ): Unit = {
-    this := factory.apply(that.asUInt)
-  }
+  override private[chisel3] def _fromUInt(that: UInt)(implicit sourceInfo: SourceInfo): Data =
+    factory.apply(that.asUInt)
 
   final def ===(that: EnumType): Bool = macro SourceInfoTransform.thatArg
   final def =/=(that: EnumType): Bool = macro SourceInfoTransform.thatArg
@@ -82,6 +77,27 @@ abstract class EnumType(private[chisel3] val factory: ChiselEnum, selfAnnotating
   def do_>=(that: EnumType)(implicit sourceInfo: SourceInfo): Bool =
     compop(sourceInfo, GreaterEqOp, that)
 
+  // This is a hack to preserve the _workaround_ for https://github.com/chipsalliance/chisel/issues/4159
+  // Note that #4159 is due to _asUIntImpl below not actually padding the UInt
+  // This override just ensures that if `that` has a known width, the result actually has that width
+  // Put another way, this is preserving a case where #4159 does **not** occur
+  override def do_asTypeOf[T <: Data](that: T)(implicit sourceInfo: SourceInfo): T = {
+    that.widthOption match {
+      // Note that default case will handle literals just fine
+      case Some(w) =>
+        val _padded = this.litOption match {
+          case Some(value) =>
+            value.U(w.W)
+          case None =>
+            val _wire = Wire(UInt(w.W))
+            _wire := this.asUInt
+            _wire
+        }
+        _padded.do_asTypeOf(that)
+      case None => super.do_asTypeOf(that)
+    }
+  }
+
   override private[chisel3] def _asUIntImpl(first: Boolean)(implicit sourceInfo: SourceInfo): UInt =
     pushOp(DefPrim(sourceInfo, UInt(width), AsUIntOp, ref))
 

diff --git a/core/src/main/scala/chisel3/Clock.scala b/core/src/main/scala/chisel3/Clock.scala
@@ -38,11 +38,6 @@ sealed class Clock(private[chisel3] val width: Width = Width(1)) extends Element
   override private[chisel3] def _asUIntImpl(first: Boolean)(implicit sourceInfo: SourceInfo): UInt = pushOp(
     DefPrim(sourceInfo, UInt(this.width), AsUIntOp, ref)
   )
-  private[chisel3] override def connectFromBits(
-    that: Bits
-  )(
-    implicit sourceInfo: SourceInfo
-  ): Unit = {
-    this := that.asBool.asClock
-  }
+
+  override private[chisel3] def _fromUInt(that: UInt)(implicit sourceInfo: SourceInfo): Data = that.asBool.asClock
 }
diff --git a/core/src/main/scala/chisel3/Data.scala b/core/src/main/scala/chisel3/Data.scala
@@ -809,18 +809,12 @@ abstract class Data extends HasId with NamedComponent with SourceInfoDoc {
 
   /** @group SourceInfoTransformMacro */
   def do_asTypeOf[T <: Data](that: T)(implicit sourceInfo: SourceInfo): T = {
-    val thatCloned = Wire(that.cloneTypeFull)
-    thatCloned.connectFromBits(this.asUInt)
-    thatCloned
+    that._fromUInt(this.asUInt).asInstanceOf[T]
   }
 
-  /** Assigns this node from Bits type. Internal implementation for asTypeOf.
+  /** Return a value of this type from a UInt type. Internal implementation for asTypeOf.
     */
-  private[chisel3] def connectFromBits(
-    that: Bits
-  )(
-    implicit sourceInfo: SourceInfo
-  ): Unit
+  private[chisel3] def _fromUInt(that: UInt)(implicit sourceInfo: SourceInfo): Data
 
   /** Reinterpret cast to UInt.
     *
@@ -1195,12 +1189,9 @@ final case object DontCare extends Element with connectable.ConnectableDocs {
 
   def toPrintable: Printable = PString("DONTCARE")
 
-  private[chisel3] def connectFromBits(
-    that: Bits
-  )(
-    implicit sourceInfo: SourceInfo
-  ): Unit = {
-    Builder.error("connectFromBits: DontCare cannot be a connection sink (LHS)")
+  private[chisel3] def _fromUInt(that: UInt)(implicit sourceInfo: SourceInfo): Data = {
+    Builder.error("DontCare cannot be a connection sink (LHS)")
+    this
   }
 
   override private[chisel3] def _asUIntImpl(first: Boolean)(implicit sourceInfo: SourceInfo): UInt = {

diff --git a/core/src/main/scala/chisel3/experimental/Analog.scala b/core/src/main/scala/chisel3/experimental/Analog.scala
@@ -3,7 +3,7 @@
 package chisel3.experimental
 
 import chisel3.internal._
-import chisel3.{ActualDirection, Bits, Data, Element, PString, Printable, RawModule, SpecifiedDirection, UInt, Width}
+import chisel3._
 
 import scala.collection.mutable
 
@@ -68,12 +68,9 @@ final class Analog private (private[chisel3] val width: Width) extends Element {
   override private[chisel3] def _asUIntImpl(first: Boolean)(implicit sourceInfo: SourceInfo): UInt =
     throwException("Analog does not support asUInt")
 
-  private[chisel3] override def connectFromBits(
-    that: Bits
-  )(
-    implicit sourceInfo: SourceInfo
-  ): Unit = {
-    throwException("Analog does not support connectFromBits")
+  override private[chisel3] def _fromUInt(that: UInt)(implicit sourceInfo: SourceInfo): Data = {
+    Builder.error("Analog does not support fromUInt")
+    Wire(Analog(that.width))
   }
 
   def toPrintable: Printable = PString("Analog")

diff --git a/core/src/main/scala/chisel3/internal/package.scala b/core/src/main/scala/chisel3/internal/package.scala
@@ -13,6 +13,8 @@ import scala.annotation.implicitNotFound
 import scala.collection.mutable
 import chisel3.ChiselException
 
+import scala.reflect.runtime.universe.{typeTag, TypeTag}
+
 package object internal {
 
   @implicitNotFound("You are trying to access a macro-only API. Please use the @public annotation instead.")
@@ -66,6 +68,40 @@ package object internal {
     if (headOk) res else s"_$res"
   }
 
+  // Hack to work around https://github.com/chipsalliance/chisel/issues/4162
+  // We can't use the .asTypeOf workaround because this is used to implement .asTypeOf
+  private[chisel3] def _padHandleBool[A <: Bits](
+    x:     A,
+    width: Int
+  )(
+    implicit sourceInfo: SourceInfo,
+    tag:                 TypeTag[A]
+  ): A = x match {
+    case b: Bool if !b.isLit && width > 1 && tag.tpe =:= typeTag[UInt].tpe =>
+      val _pad = Wire(UInt(width.W))
+      _pad := b
+      _pad.asInstanceOf[A] // This cast is safe because we know A is UInt on this path
+    case u => u.pad(width)
+  }
+
+  // Resize that to this width (if known)
+  private[chisel3] def _resizeToWidth[A <: Bits: TypeTag](
+    that:           A,
+    targetWidthOpt: Option[Int]
+  )(fromUInt:       UInt => A
+  )(
+    implicit sourceInfo: SourceInfo
+  ): A =
+    (targetWidthOpt, that.widthOption) match {
+      case (Some(targetWidth), Some(thatWidth)) =>
+        if (targetWidth == thatWidth) that
+        else if (targetWidth > thatWidth) _padHandleBool(that, targetWidth)
+        else fromUInt(that.take(targetWidth))
+      case (Some(targetWidth), None) => fromUInt(_padHandleBool(that, targetWidth).take(targetWidth))
+      case (None, Some(thatWidth))   => that
+      case (None, None)              => that
+    }
+
   /** Internal API for [[ViewParent]] */
   sealed private[chisel3] class ViewParentAPI extends RawModule() with PseudoModule {
     // We must provide `absoluteTarget` but not `toTarget` because otherwise they would be exactly

diff --git a/core/src/main/scala/chisel3/properties/Property.scala b/core/src/main/scala/chisel3/properties/Property.scala
@@ -235,9 +235,10 @@ sealed trait Property[T] extends Element { self =>
     Builder.error(s"${this._localErrorContext} does not support .asUInt.")
     0.U
   }
-  private[chisel3] def connectFromBits(that: Bits)(implicit sourceInfo: SourceInfo): Unit = {
-    Builder.error(s"${this._localErrorContext} cannot be driven by Bits")
+  override private[chisel3] def _fromUInt(that: UInt)(implicit sourceInfo: SourceInfo): Data = {
+    Builder.exception(s"${this._localErrorContext} cannot be driven by UInt")
   }
+
   override private[chisel3] def firrtlConnect(that: Data)(implicit sourceInfo: SourceInfo): Unit = {
     that match {
       case pthat: Property[_] => MonoConnect.propConnect(sourceInfo, this, pthat, Builder.forcedUserModule)

diff --git a/src/test/scala/chiselTests/AsTypeOfTester.scala b/src/test/scala/chiselTests/AsTypeOfTester.scala
@@ -477,12 +477,15 @@ class AsTypeOfSpec extends ChiselFunSpec {
   describe("Analogs") {
     describe("as the target type") {
       they("should error") {
-        val e = the[ChiselException] thrownBy ChiselStage.emitSystemVerilog(new RawModule {
-          val in = IO(Input(UInt(8.W)))
-          val out = IO(Analog(8.W))
-          out := in.asTypeOf(out)
-        })
-        e.getMessage should include("Analog does not support connectFromBits")
+        val e = the[ChiselException] thrownBy ChiselStage.emitSystemVerilog(
+          new RawModule {
+            val in = IO(Input(UInt(8.W)))
+            val out = IO(Analog(8.W))
+            out := in.asTypeOf(out)
+          },
+          Array("--throw-on-first-error")
+        )
+        e.getMessage should include("Analog does not support fromUInt")
       }
     }
     describe("as the source type") {

diff --git a/src/test/scala/chiselTests/BundleLiteralSpec.scala b/src/test/scala/chiselTests/BundleLiteralSpec.scala
@@ -463,4 +463,38 @@ class BundleLiteralSpec extends ChiselFlatSpec with Utils {
     }
   }
 
+  "Casting a Bundle literal to a complex Bundle type" should "maintain the literal value" in {
+    class OtherBundle extends Bundle {
+      val a = UInt(2.W)
+      val b = Vec(
+        2,
+        new Bundle {
+          val foo = UInt(1.W)
+          val bar = SInt(3.W)
+        }
+      )
+    }
+    val blit = (new MyBundle).Lit(_.a -> 43.U, _.b -> true.B, _.c -> MyEnum.sB)
+    val olit = blit.asTypeOf(new OtherBundle)
+    olit.litOption should be(Some(0xaf))
+    olit.a.litValue should be(0)
+    olit.b.litValue should be(0xaf)
+    olit.b(0).litValue should be(0xf)
+    olit.b(0).foo.litValue should be(1)
+    olit.b(0).bar.litValue should be(-1)
+    olit.b(1).litValue should be(0xa)
+    olit.b(1).foo.litValue should be(1)
+    olit.b(1).bar.litValue should be(2)
+
+    assertTesterPasses {
+      new BasicTester {
+        // Check that it gives the same value as the generated hardware.
+        val wire = WireInit(blit).asTypeOf(new OtherBundle)
+        // ScalaTest has its own multiversal === which overrules extension method.
+        // Manually instantiate extension method to get around it.
+        chisel3.assert(new Data.DataEquality(olit).===(wire))
+        stop()
+      }
+    }
+  }
 }