Merge pull request #296 from ucb-bar/fix-bits-type-take3

Fix bits type

5 files changed, 64 insertions, 34 deletions

diff --git a/src/main/scala/firrtl/LoweringCompilers.scala b/src/main/scala/firrtl/LoweringCompilers.scala
index c8430d2b..4d40d9a8 100644
--- a/src/main/scala/firrtl/LoweringCompilers.scala
+++ b/src/main/scala/firrtl/LoweringCompilers.scala
@@ -124,7 +124,6 @@ class HighFirrtlToMiddleFirrtl () extends Transform with SimpleRun {
 // TODO(izraelevitz): Create RenameMap from RemoveCHIRRTL
 class MiddleFirrtlToLowFirrtl () extends Transform with SimpleRun {
    val passSeq = Seq(
-      passes.Legalize,
       passes.LowerTypes,
       passes.ResolveKinds,
       passes.InferTypes,
@@ -146,6 +145,7 @@ class EmitVerilogFromLowFirrtl (val writer: Writer) extends Transform with Simpl
       passes.ConstProp,
       passes.PadWidths,
       passes.ConstProp,
+      passes.Legalize,
       passes.VerilogWrap,
       passes.SplitExpressions,
       passes.CommonSubexpressionElimination,
diff --git a/src/main/scala/firrtl/ir/IR.scala b/src/main/scala/firrtl/ir/IR.scala
index 8f1cabf6..cdbf2fd7 100644
--- a/src/main/scala/firrtl/ir/IR.scala
+++ b/src/main/scala/firrtl/ir/IR.scala
@@ -329,6 +329,9 @@ abstract class GroundType extends Type {
   val width: Width
   def mapType(f: Type => Type): Type = this
 }
+object GroundType {
+  def unapply(ground: GroundType): Option[Width] = Some(ground.width)
+}
 abstract class AggregateType extends Type {
   def mapWidth(f: Width => Width): Type = this
 }
diff --git a/src/main/scala/firrtl/passes/MemUtils.scala b/src/main/scala/firrtl/passes/MemUtils.scala
index b235213a..adbf23e5 100644
--- a/src/main/scala/firrtl/passes/MemUtils.scala
+++ b/src/main/scala/firrtl/passes/MemUtils.scala
@@ -83,8 +83,7 @@ object bitWidth {
   def widthOf(dt: Type): BigInt = dt match {
     case t: VectorType => t.size * bitWidth(t.tpe)
     case t: BundleType => t.fields.map(f => bitWidth(f.tpe)).foldLeft(BigInt(0))(_+_)
-    case UIntType(IntWidth(width)) => width
-    case SIntType(IntWidth(width)) => width
+    case GroundType(IntWidth(width)) => width
     case t => error("Unknown type encountered in bitWidth!")
   }
 }
diff --git a/src/main/scala/firrtl/passes/PadWidths.scala b/src/main/scala/firrtl/passes/PadWidths.scala
index 4cdcae59..1a134d11 100644
--- a/src/main/scala/firrtl/passes/PadWidths.scala
+++ b/src/main/scala/firrtl/passes/PadWidths.scala
@@ -18,11 +18,18 @@ object PadWidths extends Pass {
          case t: SIntType => SIntType(IntWidth(i))
          // default case should never be reached
       }
-      if (i > width(e))
+      if (i > width(e)) {
          DoPrim(Pad, Seq(e), Seq(i), tx)
-      else if (i < width(e))
-         DoPrim(Bits, Seq(e), Seq(i - 1, 0), tx)
-      else e
+      } else if (i < width(e)) {
+         val e2 = DoPrim(Bits, Seq(e), Seq(i - 1, 0), UIntType(IntWidth(i)))
+         // Bit Select always returns UInt, cast if selecting from SInt
+         e.tpe match {
+            case UIntType(_) => e2
+            case SIntType(_) => DoPrim(AsSInt, Seq(e2), Seq.empty, SIntType(IntWidth(i)))
+         }
+      } else {
+        e
+      }
    }
    // Recursive, updates expression so children exp's have correct widths
    private def onExp(e: Expression): Expression = {
diff --git a/src/main/scala/firrtl/passes/Passes.scala b/src/main/scala/firrtl/passes/Passes.scala
index b9808485..d5d9a3b6 100644
--- a/src/main/scala/firrtl/passes/Passes.scala
+++ b/src/main/scala/firrtl/passes/Passes.scala
@@ -211,42 +211,63 @@ object ExpandConnects extends Pass {
 // TODO replace UInt with zero-width wire instead
 object Legalize extends Pass {
   def name = "Legalize"
-  def legalizeShiftRight (e: DoPrim): Expression = e.op match {
-    case Shr => {
-      val amount = e.consts.head.toInt
-      val width = long_BANG(e.args.head.tpe)
-      lazy val msb = width - 1
-      if (amount >= width) {
-        e.tpe match {
-          case t: UIntType => UIntLiteral(0, IntWidth(1))
-          case t: SIntType =>
-            DoPrim(Bits, e.args, Seq(msb, msb), SIntType(IntWidth(1)))
-          case t => error(s"Unsupported type ${t} for Primop Shift Right")
-        }
-      } else {
-        e
+  private def legalizeShiftRight(e: DoPrim): Expression = {
+    require(e.op == Shr)
+    val amount = e.consts.head.toInt
+    val width = bitWidth(e.args.head.tpe)
+    lazy val msb = width - 1
+    if (amount >= width) {
+      e.tpe match {
+        case UIntType(_) => UIntLiteral(0, IntWidth(1))
+        case SIntType(_) =>
+          val bits = DoPrim(Bits, e.args, Seq(msb, msb), UIntType(IntWidth(1)))
+          DoPrim(AsSInt, Seq(bits), Seq.empty, SIntType(IntWidth(1)))
+        case t => error(s"Unsupported type ${t} for Primop Shift Right")
       }
+    } else {
+      e
+    }
+  }
+  private def legalizeBits(expr: DoPrim): Expression = {
+    lazy val (hi, low) = (expr.consts(0), expr.consts(1))
+    lazy val mask = (BigInt(1) << (hi - low + 1).toInt) - 1
+    lazy val width = IntWidth(hi - low + 1)
+    expr.args.head match {
+      case UIntLiteral(value, _) => UIntLiteral((value >> low.toInt) & mask, width)
+      case SIntLiteral(value, _) => SIntLiteral((value >> low.toInt) & mask, width)
+      case _ => expr
     }
-    case _ => e
   }
-  def legalizeConnect(c: Connect): Statement = {
+  private def legalizePad(expr: DoPrim): Expression = expr.args.head match {
+    case UIntLiteral(value, IntWidth(width)) if (width < expr.consts.head) =>
+      UIntLiteral(value, IntWidth(expr.consts.head))
+    case SIntLiteral(value, IntWidth(width)) if (width < expr.consts.head) =>
+      SIntLiteral(value, IntWidth(expr.consts.head))
+    case _ => expr
+  }
+  private def legalizeConnect(c: Connect): Statement = {
     val t = c.loc.tpe
-    val w = long_BANG(t)
-    if (w >= long_BANG(c.expr.tpe)) c
-    else {
-      val newType = t match {
-        case _: UIntType => UIntType(IntWidth(w))
-        case _: SIntType => SIntType(IntWidth(w))
+    val w = bitWidth(t)
+    if (w >= bitWidth(c.expr.tpe)) {
+      c
+    } else {
+      val bits = DoPrim(Bits, Seq(c.expr), Seq(w - 1, 0), UIntType(IntWidth(w)))
+      val expr = t match {
+        case UIntType(_) => bits
+        case SIntType(_) => DoPrim(AsSInt, Seq(bits), Seq(), SIntType(IntWidth(w)))
       }
-      Connect(c.info, c.loc, DoPrim(Bits, Seq(c.expr), Seq(w-1, 0), newType))
+      Connect(c.info, c.loc, expr)
     }
   }
   def run (c: Circuit): Circuit = {
-    def legalizeE (e: Expression): Expression = {
-      e map (legalizeE) match {
-        case e: DoPrim => legalizeShiftRight(e)
-        case e => e
+    def legalizeE(expr: Expression): Expression = expr map legalizeE match {
+      case prim: DoPrim => prim.op match {
+        case Shr => legalizeShiftRight(prim)
+        case Pad => legalizePad(prim)
+        case Bits => legalizeBits(prim)
+        case _ => prim
       }
+      case e => e // respect pre-order traversal
     }
     def legalizeS (s: Statement): Statement = {
       val legalizedStmt = s match {