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


public defn lower-and-type-primop (e:DoPrim) -> DoPrim :
   defn u () : UIntType(UnknownWidth())
   defn s () : SIntType(UnknownWidth())
   defn u-and (op1:Expression,op2:Expression) : 
      match(type(op1), type(op2)) :
         (t1:UIntType, t2:UIntType) : u()
         (t1:SIntType, t2) : s()
         (t1, t2:SIntType) : s()
         (t1, t2) : UnknownType()
        
   defn of-type (op:Expression) :
      match(type(op)) :
         (t:UIntType) : u()
         (t:SIntType) : s()
         (t) : UnknownType()

   ;println-all(["Inferencing primop type: " e])
   switch {op(e) == _} :
      ADD-OP             : DoPrim(ADD-OP,args(e),consts(e),u-and(args(e)[0],args(e)[1])) 
      SUB-OP             : DoPrim(SUB-OP,args(e),consts(e),s())
      MUL-OP             : DoPrim(MUL-OP,args(e),consts(e),u-and(args(e)[0],args(e)[1]))
      DIV-OP             : DoPrim(DIV-OP,args(e),consts(e),u-and(args(e)[0],args(e)[1]))
      MOD-OP             : DoPrim(MOD-OP,args(e),consts(e),of-type(args(e)[0]))
      QUO-OP             : DoPrim(QUO-OP,args(e),consts(e),u-and(args(e)[0],args(e)[1]))
      REM-OP             : DoPrim(REM-OP,args(e),consts(e),of-type(args(e)[1]))
      ADD-WRAP-OP        : DoPrim(ADD-WRAP-OP,args(e),consts(e),u-and(args(e)[0],args(e)[1]))
      SUB-WRAP-OP        : DoPrim(SUB-WRAP-OP,args(e),consts(e),u-and(args(e)[0],args(e)[1]))
      LESS-OP            : DoPrim(LESS-OP,args(e),consts(e),u())
      LESS-EQ-OP         : DoPrim(LESS-EQ-OP,args(e),consts(e),u())
      GREATER-OP         : DoPrim(GREATER-OP,args(e),consts(e),u())
      GREATER-EQ-OP      : DoPrim(GREATER-EQ-OP,args(e),consts(e),u())
      EQUAL-OP           : DoPrim(EQUAL-OP,args(e),consts(e),u())
      NEQUAL-OP          : DoPrim(NEQUAL-OP,args(e),consts(e),u())
      MUX-OP             : DoPrim(MUX-OP,args(e),consts(e),of-type(args(e)[1]))
      PAD-OP             : DoPrim(PAD-OP,args(e),consts(e),of-type(args(e)[0]))
      AS-UINT-OP         : DoPrim(AS-UINT-OP,args(e),consts(e),u())
      AS-SINT-OP         : DoPrim(AS-SINT-OP,args(e),consts(e),s())
      DYN-SHIFT-LEFT-OP  : DoPrim(DYN-SHIFT-LEFT-OP,args(e),consts(e),of-type(args(e)[0]))
      DYN-SHIFT-RIGHT-OP : DoPrim(DYN-SHIFT-RIGHT-OP,args(e),consts(e),of-type(args(e)[0]))
      SHIFT-LEFT-OP      : DoPrim(SHIFT-LEFT-OP,args(e),consts(e),of-type(args(e)[0]))
      SHIFT-RIGHT-OP     : DoPrim(SHIFT-RIGHT-OP,args(e),consts(e),of-type(args(e)[0]))
      CONVERT-OP         : DoPrim(CONVERT-OP,args(e),consts(e),s())
      NEG-OP             : DoPrim(NEG-OP,args(e),consts(e),u())
      BIT-NOT-OP         : DoPrim(op(e),args(e),consts(e),u())
      BIT-AND-OP         : DoPrim(op(e),args(e),consts(e),u())
      BIT-OR-OP          : DoPrim(op(e),args(e),consts(e),u())
      BIT-XOR-OP         : DoPrim(op(e),args(e),consts(e),u())
      CONCAT-OP          : DoPrim(op(e),args(e),consts(e),u())
      BIT-SELECT-OP      : DoPrim(op(e),args(e),consts(e),u())
      BITS-SELECT-OP     : DoPrim(op(e),args(e),consts(e),u())
      BIT-AND-REDUCE-OP  : DoPrim(op(e),args(e),consts(e),u())
      BIT-OR-REDUCE-OP   : DoPrim(op(e),args(e),consts(e),u())
      BIT-XOR-REDUCE-OP  : DoPrim(op(e),args(e),consts(e),u())

