;;;; Instruction definition ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

(type OperandSize extern (enum Size32 Size64))

(decl ty_to_operand_size (Type) OperandSize)
(rule (ty_to_operand_size $I32) (OperandSize.Size32))
(rule (ty_to_operand_size $I64) (OperandSize.Size64))

;; Note: in the instructions below, we order destination registers first and
;; then source registers afterwards.
(type MInst
  (enum
    ;;;; Pseudo-Instructions ;;;;

    ;; A pseudo-instruction that captures register arguments in vregs.
    (Args (args VecArgPair))

    ;; A pseudo-instruction that moves vregs to return registers.
    (Rets (rets VecRetPair))

    ;; Implementation of `br_table`, uses `idx` to jump to one of `targets` or
    ;; jumps to `default` if it's out-of-bounds.
    (BrTable
      (idx XReg)
      (default MachLabel)
      (targets BoxVecMachLabel))

    ;;;; Actual Instructions ;;;;

    ;; Trap if `cond` is true.
    (TrapIf (cond Cond) (code TrapCode))

    ;; Nothing.
    (Nop)

    ;; Move a special register (e.g. sp, fp, lr, etc) in to a general-purpose
    ;; register.
    (GetSpecial
      (dst WritableXReg)
      (reg XReg))

    ;; Load an external symbol's address into a register.
    (LoadExtName (dst WritableXReg)
                 (name BoxExternalName)
                 (offset i64))

    ;; A direct call to a known callee.
    (Call (info BoxCallInfo))

    ;; An indirect call to an unknown callee.
    (IndirectCall (info BoxCallIndInfo))

    ;; A direct return-call macro instruction.
    (ReturnCall (info BoxReturnCallInfo))

    ;; An indirect return-call macro instruction.
    (ReturnIndirectCall (info BoxReturnCallIndInfo))

    ;; An indirect call out to a host-defined function. The host function
    ;; pointer is the first "argument" of this function call.
    (IndirectCallHost (info BoxCallIndirectHostInfo))

    ;; Unconditional jumps.
    (Jump (label MachLabel))

    ;; Jump to `then` if `c` is true, otherwise to `else`.
    (BrIf (cond Cond) (taken MachLabel) (not_taken MachLabel))

    ;; Load the memory address referenced by `mem` into `dst`.
    (LoadAddr (dst WritableXReg) (mem Amode))

    ;; Load `ty` bytes from memory pointed to by `mem` and store in `dst`.
    ;;
    ;; How much is written to the register is defined by `ExtKind`. The `flags`
    ;; control behavior such as endianness.
    (XLoad (dst WritableXReg) (mem Amode) (ty Type) (flags MemFlags) (ext ExtKind))
    (FLoad (dst WritableFReg) (mem Amode) (ty Type) (flags MemFlags))
    (VLoad (dst WritableVReg) (mem Amode) (ty Type) (flags MemFlags) (ext VExtKind))

    ;; Stores.
    (XStore (mem Amode) (src XReg) (ty Type) (flags MemFlags))
    (FStore (mem Amode) (src FReg) (ty Type) (flags MemFlags))
    (VStore (mem Amode) (src VReg) (ty Type) (flags MemFlags))

    ;; A raw pulley instruction generated at compile-time via Pulley's
    ;; `for_each_op!` macro. This variant has `pulley_*` constructors to
    ;; emit this instruction and auto-generated methods for other various
    ;; bits and pieces of boilerplate in the backend.
    (Raw (raw RawInst))
  )
)

