Added muxing on passive aggregate types

13 files changed, 204 insertions, 80 deletions

diff --git a/src/main/stanza/chirrtl.stanza b/src/main/stanza/chirrtl.stanza
index 2ac76a05..fdd8424b 100644
--- a/src/main/stanza/chirrtl.stanza
+++ b/src/main/stanza/chirrtl.stanza
@@ -99,6 +99,7 @@ defn infer-types (c:Circuit) -> Circuit :
             (e:SubIndex) : SubIndex(exp(e),value(e),sub-type(type(exp(e))))
             (e:SubAccess) : SubAccess(exp(e),index(e),sub-type(type(exp(e))))
             (e:DoPrim) : set-primop-type(e)
+            (e:Mux) : Mux(cond(e),tval(e),fval(e),mux-type(tval(e),tval(e)))
             (e:UIntValue|SIntValue) : e
       defn infer-types-s (s:Stmt) -> Stmt :
          match(s) :
@@ -220,14 +221,20 @@ public val MALE = new Gender
 public val FEMALE = new Gender
 
 defn create-exps (e:Expression) -> List<Expression> :
-   match(type(e)) : 
-      (t:UIntType|SIntType|ClockType) : list(e)
-      (t:BundleType) : 
-         for f in fields(t) map-append :
-            create-exps(SubField(e,name(f),type(f)))
-      (t:VectorType) :
-         for i in 0 to size(t) map-append :
-            create-exps(SubIndex(e,i,type(t)))
+   match(e) :
+      (e:Mux) :
+         for (e1 in create-exps(tval(e)), e2 in create-exps(fval(e))) map :
+            Mux(cond(e),e1,e2,mux-type(e1,e2))
+      (e) :
+      match(type(e)) : 
+         (t:UIntType|SIntType|ClockType) : list(e)
+         (t:BundleType) : 
+            for f in fields(t) map-append :
+               create-exps(SubField(e,name(f),type(f)))
+         (t:VectorType) :
+            for i in 0 to size(t) map-append :
+               create-exps(SubIndex(e,i,type(t)))
+         (t:UnknownType) : list(e)
 
 defn remove-chirrtl (c:Circuit) :
    defn remove-chirrtl-m (m:InModule) -> InModule :
diff --git a/src/main/stanza/errors.stanza b/src/main/stanza/errors.stanza
index b9f9d4a4..1ebe5299 100644
--- a/src/main/stanza/errors.stanza
+++ b/src/main/stanza/errors.stanza
@@ -204,7 +204,7 @@ public defn check-high-form (c:Circuit) -> Circuit :
          NEQUAL-OP          : correct-num(2,0)
          EQUIV-OP           : correct-num(2,0)
          NEQUIV-OP          : correct-num(2,0)
-         MUX-OP             : correct-num(3,0)
+         ;MUX-OP             : correct-num(3,0)
          PAD-OP             : correct-num(1,1)
          AS-UINT-OP         : correct-num(1,0)
          AS-SINT-OP         : correct-num(1,0)
@@ -262,7 +262,7 @@ public defn check-high-form (c:Circuit) -> Circuit :
       defn check-high-form-e (e:Expression) -> Expression :
          defn valid-subexp (e:Expression) -> Expression :
             match(e) :
-               (e:Ref|SubField|SubIndex|SubAccess) : false
+               (e:Ref|SubField|SubIndex|SubAccess|Mux) : false
                (e) : add(errors,InvalidAccess())
             e
          match(map(check-high-form-e,e)) :
@@ -270,6 +270,7 @@ public defn check-high-form (c:Circuit) -> Circuit :
                if not key?(names,name(e)) : 
                   add(errors,UndeclaredReference(name(e)))
             (e:DoPrim) : check-high-form-primop(e)
+            (e:Mux) : e
             (e:UIntValue) : false
             (e:SubAccess) : 
                valid-subexp(exp(e))
@@ -410,9 +411,21 @@ defn OpNotAllSameType (info:FileInfo, op:Symbol) :
    PassException $ string-join $
    [info ": [module " mname "]  Primop " op " requires all operands to have the same type."]
 
-defn NodeGroundType (info:FileInfo) :
+defn NodePassiveType (info:FileInfo) :
    PassException $ string-join $
