1 files changed, 12 insertions, 4 deletions

diff --git a/src/main/scala/firrtl/ir/IR.scala b/src/main/scala/firrtl/ir/IR.scala
index cd8cd975..5263d9c0 100644
--- a/src/main/scala/firrtl/ir/IR.scala
+++ b/src/main/scala/firrtl/ir/IR.scala
@@ -206,6 +206,14 @@ abstract class Expression extends FirrtlNode {
   def foreachWidth(f: Width => Unit): Unit
 }
 
+/** Represents reference-like expression nodes: SubField, SubIndex, SubAccess and Reference
+  * The following fields can be cast to RefLikeExpression in every well formed firrtl AST:
+  * - SubField.expr, SubIndex.expr, SubAccess.expr
+  * - IsInvalid.expr, Connect.loc, PartialConnect.loc
+  * - Attach.exprs
+  */
+sealed trait RefLikeExpression extends Expression { def flow: Flow }
+
 object Reference {
   /** Creates a Reference from a Wire */
   def apply(wire: DefWire): Reference = Reference(wire.name, wire.tpe, WireKind, UnknownFlow)
@@ -222,7 +230,7 @@ object Reference {
 }
 
 case class Reference(name: String, tpe: Type = UnknownType, kind: Kind = UnknownKind, flow: Flow = UnknownFlow)
-    extends Expression with HasName with UseSerializer {
+    extends Expression with HasName with UseSerializer with RefLikeExpression {
   def mapExpr(f: Expression => Expression): Expression = this
   def mapType(f: Type => Type): Expression = this.copy(tpe = f(tpe))
   def mapWidth(f: Width => Width): Expression = this
@@ -232,7 +240,7 @@ case class Reference(name: String, tpe: Type = UnknownType, kind: Kind = Unknown
 }
 
 case class SubField(expr: Expression, name: String, tpe: Type = UnknownType, flow: Flow = UnknownFlow)
-    extends Expression with HasName with UseSerializer {
+    extends Expression with HasName with UseSerializer with RefLikeExpression {
   def mapExpr(f: Expression => Expression): Expression = this.copy(expr = f(expr))
   def mapType(f: Type => Type): Expression = this.copy(tpe = f(tpe))
   def mapWidth(f: Width => Width): Expression = this
@@ -242,7 +250,7 @@ case class SubField(expr: Expression, name: String, tpe: Type = UnknownType, flo
 }
 
 case class SubIndex(expr: Expression, value: Int, tpe: Type, flow: Flow = UnknownFlow)
-    extends Expression with UseSerializer {
+    extends Expression with UseSerializer with RefLikeExpression {
   def mapExpr(f: Expression => Expression): Expression = this.copy(expr = f(expr))
   def mapType(f: Type => Type): Expression = this.copy(tpe = f(tpe))
   def mapWidth(f: Width => Width): Expression = this
@@ -252,7 +260,7 @@ case class SubIndex(expr: Expression, value: Int, tpe: Type, flow: Flow = Unknow
 }
 
 case class SubAccess(expr: Expression, index: Expression, tpe: Type, flow: Flow = UnknownFlow)
-    extends Expression with UseSerializer {
+    extends Expression with UseSerializer with RefLikeExpression {
   def mapExpr(f: Expression => Expression): Expression = this.copy(expr = f(expr), index = f(index))
   def mapType(f: Type => Type): Expression = this.copy(tpe = f(tpe))
   def mapWidth(f: Width => Width): Expression = this