defpackage firrtl/passes :
  import core
  import verse
  import firrtl/ir2
  import firrtl/ir-utils
  import firrtl/primops
  import firrtl-main

;============== EXCEPTIONS =================================
defclass PassException <: Exception
defn PassException (msg:String) :
  new PassException :
    defmethod print (o:OutputStream, this) :
      print(o, msg)

;=============== WORKING IR ================================
definterface Kind
defstruct WireKind <: Kind 
defstruct RegKind <: Kind
defstruct InstanceKind <: Kind
defstruct ReadAccessorKind <: Kind 
defstruct WriteAccessorKind <: Kind 
defstruct PortKind <: Kind
defstruct NodeKind <: Kind ; All elems except structural memory, wires

defstruct MemKind <: Kind
defstruct ModuleKind <: Kind
defstruct StructuralMemKind <: Kind ; Separate kind because need special treatment
defstruct AccessorKind <: Kind

public definterface Gender
public val MALE = new Gender
public val FEMALE = new Gender
public val UNKNOWN-GENDER = new Gender
public val BI-GENDER = new Gender

defstruct WRef <: Expression :
   name: Symbol
   type: Type with: (as-method => true)
   kind: Kind
   gender: Gender  with: (as-method => true)

defstruct WRegInit <: Expression :
   reg: Expression
   name: Symbol
   type: Type with: (as-method => true)
   gender: Gender  with: (as-method => true)

defstruct WSubfield <: Expression :
   exp: Expression
   name: Symbol
   type: Type with: (as-method => true)
   gender: Gender with: (as-method => true)

defstruct WIndex <: Expression :
   exp: Expression
   value: Int
   type: Type with: (as-method => true)
   gender: Gender with: (as-method => true)

defstruct WDefAccessor <: Stmt :
   name: Symbol
   source: Expression
   index: Expression
   gender: Gender 

defstruct ConnectToIndexed <: Stmt :
   index: Expression
   locs: List<Expression>
   exp: Expression

defstruct ConnectFromIndexed <: Stmt :
   index: Expression
   loc: Expression
   exps: List<Expression>


;================ WORKING IR UTILS =========================

defn plus (g1:Gender,g2:Gender) -> Gender :
   switch fn ([x,y]) : g1 == x and g2 == y :
      [FEMALE,MALE]   : UNKNOWN-GENDER
      [MALE,FEMALE]   : UNKNOWN-GENDER
      [MALE,MALE]     : MALE
      [FEMALE,FEMALE] : FEMALE
      [BI-GENDER,MALE] : MALE
      [BI-GENDER,FEMALE] : FEMALE
      [MALE,BI-GENDER] : MALE
      [FEMALE,BI-GENDER] : FEMALE

defn swap (g:Gender) -> Gender :
   switch {_ == g} :
      UNKNOWN-GENDER : UNKNOWN-GENDER
      MALE : FEMALE
      FEMALE : MALE
      BI-GENDER : BI-GENDER

defn swap (f:Flip) -> Flip :
   switch {_ == f} :
      DEFAULT : REVERSE
      REVERSE : DEFAULT

defn swap (d:Direction) -> Direction :
   switch {_ == d} :
      OUTPUT : INPUT
      INPUT : OUTPUT

defn times (flip:Flip,d:Direction) -> Direction : flip * d
defn times (d:Direction,flip:Flip) -> Direction :
   switch {_ == flip} :
      DEFAULT : d
      REVERSE : swap(d)

defn times (g:Gender,flip:Flip) -> Gender : flip * g
defn times (flip:Flip,g:Gender) -> Gender :
   switch {_ == flip} :
      DEFAULT : g
      REVERSE : swap(g)

defn times (f1:Flip,f2:Flip) -> Flip :
   switch {_ == f2} :
      DEFAULT : f1
      REVERSE : swap(f1)
     
defn to-field (p:Port) -> Field :
   Field(name(p),REVERSE,type(p))
   if direction(p) == OUTPUT : Field(name(p),REVERSE,type(p))
   else if direction(p) == INPUT : Field(name(p),DEFAULT,type(p))
   else : error("Shouldn't be here")

defn to-dir (g:Gender) -> Direction :
   switch {_ == g} :
      MALE : INPUT
      FEMALE : OUTPUT

defmulti gender (e:Expression) -> Gender
defmethod gender (e:Expression) :
   MALE ; TODO, why was this OUTPUT before? It makes sense as male, not female

defmethod print (o:OutputStream, g:Gender) :
   print{o, _} $
   switch {g == _} :
      MALE : "male"
      FEMALE: "female"
      BI-GENDER : "bi"
      UNKNOWN-GENDER: "unknown"

defmethod type (exp:UIntValue) -> Type : UIntType(width(exp))
defmethod type (exp:SIntValue) -> Type : SIntType(width(exp))

;============== DEBUG STUFF =============================
public var PRINT-TYPES : True|False = false
public var PRINT-KINDS : True|False = false
public var PRINT-WIDTHS : True|False = false
public var PRINT-TWIDTHS : True|False = false
public var PRINT-GENDERS : True|False = false
public var PRINT-CIRCUITS : True|False = false
;=== Printers ===
defmethod print (o:OutputStream, k:Kind) :
   print{o, _} $
   match(k) :
      (k:WireKind) : "wire"
      (k:RegKind) : "reg"
      (k:AccessorKind) : "accessor"
      (k:PortKind) : "port"
      (k:MemKind) : "mem"
      (k:NodeKind) : "n"
      (k:ModuleKind) : "module"
      (k:InstanceKind) : "inst"
      (k:StructuralMemKind) : "smem"
      (k:ReadAccessorKind) : "racc"
      (k:WriteAccessorKind) : "wacc"

defn hasGender (e:Expression|Stmt|Type|Port|Field) :
   e typeof WRef|WSubfield|WIndex|WDefAccessor|WRegInit

defn hasWidth (e:Expression|Stmt|Type|Port|Field) :
   e typeof UIntType|SIntType|UIntValue|SIntValue

defn hasType (e:Expression|Stmt|Type|Port|Field) :
   e typeof Ref|Subfield|Index|DoPrim|WritePort|ReadPort|WRef|WSubfield
      |WIndex|DefWire|DefRegister|DefMemory|Register
      |VectorType|Port|Field|WRegInit

defn hasKind (e:Expression|Stmt|Type|Port|Field) :
   e typeof WRef

defn any-debug? (e:Expression|Stmt|Type|Port|Field) : 
   (hasGender(e) and PRINT-GENDERS) or
   (hasType(e) and PRINT-TYPES) or
   (hasWidth(e) and PRINT-WIDTHS) or
   (hasKind(e) and PRINT-KINDS)

defmethod print-debug (o:OutputStream, e:Expression|Stmt|Type|Port|Field) :
   defn wipe-width (t:Type) -> Type : 
     match(t) :
        (t:UIntType) : UIntType(UnknownWidth())
        (t:SIntType) : SIntType(UnknownWidth())
        (t) : t
      
   if any-debug?(e) : print(o,"@")
   if PRINT-KINDS and hasKind(e) : print-all(o,["<k:" kind(e as ?) ">"])
   if PRINT-TYPES and hasType(e) : print-all(o,["<t:" wipe-width(type(e as ?)) ">"])
   if PRINT-TWIDTHS and hasType(e): print-all(o,["<t:" type(e as ?) ">"])
   if PRINT-WIDTHS and hasWidth(e): print-all(o,["<w:" width(e as ?) ">"])
   if PRINT-GENDERS and hasGender(e): print-all(o,["<g:" gender(e as ?) ">"])
      
defmethod print (o:OutputStream, e:WRef) :
   print(o,name(e))
   print-debug(o,e as ?)

defmethod print (o:OutputStream, e:WRegInit) :
   print-all(o,[name(e)])
   print-debug(o,e as ?)

defmethod print (o:OutputStream, e:WSubfield) :
   print-all(o,[exp(e) "." name(e)])
   print-debug(o,e as ?)

defmethod print (o:OutputStream, e:WIndex) :
   print-all(o,[exp(e) "." value(e)])
   print-debug(o,e as ?)

defmethod print (o:OutputStream, s:WDefAccessor) :
   print-all(o,["accessor " name(s) " = " source(s) "[" index(s) "]"])
   print-debug(o,s)

defmethod print (o:OutputStream, c:ConnectToIndexed) :
   print-all(o, [locs(c) "[" index(c) "] := " exp(c)])
   print-debug(o,c as ?)

defmethod print (o:OutputStream, c:ConnectFromIndexed) :
   print-all(o, [loc(c) " := " exps(c) "[" index(c) "]"])
   print-debug(o,c as ?)

defmethod map (f: Expression -> Expression, e: WRegInit) :
   WRegInit(f(reg(e)), name(e), type(e), gender(e))
defmethod map (f: Expression -> Expression, e: WSubfield) :
   WSubfield(f(exp(e)), name(e), type(e), gender(e))
defmethod map (f: Expression -> Expression, e: WIndex) :
   WIndex(f(exp(e)), value(e), type(e), gender(e))

defmethod map (f: Expression -> Expression, c:WDefAccessor) :
   WDefAccessor(name(c), f(source(c)), f(index(c)), gender(c))
defmethod map (f: Expression -> Expression, c:ConnectToIndexed) :
   ConnectToIndexed(f(index(c)), map(f, locs(c)), f(exp(c)))
defmethod map (f: Expression -> Expression, c:ConnectFromIndexed) :
   ConnectFromIndexed(f(index(c)), f(loc(c)), map(f, exps(c)))

defmethod map (f: Type -> Type, e: WRef) :
   WRef(name(e), f(type(e)), kind(e), gender(e))
defmethod map (f: Type -> Type, e: WRegInit) :
   WRegInit(reg(e), name(e), f(type(e)), gender(e))
defmethod map (f: Type -> Type, e: WSubfield) :
   WSubfield(exp(e), name(e), f(type(e)), gender(e))
defmethod map (f: Type -> Type, e: WIndex) :
   WIndex(exp(e), value(e), f(type(e)), gender(e))

;================= Temparary Variable Elimination ========================
; Returns a new Circuit where temporary variables are removed and returns
;   the resulting nested expression

defn temp-elimination (c:Circuit) :
   val h = HashTable<Symbol,Expression>(symbol-hash)
   defn is-temp? (n:Symbol) -> True|False :
      to-string(n)[0] == 'T'
   defn temp-elim-e (e:Expression) :
      match(map(temp-elim-e,e)) :
         (e:Ref) : 
            if key?(h,name(e)) : h[name(e)]
            else : e
         (e) : e
   defn temp-elim-s (s:Stmt) :
      match(map(temp-elim-e,s)) :
         (s:DefNode) : 
            if is-temp?(name(s)) :
               h[name(s)] = value(s)
               EmptyStmt()
            else : s
         (s) : map(temp-elim-s,s)

   Circuit(modules*, main(c)) where :
      val modules* =
         for m in modules(c) map :
            Module(name(m), ports(m), temp-elim-s(body(m)))

