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{_,args(e),consts(e),u-and(args(e)[0],args(e)[1])} $
            match(type(args(e)[0]),type(args(e)[1])) :
               (t1:UIntType, t2:UIntType) : ADD-UU-OP
               (t1:UIntType, t2:SIntType) : ADD-US-OP
               (t1:SIntType, t2:UIntType) : ADD-SU-OP
               (t1:SIntType, t2:SIntType) : ADD-SS-OP
      ADD-UU-OP        : DoPrim(op(e),args(e),consts(e),u())
      ADD-US-OP        : DoPrim(op(e),args(e),consts(e),s())
      ADD-SU-OP        : DoPrim(op(e),args(e),consts(e),s())
      ADD-SS-OP        : DoPrim(op(e),args(e),consts(e),s())
      SUB-OP           : 
         DoPrim{_,args(e),consts(e),s()} $
            match(type(args(e)[0]),type(args(e)[1])) :
               (t1:UIntType, t2:UIntType) : SUB-UU-OP
               (t1:UIntType, t2:SIntType) : SUB-US-OP
               (t1:SIntType, t2:UIntType) : SUB-SU-OP
               (t1:SIntType, t2:SIntType) : SUB-SS-OP
      SUB-UU-OP        : DoPrim(op(e),args(e),consts(e),s())
      SUB-US-OP        : DoPrim(op(e),args(e),consts(e),s())
      SUB-SU-OP        : DoPrim(op(e),args(e),consts(e),s())
      SUB-SS-OP        : DoPrim(op(e),args(e),consts(e),s())
      MUL-OP           : 
         DoPrim{_,args(e),consts(e),u-and(args(e)[0],args(e)[1])} $
            match(type(args(e)[0]),type(args(e)[1])) :
               (t1:UIntType, t2:UIntType) : MUL-UU-OP
               (t1:UIntType, t2:SIntType) : MUL-US-OP
               (t1:SIntType, t2:UIntType) : MUL-SU-OP
               (t1:SIntType, t2:SIntType) : MUL-SS-OP
      MUL-UU-OP        : DoPrim(op(e),args(e),consts(e),u())
      MUL-US-OP        : DoPrim(op(e),args(e),consts(e),s())
      MUL-SU-OP        : DoPrim(op(e),args(e),consts(e),s())
      MUL-SS-OP        : DoPrim(op(e),args(e),consts(e),s())
      DIV-OP           : 
         DoPrim{_,args(e),consts(e),u-and(args(e)[0],args(e)[1])} $
            match(type(args(e)[0]),type(args(e)[1])) :
               (t1:UIntType, t2:UIntType) : DIV-UU-OP
               (t1:UIntType, t2:SIntType) : DIV-US-OP
               (t1:SIntType, t2:UIntType) : DIV-SU-OP
               (t1:SIntType, t2:SIntType) : DIV-SS-OP
      DIV-UU-OP        : DoPrim(op(e),args(e),consts(e),u())
      DIV-US-OP        : DoPrim(op(e),args(e),consts(e),s())
      DIV-SU-OP        : DoPrim(op(e),args(e),consts(e),s())
      DIV-SS-OP        : DoPrim(op(e),args(e),consts(e),s())
      MOD-OP           : 
         DoPrim{_,args(e),consts(e),of-type(args(e)[0])} $
            match(type(args(e)[0]),type(args(e)[1])) :
               (t1:UIntType, t2:UIntType) : MOD-UU-OP
               (t1:UIntType, t2:SIntType) : MOD-US-OP
               (t1:SIntType, t2:UIntType) : MOD-SU-OP
               (t1:SIntType, t2:SIntType) : MOD-SS-OP
      MOD-UU-OP        : DoPrim(op(e),args(e),consts(e),u())
      MOD-US-OP        : DoPrim(op(e),args(e),consts(e),u())
      MOD-SU-OP        : DoPrim(op(e),args(e),consts(e),s())
      MOD-SS-OP        : DoPrim(op(e),args(e),consts(e),s())
      QUO-OP           : 
         DoPrim{_,args(e),consts(e),u-and(args(e)[0],args(e)[1])} $
            match(type(args(e)[0]),type(args(e)[1])) :
               (t1:UIntType, t2:UIntType) : QUO-UU-OP
               (t1:UIntType, t2:SIntType) : QUO-US-OP
               (t1:SIntType, t2:UIntType) : QUO-SU-OP
               (t1:SIntType, t2:SIntType) : QUO-SS-OP
      QUO-UU-OP        : DoPrim(op(e),args(e),consts(e),u())
      QUO-US-OP        : DoPrim(op(e),args(e),consts(e),s())
      QUO-SU-OP        : DoPrim(op(e),args(e),consts(e),s())
      QUO-SS-OP        : DoPrim(op(e),args(e),consts(e),s())
      REM-OP           : 
         DoPrim{_,args(e),consts(e),of-type(args(e)[1])} $
            match(type(args(e)[0]),type(args(e)[1])) :
               (t1:UIntType, t2:UIntType) : REM-UU-OP
               (t1:UIntType, t2:SIntType) : REM-US-OP
               (t1:SIntType, t2:UIntType) : REM-SU-OP
               (t1:SIntType, t2:SIntType) : REM-SS-OP
      REM-UU-OP        : DoPrim(op(e),args(e),consts(e),u())
      REM-US-OP        : DoPrim(op(e),args(e),consts(e),s())
      REM-SU-OP        : DoPrim(op(e),args(e),consts(e),u())
      REM-SS-OP        : DoPrim(op(e),args(e),consts(e),s())
      ADD-WRAP-OP      : 
         DoPrim{_,args(e),consts(e),u-and(args(e)[0],args(e)[1])} $
            match(type(args(e)[0]),type(args(e)[1])) :
               (t1:UIntType, t2:UIntType) : ADD-WRAP-UU-OP
               (t1:UIntType, t2:SIntType) : ADD-WRAP-US-OP
               (t1:SIntType, t2:UIntType) : ADD-WRAP-SU-OP
               (t1:SIntType, t2:SIntType) : ADD-WRAP-SS-OP
      ADD-WRAP-UU-OP   : DoPrim(op(e),args(e),consts(e),u())
      ADD-WRAP-US-OP   : DoPrim(op(e),args(e),consts(e),s())
      ADD-WRAP-SU-OP   : DoPrim(op(e),args(e),consts(e),s())
      ADD-WRAP-SS-OP   : DoPrim(op(e),args(e),consts(e),s())
      SUB-WRAP-OP      : 
         DoPrim{_,args(e),consts(e),u-and(args(e)[0],args(e)[1])} $
            match(type(args(e)[0]),type(args(e)[1])) :
               (t1:UIntType, t2:UIntType) : SUB-WRAP-UU-OP
               (t1:UIntType, t2:SIntType) : SUB-WRAP-US-OP
               (t1:SIntType, t2:UIntType) : SUB-WRAP-SU-OP
               (t1:SIntType, t2:SIntType) : SUB-WRAP-SS-OP
      SUB-WRAP-UU-OP   : DoPrim(op(e),args(e),consts(e),u())
      SUB-WRAP-US-OP   : DoPrim(op(e),args(e),consts(e),s())
      SUB-WRAP-SU-OP   : DoPrim(op(e),args(e),consts(e),s())
      SUB-WRAP-SS-OP   : DoPrim(op(e),args(e),consts(e),s())
      LESS-OP          : 
         DoPrim{_,args(e),consts(e),u()} $
            match(type(args(e)[0]),type(args(e)[1])) :
               (t1:UIntType, t2:UIntType) : LESS-UU-OP
               (t1:UIntType, t2:SIntType) : LESS-US-OP
               (t1:SIntType, t2:UIntType) : LESS-SU-OP
               (t1:SIntType, t2:SIntType) : LESS-SS-OP
      LESS-UU-OP       : DoPrim(op(e),args(e),consts(e),u())
      LESS-US-OP       : DoPrim(op(e),args(e),consts(e),u())
      LESS-SU-OP       : DoPrim(op(e),args(e),consts(e),u())
      LESS-SS-OP       : DoPrim(op(e),args(e),consts(e),u())
      LESS-EQ-OP       : 
         DoPrim{_,args(e),consts(e),u()} $
            match(type(args(e)[0]),type(args(e)[1])) :
               (t1:UIntType, t2:UIntType) : LESS-EQ-UU-OP
               (t1:UIntType, t2:SIntType) : LESS-EQ-US-OP
               (t1:SIntType, t2:UIntType) : LESS-EQ-SU-OP
               (t1:SIntType, t2:SIntType) : LESS-EQ-SS-OP
      LESS-EQ-UU-OP    : DoPrim(op(e),args(e),consts(e),u())
      LESS-EQ-US-OP    : DoPrim(op(e),args(e),consts(e),u())
      LESS-EQ-SU-OP    : DoPrim(op(e),args(e),consts(e),u())
      LESS-EQ-SS-OP    : DoPrim(op(e),args(e),consts(e),u())
      GREATER-OP       : 
         DoPrim{_,args(e),consts(e),u()} $
            match(type(args(e)[0]),type(args(e)[1])) :
               (t1:UIntType, t2:UIntType) : GREATER-UU-OP
               (t1:UIntType, t2:SIntType) : GREATER-US-OP
               (t1:SIntType, t2:UIntType) : GREATER-SU-OP
               (t1:SIntType, t2:SIntType) : GREATER-SS-OP
      GREATER-UU-OP    : DoPrim(op(e),args(e),consts(e),u())
      GREATER-US-OP    : DoPrim(op(e),args(e),consts(e),u())
      GREATER-SU-OP    : DoPrim(op(e),args(e),consts(e),u())
      GREATER-SS-OP    : DoPrim(op(e),args(e),consts(e),u())
      GREATER-EQ-OP    : 
         DoPrim{_,args(e),consts(e),u()} $
            match(type(args(e)[0]),type(args(e)[1])) :
               (t1:UIntType, t2:UIntType) : GREATER-EQ-UU-OP
               (t1:UIntType, t2:SIntType) : GREATER-EQ-US-OP
               (t1:SIntType, t2:UIntType) : GREATER-EQ-SU-OP
               (t1:SIntType, t2:SIntType) : GREATER-EQ-SS-OP
      GREATER-EQ-UU-OP : DoPrim(op(e),args(e),consts(e),u())
      GREATER-EQ-US-OP : DoPrim(op(e),args(e),consts(e),u())
      GREATER-EQ-SU-OP : DoPrim(op(e),args(e),consts(e),u())
      GREATER-EQ-SS-OP : DoPrim(op(e),args(e),consts(e),u())
      EQUAL-OP         : 
         DoPrim{_,args(e),consts(e),u()} $
            match(type(args(e)[0]),type(args(e)[1])) :
               (t1:UIntType, t2:UIntType) : EQUAL-UU-OP
               (t1:SIntType, t2:SIntType) : EQUAL-SS-OP
      EQUAL-UU-OP      : DoPrim(op(e),args(e),consts(e),u())
      EQUAL-SS-OP      : DoPrim(op(e),args(e),consts(e),u())
      NEQUAL-OP         : 
         DoPrim{_,args(e),consts(e),u()} $
            match(type(args(e)[0]),type(args(e)[1])) :
               (t1:UIntType, t2:UIntType) : NEQUAL-UU-OP
               (t1:SIntType, t2:SIntType) : NEQUAL-SS-OP
      NEQUAL-UU-OP      : DoPrim(op(e),args(e),consts(e),u())
      NEQUAL-SS-OP      : DoPrim(op(e),args(e),consts(e),u())
      MUX-OP           : 
         DoPrim{_,args(e),consts(e),of-type(args(e)[1])} $
            match(type(args(e)[1]),type(args(e)[2])) :
               (t1:UIntType, t2:UIntType) : MUX-UU-OP
               (t1:SIntType, t2:SIntType) : MUX-SS-OP
               (t1,t2) : error(to-string(args(e)))
      MUX-UU-OP        : DoPrim(op(e),args(e),consts(e),u())
      MUX-SS-OP        : DoPrim(op(e),args(e),consts(e),s())
      PAD-OP           : 
         DoPrim{_,args(e),consts(e),of-type(args(e)[0])} $
            match(type(args(e)[0])) :
               (t1:UIntType) : PAD-U-OP
               (t1:SIntType) : PAD-S-OP
      PAD-U-OP         : DoPrim(op(e),args(e),consts(e),u())
      PAD-S-OP         : DoPrim(op(e),args(e),consts(e),s())
      AS-UINT-OP       : 
         DoPrim{_,args(e),consts(e),u()} $
            match(type(args(e)[0])) :
               (t1:UIntType) : AS-UINT-U-OP
               (t1:SIntType) : AS-UINT-S-OP
      AS-UINT-U-OP     : DoPrim(op(e),args(e),consts(e),u())
      AS-UINT-S-OP     : DoPrim(op(e),args(e),consts(e),u())
      AS-SINT-OP       : 
         DoPrim{_,args(e),consts(e),s()} $
            match(type(args(e)[0])) :
               (t1:UIntType) : AS-SINT-U-OP
               (t1:SIntType) : AS-SINT-S-OP
      AS-SINT-U-OP     : DoPrim(op(e),args(e),consts(e),s())
      AS-SINT-S-OP     : DoPrim(op(e),args(e),consts(e),s())
      DYN-SHIFT-LEFT-OP    : 
         DoPrim{_,args(e),consts(e),of-type(args(e)[0])} $
            match(type(args(e)[0]),type(args(e)[1])) :
               (t1:UIntType,t2:UIntType) : DYN-SHIFT-LEFT-U-OP
               (t1:SIntType,t2:UIntType) : DYN-SHIFT-LEFT-S-OP
      DYN-SHIFT-LEFT-U-OP  : DoPrim(op(e),args(e),consts(e),u())
      DYN-SHIFT-LEFT-S-OP  : DoPrim(op(e),args(e),consts(e),s())
      DYN-SHIFT-RIGHT-OP   : 
         DoPrim{_,args(e),consts(e),of-type(args(e)[0])} $
            match(type(args(e)[0]),type(args(e)[1])) :
               (t1:UIntType,t2:UIntType) : DYN-SHIFT-RIGHT-U-OP
               (t1:SIntType,t2:UIntType) : DYN-SHIFT-RIGHT-S-OP
      DYN-SHIFT-RIGHT-U-OP : DoPrim(op(e),args(e),consts(e),u())
      DYN-SHIFT-RIGHT-S-OP : DoPrim(op(e),args(e),consts(e),s())
      SHIFT-LEFT-OP    : 
         DoPrim{_,args(e),consts(e),of-type(args(e)[0])} $
            match(type(args(e)[0])) :
               (t1:UIntType) : SHIFT-LEFT-U-OP
               (t1:SIntType) : SHIFT-LEFT-S-OP
      SHIFT-LEFT-U-OP  : DoPrim(op(e),args(e),consts(e),u())
      SHIFT-LEFT-S-OP  : DoPrim(op(e),args(e),consts(e),s())
      SHIFT-RIGHT-OP   : 
         DoPrim{_,args(e),consts(e),of-type(args(e)[0])} $
            match(type(args(e)[0])) :
               (t1:UIntType) : SHIFT-RIGHT-U-OP
               (t1:SIntType) : SHIFT-RIGHT-S-OP
      SHIFT-RIGHT-U-OP : DoPrim(op(e),args(e),consts(e),u())
      SHIFT-RIGHT-S-OP : DoPrim(op(e),args(e),consts(e),s())
      CONVERT-OP       : 
         DoPrim{_,args(e),consts(e),s()} $
            match(type(args(e)[0])) :
               (t1:UIntType) : CONVERT-U-OP
               (t1:SIntType) : CONVERT-S-OP
      CONVERT-U-OP     : DoPrim(op(e),args(e),consts(e),s())
      CONVERT-S-OP     : DoPrim(op(e),args(e),consts(e),s())
      NEG-OP         : 
         DoPrim{_,args(e),consts(e),u()} $
            match(type(args(e)[0])) :
               (t1:UIntType) : NEG-U-OP
               (t1:SIntType) : NEG-S-OP
      NEG-U-OP         : DoPrim(op(e),args(e),consts(e),u())
      NEG-S-OP         : DoPrim(op(e),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 all-equal (ls:List<Width>) -> Width : 
      if length(ls) == 1 : (ls[0])
      else :
         val m = MaxWidth(ls)
         for (l in ls) do :
            add(v,WGeq(l,m))
         m
   println-all-debug(["Looking at " op(e) " with inputs " args(e)])
   val all-args-not-equal = to-list([MUX-UU-OP,MUX-SS-OP,CONCAT-OP,BIT-AND-OP,BIT-NOT-OP,BIT-OR-OP,BIT-XOR-OP,BIT-AND-REDUCE-OP,BIT-OR-REDUCE-OP,BIT-XOR-REDUCE-OP,AS-UINT-U-OP,AS-UINT-S-OP,AS-SINT-U-OP,AS-SINT-S-OP,DYN-SHIFT-LEFT-U-OP,DYN-SHIFT-LEFT-S-OP,DYN-SHIFT-RIGHT-U-OP,DYN-SHIFT-RIGHT-S-OP])
   
   val w-var = VarWidth(gensym(`w))
   val w* = 
      if not contains?(all-args-not-equal,op(e)) :
         val max-args-w = all-equal(map(width!,args(e)))
         switch {op(e) == _} :
            ADD-UU-OP : PlusWidth(max-args-w,IntWidth(1))
            ADD-US-OP : PlusWidth(max-args-w,IntWidth(1))
            ADD-SU-OP : PlusWidth(max-args-w,IntWidth(1))
            ADD-SS-OP : PlusWidth(max-args-w,IntWidth(1))
            SUB-UU-OP : PlusWidth(max-args-w,IntWidth(1))
            SUB-US-OP : PlusWidth(max-args-w,IntWidth(1))
            SUB-SU-OP : PlusWidth(max-args-w,IntWidth(1))
            SUB-SS-OP : PlusWidth(max-args-w,IntWidth(1))
            MUL-UU-OP : PlusWidth(max-args-w,max-args-w)
            MUL-US-OP : PlusWidth(max-args-w,max-args-w)
            MUL-SU-OP : PlusWidth(max-args-w,max-args-w)
            MUL-SS-OP : PlusWidth(max-args-w,max-args-w)
            DIV-UU-OP : max-args-w
            DIV-US-OP : PlusWidth(max-args-w,IntWidth(1))
            DIV-SU-OP : max-args-w
            DIV-SS-OP : PlusWidth(max-args-w,IntWidth(1))
            ADD-WRAP-UU-OP : max-args-w
            ADD-WRAP-US-OP : max-args-w
            ADD-WRAP-SU-OP : max-args-w
            ADD-WRAP-SS-OP : max-args-w
            SUB-WRAP-UU-OP : max-args-w
            SUB-WRAP-US-OP : max-args-w
            SUB-WRAP-SU-OP : max-args-w
            SUB-WRAP-SS-OP : max-args-w
            LESS-UU-OP : IntWidth(1)
            LESS-US-OP : IntWidth(1)
            LESS-SU-OP : IntWidth(1)
            LESS-SS-OP : IntWidth(1)
            LESS-EQ-UU-OP : IntWidth(1)
            LESS-EQ-US-OP : IntWidth(1)
            LESS-EQ-SU-OP : IntWidth(1)
            LESS-EQ-SS-OP : IntWidth(1)
            GREATER-UU-OP : IntWidth(1)
            GREATER-US-OP : IntWidth(1)
            GREATER-SU-OP : IntWidth(1)
            GREATER-SS-OP : IntWidth(1)
            GREATER-EQ-UU-OP : IntWidth(1)
            GREATER-EQ-US-OP : IntWidth(1)
            GREATER-EQ-SU-OP : IntWidth(1)
            GREATER-EQ-SS-OP : IntWidth(1)
            EQUAL-UU-OP : IntWidth(1)
            EQUAL-SS-OP : IntWidth(1)
            NEQUAL-UU-OP : IntWidth(1)
            NEQUAL-SS-OP : IntWidth(1)
            PAD-U-OP : IntWidth(consts(e)[0])
            PAD-S-OP : IntWidth(consts(e)[0])
            NEG-U-OP : PlusWidth(max-args-w,IntWidth(1))
            NEG-S-OP : PlusWidth(max-args-w,IntWidth(1))
            SHIFT-LEFT-U-OP : PlusWidth(max-args-w,IntWidth(consts(e)[0]))
            SHIFT-LEFT-S-OP : PlusWidth(max-args-w,IntWidth(consts(e)[0]))
            SHIFT-RIGHT-U-OP : MinusWidth(max-args-w,IntWidth(consts(e)[0]))
            SHIFT-RIGHT-S-OP : MinusWidth(max-args-w,IntWidth(consts(e)[0]))
            CONVERT-U-OP : PlusWidth(max-args-w,IntWidth(1))
            CONVERT-S-OP : max-args-w
            BIT-SELECT-OP : IntWidth(1)
            BITS-SELECT-OP : IntWidth(consts(e)[0] - consts(e)[1])
      else : 
         switch {op(e) == _} :
            MUX-UU-OP : 
               add(v,WGeq(width!(args(e)[0]),IntWidth(1)))
               all-equal(List(w-var,tail(map(width!,args(e)))))
            MUX-SS-OP : 
               add(v,WGeq(width!(args(e)[0]),IntWidth(1)))
               all-equal(List(w-var,tail(map(width!,args(e)))))
            CONCAT-OP : 
               PlusWidth(width!(args(e)[0]),width!(args(e)[1]))
            BIT-NOT-OP : all-equal(List(w-var,map(width!,args(e))))
            BIT-AND-OP : all-equal(List(w-var,map(width!,args(e))))
            BIT-OR-OP :  all-equal(List(w-var,map(width!,args(e))))
            BIT-XOR-OP : all-equal(List(w-var,map(width!,args(e))))
            BIT-AND-REDUCE-OP : all-equal(List(w-var,map(width!,args(e))))
            BIT-OR-REDUCE-OP : all-equal(List(w-var,map(width!,args(e))))
            BIT-XOR-REDUCE-OP : all-equal(List(w-var,map(width!,args(e))))
            AS-UINT-U-OP : all-equal(List(w-var,map(width!,args(e))))
            AS-UINT-S-OP : all-equal(List(w-var,map(width!,args(e))))
            AS-SINT-U-OP : all-equal(List(w-var,map(width!,args(e))))
            AS-SINT-S-OP : all-equal(List(w-var,map(width!,args(e))))
            DYN-SHIFT-LEFT-U-OP : 
               PlusWidth(width!(args(e)[0]),ExpWidth(width!(args(e)[1])))
            DYN-SHIFT-LEFT-S-OP : 
               PlusWidth(width!(args(e)[0]),ExpWidth(width!(args(e)[1])))
            DYN-SHIFT-RIGHT-U-OP : width!(args(e)[0])
            DYN-SHIFT-RIGHT-S-OP : width!(args(e)[0])

   add(v,WGeq(w-var,w*))
   match(type(e)) :
      (t:UIntType) : UIntType(w-var)
      (t:SIntType) : SIntType(w-var)
      (t) : error("Shouldn't be here")