-   [info ": [module " mname "]  Node must be a ground type."]
+   [info ": [module " mname "]  Node must be a passive type."]
+
+defn MuxSameType (info:FileInfo) :
+   PassException $ string-join $
+   [info ": [module " mname "]  Must mux between equivalent types."]
+
+defn MuxPassiveTypes (info:FileInfo) :
+   PassException $ string-join $
+   [info ": [module " mname "]  Must mux between passive types."]
+
+defn MuxCondUInt (info:FileInfo) :
+   PassException $ string-join $
+   [info ": [module " mname "]  A mux condition must be of type UInt."]
 
 ;---------------- Helper Functions --------------
 defmethod equal? (t1:Type,t2:Type) -> True|False :
@@ -480,9 +493,9 @@ defn check-types-primop (e:DoPrim, errors:Vector<PassException>,info:FileInfo) -
       NEQUAL-OP          : false
       EQUIV-OP           : all-same-type(args(e))
       NEQUIV-OP          : all-same-type(args(e))
-      MUX-OP             : 
-         all-same-type(tail(args(e)))
-         is-uint(head(args(e)))
+      ;MUX-OP             : 
+      ;   all-same-type(tail(args(e)))
+      ;   is-uint(head(args(e)))
       PAD-OP             : false
       AS-UINT-OP         : false
       AS-SINT-OP         : false
@@ -506,6 +519,17 @@ defn check-types-primop (e:DoPrim, errors:Vector<PassException>,info:FileInfo) -
 ;----------------- Check Types Pass ---------------------
 public defn check-types (c:Circuit) -> Circuit : 
    val errors = Vector<PassException>()
+   defn passive? (t:Type) -> True|False :
+      match(t) :
+         (t:UIntType|SIntType) : true
+         (t:VectorType) : passive?(type(t))
+         (t:BundleType) : 
+            var p? = true
+            for x in fields(t) do :
+               if (flip(x) == REVERSE) : p? = false
+               if not passive?(type(x)) : p? = false
+            p?
+         (t) : true
    defn check-types-e (info:FileInfo,e:Expression) -> Expression :
       match(map(check-types-e{info,_},e)) :
          (e:WRef) : e
@@ -528,6 +552,11 @@ public defn check-types (c:Circuit) -> Circuit :
                (t:UIntType) : false
                (t) : add(errors,AccessIndexNotUInt(info))
          (e:DoPrim) : check-types-primop(e,errors,info)
+         (e:Mux) :
+            if type(tval(e)) != type(fval(e)) : add(errors,MuxSameType(info))
+            if not passive?(type(e)) : add(errors,MuxPassiveTypes(info))
+            if not passive?(type(e)) : add(errors,MuxPassiveTypes(info))
+            if not (type(cond(e)) typeof UIntType) : add(errors,MuxCondUInt(info))
          (e:UIntValue|SIntValue) : false
       e
 
@@ -568,8 +597,8 @@ public defn check-types (c:Circuit) -> Circuit :
             (s:Conditionally) :
                if type(pred(s)) != ut() : add(errors,PredNotUInt(info(s)))
             (s:DefNode) :
-               if not type(value(s)) typeof UIntType|SIntType :
-                  add(errors,NodeGroundType(info(s)))
+               if not passive?(type(value(s))) :
+                  add(errors,NodePassiveType(info(s)))
             (s) : false
          s }()
    
@@ -668,6 +697,7 @@ public defn check-genders (c:Circuit) -> Circuit :
          (e:DoPrim) : MALE
          (e:UIntValue) : MALE
          (e:SIntValue) : MALE
+         (e:Mux) : MALE
 
    defn check-genders-e (info:FileInfo,e:Expression,genders:HashTable<Symbol,Gender>) -> False :
       do(check-genders-e{info,_,genders},e)
@@ -679,6 +709,7 @@ public defn check-genders (c:Circuit) -> Circuit :
          (e:DoPrim) : 
             for e in args(e) do :
                check-gender(info,genders,e,MALE)
+         (e:Mux) : do(check-gender{info,genders,_,MALE},e)
          (e:UIntValue) : false
          (e:SIntValue) : false
      