;; Helper type on conditional branches and traps to represent what the
;; condition that is being performed is.
;;
;; Used in `BrIf` and `TrapIf` above for example.
(type Cond
  (enum
    ;; True if `reg` contains a nonzero value in the low 32-bits.
    (If32 (reg XReg))
    ;; True if `reg` contains a zero in the low 32-bits.
    (IfNot32 (reg XReg))

    ;; Conditionals for comparing the low 32-bits of two registers.
    (IfXeq32 (src1 XReg) (src2 XReg))
    (IfXneq32 (src1 XReg) (src2 XReg))
    (IfXslt32 (src1 XReg) (src2 XReg))
    (IfXslteq32 (src1 XReg) (src2 XReg))
    (IfXult32 (src1 XReg) (src2 XReg))
    (IfXulteq32 (src1 XReg) (src2 XReg))

    (IfXeq32I32 (src1 XReg) (src2 i32))
    (IfXneq32I32 (src1 XReg) (src2 i32))
    (IfXslt32I32 (src1 XReg) (src2 i32))
    (IfXslteq32I32 (src1 XReg) (src2 i32))
    (IfXult32I32 (src1 XReg) (src2 u32))
    (IfXulteq32I32 (src1 XReg) (src2 u32))
    (IfXsgt32I32 (src1 XReg) (src2 i32))
    (IfXsgteq32I32 (src1 XReg) (src2 i32))
    (IfXugt32I32 (src1 XReg) (src2 u32))
    (IfXugteq32I32 (src1 XReg) (src2 u32))

    ;; Conditionals for comparing two 64-bit registers.
    (IfXeq64 (src1 XReg) (src2 XReg))
    (IfXneq64 (src1 XReg) (src2 XReg))
    (IfXslt64 (src1 XReg) (src2 XReg))
    (IfXslteq64 (src1 XReg) (src2 XReg))
    (IfXult64 (src1 XReg) (src2 XReg))
    (IfXulteq64 (src1 XReg) (src2 XReg))

    (IfXeq64I32 (src1 XReg) (src2 i32))
    (IfXneq64I32 (src1 XReg) (src2 i32))
    (IfXslt64I32 (src1 XReg) (src2 i32))
    (IfXslteq64I32 (src1 XReg) (src2 i32))
    (IfXult64I32 (src1 XReg) (src2 u32))
    (IfXulteq64I32 (src1 XReg) (src2 u32))
    (IfXsgt64I32 (src1 XReg) (src2 i32))
    (IfXsgteq64I32 (src1 XReg) (src2 i32))
    (IfXugt64I32 (src1 XReg) (src2 u32))
    (IfXugteq64I32 (src1 XReg) (src2 u32))
  )
)

(decl cond_invert (Cond) Cond)
(extern constructor cond_invert cond_invert)

(decl raw_inst_to_inst (RawInst) MInst)
(rule (raw_inst_to_inst inst) (MInst.Raw inst))
(convert RawInst MInst raw_inst_to_inst)

(type U6 (primitive U6))
(type BoxCallInfo (primitive BoxCallInfo))
(type BoxCallIndInfo (primitive BoxCallIndInfo))
(type BoxReturnCallInfo (primitive BoxReturnCallInfo))
(type BoxReturnCallIndInfo (primitive BoxReturnCallIndInfo))
(type XRegSet (primitive XRegSet))
(type BoxCallIndirectHostInfo (primitive BoxCallIndirectHostInfo))

;;;; Address Modes ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

(type StackAMode extern (enum))

(type Amode
  (enum
    (SpOffset (offset i32))
    (RegOffset (base XReg) (offset i32))
    (Stack (amode StackAMode))
  )
)

(type ExtKind (enum None Sign32 Sign64 Zero32 Zero64))

(type VExtKind (enum None S8x8 U8x8 S16x4 U16x4 S32x2 U32x2))

;;;; Newtypes for Different Register Classes ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

(type XReg (primitive XReg))
(type WritableXReg (primitive WritableXReg))

(type FReg (primitive FReg))
(type WritableFReg (primitive WritableFReg))

(type VReg (primitive VReg))
(type WritableVReg (primitive WritableVReg))

;; Construct a new `XReg` from a `Reg`.
;;
;; Asserts that the register has an integer RegClass.
(decl xreg_new (Reg) XReg)
(extern constructor xreg_new xreg_new)
(convert Reg XReg xreg_new)

;; Construct a new `WritableXReg` from a `WritableReg`.
;;
;; Asserts that the register has an integer RegClass.
(decl writable_xreg_new (WritableReg) WritableXReg)
(extern constructor writable_xreg_new writable_xreg_new)
(convert WritableReg WritableXReg writable_xreg_new)

;; Put a value into a XReg.
;;
;; Asserts that the value goes into a XReg.
(decl put_in_xreg (Value) XReg)
(rule (put_in_xreg val) (xreg_new (put_in_reg val)))
(convert Value XReg put_in_xreg)

;; Construct an `InstOutput` out of a single XReg register.
(decl output_xreg (XReg) InstOutput)
(rule (output_xreg x) (output_reg x))
(convert XReg InstOutput output_xreg)

