Fixed zero width cat but (#651)

2 files changed, 63 insertions, 4 deletions

diff --git a/src/main/scala/firrtl/passes/ZeroWidth.scala b/src/main/scala/firrtl/passes/ZeroWidth.scala
index 6d7766d1..32b0b833 100644
--- a/src/main/scala/firrtl/passes/ZeroWidth.scala
+++ b/src/main/scala/firrtl/passes/ZeroWidth.scala
@@ -40,10 +40,25 @@ object ZeroWidth extends Transform {
     case VectorType(t, size) => removeZero(t) map (VectorType(_, size))
     case x => Some(x)
   }
-  private def onExp(e: Expression): Expression = e.tpe match {
-    case UIntType(IntWidth(ZERO)) => UIntLiteral(ZERO, IntWidth(BigInt(1)))
-    case SIntType(IntWidth(ZERO)) => SIntLiteral(ZERO, IntWidth(BigInt(1)))
-    case other => e map onExp
+  private def onExp(e: Expression): Expression = e match {
+    case DoPrim(Cat, args, consts, tpe) =>
+      val nonZeros = args.flatMap { x =>
+        x.tpe match {
+          case UIntType(IntWidth(ZERO)) => Seq.empty[Expression]
+          case SIntType(IntWidth(ZERO)) => Seq.empty[Expression]
+          case other => Seq(x)
+        }
+      }
+      nonZeros match {
+        case Nil => UIntLiteral(ZERO, IntWidth(BigInt(1)))
+        case Seq(x) => x
+        case seq => DoPrim(Cat, seq, consts, tpe) map onExp
+      }
+    case other => other.tpe match {
+      case UIntType(IntWidth(ZERO)) => UIntLiteral(ZERO, IntWidth(BigInt(1)))
+      case SIntType(IntWidth(ZERO)) => SIntLiteral(ZERO, IntWidth(BigInt(1)))
+      case _ => e map onExp
+    }
   }
   private def onStmt(renames: RenameMap)(s: Statement): Statement = s match {
     case (_: DefWire| _: DefRegister| _: DefMemory) =>
diff --git a/src/test/scala/firrtlTests/ZeroWidthTests.scala b/src/test/scala/firrtlTests/ZeroWidthTests.scala
index 8c39dc1e..50385a80 100644
--- a/src/test/scala/firrtlTests/ZeroWidthTests.scala
+++ b/src/test/scala/firrtlTests/ZeroWidthTests.scala
@@ -132,6 +132,50 @@ class ZeroWidthTests extends FirrtlFlatSpec {
         |    node z = add(x, UInt<1>(0))""".stripMargin
       (parse(exec(input)).serialize) should be (parse(check).serialize)
   }
+  "Expression in cat with type <0>" should "be removed" in {
+    val input =
+      """circuit Top :
+        |  module Top :
+        |    input x: UInt<1>
+        |    input y: UInt<0>
+        |    node z = cat(x, y)""".stripMargin
+    val check =
+      """circuit Top :
+        |  module Top :
+        |    input x: UInt<1>
+        |    node z = x""".stripMargin
+      (parse(exec(input)).serialize) should be (parse(check).serialize)
+  }
+  "Nested cats with type <0>" should "be removed" in {
+    val input =
+      """circuit Top :
+        |  module Top :
+        |    input x: UInt<0>
+        |    input y: UInt<0>
+        |    input z: UInt<0>
+        |    node a = cat(cat(x, y), z)""".stripMargin
+    val check =
+      """circuit Top :
+        |  module Top :
+        |    skip""".stripMargin
+      (parse(exec(input)).serialize) should be (parse(check).serialize)
+  }
+  "Nested cats where one has type <0>" should "be unaffected" in {
+    val input =
+      """circuit Top :
+        |  module Top :
+        |    input x: UInt<1>
+        |    input y: UInt<0>
+        |    input z: UInt<1>
+        |    node a = cat(cat(x, y), z)""".stripMargin
+    val check =
+      """circuit Top :
+        |  module Top :
+        |    input x: UInt<1>
+        |    input z: UInt<1>
+        |    node a = cat(x, z)""".stripMargin
+      (parse(exec(input)).serialize) should be (parse(check).serialize)
+  }
 }
 
 class ZeroWidthVerilog extends FirrtlFlatSpec {