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

Fix bits type

7 files changed, 143 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 {
diff --git a/src/test/resources/passes/Legalize/Legalize.fir b/src/test/resources/passes/Legalize/Legalize.fir
new file mode 100644
index 00000000..6e27a0bf
--- /dev/null
+++ b/src/test/resources/passes/Legalize/Legalize.fir
@@ -0,0 +1,41 @@
+
+circuit Legalize :
+  module Legalize :
+    input clk : Clock
+    input reset : UInt<1>
+
+    ; Count till done
+    node done = UInt(6)
+    reg count : UInt<16>, clk with :
+      reset => (reset, UInt(0))
+    when neq(count, done) :
+      count <= add(count, UInt(1))
+    when not(reset) :
+      when eq(count, done) :
+        stop(clk, UInt(1), 0)
+
+    ; Begin Test
+    ; Check assignment to smaller width
+    node x = UInt<32>("hdeadbeef")
+    wire y : UInt<16>
+    y <- x
+    when neq(y, UInt("hbeef")) :
+      printf(clk, UInt(1), "Assertion failed!\n y != beef\n")
+      stop(clk, UInt(1), 1)
+
+    ; Check bit select of literal
+    node b = bits(UInt("hd0"), 7, 5)
+    node b2 = bits(UInt("h9"), 3, 3)
+    when neq(b, UInt(6)) :
+      printf(clk, UInt(1), "Assertion failed!\n b != 6\n")
+      stop(clk, UInt(1), 1)
+    when neq(b2, UInt(1)) :
+      printf(clk, UInt(1), "Assertion failed!\n b2 != 1\n")
+      stop(clk, UInt(1), 1)
+
+    ; Check padding of literal
+    node bar = pad(SInt(-1), 16)
+    node bar_15 = bits(bar, 15, 15)
+    when neq(bar_15, UInt(1)) :
+      printf(clk, UInt(1), "Assertion failed!\n bar_15 != 0\n")
+      stop(clk, UInt(1), 1)
diff --git a/src/test/scala/firrtlTests/LegalizeSpec.scala b/src/test/scala/firrtlTests/LegalizeSpec.scala
new file mode 100644
index 00000000..781f93d7
--- /dev/null
+++ b/src/test/scala/firrtlTests/LegalizeSpec.scala
@@ -0,0 +1,38 @@
+/*
+Copyright (c) 2014 - 2016 The Regents of the University of
+California (Regents). All Rights Reserved.  Redistribution and use in
+source and binary forms, with or without modification, are permitted
+provided that the following conditions are met:
+   * Redistributions of source code must retain the above
+     copyright notice, this list of conditions and the following
+     two paragraphs of disclaimer.
+   * Redistributions in binary form must reproduce the above
+     copyright notice, this list of conditions and the following
+     two paragraphs of disclaimer in the documentation and/or other materials
+     provided with the distribution.
+   * Neither the name of the Regents nor the names of its contributors
+     may be used to endorse or promote products derived from this
+     software without specific prior written permission.
+IN NO EVENT SHALL REGENTS BE LIABLE TO ANY PARTY FOR DIRECT, INDIRECT,
+SPECIAL, INCIDENTAL, OR CONSEQUENTIAL DAMAGES, INCLUDING LOST PROFITS,
+ARISING OUT OF THE USE OF THIS SOFTWARE AND ITS DOCUMENTATION, EVEN IF
+REGENTS HAS BEEN ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+REGENTS SPECIFICALLY DISCLAIMS ANY WARRANTIES, INCLUDING, BUT NOT
+LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
+A PARTICULAR PURPOSE. THE SOFTWARE AND ACCOMPANYING DOCUMENTATION, IF
+ANY, PROVIDED HEREUNDER IS PROVIDED "AS IS". REGENTS HAS NO OBLIGATION
+TO PROVIDE MAINTENANCE, SUPPORT, UPDATES, ENHANCEMENTS, OR
+MODIFICATIONS.
+*/
+
+package firrtlTests
+
+import firrtl._
+
+class LegalizeSpec extends FirrtlFlatSpec {
+  behavior of "Legalize"
+
+  it should "compile and run" in {
+    runFirrtlTest("Legalize", "/passes/Legalize")
+  }
+}