defpackage firrtl/flo :
   import core
   import verse
   import firrtl/ir-utils
   import firrtl/ir2
   import firrtl/passes
   import bigint2

;============= Flo Backend ================

public defstruct Flo <: Pass :
   with-output: (() -> False) -> False 
public defmethod pass (b:Flo) -> (Circuit -> Circuit) : emit-flo{with-output(b),_}
public defmethod name (b:Flo) -> String : "To Flo"

definterface FloKind
defstruct FRegKind <: FloKind
defstruct FWritePortKind <: FloKind :
   mem: Expression
   index: Expression
defstruct FOutputPortKind <: FloKind 
  
defn is-sint? (arg:Expression) -> True|False : type(arg) typeof SIntType
defn type (s:DefAccessor) -> Type : type(type(source(s)) as VectorType)

defn flo-op-name (op:PrimOp, args:List<Expression>) -> String :
   switch {op == _ } :
      ADD-OP :             "add"
      ADD-WRAP-OP :        "add"
      SUB-OP :             "sub"
      SUB-WRAP-OP :        "sub"
      MUL-OP :             "mul" ;; todo: signed version
      DIV-OP :             "div" ;; todo: signed version
      MOD-OP :             "mod" ;; todo: signed version
      QUO-OP :             "div" ;; todo: signed version
      REM-OP :             "mod" ;; todo: signed version
      LESS-OP :            "lt"  ;; todo: signed version
      LESS-EQ-OP :         "lte" ;; todo: swap args
      GREATER-OP :         "lt"  ;; todo: swap args
      GREATER-EQ-OP :      "lte" ;; todo: signed version
      NEQUIV-OP :          "neq"
      EQUIV-OP :           "eq"
      NEQUAL-OP :          "neq"
      EQUAL-OP :           "eq"
      MUX-OP :             "mux"
      NEG-OP :             "neg"
      AS-UINT-OP :         "mov"
      AS-SINT-OP :         "mov"
      SHIFT-LEFT-OP :      "lsh" 
      SHIFT-RIGHT-OP :     if is-sint?(args[0]): "arsh" else: "rsh"
      DYN-SHIFT-LEFT-OP :  "lsh" 
      DYN-SHIFT-RIGHT-OP : if is-sint?(args[0]): "arsh" else: "rsh"
      PAD-OP :             if is-sint?(args[0]): "arsh" else: "rsh"
      CONVERT-OP :         if is-sint?(args[0]): "arsh" else: "rsh"
      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"
      BIT-XOR-REDUCE-OP :  "xorr"
      else :
        error $ string-join $ ["Unable to print Primop: " op]

defn sane-width (wd:Width) -> Int|Long :
  match(wd) :
    (w:LongWidth) : max(to-long(1), width(w))
    (w)          : error(string-join(["Unknown width: " w]))

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

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