public defn primop-gen-constraints (e:DoPrim,v:Vector<WGeq>) -> Type :
   defn get-max (i0:Int,i1:Int) -> Width : get-max(list(i0,i1))
   defn get-max (ls:List<Int>) -> Width : 
      MaxWidth $ for i in ls map : width!(args(e)[i])
   defn all-max () -> Width :
      MaxWidth $ for x in args(e) map : width!(x)

   println-all-debug(["Looking at " op(e) " with inputs " args(e)])
   val w* = switch {op(e) == _} :
      ADD-OP : PlusWidth(get-max(0,1),IntWidth(1))
      SUB-OP : PlusWidth(get-max(0,1),IntWidth(1))
      MUL-OP : PlusWidth(get-max(0,1),get-max(0,1))
      DIV-OP : 
         match(type(args(e)[0]),type(args(e)[1])) :
            (t0:UIntType,t1:SIntType) : PlusWidth(width!(args(e)[0]),IntWidth(1))
            (t0:SIntType,t1:SIntType) : PlusWidth(width!(args(e)[0]),IntWidth(1))
            (t0,t1) : width!(args(e)[0])
      MOD-OP : 
         match(type(args(e)[0]),type(args(e)[1])) :
            (t0:SIntType,t1:UIntType) : PlusWidth(width!(args(e)[1]),IntWidth(1))
            (t0,t1) : width!(args(e)[1])
      QUO-OP : 
         match(type(args(e)[0]),type(args(e)[1])) :
            (t0:UIntType,t1:SIntType) : PlusWidth(width!(args(e)[0]),IntWidth(1))
            (t0:SIntType,t1:SIntType) : PlusWidth(width!(args(e)[0]),IntWidth(1))
            (t0,t1) : width!(args(e)[0])
      REM-OP : 
         match(type(args(e)[0]),type(args(e)[1])) :
            (t0:SIntType,t1:UIntType) : PlusWidth(width!(args(e)[1]),IntWidth(1))
            (t0,t1) : width!(args(e)[1])
      ADD-WRAP-OP : get-max(0,1)
      SUB-WRAP-OP : get-max(0,1)
      LESS-OP : IntWidth(1)
      LESS-EQ-OP : IntWidth(1)
      GREATER-OP : IntWidth(1)
      GREATER-EQ-OP : IntWidth(1)
      EQUAL-OP : IntWidth(1)
      NEQUAL-OP : IntWidth(1)
      MUX-OP : 
         add(v,WGeq(IntWidth(1),width!(args(e)[0])))
         add(v,WGeq(width!(args(e)[0]),IntWidth(1)))
         get-max(1,2)
      PAD-OP : IntWidth(consts(e)[0])
      AS-UINT-OP : width!(args(e)[0])
      AS-SINT-OP : width!(args(e)[0])
      SHIFT-LEFT-OP : PlusWidth(width!(args(e)[0]),IntWidth(consts(e)[0]))
      SHIFT-RIGHT-OP : MinusWidth(width!(args(e)[0]),IntWidth(consts(e)[0]))
      DYN-SHIFT-LEFT-OP : PlusWidth(width!(args(e)[0]),ExpWidth(width!(args(e)[1])))
      DYN-SHIFT-RIGHT-OP : width!(args(e)[0])
      CONVERT-OP : 
         match(type(args(e)[0])) :
            (t0:UIntType) : PlusWidth(width!(args(e)[0]),IntWidth(1))
            (t0:SIntType) : width!(args(e)[0])
      NEG-OP : PlusWidth(width!(args(e)[0]),IntWidth(1))
      BIT-NOT-OP : width!(args(e)[0])
      BIT-AND-OP : get-max(0,1)
      BIT-OR-OP :  get-max(0,1)
      BIT-XOR-OP : get-max(0,1)
      BIT-AND-REDUCE-OP : all-max()
      BIT-OR-REDUCE-OP  : all-max()
      BIT-XOR-REDUCE-OP : all-max()
      CONCAT-OP : PlusWidth(width!(args(e)[0]),width!(args(e)[1]))
      BIT-SELECT-OP  : IntWidth(1)
      BITS-SELECT-OP : IntWidth(consts(e)[0] - consts(e)[1] + 1)

   match(type(e)) :
      (t:UIntType) : UIntType(w*)
      (t:SIntType) : SIntType(w*)
      (t) : error("Shouldn't be here")