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

;========== Pad Widths ==================

public defstruct Pad <: Pass
public defmethod pass (b:Pad) -> (Circuit -> Circuit) : pad-widths
public defmethod name (b:Pad) -> String : "Pad Widths"

defn int-width! (t:Type) -> Int :
   match(width!(t)) :
      (w:IntWidth) : width(w)
      (w) : error("Non-int width")

defn set-width (desired:Int,t:Type) -> Type :
   match(t) :
      (t:UIntType) : UIntType(IntWidth(desired))
      (t:SIntType) : SIntType(IntWidth(desired))
      (t) : error("Non-ground type")

defn pad-widths-e (desired:Int,e:Expression) -> Expression :
   defn trim (desired:Int, e:Expression) :
     ;; println-all(["TRIM " desired " e " e])
     DoPrim(BITS-SELECT-OP,list(e),list(desired - 1, 0),set-width(desired,type(e)))
   defn pad (desired:Int, e:Expression) :
     ;; println-all(["PAD " desired " e " e])
     DoPrim(PAD-OP,list(e),list(desired),set-width(desired,type(e)))
   defn trim-pad (desired:Int, e:Expression) :
     val i = int-width!(type(e))
     if i > desired : trim(desired, e)
     else if i == desired : e
     else : pad(desired, e)
   defn self-pad-widths-e (e:Expression) -> Expression :
     pad-widths-e(int-width!(type(e)), e)
   ;; println-all(["PAD-E " desired " " e])
   match(e) :
      (e:DoPrim) :
         val new-desired = reduce(max, 0, map(int-width!{type(_)}, args(e)))
         ;; println-all(["  NEW DESIRED " new-desired])
         val e* =
           if contains?([CONCAT-OP, DYN-SHIFT-RIGHT-OP, DYN-SHIFT-LEFT-OP], op(e)) :
             DoPrim(op(e), map(self-pad-widths-e, args(e)), consts(e), type(e))
           else if contains?([MUX-OP], op(e)) :
             DoPrim(op(e), list(pad-widths-e(1, args(e)[0]), pad-widths-e(new-desired, args(e)[1]), pad-widths-e(new-desired, args(e)[2])), consts(e), type(e))
           else :
             map(pad-widths-e{new-desired,_},e)
         trim-pad(desired, e*)
      (e:WRef|WSubfield|WIndex) : 
         trim-pad(desired, e)
      (e:UIntValue) : 
         val i = int-width!(type(e))
         if i > desired : trim(desired, e)
         else : UIntValue(value(e),IntWidth(desired))
      (e:SIntValue) : 
         val i = int-width!(type(e))
         if i > desired : trim(desired, e)
         else : SIntValue(value(e),IntWidth(desired))
      (e:Register) :
         trim-pad(desired, Register(type(e), pad-widths-e(int-width!(type(e)), value(e)), pad-widths-e(1, enable(e))))
      (e:ReadPort) :
         trim-pad(desired, ReadPort(mem(e), self-pad-widths-e(index(e)), type(e), pad-widths-e(1, enable(e))))
      (e:WritePort) :
         trim-pad(desired, WritePort(mem(e), self-pad-widths-e(index(e)), type(e), pad-widths-e(1, enable(e))))
      (e) : error(to-string $ e)

defn pad-widths-s (s:Stmt) -> Stmt :
   ;; println-all(["PAD-S " s])
   match(map(pad-widths-s,s)) :
      (s:Connect) : 
         val i = int-width!(type(loc(s)))
         val loc* = pad-widths-e(i,loc(s))
         val exp* = pad-widths-e(i,exp(s))
         Connect(info(s),loc*,exp*)
      (s:DefNode) : 
         val i = int-width!(type(value(s)))
         val exp* = pad-widths-e(i,value(s))
         DefNode(info(s),name(s),exp*)
      (s) :
         s

public defn pad-widths (c:Circuit) -> Circuit :
   Circuit{info(c),_,main(c)} $
      for m in modules(c) map :
         match(m) :
            (m:ExModule) : error("Cannot use flo backend with external modules")
            (m:InModule) : InModule(info(m),name(m),ports(m),pad-widths-s(body(m)))

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

public defstruct Flo <: Pass :
   file : String
public defmethod pass (b:Flo) -> (Circuit -> Circuit) : emit-flo{file(b),_}
public defmethod name (b:Flo) -> String : "To Flo"
  
defn is-sint? (arg:Expression) -> True|False : type(arg) typeof SIntType

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
      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 :
  match(wd) :
    (w:IntWidth) : max(1, width(w))
    (w)          : error(string-join(["Unknown width: " w]))

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 :
  max(1, ceil-log2(in + 1))

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:Pad) : 
         ;emit-all(["rsh'" prim-width(type(e)) " " value(e) " 0"], top)
      (e:Register) : 
         emit-all(["reg'" prim-width(type(e)) " " enable(e) " " value(e)], top) 
      (e:ReadPort) : 
         emit-all(["rd'" prim-width(type(e)) " " "1" " " mem(e) " " index(e)], top) ;; enable(e)
      (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 "'" sizeof(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, v:List<Symbol>, top:Symbol,sh:HashTable<Symbol,Int>) :
   match(s) :
      (s:DefWire) : ""
      (s:DefInstance) : error("Shouldn't be here")
      (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{_, v, top,sh}, body(s))
      (s:Connect) : 
         match(loc(s)) :
           (r:Ref) :
             val n = name(r)
             if contains?(v,n) :
                emit-all([top "::" n " = out'" prim-width(type(r)) " " exp(s) "\n"], top)
             else :
                emit-all([top "::" n " = " maybe-mov(exp(s)) exp(s) "\n"], top)
           (w:WritePort) :
             val n = firrtl-gensym(`F,sh)
             emit-all([top "::" n " = wr'" prim-width(type(w)) " " enable(w) " " mem(w) " " index(w) " " exp(s) "\n"], top)
           (o) :
             println-all(["CONNEcT LOC " loc(s)])
             error("Unknown Connect")
      (s) : s 

defn emit-module (m:InModule,sh:HashTable<Symbol,Int>) :
   val v = Vector<Symbol>()
   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 : add(v,name(port))
   emit-s(body(m), to-list(v), name(m),sh)

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