defn emit! (e:Expression,top:Symbol) :
   defn greater-op? (op: PrimOp) -> True|False :
      contains?([GREATER-OP], op)
   defn greater-eq-op? (op: PrimOp) -> True|False :
      contains?([GREATER-EQ-OP], op)     
   defn less-eq-op? (op: PrimOp) -> True|False :
      contains?([LESS-EQ-OP], op)
   defn less-op? (op: PrimOp) -> True|False :
      contains?([LESS-OP], op)
   defn cmp-op? (op: PrimOp) -> True|False :
      greater-op?(op) or greater-eq-op?(op) or less-op?(op) or less-eq-op?(op) or 
      contains?([EQUAL-OP NEQUAL-OP] op)
   match(e) :
      (e:Ref) :       emit-all([top "::" name(e)], top)
      (e:UIntValue) : emit-all([value(e) "'" sane-width(width(e))], top)
      (e:SIntValue) : emit-all([value(e) "'" sane-width(width(e))], top)
      (e:Subfield) :  emit-all([exp(e) "/" name(e)], top)
      (e:Index) :     emit-all([exp(e) "/" value(e)], top)
      (e:DoPrim) : 
         if cmp-op?(op(e)) :
           emit-all([flo-op-name(op(e), args(e)) "'" prim-width(type(args(e)[0]))], top)
           if greater-op?(op(e)) or greater-eq-op?(op(e)) :
             emit-all([" " args(e)[1] " " args(e)[0]], top)
           else :
             emit-all([" " args(e)[0] " " args(e)[1]], top)
         else if op(e) == BIT-SELECT-OP :
           emit-all([flo-op-name(op(e), args(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([flo-op-name(op(e), args(e)) "'" w " " args(e)[0] " " consts(e)[1]], top)
         else if op(e) == CONCAT-OP :
           val w = prim-width(type(args(e)[1]))
           emit-all([flo-op-name(op(e), args(e)) "'" w " " args(e)[0] " " args(e)[1]], top)
         else if op(e) == PAD-OP or op(e) == CONVERT-OP :
           emit-all([flo-op-name(op(e), args(e)) "'" prim-width(type(e)) " " args(e)[0] " 0"], top)
         else :
           emit-all([flo-op-name(op(e), args(e)) "'" prim-width(type(e))], top)
           for arg in args(e) do :
             print(" ")
             emit!(arg, top)
           for const in consts(e) do :
             print-all([" " const "'" req-num-bits(const)])
      (e) : error("SHOULDN'T EMIT THIS") ;; print-all(["EMIT(" e ")"])
      ;(e) : emit-all(["mov'" prim-width(type(e)) " " e], top) ;TODO, not sure which one is right

defn maybe-mov (e:Expression) -> String :
   val need-mov? = match(e) :
      (e:Ref) :       true
      (e:UIntValue) : true
      (e:SIntValue) : true
      (e:Subfield) :  true
      (e:Index) :     true
      (e) :           false
   if need-mov?: "mov " else: ""

defn emit-s (s:Stmt, flokinds:HashTable<Symbol,FloKind>, top:Symbol,sh:HashTable<Symbol,Int>) :
   defn emit-connect (s:Connect, en:Expression) :
      match(loc(s)) :
        (r:Ref) :
           val n = name(r)
           if key?(flokinds,n) :
              match(flokinds[n]) :
                 (k:FRegKind) : 
                    emit-all(["reg'" prim-width(type(r)) " " en " " exp(s)], top) 
                 (k:FWritePortKind) :
                    emit-all([top "::" n " = wr'" prim-width(type(r)) " " en " " mem(k) " " index(k) " " exp(s) "\n"], top)
                 (k:FOutputPortKind) : 
                    emit-all([top "::" n " = out'" prim-width(type(r)) " " exp(s) "\n"], top)
                 (k) : error("Shouldn't be here")
           else :
              emit-all([top "::" n " = " maybe-mov(exp(s)) exp(s) "\n"], top)
        (o) :
           println-all(["CONNEcT LOC " loc(s)])
           error("Unknown Connect")
   match(s) :
      (s:DefWire) : ""
      (s:DefPoison) : ""
      (s:DefInstance) : error("Shouldn't be here")
      (e:DefAccessor) : 
         if acc-dir == READ :
            emit-all(["rd'" prim-width(type(e)) " " "1" " " source(e) " " index(e)], top)
      (s:DefMemory) :
         val vtype = type(s) as VectorType
         emit-all([top "::" name(s) " = mem'" prim-width(type(vtype)) " " size(vtype) "\n"], top)
      (s:DefNode) :
         emit-all([top "::" name(s) " = " maybe-mov(value(s)) value(s) "\n"], top)
      (s:Begin) : do(emit-s{_, flokinds, top,sh}, body(s))
      (s:Connect) : emit-connect(s,UIntValue(BigIntLit("1"),LongWidth(1)))
      (s:Conditionally) : emit-connect(conseq(s) as Connect,pred(s))
      (s) : s 

defn emit-module (m:InModule,sh:HashTable<Symbol,Int>) :
   val flokinds = HashTable<Symbol,FloKind>(symbol-hash)
   defn build-table (s:Stmt) -> False :
      do(build-table,s)
      match(s) : 
         (s:DefRegister) : flokinds[name(s)] = FRegKind()
         (s:DefAccessor) : 
            switch {_ == acc-dir(s)} :
               WRITE : flokinds[name(s)] = FWritePortKind(source(s),index(s))
               else : false
         (s) : false
      
   for port in ports(m) do :
      if name(port) ==`reset :
         emit-all([name(m) "::" name(port) " = rst'1\n"], name(m))
      else : switch {_ == direction(port)} :
         INPUT : print-all([name(m) "::" name(port) " = " "in'" prim-width(type(port)) "\n"])
         OUTPUT : flokinds[name(port)] = FOutputPortKind()
   emit-s(body(m), flokinds, name(m),sh)

public defn emit-flo (with-output:(() -> False) -> False, c:Circuit) :
   with-output $ {
      emit-module(modules(c)[0] as InModule,get-sym-hash(modules(c)[0] as InModule))
      false
      }
   c