;; Convert a `WritableXReg` to an `XReg`.
(decl pure writable_xreg_to_xreg (WritableXReg) XReg)
(extern constructor writable_xreg_to_xreg writable_xreg_to_xreg)
(convert WritableXReg XReg writable_xreg_to_xreg)

;; Convert a `WritableXReg` to an `WritableReg`.
(decl pure writable_xreg_to_writable_reg (WritableXReg) WritableReg)
(extern constructor writable_xreg_to_writable_reg writable_xreg_to_writable_reg)
(convert WritableXReg WritableReg writable_xreg_to_writable_reg)

;; Convert a `WritableXReg` to an `Reg`.
(decl pure writable_xreg_to_reg (WritableXReg) Reg)
(rule (writable_xreg_to_reg x) (writable_xreg_to_writable_reg x))
(convert WritableXReg Reg writable_xreg_to_reg)

;; Convert an `XReg` to a `Reg`.
(decl pure xreg_to_reg (XReg) Reg)
(extern constructor xreg_to_reg xreg_to_reg)
(convert XReg Reg xreg_to_reg)

;; Convert a `XReg` to a `ValueRegs`.
(decl xreg_to_value_regs (XReg) ValueRegs)
(rule (xreg_to_value_regs x) (value_reg x))
(convert XReg ValueRegs xreg_to_reg)

;; Convert a `WritableXReg` to a `ValueRegs`.
(decl writable_xreg_to_value_regs (WritableXReg) ValueRegs)
(rule (writable_xreg_to_value_regs x) (value_reg x))
(convert WritableXReg ValueRegs writable_xreg_to_value_regs)

;; Allocates a new `WritableXReg`.
(decl temp_writable_xreg () WritableXReg)
(rule (temp_writable_xreg) (temp_writable_reg $I64))

;; Construct a new `FReg` from a `Reg`.
;;
;; Asserts that the register has a Float RegClass.
(decl freg_new (Reg) FReg)
(extern constructor freg_new freg_new)
(convert Reg FReg freg_new)

;; Construct a new `WritableFReg` from a `WritableReg`.
;;
;; Asserts that the register has a Float RegClass.
(decl writable_freg_new (WritableReg) WritableFReg)
(extern constructor writable_freg_new writable_freg_new)
(convert WritableReg WritableFReg writable_freg_new)

;; Put a value into a FReg.
;;
;; Asserts that the value goes into a FReg.
(decl put_in_freg (Value) FReg)
(rule (put_in_freg val) (freg_new (put_in_reg val)))
(convert Value FReg put_in_freg)

;; Construct an `InstOutput` out of a single FReg register.
(decl output_freg (FReg) InstOutput)
(rule (output_freg x) (output_reg x))
(convert FReg InstOutput output_freg)

;; Convert a `WritableFReg` to an `FReg`.
(decl pure writable_freg_to_freg (WritableFReg) FReg)
(extern constructor writable_freg_to_freg writable_freg_to_freg)
(convert WritableFReg FReg writable_freg_to_freg)

;; Convert a `WritableFReg` to an `WritableReg`.
(decl pure writable_freg_to_writable_reg (WritableFReg) WritableReg)
(extern constructor writable_freg_to_writable_reg writable_freg_to_writable_reg)
(convert WritableFReg WritableReg writable_freg_to_writable_reg)

;; Convert a `WritableFReg` to an `Reg`.
(decl pure writable_freg_to_reg (WritableFReg) Reg)
(rule (writable_freg_to_reg x) (writable_freg_to_writable_reg x))
(convert WritableFReg Reg writable_freg_to_reg)

;; Convert an `FReg` to a `Reg`.
(decl pure freg_to_reg (FReg) Reg)
(extern constructor freg_to_reg freg_to_reg)
(convert FReg Reg freg_to_reg)

;; Convert a `FReg` to a `ValueRegs`.
(decl freg_to_value_regs (FReg) ValueRegs)
(rule (freg_to_value_regs x) (value_reg x))
(convert FReg ValueRegs xreg_to_reg)

;; Convert a `WritableFReg` to a `ValueRegs`.
(decl writable_freg_to_value_regs (WritableFReg) ValueRegs)
(rule (writable_freg_to_value_regs x) (value_reg x))
(convert WritableFReg ValueRegs writable_freg_to_value_regs)

;; Allocates a new `WritableFReg`.
(decl temp_writable_freg () WritableFReg)
(rule (temp_writable_freg) (temp_writable_reg $F64))