@@ -713,7 +744,6 @@ public defn check-genders (c:Circuit) -> Circuit :
             check-gender(info(s),genders,clk(s),MALE)
          (s:Begin) : false
 
-
    for m in modules(c) do :
       mname = name(m)
       val genders = HashTable<Symbol,Gender>(symbol-hash)
diff --git a/src/main/stanza/firrtl-ir.stanza b/src/main/stanza/firrtl-ir.stanza
index ace1e76f..6a69b9ac 100644
--- a/src/main/stanza/firrtl-ir.stanza
+++ b/src/main/stanza/firrtl-ir.stanza
@@ -45,7 +45,7 @@ public val NEQUAL-OP = new PrimOp
 public val EQUAL-OP = new PrimOp
 public val NEQUIV-OP = new PrimOp
 public val EQUIV-OP = new PrimOp
-public val MUX-OP = new PrimOp
+;public val MUX-OP = new PrimOp
 public val PAD-OP = new PrimOp
 public val AS-UINT-OP = new PrimOp
 public val AS-SINT-OP = new PrimOp
@@ -84,6 +84,11 @@ public defstruct SubAccess <: Expression :
    exp: Expression
    index: Expression
    type: Type with: (as-method => true)
+public defstruct Mux <: Expression :
+   cond: Expression
+   tval: Expression
+   fval: Expression
+   type: Type with: (as-method => true)
 public defstruct UIntValue <: Expression :
    value: BigInt
    width: Width
diff --git a/src/main/stanza/ir-parser.stanza b/src/main/stanza/ir-parser.stanza
index 7c457c19..5d5d56ec 100644
--- a/src/main/stanza/ir-parser.stanza
+++ b/src/main/stanza/ir-parser.stanza
@@ -76,7 +76,7 @@ OPERATORS[`eq] = EQUAL-OP
 OPERATORS[`neq] = NEQUAL-OP
 OPERATORS[`eqv] = EQUIV-OP
 OPERATORS[`neqv] = NEQUIV-OP