;================= Bring to Working IR ========================
; Returns a new Circuit with Refs, Subfields, Indexes and DefAccessors 
;   replaced with IR-internal nodes that contain additional 
;   information (kind, gender)

defn to-working-ir (c:Circuit) :
   defn to-exp (e:Expression) :
      match(map(to-exp,e)) :
         (e:Ref) : WRef(name(e), type(e), NodeKind(), UNKNOWN-GENDER)
         (e:Subfield) : 
            if name(e) == `init : WRegInit(exp(e), to-symbol("~.init" % [name(exp(e) as WRef)]), type(e), UNKNOWN-GENDER)
            else : WSubfield(exp(e), name(e), type(e), UNKNOWN-GENDER)
         (e:Index) : WIndex(exp(e), value(e), type(e), UNKNOWN-GENDER)
         (e) : e
   defn to-stmt (s:Stmt) :
      match(map(to-exp,s)) :
         (s:DefAccessor) : WDefAccessor(name(s),source(s),index(s), UNKNOWN-GENDER)
         (s) : map(to-stmt,s)

   Circuit(modules*, main(c)) where :
      val modules* =
         for m in modules(c) map :
            Module(name(m), ports(m), to-stmt(body(m)))

;=============== Resolve Kinds =============================
; It is useful for the compiler to know information about 
;   objects referenced. This information is stored in the kind
;   field in WRef. This pass walks the graph and returns a new
;   Circuit where all WRef kinds are resolved

defn resolve-kinds (c:Circuit) :
   defn resolve (body:Stmt, kinds:HashTable<Symbol,Kind>) :
      defn resolve-stmt (s:Stmt) -> Stmt :
         map{resolve-expr,_} $
         map(resolve-stmt,s)

      defn resolve-expr (e:Expression) -> Expression :
         match(e) :
            (e:WRef) : WRef(name(e),type(e),kinds[name(e)],gender(e))
            (e) : map(resolve-expr,e)

      resolve-stmt(body)

   defn find (m:Module, kinds:HashTable<Symbol,Kind>) :
      defn find-stmt (s:Stmt) -> Stmt :
         match(s) :
            (s:DefWire) : kinds[name(s)] = NodeKind()
            (s:DefNode) : kinds[name(s)] = NodeKind()
            (s:DefRegister) : kinds[name(s)] = RegKind()
            (s:DefInstance) : kinds[name(s)] = InstanceKind()
            (s:DefMemory) : kinds[name(s)] = MemKind()
            (s:WDefAccessor) : kinds[name(s)] = AccessorKind()
            (s) : false
         map(find-stmt,s)

      kinds[name(m)] = ModuleKind()
      for p in ports(m) do :
         kinds[name(p)] = PortKind()
      find-stmt(body(m))
    
   defn resolve-kinds (m:Module, c:Circuit) -> Module :
      val kinds = HashTable<Symbol,Kind>(symbol-hash)
      for m in modules(c) do :
         kinds[name(m)] = ModuleKind()
      find(m,kinds)   
      val body! = resolve(body(m),kinds)
      Module(name(m),ports(m),body!)

   Circuit(modules*, main(c)) where :
      val modules* = 
         for m in modules(c) map :
            resolve-kinds(m,c)

;=============== MAKE EXPLICIT RESET =======================
; All modules have an implicit reset signal - however, the 
;   programmer can explicitly reference this signal if desired.
;   This pass makes all implicit resets explicit while
;   preserving any previously explicit resets
; If reset is not explicitly passed to instantiations, then this
;   pass autmatically connects the parent module's reset to the
;   instantiation's reset

defn make-explicit-reset (c:Circuit) :
   defn find-explicit (c:Circuit) -> List<Symbol> :
      defn explicit? (m:Module) -> True|False :
         for p in ports(m) any? :
            name(p) == `reset
      val explicit-reset = Vector<Symbol>()
      for m in modules(c) do:
         if explicit?(m) : add(explicit-reset,name(m))
      to-list(explicit-reset)

   defn make-explicit (m:Module, explicit-reset:List<Symbol>) -> Module :
      defn route-reset (s:Stmt) -> Stmt :
         match(s) :
            (s:DefInstance) : 
               val iref = WSubfield(WRef(name(s), UnknownType(), InstanceKind(), UNKNOWN-GENDER),`reset,UnknownType(),UNKNOWN-GENDER)
               val pref = WRef(`reset, UnknownType(), PortKind(), MALE)
               Begin(to-list([s,Connect(iref,pref)]))
            (s) : map(route-reset,s)

      var ports! = ports(m)
      if not contains?(explicit-reset,name(m)) :
         ports! = append(ports(m),list(Port(`reset,INPUT,UIntType(IntWidth(1)))))
      val body! = route-reset(body(m))
      Module(name(m),ports!,body!)
      
   defn make-explicit-reset (m:Module, c:Circuit) -> Module :
      val explicit-reset = find-explicit(c)
      make-explicit(m,explicit-reset)

   Circuit(modules*, main(c)) where :
      val modules* = 
         for m in modules(c) map :
            make-explicit-reset(m,c)


;============== INFER TYPES ================================
; This pass infers the type field in all IR nodes by updating
;   and passing an environment to all statements in pre-order
;   traversal, and resolving types in expressions in post-
;   order traversal.
; Type propagation for primary ops are defined here.
; Notable cases: LetRec requires updating environment before
;   resolving the subexpressions in its elements.
; Type errors are not checked in this pass, as this is
;   postponed for a later/earlier pass.

defn type (m:Module) -> Type : 
   BundleType(for p in ports(m) map : to-field(p))

defn get-type (b:Symbol,l:List<KeyValue<Symbol,Type>>) -> Type :
   val ma  = for kv in l find : b == key(kv)
   if ma != false : 
      val ret = value(ma as KeyValue<Symbol,Type>)
      ret
   else :
      UnknownType()

defn bundle-field-type (v:Type,s:Symbol) -> Type :
   match(v) :
      (v:BundleType) : 
         val ft = for p in fields(v) find : name(p) == s
         if ft != false : type(ft as Field)
         else : UnknownType()
      (v) : error(string-join(["Accessing subfield " s " on a non-Bundle type."]))

defn get-vector-subtype (v:Type) -> Type :
   match(v) :
      (v:VectorType) : type(v)
      (v) : UnknownType()

defn infer-exp-types (e:Expression, l:List<KeyValue<Symbol,Type>>) -> Expression :
   val r = map(infer-exp-types{_,l},e)
   match(r) :
      (e:WRef) : WRef(name(e), get-type(name(e),l),kind(e),gender(e))
      (e:WSubfield) : WSubfield(exp(e),name(e), bundle-field-type(type(exp(e)),name(e)),gender(e))
      (e:WRegInit) : WRegInit(reg(e),name(e),get-type(name(reg(e) as WRef),l),gender(e))
      (e:WIndex) : WIndex(exp(e),value(e), get-vector-subtype(type(exp(e))),gender(e))
      (e:DoPrim) : lower-and-type-primop(e)
         ;DoPrim(op(e),args(e),consts(e),get-primop-rettype(e))
      (e:ReadPort) : ReadPort(mem(e),index(e),get-vector-subtype(type(mem(e))),enable(e))
      (e:WritePort) : WritePort(mem(e),index(e),get-vector-subtype(type(mem(e))),enable(e))
      (e:UIntValue|SIntValue) : e

defn infer-types (s:Stmt, l:List<KeyValue<Symbol,Type>>) -> [Stmt List<KeyValue<Symbol,Type>>] :
   match(map(infer-exp-types{_,l},s)) :
      (s:Begin) : 
         var env = l
         val body* = 
            for s in body(s) map :
               val [s*,l*] = infer-types(s,env)
               env = l*
               s*
         [Begin(body*),env]
      (s:DefWire) : [s,List(name(s) => type(s),l)]
      (s:DefRegister) : [s,List(name(s) => type(s),l)]
      (s:DefMemory) : [s,List(name(s) => type(s),l)]
      (s:DefInstance) : [s, List(name(s) => type(module(s)),l)]
      (s:DefNode) : [s, List(name(s) => type(value(s)),l)]
      (s:WDefAccessor) : [s, List(name(s) => get-vector-subtype(type(source(s))),l)]
      (s:Conditionally) : 
        val [s*,l*] = infer-types(conseq(s),l)
        val [s**,l**] = infer-types(alt(s),l)
        [Conditionally(pred(s),s*,s**),l]
      (s:Connect|EmptyStmt) : [s,l]

defn infer-types (m:Module, l:List<KeyValue<Symbol,Type>>) -> Module :
   val ptypes = 
      for p in ports(m) map :
         name(p) => type(p)
   ;println-all(append(ptypes,l))
   val [s,l*] = infer-types(body(m),append(ptypes, l))
   Module(name(m),ports(m),s)

defn infer-types (c:Circuit) -> Circuit :
   val l = 
      for m in modules(c) map :
         name(m) => BundleType(map(to-field,ports(m)))
   ;println-all(l)
   Circuit{ _, main(c) } $ 
      for m in modules(c) map :
         infer-types(m,l)
    
;============= RESOLVE ACCESSOR GENDER  ============================
; To ensure a proper circuit, we must ensure that assignments
;   only work on expressions that can be assigned to. Similarly,
;   we must ensure that only expressions that can be read from
;   are used to assign from. This invariant requires each 
;   expression's gender to be inferred.
; Various elements can be bi-gender (e.g. wires) and can 
;   thus be treated as either female or male. Conversely, some
;   elements are single-gender (e.g. accessors, ports). 
; Because accessor gender is not known during declaration, 
;   this pass requires iterating until a fixed point is reached.

defn bundle-field-flip (n:Symbol,t:Type) -> Flip :
   match(t) :
      (b:BundleType) : 
          val field = for f in fields(b) find :
                         name(f) == n
          match(field):
                   (f:Field) : flip(f)
                   (f) : error(string-join(["Could not find " n " in bundle "]))
      (b) : error(string-join(["Accessing subfield " n " on a non-Bundle type."]))

defn resolve-genders (c:Circuit) :
   defn resolve-module (m:Module, genders:HashTable<Symbol,Gender>) -> Module :
      var done? = true

      defn resolve-iter (m:Module) -> Module : 
         val body* = resolve-stmt(body(m))
         Module(name(m),ports(m),body*)

      defn get-gender (n:Symbol,g:Gender) -> Gender :
         defn force-gender (n:Symbol,g:Gender) -> Gender : 
            genders[n] = g
            done? = false
            g
         val entry = for kv in genders find :
            key(kv) == n
         match(entry) :
            (e:KeyValue<Symbol,Gender>) :
               val value = value(e)
               if      value == UNKNOWN-GENDER and g == UNKNOWN-GENDER : g
               else if value != UNKNOWN-GENDER and g == UNKNOWN-GENDER : value
               else if value == UNKNOWN-GENDER and g != UNKNOWN-GENDER : force-gender(n,g)
               else : value
            (e:False) : force-gender(n,g)
            
      defn resolve-stmt (s:Stmt) -> Stmt :
         match(s) :
            (s:DefWire) :
               get-gender(name(s),BI-GENDER)
               s
            (s:DefRegister) :
               get-gender(name(s),BI-GENDER)
               s
            (s:DefMemory) :
               get-gender(name(s),BI-GENDER)
               s
            (s:DefNode) : 
               DefNode(name(s),resolve-expr(value(s),get-gender(name(s),MALE)))
            (s:DefInstance) :
               get-gender(name(s),FEMALE)
               DefInstance(name(s),resolve-expr(module(s),FEMALE))
            (s:WDefAccessor) : 
               val gender* = get-gender(name(s),UNKNOWN-GENDER)
               val index* = resolve-expr(index(s),MALE)
               val source* = resolve-expr(source(s),gender*)
               WDefAccessor(name(s),source*,index*,gender*)
            (s:Connect) : 
               Connect(resolve-expr(loc(s),FEMALE),resolve-expr(exp(s),MALE))
            (s:Conditionally) :
               val pred* = resolve-expr(pred(s),MALE)
               val conseq* = resolve-stmt(conseq(s))
               val alt* = resolve-stmt(alt(s))
               Conditionally(pred*,conseq*,alt*)
            (s) : map(resolve-stmt,s)

      defn resolve-expr (e:Expression,desired:Gender) -> Expression :
         match(e) :
            (e:WRef) : 
               val gender = get-gender(name(e),desired)
               WRef{name(e),type(e),kind(e),_} $
                  if gender == BI-GENDER : desired
                  else : gender
            (e:WRegInit) : 
               val gender = get-gender(name(reg(e) as WRef),desired)
               WRegInit{reg(e),name(e),type(e),_} $
                  if gender == BI-GENDER : desired
                  else : gender
            (e:WSubfield) : 
               val field-flip = bundle-field-flip(name(e),type(exp(e)))
               val exp* = resolve-expr(exp(e),field-flip * desired)
               val gender* = field-flip * gender(exp*)
               WSubfield(exp*,name(e),type(e),gender*)
            (e:WIndex) : 
               val exp* = resolve-expr(exp(e),desired)
               val gender* = gender(exp*)
               WIndex(exp*,value(e),type(e),gender*)
            (e) : map(resolve-expr{_,MALE},e)

      var module* = resolve-iter(m)
      ;println(genders)
      while not done? : 
        done? = true
        module* = resolve-iter(m)
        ;println(genders)
      module*
    
   defn resolve-genders (m:Module, c:Circuit) -> Module :
      val genders = HashTable<Symbol,Gender>(symbol-hash)
      resolve-module(m,genders)

   Circuit(modules*, main(c)) where :
      val modules* = 
         for m in modules(c) map :
            resolve-genders(m,c)

;;============== EXPAND ACCESSORS ================================
; This pass expands non-memory accessors into ConnectToIndexed or
;   ConnectFromIndexed. All elements of the vector are 
;   explicitly written out, then indexed. Depending on the gender
;   of the accessor, it is transformed into ConnectToIndexed (male) or
;   ConnectFromIndexed (female)
; Eg:

defn expand-vector (e:Expression) -> List<Expression> :
   val t = type(e) as VectorType
   for i in 0 to size(t) map-append :
      list(WIndex(e,i,type(t),gender(e as ?))) ;always be WRef|WSubfield|WIndex

defn expand-stmt (s:Stmt) -> Stmt :
   match(s) :
      (s:WDefAccessor) : 
         ;println-all(["Matched WDefAcc with " name(s)])
         val mem? = match(source(s)) :
            (e:WRef) : kind(e) typeof MemKind
            (e) : false
         if mem? : s
         else :
            val vtype = type(type(source(s)) as VectorType)
            val wire = DefWire(name(s),vtype)
            switch {gender(s) == _} :
               MALE : Begin{list(wire,_)} $ ConnectFromIndexed(
                  index(s),
                  WRef(name(wire),vtype,NodeKind(),FEMALE),
                  expand-vector(source(s)))
               FEMALE: Begin{list(wire,_)} $ ConnectToIndexed(
                  index(s),
                  expand-vector(source(s)),
                  WRef(name(wire),vtype,NodeKind(),MALE))
      (s) : map(expand-stmt,s)

defn expand-accessors (c:Circuit) :
   Circuit(modules*, main(c)) where :
      val modules* = 
         for m in modules(c) map :
            Module(name(m),ports(m),expand-stmt(body(m)))

;;=============== LOWERING TO GROUND TYPES =============================
; All non-ground (elevated) types (Vectors, Bundles) are expanded out to
;   individual ground types.
; This pass involves filling a table mapping the name of elevated types
;   to the lowered ground expression names and genders. This allows
;   references to be resolved.

defn num-elems (t:Type) -> Int :
   match(t) :
      (t:BundleType) : 
         var sum = 0
         for f in fields(t) do : 
            sum = sum + num-elems(type(f))
         sum
      (t:VectorType) : size(t) * num-elems(type(t))
      (t) : 1

defn index-of-elem (t:BundleType, s:Symbol) -> Int :
   var sum = 0
   label<Int> ret : 
      for f in fields(t) do :
         if s == name(f) : ret(sum)
         else : sum = sum + num-elems(type(f))
      error("Shouldn't be here")

defn lower-ports (m:Module, table:HashTable<Symbol,List<KeyValue<Expression,Flip>>>) -> List<Port> :
   val entries = table[name(m)]
   val directions = 
      for p in ports(m) map-append :
         to-list(for i in 0 to num-elems(type(p)) stream : direction(p))
   for (kv in entries, d in directions) map :
      val exp = key(kv) as WRef
      val dir* = d * value(kv)
      Port(name(exp),dir*,type(exp))

defn lower (body:Stmt, table:HashTable<Symbol,List<KeyValue<Expression,Flip>>>) -> Stmt :
   defn lower-stmt (s:Stmt) -> Stmt :
      defn add-to-table (y:Symbol,k:KeyValue<Expression,Flip>,ctable:HashTable<Symbol,List<KeyValue<Expression,Flip>>>) :
         val contains? = for x in ctable any? :
                            key(x) == y
         if contains? : ctable[y] = append(ctable[y],list(k))
         else : ctable[y] = list(k)
      defn is-instance (e:Expression) -> True|False :
         match(e) :
            (e:WRef) : kind(e) == InstanceKind()
            (e) : false
      defn calc-gender (g:Gender, e:Expression) -> Gender :
         match(e) :
            (e:WRef) : gender(e)
            (e:WRegInit) : gender(e)
            (e:WSubfield) : 
               if is-instance(exp(e)) : gender(e)
               else : calc-gender(bundle-field-flip(name(e),type(exp(e))) * g,exp(e))
            (e:WIndex) : gender(e)
            (e) : g

      match(s) :
         (s:DefWire) : Begin{_} $
            for t in table[name(s)] map :
               DefWire(name(key(t) as WRef),type(key(t)))
         (s:DefRegister) : Begin{_} $
            for t in table[name(s)] map :
               DefRegister(name(key(t) as WRef),type(key(t)))
         (s:DefInstance) : s
         (s:DefNode) :
            val s* = Begin $ list(
               DefWire(name(s),type(value(s))),
               Connect(WRef(name(s),type(value(s)),NodeKind(),FEMALE),value(s)))
            lower-stmt(s*)
         (s:Connect) : Begin{_} $
            for (l in expand-expr(loc(s)), r in expand-expr(exp(s))) map :
               val lgender = calc-gender(FEMALE,loc(s)) * value(l)
               val rgender = calc-gender(MALE,exp(s)) * value(r)
               switch fn ([x,y]) : lgender == x and rgender == y :
                  [FEMALE,MALE]   : Connect(key(l),key(r))
                  [MALE,FEMALE]   : Connect(key(r),key(l))
         (s:WDefAccessor) : Begin{_} $
            for (l in table[name(s)], r in expand-expr(source(s))) map:
               WDefAccessor(name(key(l) as WRef),key(r),index(s),value(r) * gender(s))
         (s:ConnectFromIndexed) : Begin(ls) where :
            val ctable = HashTable<Symbol,List<KeyValue<Expression,Flip>>>(symbol-hash)
            for e in exps(s) do :
               for (r in expand-expr(e),l in expand-expr(loc(s))) do :
                  add-to-table(name(key(l) as WRef),r,ctable)
            val ls = 
               for l in expand-expr(loc(s)) map :
                  val cg = calc-gender(FEMALE,loc(s))
                  val lgender = cg * value(l)
                  var rgender = BI-GENDER
                  val exps = for e in ctable[name(key(l) as WRef)] map :
                                rgender = rgender + (swap(cg) * value(e))
                                key(e)
                  switch fn ([x,y]) : lgender == x and rgender == y :
                     [FEMALE,MALE]   : ConnectFromIndexed(index(s),key(l),exps)
                     [MALE,FEMALE]   : ConnectToIndexed(index(s),exps,key(l))
         (s:ConnectToIndexed) : Begin(ls) where :
            val ctable = HashTable<Symbol,List<KeyValue<Expression,Flip>>>(symbol-hash)
            for ls in locs(s) do :
               for (l in expand-expr(ls),r in expand-expr(exp(s))) do :
                  add-to-table(name(key(r) as WRef),l,ctable)
            val ls = 
               for r in expand-expr(exp(s)) map :
                  val n = name(key(r) as WRef)
                  val cg = calc-gender(MALE,exp(s))
                  val rgender = cg * value(r)
                  var lgender = BI-GENDER
                  val locs = for l in ctable[n] map :
                                lgender = lgender + (swap(cg) * value(l))
                                key(l)
                  switch fn ([x,y]) : lgender == x and rgender == y :
                     [FEMALE,MALE]   : ConnectToIndexed(index(s),locs,key(r))
                     [MALE,FEMALE]   : ConnectFromIndexed(index(s),key(r),locs)
         (s:DefMemory) : Begin{_} $ 
            for t in table[name(s)] map :
               DefMemory(name(key(t) as WRef),type(key(t)) as VectorType)
         (s) : map(lower-stmt,s)

   defn expand-expr (e:Expression) -> List<KeyValue<Expression,Flip>> :
      match(e) :
         (e:WRef) : table[name(e)]
         (e:WRegInit) : table[name(e)]
         (e:WSubfield) : 
            val exps = expand-expr(exp(e))
            val begin = index-of-elem(type(exp(e)) as BundleType,name(e))
            val len = num-elems(type(e))
            headn(tailn(exps,begin),len)
         (e:WIndex) :
            val exps = expand-expr(exp(e))
            val len = num-elems(type(e))
            headn(tailn(exps,len * value(e)),len)
         (e) : list(KeyValue(e, DEFAULT))

   ;println(table)
   lower-stmt(body)

defn get-entries (n:Symbol,t:Type) -> List<KeyValue<WRef,Flip>> :
   defn uniquify (w:WRef) -> WRef :
      val name* = symbol-join([n "$" name(w)])
      WRef(name*,type(w),kind(w),gender(w))
   match(t) :
      (t:BundleType) : 
         for f in fields(t) map-append :
            val es = get-entries(name(f),type(f))
            for e in es map :
               uniquify(key(e)) => value(e) * flip(f)
      (t:VectorType) :
         for i in 0 to size(t) map-append :
            val es = get-entries(to-symbol(i),type(t))
            for e in es map :
               uniquify(key(e)) => value(e)
      (t) : list(KeyValue(WRef(n,t,NodeKind(),UNKNOWN-GENDER),DEFAULT))

defn lower-module (m:Module,table:HashTable<Symbol,List<KeyValue<Expression,Flip>>>) -> Module :
   defn build-table-ports (ports:List<Port>) :
      for p in ports do :
         table[name(p)] = get-entries(name(p),type(p))

   defn build-table-stmt (stmt:Stmt) -> Stmt:
      match(stmt) :
         (s:DefWire) : table[name(s)] = get-entries(name(s),type(s))
         (s:DefRegister) : 
            val regs = get-entries(name(s),type(s))
            val init-sym = symbol-join([name(s),`\|.init|])
            val init-regs = for r in regs map : 
               val [e f] = [key(r) value(r)]
               WRegInit(e,symbol-join([name(e),`\|.init|]),type(e),gender(e)) => f
            table[name(s)] = regs
            table[init-sym] = init-regs
         (s:DefInstance) : 
            val r = WRef(name(s),type(module(s)),InstanceKind(),FEMALE)
            val ports = table[name(module(s) as WRef)]
            table[name(s)] =
               for w in ports map-append :
                  list(KeyValue(WSubfield(r,name(key(w) as WRef),type(key(w) as WRef),UNKNOWN-GENDER), value(w)))
         (s:DefMemory) : table[name(s)] = 
            for x in get-entries(name(s),type(type(s) as VectorType)) map :
               val [w f] = [key(x) value(x)]
               WRef(name(w),VectorType(type(w),size(type(s) as VectorType)),kind(w),gender(w)) => f
         (s:DefNode) : table[name(s)] = get-entries(name(s),type(value(s)))
         (s:WDefAccessor) : table[name(s)] = get-entries(name(s),type(type(source(s)) as VectorType))
         (s) : map(build-table-stmt,s)
      stmt

   build-table-ports(ports(m))
   build-table-stmt(body(m)) 
   Module(name(m),ports*,body*) where :
      val body* = lower(body(m),table)
      val ports* = lower-ports(m,table)

defn lower-to-ground (c:Circuit) -> Circuit :
   val table = HashTable<Symbol,List<KeyValue<Expression,Flip>>>(symbol-hash)
   defn build-table-module (m:Module) -> ? :
      table[name(m)] = for p in ports(m) map-append : get-entries(name(p),type(p))

   for m in modules(c) map :
      build-table-module(m)

   Circuit(modules*, main(c)) where :
      val modules* = 
         for m in modules(c) map :
            lower-module(m,table)


;;=========== CONVERT MULTI CONNECTS to WHEN ================
; This pass converts ConnectToIndexed and ConnectFromIndexed
;   into a series of when statements. TODO what about initial
;   values?

defn expand-connect-indexed-stmt (s: Stmt) -> Stmt :
   defn equality (e1:Expression,e2:Expression) -> Expression :
      DoPrim(EQUAL-UU-OP,list(e1,e2),List(),UIntType(UnknownWidth()))
   match(s) :
      (s:ConnectToIndexed) : Begin $
         if length(locs(s)) == 0 : list(EmptyStmt())
         else :
            List(Connect(head(locs(s)),exp(s)), to-list $
               for (i in 1 to false, l in tail(locs(s))) stream : Conditionally(
                  equality(index(s),UIntValue(i,UnknownWidth())),
                  Connect(l,exp(s)),
                  EmptyStmt())
            )
      (s:ConnectFromIndexed) : Begin $
         if length(exps(s)) == 0 : list(EmptyStmt())
         else :
            List(Connect(loc(s),head(exps(s))), to-list $
               for (i in 1 to false, e in tail(exps(s))) stream : Conditionally(
                  equality(index(s),UIntValue(i,UnknownWidth())),
                  Connect(loc(s),e),
                  EmptyStmt())
            )
      (s) : map(expand-connect-indexed-stmt,s)
  
defn expand-connect-indexed (m: Module) -> Module :
   Module(name(m),ports(m),expand-connect-indexed-stmt(body(m)))

defn expand-connect-indexed (c: Circuit) -> Circuit :
  Circuit(modules*, main(c)) where :
     val modules* = 
        for m in modules(c) map :
           expand-connect-indexed(m)

;======= MAKE EXPLICIT REGISTER INITIALIZATION =============
; This pass replaces the reg.init construct by creating a new
;   wire that holds the value at initialization. This wire
;   is then connected to the register conditionally on reset,
;   at the end of the scope containing the  register 
;   declaration
; If a register has no inital value, the wire is connected to
;   a NULL node. Later passes will remove these with the base
;   case Mux(reset,NULL,a) -> a, and Mux(reset,a,NULL) -> a.
;   This ensures proper behavior if this pass is run multiple 
;   times.

defn initialize-registers (c:Circuit) :
   defn to-wire-name (y:Symbol) : symbol-join([ y "$init"])
   defn add-when (s:Stmt,h:HashTable<Symbol,Type>) -> Stmt :
      var inits = List<Stmt>()
      for kv in h do :
         val refreg = WRef(key(kv),value(kv),RegKind(),FEMALE)
         val refwire = WRef(to-wire-name(key(kv)),value(kv),NodeKind(),MALE)
         val connect = Connect(refreg,refwire)
         inits = append(inits,list(connect))
      if empty?(inits) : s
      else :
         val pred = WRef(`reset, UIntType(IntWidth(1)), PortKind(), MALE)
         val when-reset = Conditionally(pred,Begin(inits),Begin(List<Stmt>()))
         Begin(list(s,when-reset))

   defn rename (s:Stmt,h:HashTable<Symbol,True|False>) -> [Stmt HashTable<Symbol,Type>] :
      val t = HashTable<Symbol,Type>(symbol-hash)
      defn rename-expr (e:Expression) -> Expression :
         match(map(rename-expr,e)) :
            (e:WRegInit) : 
               val new-name = to-wire-name(name(reg(e) as WRef))
               WRef(new-name,type(reg(e)),RegKind(),gender(e))
            (e) : e
      defn rename-stmt (s:Stmt) -> Stmt :
         match(map(rename-stmt,s)) :
            (s:DefRegister) : 
               if h[name(s)] :
                  t[name(s)] = type(s)
                  Begin(list(s,DefWire(to-wire-name(name(s)),type(s))))
               else : s
            (s) : map(rename-expr,s)
      [rename-stmt(s) t]

   defn init? (y:Symbol,s:Stmt) -> True|False :
      var used? = false
      defn has? (e:Expression) -> Expression :
         match(map(has?,e)) :
            (e:WRegInit) : 
               if name(reg(e) as WRef) == y : used? = true
            (e) : map(has?,e)
         e
      map(has?,s)
      used?

   defn using-init (s:Stmt,h:HashTable<Symbol,True|False>) -> Stmt :
      match(s) :
         (s:DefRegister) : h[name(s)] = false
         (s) : 
            for x in h do :
               h[key(x)] = value(x) or init?(key(x),s)
      map(using-init{_,h},s)

   defn explicit-init-scope (s:Stmt) -> Stmt :
      val h = HashTable<Symbol,True|False>(symbol-hash)
      using-init(s,h)
      ;println(h)
      val [s* t] = rename(s,h)
      add-when(s*,t)
     
   Circuit(modules*, main(c)) where :
      val modules* = 
         for m in modules(c) map :
            Module(name(m), ports(m), body*) where :
               val body* = explicit-init-scope(body(m))

;;================ EXPAND WHENS =============================
; This pass does three things: remove last connect semantics,
;   remove conditional blocks, and eliminate concept of scoping.
; First, we scan the circuit to build a table mapping references
;   to the final assigned value, represented with SymbolicValues.
;   Within a scope, we remove the last connect symantics to get
;   the final value. When leaving a scope, the resulting table
;   is merged with the parent scope by using the SVMux.
;   We also collect the kind of reference to know how to declare
;   it in a following stage.
; Second, we use the table to declare each reference, then
;   assign to each once. This is relatively straightforward 
;   except calculating the WritePort/ReadPort enables.
; Finally, we scan the table to remove redundant values
; The WritePort enable is calculated by returning 1 for all conditions
;   for which the corresponding symbolic value is not SVNul.
; The ReadPort enable is calcuated by scanning all entries in
;   the table for when this is referenced (a read). All conditions
;   are accumulated and OR'ed together.

; ======== Expression Computation Library ===========

val zero = UIntValue(0,IntWidth(1))
val one = UIntValue(1,IntWidth(1))

defmethod equal? (e1:Expression,e2:Expression) -> True|False :
   match(e1,e2) :
      (e1:UIntValue,e2:UIntValue) : 
         if value(e1) == value(e2) : 
            match(width(e1), width(e2)) :
               (w1:IntWidth,w2:IntWidth) : width(w1) == width(w2)
               (w1,w2) : false
         else : false
      (e1:SIntValue,e2:SIntValue) : 
         if value(e1) == value(e2) : 
            match(width(e1), width(e2)) :
               (w1:IntWidth,w2:IntWidth) : width(w1) == width(w2)
         else : false
      (e1:WRef,e2:WRef) : name(e1) == name(e2)
      ;(e1:DoPrim,e2:DoPrim) :  TODO
      (e1:WRegInit,e2:WRegInit) : reg(e1) == reg(e2) and name(e1) == name(e2)
      (e1:WSubfield,e2:WSubfield) : name(e1) == name(e2)
      (e1,e2) : false

defn AND (e1:Expression,e2:Expression) -> Expression :
   if e1 == e2 : e1
   else if e1 == zero or e2 == zero : zero
   else if e1 == one : e2
   else if e2 == one : e1
   else : DoPrim(BIT-AND-OP,list(e1,e2),list(),UIntType(IntWidth(1)))

defn OR (e1:Expression,e2:Expression) -> Expression :
   if e1 == e2 : e1
   else if e1 == one or e2 == one : one
   else if e1 == zero : e2
   else if e2 == zero : e1
   else : DoPrim(BIT-OR-OP,list(e1,e2),list(),UIntType(IntWidth(1)))

defn NOT (e1:Expression) -> Expression :
   if e1 == one : zero
   else if e1 == zero : one
   else : DoPrim(EQUAL-UU-OP,list(e1,zero),list(),UIntType(IntWidth(1)))

defn children (e:Expression) -> List<Expression> :
   val es = Vector<Expression>()
   defn f (e:Expression) :
      add(es,e)
      e
   map(f,e)
   to-list(es)
   



; ======= Symbolic Value Library ==========
public definterface SymbolicValue
public defstruct SVExp <: SymbolicValue :
   exp : Expression
public defstruct SVMux <: SymbolicValue :
   pred : Expression
   conseq : SymbolicValue
   alt : SymbolicValue
public defstruct SVNul <: SymbolicValue

defmethod print (o:OutputStream, sv:SymbolicValue) :
   match(sv) :
      (sv: SVExp) : print(o, exp(sv))
      (sv: SVMux) : print-all(o, ["(" pred(sv) " ? " conseq(sv) " : " alt(sv) ")"])
      (sv: SVNul) : print(o, "SVNUL")

defmulti map<?T> (f: SymbolicValue -> SymbolicValue, sv:?T&SymbolicValue) -> T 
defmethod map (f: SymbolicValue -> SymbolicValue, sv:SymbolicValue) -> SymbolicValue :
   match(sv) :
      (sv: SVMux) : SVMux(pred(sv),f(conseq(sv)),f(alt(sv)))
      (sv) : sv

defn do (f:SymbolicValue -> ?, s:SymbolicValue) -> False :
   for x in s map :
      f(x)
      x
   false
defn dor (f:SymbolicValue -> ?, e:SymbolicValue) -> False :
   do(f,e)
   for x in e map :
      dor(f,x)
      x
   false

defmethod equal? (a:SymbolicValue,b:SymbolicValue) -> True|False :
   match(a,b) :
      (a:SVNul,b:SVNul) : true
      (a:SVExp,b:SVExp) : exp(a) == exp(b)
      (a:SVMux,b:SVMux) : pred(a) == pred(b) and conseq(a) == conseq(b) and alt(a) == alt(b)
      (a,b) : false

;TODO add invert to primop
defn optimize (sv:SymbolicValue) -> SymbolicValue :
   match(map(optimize,sv)) :
      (sv:SVMux) : 
         if conseq(sv) == alt(sv) : conseq(sv)
         else : 
            match(conseq(sv),alt(sv)) :
               (c:SVExp,a:SVExp) : 
                  if exp(c) == one and exp(a) == zero : SVExp(pred(sv))
                  else if exp(c) == zero and exp(a) == one : SVExp(NOT(pred(sv)))
                  else if exp(c) == exp(a) : c
                  else : sv
               (c,a) : sv
      (sv) : sv

; ========== Expand When Utilz ==========

defn deepcopy (t:HashTable<Symbol,SymbolicValue>) -> HashTable<Symbol,SymbolicValue> :
   t0 where :
      val t0 = HashTable<Symbol,SymbolicValue>(symbol-hash)
      for x in t do :
         t0[key(x)] = value(x)
defn get-unique-keys (ts:List<HashTable<Symbol,SymbolicValue>>) -> Vector<Symbol> :
   t0 where :
      val t0 = Vector<Symbol>()
      for v in ts do : 
         for t in v do :
            val duplicate? = for x in t0 any? : x == key(t)
            if not duplicate? : add(t0,key(t))
defn has-nul? (sv:SymbolicValue) -> True|False :
   var has? = false
   if sv typeof SVNul : has? = true
   for x in sv dor :
      if x typeof SVNul : has? = true
   has?
defn remove-nul (sv:SymbolicValue) -> SymbolicValue :
   match(map(remove-nul,sv)) :
      (sv:SVMux) : 
         match(conseq(sv),alt(sv)) :
            (c,a:SVNul) : c
            (c:SVNul,a) : a
            (c,a) : sv
      (sv) : sv
defn to-exp (sv:SymbolicValue) -> Expression :
   match(remove-nul(sv)) :
      (sv:SVMux) : 
         DoPrim(MUX-UU-OP,
                list(pred(sv),to-exp(conseq(sv)),to-exp(alt(sv))),
                list(),
                UIntType(IntWidth(1)))
      (sv:SVExp) : exp(sv)
      (sv) : error("Shouldn't be here")
defn reduce-or (l:List<True|False>) -> True|False :
   if length(l) == 0 : false
   else : head(l) or reduce-or(tail(l))
defn reduce-or (l:List<Expression>) -> Expression :
   if length(l) == 0 : zero
   else : OR(head(l) reduce-or(tail(l)))

; ========= Expand When Pass ===========
; TODO: replace stmt with wr (WRefs). The KIND of wref will help figure out what to emit as far as 
;   declarations, especially with not declaring anything for ports. We need WRefs, and not just Kinds,
;   because we need the name of the symbolic expression. I think? Or maybe we can use the key?

; 1) Build Table, Build Declaration List

defn expand-whens (assign:HashTable<Symbol,SymbolicValue>,
                   kinds:HashTable<Symbol,Kind>,
                   stmts:HashTable<Symbol,Stmt>,
                   decs:Vector<Stmt>,
                   enables:HashTable<Symbol,SymbolicValue>) -> Stmt :
           
   for x in assign do :
      val [n sv] = [key(x) value(x)]
      match(kinds[n]) : 
         (k:WriteAccessorKind) :
            ;First create WritePort and assign from accessor-turned-wire
            val s = stmts[n] as WDefAccessor
            val t = type(type(source(s)) as VectorType)
            val ref = WRef(n,t,k,MALE)
            val wp = WritePort(source(s),index(s),t,to-exp(enables[n]))
            add(decs,Connect(wp,ref))
            ;If initialized, assign input to accessor-turned-wire
            val sv = remove-nul(assign[n])
            if sv == SVNul : println("Uninitialized: ~" % [to-string(n)]) ;TODO actually collect error
            else : add(decs,Connect(ref,to-exp(sv)))
         (k:ReadAccessorKind) :
            val s = stmts[n] as WDefAccessor
            val t = type(type(source(s)) as VectorType)
            val ref = WRef(n,t,k,FEMALE)
            val rp = ReadPort(source(s),index(s),t,to-exp(enables[n]))
            add(decs,Connect(ref,rp))
         (k:RegKind) :
            val s = stmts[n] as DefRegister
            val ref = WRef(n,type(s),k,FEMALE)
            val sv = remove-nul(assign[n])
            val reg = 
               if sv typeof SVNul : Register(type(s),UIntValue(0,width(type(s) as ?)),zero)
               else : Register(type(s),to-exp(sv),to-exp(enables[n]))
            add(decs,Connect(ref,reg))
         (k:InstanceKind) :
            val s = stmts[n] as DefInstance
            val x = to-symbol(split(to-string(n),'.')[0])
            val f = to-symbol(split(to-string(n),'.')[1])
            val ref = WRef(x,type(module(s)),k,FEMALE)
            val sref = WSubfield(ref,f,bundle-field-type(type(module(s)),f),FEMALE)
            if has-nul?(assign[n]) : println("Uninitialized: ~" % [to-string(n)]);TODO actually collect error
            else : add(decs,Connect(sref,to-exp(assign[n])))
         (k) : 
            val s = stmts[n] as DefWire
            val ref = WRef(n,type(s),k,FEMALE)
            if has-nul?(assign[n]) : println("Uninitialized: ~" % [to-string(n)]);TODO actually collect error
            else : add(decs,Connect(ref,to-exp(assign[n])))
   Begin(to-list(decs))

defn get-enables (assign:HashTable<Symbol,SymbolicValue>,
                  kinds:HashTable<Symbol,Kind>) -> HashTable<Symbol,SymbolicValue> :
   defn get-read-enable (sym:Symbol,sv:SymbolicValue) -> Expression :
      defn active (e:Expression) -> True|False : 
         match(e) :
            (e:WRef) : name(e) == sym
            (e) : reduce-or{_} $ map(active,children(e))
            (e) : false
      match(sv) : 
         (sv: SVNul) : zero
         (sv: SVExp) : 
            if active(exp(sv)) : one 
            else : zero
         (sv: SVMux) : 
            val e0 = get-read-enable(sym,SVExp(pred(sv)))
            val e1 = get-read-enable(sym,conseq(sv))
            val e2 = get-read-enable(sym,alt(sv))
            if e1 == e2 : OR(e0,e1)
            else : OR(e0,OR(AND(pred(sv),e1),AND(NOT(pred(sv)),e2)))

   defn get-write-enable (sv:SymbolicValue) -> SymbolicValue :
      match(map(get-write-enable,sv)) :
         (sv: SVExp) : SVExp(one)
         (sv: SVNul) : SVExp(zero)
         (sv) : sv

   val enables = HashTable<Symbol,SymbolicValue>(symbol-hash)
   for x in assign do :
      val sym = key(x)
      match(kinds[sym]) :
         (k:ReadAccessorKind) : 
            enables[sym] = SVExp{_} $ reduce-or{_} $ to-list{_} $
               for y in assign stream :
                  get-read-enable(sym,value(y))
         (k:WriteAccessorKind) : enables[sym] = get-write-enable(value(x))
         (k:RegKind) : enables[sym] = get-write-enable(value(x))
         (k) : k
   enables

defn build-tables (s:Stmt, 
                   assign:HashTable<Symbol,SymbolicValue>,
                   kinds:HashTable<Symbol,Kind>,
                   decs:Vector<Stmt>,
                   stmts:HashTable<Symbol,Stmt>) -> False :
   match(s) :
      (s:DefWire) : 
         add(decs,s)
         kinds[name(s)] = WireKind()
         assign[name(s)] = SVNul()
         stmts[name(s)] = s
      (s:DefNode) : add(decs,s)
      (s:DefRegister) :
         add(decs,DefWire(name(s),type(s)))
         kinds[name(s)] = RegKind()
         assign[name(s)] = SVNul()
         stmts[name(s)] = s
      (s:WDefAccessor) : 
         add(decs,DefWire(name(s),type(type(source(s)) as VectorType)))
         assign[name(s)] = SVNul()
         kinds[name(s)] = switch {_ == gender(s)} :
            MALE : ReadAccessorKind()
            FEMALE : WriteAccessorKind()
         stmts[name(s)] = s
      (s:DefInstance) : 
         add(decs,s)
         for f in fields(type(module(s)) as BundleType) do :
            val n = to-symbol("~.~" % [name(s),name(f)]) ; only on inputs
            ;println-all(["In DefInst adding: " n])
            kinds[n] = InstanceKind()
            assign[n] = SVNul()
            stmts[n] = s
      (s:DefMemory) : add(decs,s)
      (s:Conditionally) :
         val assign-c = deepcopy(assign)
         val assign-a = deepcopy(assign)
         build-tables(conseq(s),assign-c,kinds,decs,stmts)
         build-tables(alt(s),assign-a,kinds,decs,stmts)
         for i in get-unique-keys(list(assign-c,assign-a)) do :
            assign[i] = match(get?(assign-c,i,false),get?(assign-a,i,false)) : ;TODO add to syntax highlighting
                  (c:SymbolicValue,a:SymbolicValue) : SVMux(pred(s),c,a)
                  (c:SymbolicValue,a:False) :
                     if kinds[i] typeof WireKind|InstanceKind|NodeKind : c
                     else : SVMux(pred(s),c,SVNul())
                  (c:False,a:SymbolicValue) : 
                     if kinds[i] typeof WireKind|InstanceKind|NodeKind : a
                     else : SVMux(pred(s),SVNul(),a)
                  (c:False,a:False) : error("Shouldn't be here")
         ;println("TABLE-C")
         ;for x in assign-c do : println(x)
         ;println("TABLE-A")
         ;for x in assign-a do : println(x)
         ;println("TABLE")
         ;for x in assign do : println(x)
      (s:Connect) : 
         val key* = match(loc(s)) :
            (e:WRef) : name(e)
            (e:WSubfield) : symbol-join([name(exp(e) as ?) `. name(e)])
            (e) : error("Shouldn't be here with ~" % [e])
         assign[key*] = SVExp(exp(s)); TODO, need to check all references are declared before this point
      (s:Begin) : for s* in body(s) do: build-tables(s*,assign,kinds,decs,stmts)
      (s) : false
         
defn expand-whens (m:Module) -> Module :
   val assign = HashTable<Symbol,SymbolicValue>(symbol-hash)
   val decs = Vector<Stmt>()
   val kinds = HashTable<Symbol,Kind>(symbol-hash)
   val stmts = HashTable<Symbol,Stmt>(symbol-hash)

   for p in ports(m) do :
      if direction(p) == OUTPUT :
         assign[name(p)] = SVNul()
         kinds[name(p)] = PortKind()
         stmts[name(p)] = DefWire(name(p),type(p))
         
   build-tables(body(m),assign,kinds,decs,stmts)
   for x in assign do : assign[key(x)] = optimize(value(x))
   val enables = get-enables(assign,kinds)
   for x in enables do : enables[key(x)] = optimize(value(x))
         
   ;println("Assigns")
   ;for x in assign do : println(x)
   ;println("Kinds")
   ;for x in kinds do : println(x)
   ;println("Decs")
   ;for x in decs do : println(x)
   ;println("Enables")
   ;for x in enables do : println(x)

   Module(name(m),ports(m),expand-whens(assign,kinds,stmts,decs,enables))

defn expand-whens (c:Circuit) -> Circuit :
   Circuit(modules*, main(c)) where :
      val modules* = 
         for m in modules(c) map :
            expand-whens(m)


;;================ INFER WIDTHS =============================
; First, you replace all unknown widths with a unique width
;   variable.
; Then, you collect all width constraints.
; Then, you solve width constraints.
; Finally, you replace all width variables with the solved
;   widths.
; Low FIRRTL Pass.

public defstruct VarWidth <: Width :
   name: Symbol
public defstruct PlusWidth <: Width :
   arg1 : Width
   arg2 : Width
public defstruct MinusWidth <: Width :
   arg1 : Width
   arg2 : Width
public defstruct MaxWidth <: Width :
   args : List<Width>

public defmulti map<?T> (f: Width -> Width, w:?T&Width) -> T
defmethod map (f: Width -> Width, w:Width) -> Width :
   match(w) :
      (w:MaxWidth) : MaxWidth(map(f,args(w)))
      (w:PlusWidth) : PlusWidth(f(arg1(w)),f(arg2(w)))
      (w:MinusWidth) : MinusWidth(f(arg1(w)),f(arg2(w)))
      (w) : w

public defmethod print (o:OutputStream, w:VarWidth) :
   print(o,name(w))
public defmethod print (o:OutputStream, w:MaxWidth) :
   print-all(o,["max" args(w)])
public defmethod print (o:OutputStream, w:PlusWidth) :
   print-all(o,[ arg1(w) " + " arg2(w)])
public defmethod print (o:OutputStream, w:MinusWidth) :
   print-all(o,[ arg1(w) " - " arg2(w)])

public definterface Constraint
public defstruct WGeq <: Constraint :
   loc : Width
   exp : Width
public defmethod print (o:OutputStream, c:WGeq) :
   print-all(o,[ loc(c) " >= " exp(c)])
defmethod equal? (w1:Width,w2:Width) -> True|False :
   match(w1,w2) :
      (w1:VarWidth,w2:VarWidth) : name(w1) == name(w2)
      (w1:MaxWidth,w2:MaxWidth) : 
         label<True|False> ret :
            if not length(args(w1)) == length(args(w2)) : ret(false)
            else :
               for w in args(w1) do :
                  if not contains?(args(w2),w) : ret(false)
               ret(true)
      (w1:IntWidth,w2:IntWidth) : width(w1) == width(w2)
      (w1,w2) : false
defn apply (a:Int|False,b:Int|False, f: (Int,Int) -> Int) -> Int|False :
   if a typeof Int and b typeof Int : f(a as Int, b as Int)
   else : false

; TODO: I should make MaxWidth take a variable list of arguments, which would make it easier to write the simplify function. It looks like there isn't a bug in the algorithm, but simplification reallllly speeds it up.

   
defn solve-constraints (l:List<WGeq>) -> HashTable<Symbol,Int> :
   defn contains? (n:Symbol,h:HashTable<Symbol,?>) -> True|False : key?(h,n)
   defn make-unique (ls:List<WGeq>) -> HashTable<Symbol,Width> :
      val h = HashTable<Symbol,Width>(symbol-hash)
      for g in ls do :
         match(loc(g)) :
            (w:VarWidth) :
               val n = name(w)
               if contains?(n,h) : h[n] = MaxWidth(list(exp(g),h[n]))
               else : h[n] = exp(g)
            (w) : w
      h
   defn simplify (w:Width) -> Width :
      match(map(simplify,w)) :
         (w:MaxWidth) : 
            val v = Vector<Width>()
            for w* in args(w) do :
               match(w*) :
                  (w:MaxWidth) : 
                     for x in args(w) do : add(v,x)
                  (w) : add(v,w)
            MaxWidth(unique(v))
         (w) : w
   defn substitute (w:Width,h:HashTable<Symbol,Width>) -> Width :
      ;println-all(["Substituting for [" w "]"])
      val w* = simplify(w)
      ;println-all(["After Simplify: [" w "]"])
      match(map(substitute{_,h},simplify(w))) :
         (w:VarWidth) : 
            ;println("matched varwidth!")
            if contains?(name(w),h) : 
               ;println("Contained!")
               ;println-all(["Width: " w])
               ;println-all(["Accessed: " h[name(w)]])
               val t = simplify(substitute(h[name(w)],h))
               ;val t = h[name(w)]
               ;println-all(["Width after sub: " t])
               h[name(w)] = t
               t
            else : w
         (w): 
            ;println-all(["not varwidth!" w])
            w
   defn b-sub (w:Width,h:HashTable<Symbol,Width>) -> Width:
      match(map(b-sub{_,h},w)) :
         (w:VarWidth) : 
            if key?(h,name(w)) : h[name(w)]
            else : w
         (w) : w
   defn remove-cycle (n:Symbol,w:Width) -> Width :
      ;println-all(["Removing cycle for " n " inside " w])
      val w* = match(map(remove-cycle{n,_},w)) :
         (w:MaxWidth) : MaxWidth(to-list(filter({_ != VarWidth(n)},args(w))))
         (w) : w
      ;println-all(["After removing cycle for " n ", returning " w*])
      w*
   defn self-rec? (n:Symbol,w:Width) -> True|False :
      var has? = false
      defn look (w:Width) -> Width :
         match(map(look,w)) :
            (w:VarWidth) : if name(w) == n : has? = true
            (w) : w
         w
      look(w)
      has?
   defn evaluate (h:HashTable<Symbol,Width>) -> HashTable<Symbol,Int> :
      defn apply (a:Int|False,b:Int|False, f: (Int,Int) -> Int) -> Int|False :
         if a typeof Int and b typeof Int : f(a as Int, b as Int)
         else : false
      defn apply-l (l:List<Int|False>,f:(Int,Int) -> Int) -> Int|False :
         if length(l) == 0 : 0
         else : apply(head(l),apply-l(tail(l),f),f)
      defn max (a:Int,b:Int) -> Int :
         if a > b : a
         else : b
      defn solve (w:Width) -> Int|False :
         match(w) :
            (w:VarWidth) : false
            (w:MaxWidth) : apply-l(map(solve,args(w)),max)
            (w:PlusWidth) : apply(solve(arg1(w)),solve(arg2(w)),{_ + _})
            (w:MinusWidth) : apply(solve(arg1(w)),solve(arg2(w)),{_ - _})
            (w:IntWidth) : width(w)
            (w) : error("Shouldn't be here")

      val i = HashTable<Symbol,Int>(symbol-hash)
      for x in h do :
         val s = solve(value(x))
         if s typeof Int : i[key(x)] = s as Int
      i
       
   ; Forward solve
   ; Returns a solved list where each constraint undergoes:
   ;  1) Continuous Solving (using triangular solving)
   ;  2) Remove Cycles
   ;  3) Move to solved if not self-recursive
   val u = make-unique(l)
   ;println("======== UNIQUE CONSTRAINTS ========")
   ;for x in u do : println(x)
   ;println("====================================")

   val f = HashTable<Symbol,Width>(symbol-hash)
   val o = Vector<Symbol>()
   for x in u do :
      ;println("==== SOLUTIONS TABLE ====")
      ;for x in f do : println(x)
      ;println("=========================")

      val [n e] = [key(x) value(x)]

      val e-sub = substitute(e,f)
      ;println(["Solving " n " => " e])
      ;println(["After Substitute: " n " => " e-sub])
      ;println("==== SOLUTIONS TABLE (Post Substitute) ====")
      ;for x in f do : println(x)
      ;println("=========================")
      val e* = remove-cycle{n,_} $ e-sub
      ;println(["After Remove Cycle: " n " => " e*])
      if not self-rec?(n,e*) : 
         ;println-all(["Not rec!: " n " => " e*])
         ;println-all(["Adding [" n "=>" e* "] to Solutions Table"])
         add(o,n)
         f[n] = e*

   ;println("Forward Solved Constraints")
   ;for x in f do : println(x)

   ; Backwards Solve
   val b = HashTable<Symbol,Width>(symbol-hash)
   for i in (length(o) - 1) through 0 by -1 do :
      val n = o[i]
      ;println-all(["SOLVE BACK: [" n " => " f[n] "]"])
      ;println("==== SOLUTIONS TABLE ====")
      ;for x in b do : println(x)
      ;println("=========================")
      val e* = simplify(b-sub(f[n],b))
      ;println-all(["BACK RETURN: [" n " => " e* "]"])
      b[n] = e*
      ;println("==== SOLUTIONS TABLE (Post backsolve) ====")
      ;for x in b do : println(x)
      ;println("=========================")

   ; Evaluate
   val e = evaluate(b)
   ;println("Evaluated Constraints")
   ;for x in e do : println(x)
   e
   
public defn width! (t:Type) -> Width :
   match(t) :
      (t:UIntType) : width(t)
      (t:SIntType) : width(t)
      (t) : error("No width!")
public defn width! (e:Expression) -> Width : width!(type(e))

defn gen-constraints (m:Module, h:HashTable<Symbol,Type>, v:Vector<WGeq>) -> Module:

   defn gen-constraints-s (s:Stmt) -> Stmt :
      match(map(gen-constraints-s,s)) :
         (s:DefWire) : DefWire(name(s),h[name(s)])
         (s:DefInstance) : DefInstance(name(s),gen-constraints(module(s)))
         (s:DefMemory) : DefMemory(name(s),h[name(s)] as VectorType)
         (s:DefNode) : DefNode(name(s),gen-constraints(value(s)))
         (s:Connect) : 
            val l = gen-constraints(loc(s))
            val e = gen-constraints(exp(s))
            add(v,WGeq(width!(type(l)),width!(type(e))))
            add(v,WGeq(width!(type(e)),width!(type(l))))
            Connect(l,e)
         (s) : s

   defn gen-constraints (e:Expression) -> Expression :
      match(map(gen-constraints,e)) :
         (e:WRef) : WRef(name(e),h[name(e)],kind(e),gender(e))
         (e:WSubfield) : WSubfield(exp(e),name(e),bundle-field-type(type(exp(e)),name(e)),gender(e))
         (e:WIndex) : error("Shouldn't be here")
         (e:DoPrim) : DoPrim(op(e),args(e),consts(e),primop-gen-constraints(e,v))
         (e:ReadPort) : ReadPort(mem(e),index(e),type(type(mem(e)) as VectorType),enable(e))
         (e:WritePort) : WritePort(mem(e),index(e),type(type(mem(e)) as VectorType),enable(e))
         (e:Register) : Register(type(value(e)),value(e),enable(e))
         (e:UIntValue) : 
            match(width(e)) :
               (w:UnknownWidth) : 
                  val w* = VarWidth(gensym(`w))
                  add(v,WGeq(w*,IntWidth(ceil-log2(value(e)))))
                  UIntValue(value(e),w*)
               (w) : e
         (e:SIntValue) : 
            match(width(e)) :
               (w:UnknownWidth) : 
                  val w* = VarWidth(gensym(`w))
                  add(v,WGeq(w*,IntWidth(1 + ceil-log2(abs(value(e))))))
                  SIntValue(value(e),w*)
               (w) : e
         (e) : e

   val ports* = 
      for p in ports(m) map : Port(name(p),direction(p),h[name(p)])
         
   Module(name(m),ports*,gen-constraints-s(body(m)))

defn build-environment (c:Circuit,m:Module,h:HashTable<Symbol,Type>) -> HashTable<Symbol,Type> :
   defn build-environment (s:Stmt) -> False :
      match(s) :
         (s:DefWire) : h[name(s)] = remove-unknowns(type(s))
         (s:DefInstance) : h[name(s)] = h[name(module(s) as WRef)]
         (s:DefMemory) : h[name(s)] = remove-unknowns(type(s))
         (s:DefNode) : h[name(s)] = remove-unknowns(type(value(s)))
         (s) : false
      do(build-environment,s)
   for p in ports(m) do :
      h[name(p)] = bundle-field-type(h[name(m)],name(p))
   build-environment(body(m))
   h
         
defn replace-var-widths (c:Circuit,h:HashTable<Symbol,Int>) -> Circuit :
   defn replace-var-widths-w (w:Width) -> Width :
      ;println-all(["REPLACE: " w])
      val w* = match(w) :
         (w:VarWidth) : 
            if key?(h,name(w)) : IntWidth(h[name(w)])
            else: w
         (w) : w
      ;println-all(["WITH: " w*])
      w*

   val modules* = for m in modules(c) map :
      Module{name(m),_,mapr(replace-var-widths-w,body(m))} $
         for p in ports(m) map :
            Port(name(p),direction(p),mapr(replace-var-widths-w,type(p)))

   Circuit(modules*,main(c))

;defn remove-unknown-widths (c:Circuit) -> Circuit :
;   defn remove-unknown-widths-w (w:Width) -> Width :
;      match(w) :
;         (w:UnknownWidth) : VarWidth(gensym(`w)) 
;         (w) : w
;   val modules* = for m in modules(c) map :
;      Module{name(m),_,body(m)} $
;         for p in ports(m) map :
;            Port(name(p),direction(p),mapr(remove-unknown-widths-w,type(p)))
;
;   val modules** = for m in modules* map :
;      Module(name(m),ports(m),mapr(remove-unknown-widths-w,body(m)))
;   Circuit(modules**,main(c))
   
defn remove-unknowns-w (w:Width) -> Width :
   match(w) :
      (w:UnknownWidth) : VarWidth(gensym(`w)) 
      (w) : w
defn remove-unknowns (t:Type) -> Type : mapr(remove-unknowns-w,t)

defn infer-widths (c:Circuit) -> Circuit :
   defn deepcopy (t:HashTable<Symbol,Type>) -> HashTable<Symbol,Type> :
      t0 where :
         val t0 = HashTable<Symbol,Type>(symbol-hash)
         for x in t do :
            t0[key(x)] = value(x)
   
   ;val c* = remove-unknown-widths(c)
   ;println(c*)
   val v = Vector<WGeq>()
   val ports* = HashTable<Symbol,Type>(symbol-hash)
   for m in modules(c) do : 
      ports*[name(m)] = remove-unknowns(BundleType(map(to-field,ports(m))))
   val modules* = for m in modules(c) map :
      ;println-all(["====== MODULE(" name(m) ") ENV ======"])
      val h = build-environment(c,m,deepcopy(ports*))
      ;for x in h do: println(x)
      ;println-all(["====================================="])
      val m* = gen-constraints(m,h,v)
      ;println-all(["====== MODULE(" name(m) ") ======"])
      ;println(m*)
      ;println-all(["====================================="])
      m*
   ;println("======== ALL CONSTRAINTS ========")
   ;for x in v do : println(x)
   ;println("=================================")
   val h = solve-constraints(to-list(v))
   ;println("======== SOLVED CONSTRAINTS ========")
   ;for x in h do : println(x)
   ;println("====================================")
   replace-var-widths(Circuit(modules*,main(c)),h)


;================= Inline Instances ========================
; Inlines instances. Assumes module with same name as the
;   Circuit is the top level module

defn inline-instances (c:Circuit) :
   val h = HashTable<Symbol,Module>(symbol-hash)
   val h-s = HashTable<Symbol,Stmt>(symbol-hash)
   defn inline-inst (s:Stmt) -> Stmt :
      match(map(inline-inst,s)) :
        (s:DefInstance) : 
           val n = name(module(s) as WRef)
           val m = h[n]
           val body* = 
              if key?(h-s,n) : h-s[n]
              else : 
                 val v = Vector<Stmt>()
                 for p in ports(m) do :
                    add(v,DefWire(name(p),type(p)))
                 add(v,inline-inst(body(m)))
                 Begin(to-list(v))
           h-s[n] = body*
           rename-s(body*,name(s))
        (s) : map(inline-inst-e,s)
   defn inline-inst-e (e:Expression) -> Expression :
      match(map(inline-inst-e,e)) :
         (e:WSubfield) : 
            match(kind(exp(e) as WRef)) :
               (k:InstanceKind) : 
                  WRef(symbol-join([name(exp(e) as WRef) "$" name(e)]),type(e),k,gender(e))
               (k:MemKind) : e
         (e) : e
   defn rename (ref:Symbol,n:Symbol) -> Symbol : symbol-join([n "$" ref])
   defn rename-e (e:Expression,n:Symbol) -> Expression :
      match(map(rename-e{_,n},e)) :
         (e:WRef) : WRef(rename(name(e),n),type(e),kind(e),gender(e))
         (e:WSubfield) : 
            match(kind(exp(e) as WRef)) :
               (k:InstanceKind) : 
                  WRef(symbol-join([name(exp(e) as WRef) "$" name(e)]),type(e),k,gender(e))
               (k:MemKind) : e
         (e) : e
   defn rename-s (s:Stmt,n:Symbol) -> Stmt :
      map{rename-e{_,n},_} $ match(map(rename-s{_,n},s)) :
         (s:DefWire) : DefWire(rename(name(s),n),type(s))
         (s:DefInstance) : error("Shouldn't be here")
         (s:DefMemory) : DefMemory(rename(name(s),n),type(s))
         (s:DefNode) : DefNode(rename(name(s),n),value(s))
         (s) : s
   for m in modules(c) do :
      h[name(m)] = m
   val top = (for m in modules(c) find : name(m) == main(c)) as Module
   Circuit(list(Module(name(top),ports(top),inline-inst(body(top)))),main(c))


;================= Split Expressions ========================
; Intended to only work on low firrtl

defn split-exp (c:Circuit) :
   defn split-exp-s (s:Stmt,v:Vector<Stmt>) -> False :
      match(s) :
         (s:Begin) : 
            defn f (s:Stmt) -> False: split-exp-s(s,v)
            do(f,s)
         (s:Conditionally) : error("Shouldn't be here")
         (s) : add(v,map(split-exp-e{_,v},s))
      false
   defn split-exp-e (e:Expression,v:Vector<Stmt>) -> Expression :
      match(map(split-exp-e{_,v},e)):
         (e:Subfield|DoPrim|ReadPort|Register|WritePort) : 
            val n = gensym(`T)
            add(v,DefNode(n,e))
            WRef(n,type(e),NodeKind(),UNKNOWN-GENDER)
         (e) : e

   Circuit{_,main(c)} $
      for m in modules(c) map :
         val v = Vector<Stmt>()
         split-exp-s(body(m),v)
         Module(name(m),ports(m),Begin(to-list(v)))

;================= Bring to Real IR ========================
; Returns a new Circuit with only real IR nodes.

defn to-real-ir (c:Circuit) :
   defn to-exp (e:Expression) :
      match(map(to-exp,e)) :
         (e:WRef) : Ref(name(e), type(e))
         (e:WSubfield) : Subfield(exp(e),name(e),type(e))
         (e:WRegInit) : error("Shouldn't be here")
         (e:WIndex) : error("Shouldn't be here")
         (e) : e
   defn to-stmt (s:Stmt) :
     match(map(to-exp,s)) :
        (e:WDefAccessor) : error("Shouldn't be here")
        (e:ConnectToIndexed) : error("Shouldn't be here")
        (e:ConnectFromIndexed) : error("Shouldn't be here")
        (e) : map(to-stmt,e)

   Circuit(modules*, main(c)) where :
      val modules* =
         for m in modules(c) map :
            Module(name(m), ports(m), to-stmt(body(m)))

;============= FLO PRINTER ======================================
; Emit 

defn flo-op-name (op:PrimOp) -> String :
   switch {op == _ } :
      ;NEG-OP :         "neg"
      ;NEG-OP :         "neg"
      ;NEG-OP :         "neg"
      ;NEG-OP :         "neg"
      ADD-UU-OP :      "add"
      ADD-US-OP :      "add"
      ADD-SU-OP :      "add"
      ADD-SS-OP :      "add"
      ADD-WRAP-UU-OP : "add"
      ADD-WRAP-US-OP : "add"
      ADD-WRAP-SU-OP : "add"
      ADD-WRAP-SS-OP : "add"
      SUB-UU-OP :      "sub"
      SUB-US-OP :      "sub"
      SUB-SU-OP :      "sub"
      SUB-SS-OP :      "sub"
      SUB-WRAP-UU-OP : "sub"
      SUB-WRAP-US-OP : "sub"
      SUB-WRAP-SU-OP : "sub"
      SUB-WRAP-SS-OP : "sub"
      MUL-UU-OP :      "mul" ;; todo: signed version
      MUL-US-OP :      "mul" ;; todo: signed version
      MUL-SU-OP :      "mul" ;; todo: signed version
      MUL-SS-OP :      "mul" ;; todo: signed version
      DIV-UU-OP :      "div" ;; todo: signed version
      DIV-US-OP :      "div" ;; todo: signed version
      DIV-SU-OP :      "div" ;; todo: signed version
      DIV-SS-OP :      "div" ;; todo: signed version
      MOD-UU-OP :      "mod" ;; todo: signed version
      MOD-US-OP :      "mod" ;; todo: signed version
      MOD-SU-OP :      "mod" ;; todo: signed version
      MOD-SS-OP :      "mod" ;; todo: signed version
      LESS-UU-OP :        "lt"  ;; todo: signed version
      LESS-US-OP :        "lt"  ;; todo: signed version
      LESS-SU-OP :        "lt"  ;; todo: signed version
      LESS-SS-OP :        "lt"  ;; todo: signed version
      LESS-EQ-UU-OP :     "lte" ;; todo: swap args
      LESS-EQ-US-OP :     "lte" ;; todo: swap args
      LESS-EQ-SU-OP :     "lte" ;; todo: swap args
      LESS-EQ-SS-OP :     "lte" ;; todo: swap args
      GREATER-UU-OP :     "lte" ;; todo: swap args BROKEN
      GREATER-US-OP :     "lte" ;; todo: swap args BROKEN
      GREATER-SU-OP :     "lte" ;; todo: swap args BROKEN
      GREATER-SS-OP :     "lte" ;; todo: swap args BROKEN
      GREATER-EQ-UU-OP :  "lt"  ;; todo: signed version
      GREATER-EQ-US-OP :  "lt"  ;; todo: signed version
      GREATER-EQ-SU-OP :  "lt"  ;; todo: signed version
      GREATER-EQ-SS-OP :  "lt"  ;; todo: signed version
      NEQUAL-UU-OP :      "neq"
      NEQUAL-SS-OP :      "neq"
      EQUAL-UU-OP :       "eq"
      EQUAL-SS-OP :       "eq"
      MUX-UU-OP :         "mux"
      MUX-SS-OP :         "mux"
      PAD-U-OP :         "rsh" ;; todo: signed version
      PAD-S-OP :         "rsh" ;; todo: signed version
      ;AS-UINT-U-OP :
      ;AS-UINT-S-OP :
      ;AS-SINT-U-OP :
      ;AS-SINT-S-OP :
      SHIFT-LEFT-U-OP :  "lsh" ;; todo: signed version
      SHIFT-LEFT-S-OP :  "lsh" ;; todo: signed version
      SHIFT-RIGHT-U-OP : "rsh"
      SHIFT-RIGHT-S-OP : "rsh"
      ;CONVERT-U-OP : 
      ;CONVERT-S-OP : 
      BIT-AND-OP :     "and"
      BIT-NOT-OP :     "not"
      BIT-OR-OP :      "or"
      BIT-XOR-OP :     "xor"
      CONCAT-OP :      "cat"
      BIT-SELECT-OP :  "rsh"
      BITS-SELECT-OP : "rsh"
      else : error $ string-join $ ["Unable to print Primop: " op]

defn sane-width (wd:Width) -> Int :
  match(wd) :
    (w:IntWidth) : max(1, width(w))
    (w)          : error("Unknown width")

defn prim-width (type:Type) -> Int :
   match(type) :
      (t:UIntType) : sane-width(width(t))
      (t:SIntType) : sane-width(width(t))
      (t) :          error("Bad prim width type")

defn sizeof (in: Int) -> Int :
  ;; if in == 1: 1 else: to-int(ceil(log(in)/log(2)))
  max(1, ceil-log2(in))

defn emit-all (o:OutputStream, es:Streamable, top:Symbol) :
  for e in es do :
     match(e) :
       (ex:Expression) : emit!(o,e,top)
       (ex:String) :     print(o, ex)
       (ex:Symbol) :     print(o, ex)
       ;; (ex:Int) :        print-all(o, [ex "'" sizeof(ex)])
       (ex:Int) :        print(o, ex)
       (ex) :            print(o, ex)

defn emit! (o:OutputStream, e:Expression,top:Symbol) :
   defn cmp-op? (op: PrimOp) -> True|False :
      contains?([EQUAL-OP, NEQUAL-OP, GREATER-OP, LESS-EQ-OP, LESS-OP, GREATER-EQ-OP], op)
   match(e) :
      (e:Ref) :       emit-all(o,[top "::" name(e)], top)
      (e:UIntValue) : emit-all(o,[value(e) "'" sane-width(width(e))], top)
      (e:SIntValue) : emit-all(o,[value(e) "'" sane-width(width(e))], top)
      (e:Subfield) :  emit-all(o,[exp(e) "/" name(e)], top)
      (e:Index) :     emit-all(o,[exp(e) "/" value(e)], top)
      (e:Register) : 
         emit-all(o,["reg'" prim-width(type(e)) " " enable(e) " " value(e)], top)
      (e:ReadPort) : 
         emit-all(o,["rd'" prim-width(type(e)) " " enable(e) " " mem(e) " " index(e)], top)
      (e:DoPrim) : 
         if cmp-op?(op(e)) :
           emit-all(o, [flo-op-name(op(e)) "'" prim-width(type(args(e)[0]))], top)
           if op(e) == GREATER-OP or op(e) == LESS-EQ-OP :
             emit-all(o, ["  " args(e)[1] " " args(e)[0]], top)
           else :
             emit-all(o, ["  " args(e)[0] " " args(e)[1]], top)
         else if op(e) == BIT-SELECT-OP :
           emit-all(o, [flo-op-name(op(e)) "'1 " args(e)[0] " " consts(e)[0]], top)
         else if op(e) == BITS-SELECT-OP :
           val w = consts(e)[0] - consts(e)[1] + 1
           emit-all(o, [flo-op-name(op(e)) "'" w " " args(e)[0] " " consts(e)[1]], top)
         ;; else if op(e) == CONCAT-OP :
         ;;   val w = consts(e)[0] - consts(e)[1] + 1
         ;;   emit-all(o, [flo-op-name(op(e)) "'" w " " args(e)[0] " " consts(e)[1]], top)
         else :
           emit-all(o, [flo-op-name(op(e)) "'" prim-width(type(e))], top)
           if op(e) == PAD-OP :
             emit-all(o, ["  " args(e)[0] " 0"], top)
           else :
             for arg in args(e) do :
               print(o, " ")
               emit!(o, arg, top)
             for const in consts(e) do :
               print(o, " ")
               print(o, const)
      (e) : print-all(o, ["EMIT(" e ")"])
      ;(e) : emit-all(o, ["mov'" prim-width(type(e)) " " e], top) ;TODO, not sure which one is right

defn emit-s (o:OutputStream, s:Stmt, v:List<Symbol>, top:Symbol) :
   match(s) :
      (s:DefWire) : ""
      (s:DefInstance) : error("Shouldn't be here")
      (s:DefMemory) :
           val vtype = type(s) as VectorType
           emit-all(o, [top "::" name(s) " = mem'" prim-width(type(vtype)) " " size(vtype) "\n"], top)
      (s:DefNode) : emit-all(o, [top "::" name(s) " = " value(s) "\n"], top)
      (s:Begin) : do(emit-s{o, _, v, top}, body(s))
      (s:Connect) : 
         if loc(s) typeof WritePort :
            val e = loc(s) as WritePort
            val name = gensym(`F)
            emit-all(o, [top "::" name " = wr'" prim-width(type(e)) " " enable(e) " " top "::" mem(e) " " index(e) " " exp(s) "\n"], top)
         else :
            val n = name(loc(s) as Ref)
            if contains?(v,n) :
               emit-all(o, [n " = out'" prim-width(type(loc(s))) " " exp(s) "\n"], top)
            else :
               emit-all(o, [top "::" n " = " exp(s) "\n"], top)
      (s) : s 

defn emit-module (o:OutputStream, m:Module) :
   val v = Vector<Symbol>()
   for port in ports(m) do :
      if name(port) ==`reset :
         emit-all(o, [name(m) "::" name(port) " = rst'1\n"], name(m))
      else : switch {_ == direction(port)} :
         INPUT : print-all(o, [name(m) "::" name(port) " = " "in'" prim-width(type(port)) "\n"])
         OUTPUT : add(v,name(port))
   emit-s(o, body(m), to-list(v), name(m))

public defn emit-flo (o:OutputStream, c:Circuit) :
   emit-module(o, modules(c)[0])
   c

public defn emit-flo (pathname:String, c:Circuit) :
   val out = FileOutputStream(pathname)
   emit-flo(out, c)
   close(out)
   c

;============= DRIVER ======================================
public defn run-passes (c: Circuit, p: List<Char>, pathname:String) :
   var c*:Circuit = c
   println("Compiling!")
   if PRINT-CIRCUITS : println("Original Circuit")
   if PRINT-CIRCUITS : print(c)
   defn do-stage (name:String, f: Circuit -> Circuit) :
      if PRINT-CIRCUITS : println(name)
      c* = f(c*)
      if PRINT-CIRCUITS : print(c*)
      if PRINT-CIRCUITS : println-all(["Finished " name "\n"])

  ; Early passes:
  ;  If modules have a reset defined, must be an INPUT and UInt(1)
   if contains(p,'q') : do-stage("Temp Elimination", temp-elimination)
   if contains(p,'a') : do-stage("Working IR", to-working-ir)
   if contains(p,'b') : do-stage("Resolve Kinds", resolve-kinds)
   if contains(p,'c') : do-stage("Make Explicit Reset", make-explicit-reset)
   if contains(p,'e') : do-stage("Infer Types", infer-types)
   if contains(p,'f') : do-stage("Resolve Genders", resolve-genders)
   if contains(p,'g') : do-stage("Expand Accessors", expand-accessors)
   if contains(p,'h') : do-stage("Lower To Ground", lower-to-ground)
   if contains(p,'i') : do-stage("Expand Indexed Connects", expand-connect-indexed)
   if contains(p,'p') : do-stage("Initialize Registers", initialize-registers)
   if contains(p,'j') : do-stage("Expand Whens", expand-whens)
   if contains(p,'k') : do-stage("Infer Widths", infer-widths)
   if contains(p,'l') : do-stage("Inline Instances", inline-instances)
   if contains(p,'m') : do-stage("Split Expressions", split-exp)
   if contains(p,'n') : do-stage("Real IR", to-real-ir)
   if contains(p,'o') : do-stage("To Flo", emit-flo{pathname,_})
   println("Done!")