;; Construct a new `VReg` from a `Reg`.
;;
;; Asserts that the register has a Vector RegClass.
(decl vreg_new (Reg) VReg)
(extern constructor vreg_new vreg_new)
(convert Reg VReg vreg_new)

;; Construct a new `WritableVReg` from a `WritableReg`.
;;
;; Asserts that the register has a Vector RegClass.
(decl writable_vreg_new (WritableReg) WritableVReg)
(extern constructor writable_vreg_new writable_vreg_new)
(convert WritableReg WritableVReg writable_vreg_new)

;; Put a value into a VReg.
;;
;; Asserts that the value goes into a VReg.
(decl put_in_vreg (Value) VReg)
(rule (put_in_vreg val) (vreg_new (put_in_reg val)))
(convert Value VReg put_in_vreg)

;; Construct an `InstOutput` out of a single VReg register.
(decl output_vreg (VReg) InstOutput)
(rule (output_vreg x) (output_reg x))
(convert VReg InstOutput output_vreg)

;; Convert a `WritableVReg` to an `VReg`.
(decl pure writable_vreg_to_vreg (WritableVReg) VReg)
(extern constructor writable_vreg_to_vreg writable_vreg_to_vreg)
(convert WritableVReg VReg writable_vreg_to_vreg)

;; Convert a `WritableVReg` to an `WritableReg`.
(decl pure writable_vreg_to_writable_reg (WritableVReg) WritableReg)
(extern constructor writable_vreg_to_writable_reg writable_vreg_to_writable_reg)
(convert WritableVReg WritableReg writable_vreg_to_writable_reg)

;; Convert a `WritableVReg` to an `Reg`.
(decl pure writable_vreg_to_reg (WritableVReg) Reg)
(rule (writable_vreg_to_reg x) (writable_vreg_to_writable_reg x))
(convert WritableVReg Reg writable_vreg_to_reg)

;; Convert an `VReg` to a `Reg`.
(decl pure vreg_to_reg (VReg) Reg)
(extern constructor vreg_to_reg vreg_to_reg)
(convert VReg Reg vreg_to_reg)

;; Convert a `VReg` to a `ValueRegs`.
(decl vreg_to_value_regs (VReg) ValueRegs)
(rule (vreg_to_value_regs x) (value_reg x))
(convert VReg ValueRegs xreg_to_reg)

;; Convert a `WritableVReg` to a `ValueRegs`.
(decl writable_vreg_to_value_regs (WritableVReg) ValueRegs)
(rule (writable_vreg_to_value_regs x) (value_reg x))
(convert WritableVReg ValueRegs writable_vreg_to_value_regs)

;; Allocates a new `WritableVReg`.
(decl temp_writable_vreg () WritableVReg)
(rule (temp_writable_vreg) (temp_writable_reg $I8X16))

;;;; Materializing Constants ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

;; Lower a constant into a register.
(decl imm (Type u64) Reg)

;; If a value can fit into 8 bits, then prioritize that.
(rule 3 (imm (ty_int _) x)
      (if-let y (i8_try_from_u64 x))
      (pulley_xconst8 y))

;; If a value can fit into 16 bits, then prioritize that.
(rule 2 (imm (ty_int _) x)
      (if-let y (i16_try_from_u64 x))
      (pulley_xconst16 y))

;; If a value can fit into 32 bits, then prioritize that.
(rule 1 (imm (ty_int _) x)
      (if-let y (i32_try_from_u64 x))
      (pulley_xconst32 y))

;; Base cases for integers.
(rule 0 (imm $I8 x) (pulley_xconst8 (u8_as_i8 (u64_as_u8 x))))
(rule 0 (imm $I16 x) (pulley_xconst16 (u16_as_i16 (u64_as_u16 x))))
(rule 0 (imm $I32 x) (pulley_xconst32 (u64_as_i32 x)))
(rule 0 (imm $I64 x) (pulley_xconst64 (u64_as_i64 x)))

;; Base cases for floats.
(rule 0 (imm $F32 (u64_as_u32 c)) (pulley_fconst32 c))
(rule 0 (imm $F64 c) (pulley_fconst64 c))

;;;; Instruction Constructors ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

(decl pulley_trap_if (Cond TrapCode) SideEffectNoResult)
(rule (pulley_trap_if cond code)
      (SideEffectNoResult.Inst (MInst.TrapIf cond code)))