-OPERATORS[`mux] = MUX-OP
+;OPERATORS[`mux] = MUX-OP
 OPERATORS[`pad] = PAD-OP
 OPERATORS[`neg] = NEG-OP
 OPERATORS[`asUInt] = AS-UINT-OP
@@ -365,6 +365,8 @@ defsyntax firrtl :
                      val num-bits = req-num-bits(v as Int)
                      SIntValue(BigIntLit(v as Int,num-bits), IntWidth(num-bits))
 
+      expterm = (mux(?cond:#exp ?tval:#exp ?fval:#exp)) :
+         Mux(cond,tval,fval,UnknownType())
       expterm = (?op:#sym(?es:#exp ... ?ints:#int ... ?rest ...)) :
          if not empty?(rest) :
             FPE(rest, "Illegal operands to primitive operator.")
diff --git a/src/main/stanza/ir-utils.stanza b/src/main/stanza/ir-utils.stanza
index 59a4b659..ac3946c2 100644
--- a/src/main/stanza/ir-utils.stanza
+++ b/src/main/stanza/ir-utils.stanza
@@ -113,7 +113,20 @@ public defn EQV (e1:Expression,e2:Expression) -> Expression :
    DoPrim(EQUIV-OP,list(e1,e2),list(),type(e1))
 
 public defn MUX (p:Expression,e1:Expression,e2:Expression) -> Expression :
-   DoPrim(MUX-OP,list(p,e1,e2),list(),type(e1))
+   Mux(p,e1,e2,mux-type(type(e1),type(e2)))
+
+public defn mux-type (e1:Expression,e2:Expression) -> Type :
+   mux-type(type(e1),type(e2))
+public defn mux-type (t1:Type,t2:Type) -> Type :
+   if t1 == t2 :
+      match(t1,t2) :
+         (t1:UIntType,t2:UIntType) : UIntType(UnknownWidth())
+         (t1:SIntType,t2:SIntType) : SIntType(UnknownWidth())
+         (t1:VectorType,t2:VectorType) : VectorType(mux-type(type(t1),type(t2)),size(t1))
+         (t1:BundleType,t2:BundleType) : 
+            BundleType $ for (f1 in fields(t1),f2 in fields(t2)) map :
+               Field(name(f1),flip(f1),mux-type(type(f1),type(f2)))
+   else : UnknownType()
 
 public defn CAT (e1:Expression,e2:Expression) -> Expression :
    DoPrim(CONCAT-OP,list(e1,e2),list(),type(e1))
@@ -365,7 +378,7 @@ defmethod print (o:OutputStream, op:PrimOp) :
       NEQUIV-OP : "neqv"
       EQUAL-OP : "eq"
       NEQUAL-OP : "neq"
-      MUX-OP : "mux"
+      ;MUX-OP : "mux"
       PAD-OP : "pad"
       AS-UINT-OP : "asUInt"
       AS-SINT-OP : "asSInt"
@@ -398,6 +411,8 @@ defmethod print (o:OutputStream, e:Expression) :
          print-all(o, [op(e) "("])
          print-all(o, join(concat(args(e), consts(e)), ", "))
          print(o, ")")
+      (e:Mux) :
+         print-all(o, ["mux(" cond(e) ", " tval(e) ", " fval(e) ")"])
    print-debug(o,e)
 
 defmethod print (o:OutputStream, c:Stmt) :
@@ -529,6 +544,7 @@ defmethod map (f: Expression -> Expression, e:Expression) -> Expression :
       (e:SubIndex) : SubIndex(f(exp(e)), value(e), type(e))
       (e:SubAccess) : SubAccess(f(exp(e)), f(index(e)), type(e)) 
       (e:DoPrim) : DoPrim(op(e), map(f, args(e)), consts(e), type(e))
+      (e:Mux) : Mux(f(cond(e)),f(tval(e)),f(fval(e)),type(e))
       (e) : e
 
 public defmulti map<?T> (f: Symbol -> Symbol, c:?T&Stmt) -> T
@@ -584,6 +600,7 @@ defmethod map (f: Type -> Type, c:Expression) -> Expression :
       (c:SubIndex) : SubIndex(exp(c),value(c),f(type(c)))
       (c:SubAccess) : SubAccess(exp(c),index(c),f(type(c)))
       (c:DoPrim) : DoPrim(op(c),args(c),consts(c),f(type(c)))
+      (c:Mux) : Mux(cond(c),tval(c),fval(c),f(type(c)))
       (c) : c
 
 public defmulti map<?T> (f: Type -> Type, c:?T&Stmt) -> T
diff --git a/src/main/stanza/passes.stanza b/src/main/stanza/passes.stanza
index 47e8711e..d243ce76 100644
--- a/src/main/stanza/passes.stanza
+++ b/src/main/stanza/passes.stanza
@@ -115,6 +115,7 @@ defmethod info (stmt:Empty) -> FileInfo : FileInfo()
 
 defmethod type (exp:UIntValue) -> Type : UIntType(width(exp))
 defmethod type (exp:SIntValue) -> Type : SIntType(width(exp))
+defmethod type (exp:WVoid) -> Type : UnknownType()
 
 defmethod get-type (s:WDefInstance) -> Type : type(s)
 
@@ -148,6 +149,10 @@ defmethod equal? (e1:Expression,e2:Expression) -> True|False :
             if not x == y : 
                are-equal? = false
          are-equal?
+      (e1:Mux,e2:Mux) : 
+         (cond(e1) == cond(e2)) and 
+         (tval(e1) == tval(e2)) and 
+         (fval(e1) == fval(e2))
       (e1,e2) : false
 
 ; ================= PRINTERS ===================
@@ -439,6 +444,23 @@ defn to-gender (d:Direction) -> Gender :
       INPUT: MALE
       OUTPUT: FEMALE
 
+public defn mux-type-and-widths (e1:Expression,e2:Expression) -> Type :
+   mux-type-and-widths(type(e1),type(e2))
+public defn mux-type-and-widths (t1:Type,t2:Type) -> Type :
+   defn wmax (w1:Width,w2:Width) -> Width :
+      match(w1,w2) :
+         (w1:IntWidth,w2:IntWidth) : IntWidth(max(width(w1),width(w2)))
+         (w1,w2) : MaxWidth(list(w1,w2))
+   if t1 == t2 :
+      match(t1,t2) :
+         (t1:UIntType,t2:UIntType) : UIntType(wmax(width(t1),width(t2)))
+         (t1:SIntType,t2:SIntType) : SIntType(wmax(width(t1),width(t2)))
+         (t1:VectorType,t2:VectorType) : VectorType(mux-type-and-widths(type(t1),type(t2)),size(t1))
+         (t1:BundleType,t2:BundleType) : 
+            BundleType $ for (f1 in fields(t1),f2 in fields(t2)) map :
+               Field(name(f1),flip(f1),mux-type-and-widths(type(f1),type(f2)))
+   else : UnknownType()
+
 ;================= Remove Special Characters ========================
 ; Returns a new Circuit where all names have all special characters 
 ;    removed, except _.
@@ -675,6 +697,7 @@ defn infer-types (c:Circuit) -> Circuit :
             (e:WSubIndex) : WSubIndex(exp(e),value(e),sub-type(type(exp(e))),gender(e))
             (e:WSubAccess) : WSubAccess(exp(e),index(e),sub-type(type(exp(e))),gender(e))
             (e:DoPrim) : set-primop-type(e)
+            (e:Mux) : Mux(cond(e),tval(e),fval(e),mux-type-and-widths(tval(e),tval(e)))
             (e:UIntValue|SIntValue) : e
       defn infer-types-s (s:Stmt) -> Stmt :
          match(s) :
@@ -846,14 +869,19 @@ defn get-point (e:Expression) -> Int :
 defn create-exps (n:Symbol, t:Type) -> List<Expression> :
    create-exps(WRef(n,t,ExpKind(),UNKNOWN-GENDER))
 defn create-exps (e:Expression) -> List<Expression> :
-   match(type(e)) : 
-      (t:UIntType|SIntType|ClockType) : list(e)
-      (t:BundleType) : 
-         for f in fields(t) map-append :
-            create-exps(WSubField(e,name(f),type(f),gender(e) * flip(f)))
-      (t:VectorType) :
-         for i in 0 to size(t) map-append :
-            create-exps(WSubIndex(e,i,type(t),gender(e)))
+   match(e) :
+      (e:Mux) :
+         for (e1 in create-exps(tval(e)), e2 in create-exps(fval(e))) map :
+            Mux(cond(e),e1,e2,mux-type-and-widths(e1,e2))
+      (e) :
+         match(type(e)) : 
+            (t:UIntType|SIntType|ClockType) : list(e)
+            (t:BundleType) : 
+               for f in fields(t) map-append :
+                  create-exps(WSubField(e,name(f),type(f),gender(e) * flip(f)))
+            (t:VectorType) :
+               for i in 0 to size(t) map-append :
+                  create-exps(WSubIndex(e,i,type(t),gender(e)))
 
 defn gexp-hash (e:Expression) -> Int :
    turn-off-debug(false)
@@ -870,18 +898,25 @@ defn fast-create-exps (e:Expression) -> List<Expression> :
    if key?(hashed-create-exps,e) : 
       hashed-create-exps[e]
    else :
-      val es = Vector<List<Expression>>()
-      match(type(e)) : 
-         (t:UIntType|SIntType|ClockType) : add(es,list(e))
-         (t:BundleType) : 
-            for f in fields(t) do :
-               add(es,fast-create-exps(WSubField(e,name(f),type(f),gender(e) * flip(f))))
-         (t:VectorType) :
-            for i in 0 to size(t) do :
-               add(es,fast-create-exps(WSubIndex(e,i,type(t),gender(e))))
-      val x = append-all(es)
-      hashed-create-exps[e] = x
-      x
+      match(e) :
+         (e:Mux) :
+            val x = for (e1 in create-exps(tval(e)), e2 in create-exps(fval(e))) map :
+               Mux(cond(e),e1,e2,mux-type-and-widths(e1,e2))
+            hashed-create-exps[e] = x
+            x
+         (e) : 
+            val es = Vector<List<Expression>>()
+            match(type(e)) : 
+               (t:UIntType|SIntType|ClockType) : add(es,list(e))
+               (t:BundleType) : 
+                  for f in fields(t) do :
+                     add(es,fast-create-exps(WSubField(e,name(f),type(f),gender(e) * flip(f))))
+               (t:VectorType) :
+                  for i in 0 to size(t) do :
+                     add(es,fast-create-exps(WSubIndex(e,i,type(t),gender(e))))
+            val x = append-all(es)
+            hashed-create-exps[e] = x
+            x
 
 ;---------------- Pass ---------------------
 
@@ -995,6 +1030,7 @@ defn remove-access (c:Circuit) :
          defn remove-e (e:Expression) -> Expression : ;NOT RECURSIVE (except primops) INTENTIONALLY!
             match(e) :
                (e:DoPrim) : map(remove-e,e)
+               (e:Mux) : map(remove-e,e)
                (e:UIntValue|SIntValue) : e
                (e) :
                   if has-access?(e) :
@@ -1115,7 +1151,9 @@ defn expand-whens (c:Circuit) -> Circuit :
                   val value = get?(netlist,lvalue,false)
                   match(value) :
                      (value:Expression) :
-                        netlist[lvalue] = MUX(pred(s),c-netlist[lvalue],value)
+                        val tv = c-netlist[lvalue]
+                        val fv = value
+                        netlist[lvalue] = Mux(pred(s),tv,fv,mux-type-and-widths(tv,fv))
                      (value:False) :
                         netlist[lvalue] = c-netlist[lvalue]
             (s:Print) : 
@@ -1479,10 +1517,9 @@ defn infer-widths (c:Circuit) -> Circuit :
             get-constraints(type(t1),type(t2),f)
    defn get-constraints-e (e:Expression) -> Expression :
       match(map(get-constraints-e,e)) :
-         (e:DoPrim) : 
-            if op(e) == MUX-OP : 
-               constrain(width!(args(e)[0]),ONE)
-               constrain(ONE,width!(args(e)[0]))
+         (e:Mux) : 
+            constrain(width!(cond(e)),ONE)
+            constrain(ONE,width!(cond(e)))
             e
          (e) : e
    defn get-constraints (s:Stmt) -> Stmt :
@@ -1638,7 +1675,7 @@ defn split-exp (m:InModule) -> InModule :
       defn split-exp-e (e:Expression,i:Int) -> Expression :
          match(map(split-exp-e{_,i + 1},e)) :
             (e:DoPrim) :
-               if i > 0 and op(e) != MUX-OP : split(e)
+               if i > 0 : split(e)
                else : e
             (e) : e
       match(map(split-exp-e{_,0},s)) :
@@ -2115,6 +2152,7 @@ defn lower-types (m:Module) -> Module :
                      else : e
                   (k) : WRef(lowered-name(e),type(e),kind(e),gender(e))
             (e:DoPrim) : map(lower-types-e,e)
+            (e:Mux) : map(lower-types-e,e)
       match(map(lower-types-e,s)) :
          (s:DefWire|DefPoison) :
             if is-ground?(type(s)) : s
@@ -2216,11 +2254,17 @@ defn rand-string (t:Type) -> Streamable :
    ["{" w* "{" ran "}}"]
 defn emit (x:?) : emit(x,0)
 defn emit (x:?, top:Int) :
+   defn cast (e:Expression) -> ? :
+      match(type(e)) :
+         (t:UIntType) : e
+         (t:SIntType) : ["$signed(" e ")"]
    match(x) :
       (e:Expression) :
          turn-off-debug(false)
          match(e) :
             (e:DoPrim) : emit(op-stream(e), top + 1)
+            (e:Mux) : 
+               emit([cond(e) " ? " cast(tval(e)) " : " cast(fval(e))],top + 1)
             (e:WRef) : print(e)
             (e:WSubField) : print(lowered-name(e))
             (e:WSubAccess) : print-all([lowered-name(exp(e)) "[" lowered-name(index(e)) "]"])
@@ -2296,7 +2340,7 @@ defn op-stream (doprim:DoPrim) -> Streamable :
       EQUIV-OP : [cast-if(a0()) " == " cast-if(a1())]
       NEQUAL-OP : [cast-if(a0()) " != " cast-if(a1())]
       EQUAL-OP : [cast-if(a0()) " == " cast-if(a1())]
-      MUX-OP : [a0() " ? " cast(a1()) " : " cast(a2())]
+      ;MUX-OP : [a0() " ? " cast(a1()) " : " cast(a2())]
       PAD-OP : 
          val w = long!(type(a0()))
          val diff = (to-long(c0()) - w)
@@ -2371,27 +2415,20 @@ defn emit-verilog (m:InModule) -> Module :
       add(assigns,["assign " e " = " value ";"])
    defn update-and-reset (r:Expression,clk:Expression,reset?:Expression,init:Expression) :
       if not key?(at-clock,clk) : at-clock[clk] = Vector<Streamable>()
-
-      defn mux? (e:Expression) -> True|False :
-         match(e) :
-            (e:DoPrim) : 
-               if (op(e) == MUX-OP) : true
-               else : false
-            (e) : false
-
       defn add-update (e:Expression,tabs:String) :
-         if mux?(e) :
-            val e* = e as DoPrim
-            add(at-clock[clk],[tabs "if(" args(e*)[0] ") begin"])
-            add-update(args(e*)[1],string-join([tabs tab]))
-            add(at-clock[clk],[tabs "end else begin"])
-            add-update(args(e*)[2],string-join([tabs tab]))
-            add(at-clock[clk],[tabs "end"])
-         else :
-            if e == r : add(at-clock[clk],[tabs ";"])
-            else : add(at-clock[clk],[tabs r " <= " e ";"])
-
-      add-update(MUX(reset?,init,netlist[r]),"");"
+         match(e) :
+            (e:Mux) :
+               add(at-clock[clk],[tabs "if(" cond(e) ") begin"])
+               add-update(tval(e),string-join([tabs tab]))
+               add(at-clock[clk],[tabs "end else begin"])
+               add-update(fval(e),string-join([tabs tab]))
+               add(at-clock[clk],[tabs "end"])
+            (e) :
+               if e == r : add(at-clock[clk],[tabs ";"])
+               else : add(at-clock[clk],[tabs r " <= " e ";"])
+      val tv = init
+      val fv = netlist[r]
+      add-update(Mux(reset?,tv,fv,mux-type-and-widths(tv,fv)),"");"
    defn update (e:Expression,value:Expression,clk:Expression,en:Expression) :
       if not key?(at-clock,clk) :
          at-clock[clk] = Vector<Streamable>()
diff --git a/src/main/stanza/primop.stanza b/src/main/stanza/primop.stanza
index 2a2d75f2..4b8d0cbc 100644
--- a/src/main/stanza/primop.stanza
+++ b/src/main/stanza/primop.stanza
@@ -95,11 +95,11 @@ public defn set-primop-type (e:DoPrim) -> DoPrim :
       NEQUAL-OP : DoPrim(o,a,c,BoolType())
       EQUIV-OP : DoPrim(o,a,c,BoolType())
       NEQUIV-OP : DoPrim(o,a,c,BoolType())
-      MUX-OP : DoPrim{o,a,c,_} $
-         match(t2(),t3()) :
-            (t2:UIntType, t3:UIntType) : UIntType(MAX(w2(),w3()))
-            (t2:SIntType, t3:SIntType) : SIntType(MAX(w2(),w3()))
-            (t2, t3) : UnknownType()
+      ;MUX-OP : DoPrim{o,a,c,_} $
+      ;   match(t2(),t3()) :
+      ;      (t2:UIntType, t3:UIntType) : UIntType(MAX(w2(),w3()))
+      ;      (t2:SIntType, t3:SIntType) : SIntType(MAX(w2(),w3()))
+      ;      (t2, t3) : UnknownType()
       PAD-OP : DoPrim{o,a,c,_} $
          match(t1()) :
             (t1:UIntType) : UIntType(c1())
diff --git a/test/errors/high-form/NumArgs.fir b/test/errors/high-form/NumArgs.fir
index 6c083905..7db78509 100644
--- a/test/errors/high-form/NumArgs.fir
+++ b/test/errors/high-form/NumArgs.fir
@@ -1,11 +1,9 @@
 ; RUN: firrtl -i %s -o %s.v -X verilog -p c 2>&1 | tee %s.out | FileCheck %s
-; CHECK: Primop mux requires 3 expression arguments.
 ; CHECK: Primop add requires 2 expression arguments.
 ; CHECK: Primop bits requires 2 integer arguments.
 
 circuit Top : 
   module Top : 
-    node x = mux(UInt(1),UInt(1))
     node y = add(SInt(1),UInt(1),UInt(1))
     node z = bits(UInt(1),1,2,3)
 
diff --git a/test/errors/type/MuxWithFlips.fir b/test/errors/type/MuxWithFlips.fir
new file mode 100644
index 00000000..9a8b2170
--- /dev/null
+++ b/test/errors/type/MuxWithFlips.fir
@@ -0,0 +1,12 @@
+; RUN: firrtl -i %s -o %s.v -X verilog -p c 2>&1 | tee %s.out | FileCheck %s
+; CHECK: [module Top]  Must mux between passive types.
+; CHECK: [module Top]  A mux condition must be of type UInt.
+
+circuit Top : 
+  module Top : 
+    input c : SInt<1>
+    input x : {x : UInt, flip y : UInt}
+    input y : {x : UInt, flip y : UInt}
+    output z : {x : UInt, flip y : UInt}
+    z <= mux(c,x,y)
+
diff --git a/test/errors/type/NodeWithFlips.fir b/test/errors/type/NodeWithFlips.fir
index 0b8b9e6b..2030955f 100644
--- a/test/errors/type/NodeWithFlips.fir
+++ b/test/errors/type/NodeWithFlips.fir
@@ -1,5 +1,5 @@
 ; RUN: firrtl -i %s -o %s.v -X verilog -p c 2>&1 | tee %s.out | FileCheck %s
-; CHECK: Node must be a ground type.
+; CHECK: Node must be a passive type.
 
 circuit Top : 
   module Top : 
diff --git a/test/errors/type/Primop.fir b/test/errors/type/Primop.fir
index f687c754..dcf19935 100644
--- a/test/errors/type/Primop.fir
+++ b/test/errors/type/Primop.fir
@@ -1,7 +1,7 @@
 ; RUN: firrtl -i %s -o %s.v -X verilog -p c 2>&1 | tee %s.out | FileCheck %s
-; CHECK: Primop mux requires all operands to have the same type.
+; CHECK: Must mux between equivalent types.
 ; CHECK: Primop add cannot operate on non-ground types.
-; CHECK: Primop mux requires argument SInt("h1") to be a UInt type.
+; CHECK: A mux condition must be of type UInt.
 
 circuit Top : 
   module Top :
diff --git a/test/features/InitAccessor.fir b/test/features/InitAccessor.fir
index 5a81a62e..6261ec01 100644
--- a/test/features/InitAccessor.fir
+++ b/test/features/InitAccessor.fir
@@ -1,4 +1,4 @@
-; RUN: firrtl -i %s -o %s.v -X verilog -p c 2>&1 | tee %s.out | FileCheck %s
+; RUN: firrtl -i %s -o %s.v -X verilog -p cT 2>&1 | tee %s.out | FileCheck %s
 
 ;CHECK: Done!
 circuit Top : 
diff --git a/test/features/MuxBundle.fir b/test/features/MuxBundle.fir
new file mode 100644
index 00000000..764078d5
--- /dev/null
+++ b/test/features/MuxBundle.fir
@@ -0,0 +1,16 @@
+; RUN: firrtl -i %s -o %s.v -X verilog -p c 2>&1 | tee %s.out | FileCheck %s
+;CHECK: Expand Connects
+circuit Top : 
+  module Top : 
+    input a: {w:UInt<42>,x:UInt<20>}
+    input b: {w:UInt<42>,x:UInt<20>}
+    input c: {w:UInt<42>,x:UInt<20>}
+    input p : UInt<1>
+    output d: {w:UInt<42>,x:UInt<20>}
+    d <= mux(p,mux(p,a,b),c)
+;CHECK: d.w <= mux(p, mux(p, a.w, b.w), c.w)
+;CHECK: d.x <= mux(p, mux(p, a.x, b.x), c.x)   
+
+;CHECK: Finished Expand Connects
+;CHECK: Done!
+