(decl sp_reg () XReg)
(extern constructor sp_reg sp_reg)

(decl pulley_get_special (XReg) XReg)
(rule (pulley_get_special reg)
      (let ((dst WritableXReg (temp_writable_xreg))
            (_ Unit (emit (MInst.GetSpecial dst reg))))
        dst))

(decl pulley_jump (MachLabel) SideEffectNoResult)
(rule (pulley_jump label)
      (SideEffectNoResult.Inst (MInst.Jump label)))

(decl pulley_br_if (Cond MachLabel MachLabel) SideEffectNoResult)
(rule (pulley_br_if cond taken not_taken)
      (SideEffectNoResult.Inst (MInst.BrIf cond taken not_taken)))

(decl pulley_xload (Amode Type MemFlags ExtKind) XReg)
(rule (pulley_xload amode ty flags ext)
      (let ((dst WritableXReg (temp_writable_xreg))
            (_ Unit (emit (MInst.XLoad dst amode ty flags ext))))
        dst))

(decl pulley_xstore (Amode XReg Type MemFlags) SideEffectNoResult)
(rule (pulley_xstore amode src ty flags)
      (SideEffectNoResult.Inst (MInst.XStore amode src ty flags)))

(decl pulley_fload (Amode Type MemFlags) FReg)
(rule (pulley_fload amode ty flags)
      (let ((dst WritableFReg (temp_writable_freg))
            (_ Unit (emit (MInst.FLoad dst amode ty flags))))
        dst))

(decl pulley_fstore (Amode FReg Type MemFlags) SideEffectNoResult)
(rule (pulley_fstore amode src ty flags)
      (SideEffectNoResult.Inst (MInst.FStore amode src ty flags)))

(decl pulley_vload (Amode Type MemFlags VExtKind) VReg)
(rule (pulley_vload amode ty flags ext)
      (let ((dst WritableVReg (temp_writable_vreg))
            (_ Unit (emit (MInst.VLoad dst amode ty flags ext))))
        dst))

(decl pulley_vstore (Amode VReg Type MemFlags) SideEffectNoResult)
(rule (pulley_vstore amode src ty flags)
      (SideEffectNoResult.Inst (MInst.VStore amode src ty flags)))

(decl gen_br_table (XReg MachLabel BoxVecMachLabel) Unit)
(rule (gen_br_table idx default labels)
      (emit (MInst.BrTable idx default labels)))

;;;; Helpers for Emitting Calls ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

(decl gen_call (SigRef ExternalName RelocDistance ValueSlice) InstOutput)
(extern constructor gen_call gen_call)

(decl gen_call_indirect (SigRef Value ValueSlice) InstOutput)
(extern constructor gen_call_indirect gen_call_indirect)

;;;; Helpers for Sign extension ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

;; Sign extend a `Value` to at least 32-bit
(decl zext32 (Value) XReg)
(rule (zext32 val @ (value_type $I8)) (pulley_zext8 val))
(rule (zext32 val @ (value_type $I16)) (pulley_zext16 val))
(rule (zext32 val @ (value_type $I32)) val)
(rule (zext32 val @ (value_type $I64)) val)

;; Same as `zext32` but for sign-extension
(decl sext32 (Value) XReg)
(rule (sext32 val @ (value_type $I8)) (pulley_sext8 val))
(rule (sext32 val @ (value_type $I16)) (pulley_sext16 val))
(rule (sext32 val @ (value_type $I32)) val)
(rule (sext32 val @ (value_type $I64)) val)

;; Sign extend a `Value` to at least 64-bit
(decl zext64 (Value) XReg)
(rule (zext64 val @ (value_type $I8)) (pulley_zext8 val))
(rule (zext64 val @ (value_type $I16)) (pulley_zext16 val))
(rule (zext64 val @ (value_type $I32)) (pulley_zext32 val))
(rule (zext64 val @ (value_type $I64)) val)

;; Same as `zext64` but for sign-extension
(decl sext64 (Value) XReg)
(rule (sext64 val @ (value_type $I8)) (pulley_sext8 val))
(rule (sext64 val @ (value_type $I16)) (pulley_sext16 val))
(rule (sext64 val @ (value_type $I32)) (pulley_sext32 val))
(rule (sext64 val @ (value_type $I64)) val)