;; Algebraic optimizations.

;; Rules here are allowed to rewrite pure expressions arbitrarily,
;; using the same inputs as the original, or fewer. In other words, we
;; cannot pull a new eclass id out of thin air and refer to it, other
;; than a piece of the input or a new node that we construct; but we
;; can freely rewrite e.g. `x+y-y` to `x`.

;; Chained `uextend` and `sextend`.
(rule (simplify (uextend ty (uextend _intermediate_ty x)))
      (uextend ty x))
(rule (simplify (sextend ty (sextend _intermediate_ty x)))
      (sextend ty x))

;; x+0 == 0+x == x.
(rule (simplify (iadd ty
                      x
                      (iconst ty (u64_from_imm64 0))))
      (subsume x))
(rule (simplify (iadd ty
                      (iconst ty (u64_from_imm64 0))
                      x))
      (subsume x))
;; x-0 == x.
(rule (simplify (isub ty
                      x
                      (iconst ty (u64_from_imm64 0))))
      (subsume x))
;; 0-x == (ineg x).
(rule (simplify (isub ty
                      (iconst ty (u64_from_imm64 0))
                      x))
      (ineg ty x))

;; x*1 == 1*x == x.
(rule (simplify (imul ty
                      x
                      (iconst ty (u64_from_imm64 1))))
      (subsume x))
(rule (simplify (imul ty
                      (iconst ty (u64_from_imm64 1))
                      x))
      (subsume x))

;; x*0 == 0*x == 0.
(rule (simplify (imul ty
                      _
                      zero @ (iconst ty (u64_from_imm64 0))))
      (subsume zero))
(rule (simplify (imul ty
                      zero @ (iconst ty (u64_from_imm64 0))
                      _))
      (subsume zero))

;; x/1 == x.
(rule (simplify (sdiv ty
                      x
                      (iconst ty (u64_from_imm64 1))))
      (subsume x))
(rule (simplify (udiv ty
                      x
                      (iconst ty (u64_from_imm64 1))))
      (subsume x))

;; x>>0 == x<<0 == x rotr 0 == x rotl 0 == x.
(rule (simplify (ishl ty
                      x
                      (iconst ty (u64_from_imm64 0))))
      (subsume x))
(rule (simplify (ushr ty
                      x
                      (iconst ty (u64_from_imm64 0))))
      (subsume x))
(rule (simplify (sshr ty
                      x
                      (iconst ty (u64_from_imm64 0))))
      (subsume x))
(rule (simplify (rotr ty
                      x
                      (iconst ty (u64_from_imm64 0))))
      (subsume x))
(rule (simplify (rotl ty
                      x
                      (iconst ty (u64_from_imm64 0))))
      (subsume x))

;; x | 0 == 0 | x == x | x == x.
(rule (simplify (bor ty
                     x
                     (iconst ty (u64_from_imm64 0))))
      (subsume x))
(rule (simplify (bor ty
                     (iconst ty (u64_from_imm64 0))
                     x))
      (subsume x))
(rule (simplify (bor ty x x))
      (subsume x))

;; x ^ 0 == 0 ^ x == x.
(rule (simplify (bxor ty
                     x
                     (iconst ty (u64_from_imm64 0))))
      (subsume x))
(rule (simplify (bxor ty
                     (iconst ty (u64_from_imm64 0))
                     x))
      (subsume x))

;; x ^ x == 0.
(rule (simplify (bxor (fits_in_64 (ty_int ty)) x x))
      (subsume (iconst ty (imm64 0))))

;; x ^ not(x) == not(x) ^ x == x | not(x) == not(x) | x == -1.
;; This identity also holds for non-integer types, vectors, and wider types.
;; But `iconst` is only valid for integers up to 64 bits wide.
(rule (simplify (bxor (fits_in_64 (ty_int ty)) x (bnot ty x))) (subsume (iconst ty (imm64 (ty_mask ty)))))
(rule (simplify (bxor (fits_in_64 (ty_int ty)) (bnot ty x) x)) (subsume (iconst ty (imm64 (ty_mask ty)))))
(rule (simplify (bor (fits_in_64 (ty_int ty)) x (bnot ty x))) (subsume (iconst ty (imm64 (ty_mask ty)))))
(rule (simplify (bor (fits_in_64 (ty_int ty)) (bnot ty x) x)) (subsume (iconst ty (imm64 (ty_mask ty)))))

;; x & -1 == -1 & x == x & x == x.
(rule (simplify (band ty x x)) (subsume x))
(rule (simplify (band ty x (iconst ty k)))
      (if-let -1 (i64_sextend_imm64 ty k))
      (subsume x))
(rule (simplify (band ty (iconst ty k) x))
      (if-let -1 (i64_sextend_imm64 ty k))
      (subsume x))

;; x & 0 == 0 & x == x & not(x) == not(x) & x == 0.
(rule (simplify (band ty _ zero @ (iconst ty (u64_from_imm64 0)))) (subsume zero))
(rule (simplify (band ty zero @ (iconst ty (u64_from_imm64 0)) _)) (subsume zero))
(rule (simplify (band (fits_in_64 (ty_int ty)) x (bnot ty x))) (subsume (iconst ty (imm64 0))))
(rule (simplify (band (fits_in_64 (ty_int ty)) (bnot ty x) x)) (subsume (iconst ty (imm64 0))))

;; not(not(x)) == x.
(rule (simplify (bnot ty (bnot ty x))) (subsume x))

;; DeMorgan's rule (two versions):
;; bnot(bor(x, y)) == band(bnot(x), bnot(y))
(rule (simplify (bnot ty (bor ty x y)))
      (band ty (bnot ty x) (bnot ty y)))
;; bnot(band(x, y)) == bor(bnot(x), bnot(y))
(rule (simplify (bnot ty (band t x y)))
      (bor ty (bnot ty x) (bnot ty y)))

;; `or(and(x, y), not(y)) == or(x, not(y))`
(rule (simplify (bor ty
                     (band ty x y)
                     z @ (bnot ty y)))
      (bor ty x z))
;; Duplicate the rule but swap the `bor` operands because `bor` is
;; commutative. We could, of course, add a `simplify` rule to do the commutative
;; swap for all `bor`s but this will bloat the e-graph with many e-nodes. It is
;; cheaper to have additional rules, rather than additional e-nodes, because we
;; amortize their cost via ISLE's smart codegen.
(rule (simplify (bor ty
                     z @ (bnot ty y)
                     (band ty x y)))
      (bor ty x z))

;; `or(and(x, y), not(y)) == or(x, not(y))` specialized for constants, since
;; otherwise we may not know that `z == not(y)` since we don't generally expand
;; constants in the e-graph.
;;
;; (No need to duplicate for commutative `bor` for this constant version because
;; we move constants to the right.)
(rule (simplify (bor ty
                     (band ty x (iconst ty (u64_from_imm64 y)))
                     z @ (iconst ty (u64_from_imm64 zk))))
      (if-let $true (u64_eq (u64_and (ty_mask ty) zk)
                            (u64_and (ty_mask ty) (u64_not y))))
      (bor ty x z))

;; x*2 == 2*x == x+x.
(rule (simplify (imul ty x (iconst _ (simm32 2))))
      (iadd ty x x))
(rule (simplify (imul ty (iconst _ (simm32 2)) x))
      (iadd ty x x))

;; x*c == x<<log2(c) when c is a power of two.
;; Note that the type of `iconst` must be the same as the type of `imul`,
;; so these rules can only fire in situations where it's safe to construct an
;; `iconst` of that type.
(rule (simplify (imul ty x (iconst _ (imm64_power_of_two c))))
      (ishl ty x (iconst ty (imm64 c))))
(rule (simplify (imul ty (iconst _ (imm64_power_of_two c)) x))
      (ishl ty x (iconst ty (imm64 c))))

;; TODO: strength reduction: div to shifts
;; TODO: div/rem by constants -> magic multiplications


;; `(x >> k) << k` is the same as masking off the bottom `k` bits (regardless if
;; this is a signed or unsigned shift right).
(rule (simplify (ishl (fits_in_64 ty)
                      (ushr ty x (iconst _ k))
                      (iconst _ k)))
      (let ((mask Imm64 (imm64_shl ty (imm64 0xFFFF_FFFF_FFFF_FFFF) k)))
        (band ty x (iconst ty mask))))
(rule (simplify (ishl (fits_in_64 ty)
                      (sshr ty x (iconst _ k))
                      (iconst _ k)))
      (let ((mask Imm64 (imm64_shl ty (imm64 0xFFFF_FFFF_FFFF_FFFF) k)))
        (band ty x (iconst ty mask))))


;; For unsigned shifts, `(x << k) >> k` is the same as masking out the top
;; `k` bits. A similar rule is valid for vectors but this `iconst` mask only
;; works for scalar integers.
(rule (simplify (ushr (fits_in_64 (ty_int ty))
                      (ishl ty x (iconst _ k))
                      (iconst _ k)))
      (band ty x (iconst ty (imm64_ushr ty (imm64 (ty_mask ty)) k))))


;; For signed shifts, `(x << k) >> k` does sign-extension from `n` bits to
;; `n+k` bits. In the special case where `x` is the result of either `sextend`
;; or `uextend` from `n` bits to `n+k` bits, we can implement this using
;; `sextend`.
(rule (simplify (sshr wide
                 (ishl wide
                  (uextend wide x @ (value_type narrow))
                  (iconst _ shift))
                 (iconst _ shift)))
      (if-let (u64_from_imm64 shift_u64) shift)
      (if-let $true (u64_eq shift_u64 (u64_sub (ty_bits_u64 wide) (ty_bits_u64 narrow))))
      (sextend wide x))

;; If `k` is smaller than the difference in bit widths of the two types, then
;; the intermediate sign bit comes from the extend op, so the final result is
;; the same as the original extend op.
(rule (simplify (sshr wide
                 (ishl wide
                  x @ (uextend wide (value_type narrow))
                  (iconst _ shift))
                 (iconst _ shift)))
      (if-let (u64_from_imm64 shift_u64) shift)
      (if-let $true (u64_lt shift_u64 (u64_sub (ty_bits_u64 wide) (ty_bits_u64 narrow))))
      x)

;; If the original extend op was `sextend`, then both of the above cases say
;; the result should also be `sextend`.
(rule (simplify (sshr wide
                 (ishl wide
                  x @ (sextend wide (value_type narrow))
                  (iconst _ shift))
                 (iconst _ shift)))
      (if-let (u64_from_imm64 shift_u64) shift)
      (if-let $true (u64_le shift_u64 (u64_sub (ty_bits_u64 wide) (ty_bits_u64 narrow))))
      x)


;; Masking out any of the top bits of the result of `uextend` is a no-op. (This
;; is like a cheap version of known-bits analysis.)
(rule (simplify (band wide x @ (uextend _ (value_type narrow)) (iconst _ (u64_from_imm64 mask))))
      ; Check that `narrow_mask` has a subset of the bits that `mask` does.
      (if-let $true (let ((narrow_mask u64 (ty_mask narrow))) (u64_eq narrow_mask (u64_and mask narrow_mask))))
      x)

;; Masking out the sign-extended bits of an `sextend` turns it into a `uextend`.
(rule (simplify (band wide (sextend _ x @ (value_type narrow)) (iconst _ (u64_from_imm64 mask))))
      (if-let $true (u64_eq mask (ty_mask narrow)))
      (uextend wide x))


;; Rematerialize ALU-op-with-imm and iconsts in each block where they're
;; used. This is neutral (add-with-imm) or positive (iconst) for
;; register pressure, and these ops are very cheap.
(rule (simplify x @ (iadd _ (iconst _ _) _))
      (remat x))
(rule (simplify x @ (iadd _ _ (iconst _ _)))
      (remat x))
(rule (simplify x @ (isub _ (iconst _ _) _))
      (remat x))
(rule (simplify x @ (isub _ _ (iconst _ _)))
      (remat x))
(rule (simplify x @ (band _ (iconst _ _) _))
      (remat x))
(rule (simplify x @ (band _ _ (iconst _ _)))
      (remat x))
(rule (simplify x @ (bor _ (iconst _ _) _))
      (remat x))
(rule (simplify x @ (bor _ _ (iconst _ _)))
      (remat x))
(rule (simplify x @ (bxor _ (iconst _ _) _))
      (remat x))
(rule (simplify x @ (bxor _ _ (iconst _ _)))
      (remat x))
(rule (simplify x @ (bnot _ _))
      (remat x))
(rule (simplify x @ (iconst _ _))
      (remat x))
(rule (simplify x @ (f32const _ _))
      (remat x))
(rule (simplify x @ (f64const _ _))
      (remat x))

;; Optimize icmp-of-icmp.
(rule (simplify (icmp ty
                      (IntCC.NotEqual)
                      (uextend _ inner @ (icmp ty _ _ _))
                      (iconst _ (u64_from_imm64 0))))
      (subsume inner))

(rule (simplify (icmp ty
                      (IntCC.Equal)
                      (uextend _ (icmp ty cc x y))
                      (iconst _ (u64_from_imm64 0))))
      (subsume (icmp ty (intcc_inverse cc) x y)))

;; Optimize select-of-uextend-of-icmp to select-of-icmp, because
;; select can take an I8 condition too.
(rule (simplify
       (select ty (uextend _ c @ (icmp _ _ _ _)) x y))
      (select ty c x y))
(rule (simplify
       (select ty (uextend _ c @ (icmp _ _ _ _)) x y))
      (select ty c x y))

;; `x == x` is always true for integers; `x != x` is false. Strict
;; inequalities are false, and loose inequalities are true.
(rule (simplify
       (icmp (fits_in_64 (ty_int ty)) (IntCC.Equal) x x))
      (iconst ty (imm64 1)))
(rule (simplify
       (icmp (fits_in_64 (ty_int ty)) (IntCC.NotEqual) x x))
      (iconst ty (imm64 0)))
(rule (simplify
       (icmp (fits_in_64 (ty_int ty)) (IntCC.UnsignedGreaterThan) x x))
      (iconst ty (imm64 0)))
(rule (simplify
       (icmp (fits_in_64 (ty_int ty)) (IntCC.UnsignedGreaterThanOrEqual) x x))
      (iconst ty (imm64 1)))
(rule (simplify
       (icmp (fits_in_64 (ty_int ty)) (IntCC.SignedGreaterThan) x x))
      (iconst ty (imm64 0)))
(rule (simplify
       (icmp (fits_in_64 (ty_int ty)) (IntCC.SignedGreaterThanOrEqual) x x))
      (iconst ty (imm64 1)))
(rule (simplify
       (icmp (fits_in_64 (ty_int ty)) (IntCC.UnsignedLessThan) x x))
      (iconst ty (imm64 0)))
(rule (simplify
       (icmp (fits_in_64 (ty_int ty)) (IntCC.UnsignedLessThanOrEqual) x x))
      (iconst ty (imm64 1)))
(rule (simplify
       (icmp (fits_in_64 (ty_int ty)) (IntCC.SignedLessThan) x x))
      (iconst ty (imm64 0)))
(rule (simplify
       (icmp (fits_in_64 (ty_int ty)) (IntCC.SignedLessThanOrEqual) x x))
      (iconst ty (imm64 1)))

;; (x ^ -1) can be replaced with the `bnot` instruction
(rule (simplify (bxor ty x (iconst ty k)))
  (if-let -1 (i64_sextend_imm64 ty k))
  (bnot ty x))


;; Masking the result of a comparison with 1 always results in the comparison
;; itself. Note that comparisons in wasm may sometimes be hidden behind
;; extensions.
(rule (simplify
       (band (ty_int _)
             cmp @ (icmp _ _ _ _)
             (iconst _ (u64_from_imm64 1))))
      cmp)
(rule (simplify
       (band (ty_int _)
             extend @ (uextend _ (icmp _ _ _ _))
             (iconst _ (u64_from_imm64 1))))
      extend)

;; `x < 0` is always false for unsigned integers, and `x >= 0` is always true
;; for unsigned integers, along with their reversals.
(rule (simplify
       (icmp (fits_in_64 (ty_int ty))
             (IntCC.UnsignedLessThan)
             _
             (iconst _ (u64_from_imm64 0))))
      (iconst ty (imm64 0)))
(rule (simplify
       (icmp (fits_in_64 (ty_int ty))
             (IntCC.UnsignedGreaterThanOrEqual)
             _
             (iconst _ (u64_from_imm64 0))))
      (iconst ty (imm64 1)))

;; 32-bit integers zero-extended to 64-bit integers are never negative
(rule (simplify
       (icmp (ty_int ty)
             (IntCC.SignedLessThan)
             (uextend $I64 x @ (value_type $I32))
             (iconst _ (u64_from_imm64 0))))
      (iconst ty (imm64 0)))
(rule (simplify
       (icmp (ty_int ty)
             (IntCC.SignedGreaterThanOrEqual)
             (uextend $I64 x @ (value_type $I32))
             (iconst _ (u64_from_imm64 0))))
      (iconst ty (imm64 1)))


;; Transform select-of-icmp into {u,s}{min,max} instructions where possible.
(rule (simplify
       (select ty (icmp _ (IntCC.SignedGreaterThan) x y) x y))
      (smax ty x y))
(rule (simplify
       (select ty (icmp _ (IntCC.SignedGreaterThanOrEqual) x y) x y))
      (smax ty x y))
(rule (simplify
       (select ty (icmp _ (IntCC.UnsignedGreaterThan) x y) x y))
      (umax ty x y))
(rule (simplify
       (select ty (icmp _ (IntCC.UnsignedGreaterThanOrEqual) x y) x y))
      (umax ty x y))
(rule (simplify
       (select ty (icmp _ (IntCC.SignedLessThan) x y) x y))
      (smin ty x y))
(rule (simplify
       (select ty (icmp _ (IntCC.SignedLessThanOrEqual) x y) x y))
      (smin ty x y))
(rule (simplify
       (select ty (icmp _ (IntCC.UnsignedLessThan) x y) x y))
      (umin ty x y))
(rule (simplify
       (select ty (icmp _ (IntCC.UnsignedLessThanOrEqual) x y) x y))
      (umin ty x y))


;; These are the same rules as above, but when the operands for select are swapped
(rule (simplify
       (select ty (icmp _ (IntCC.SignedLessThan) x y) y x))
      (smax ty x y))
(rule (simplify
       (select ty (icmp _ (IntCC.SignedLessThanOrEqual) x y) y x))
      (smax ty x y))
(rule (simplify
       (select ty (icmp _ (IntCC.UnsignedLessThan) x y) y x))
      (umax ty x y))
(rule (simplify
       (select ty (icmp _ (IntCC.UnsignedLessThanOrEqual) x y) y x))
      (umax ty x y))
(rule (simplify
       (select ty (icmp _ (IntCC.SignedGreaterThan) x y) y x))
      (smin ty x y))
(rule (simplify
       (select ty (icmp _ (IntCC.SignedGreaterThanOrEqual) x y) y x))
      (smin ty x y))
(rule (simplify
       (select ty (icmp _ (IntCC.UnsignedGreaterThan) x y) y x))
      (umin ty x y))
(rule (simplify
       (select ty (icmp _ (IntCC.UnsignedGreaterThanOrEqual) x y) y x))
      (umin ty x y))

;; Transform vselect-of-icmp into {u,s}{min,max} instructions where possible.
(rule (simplify
       (vselect ty (icmp _ (IntCC.SignedGreaterThan) x y) x y))
      (smax ty x y))
(rule (simplify
       (vselect ty (icmp _ (IntCC.SignedGreaterThanOrEqual) x y) x y))
      (smax ty x y))
(rule (simplify
       (vselect ty (icmp _ (IntCC.UnsignedGreaterThan) x y) x y))
      (umax ty x y))
(rule (simplify
       (vselect ty (icmp _ (IntCC.UnsignedGreaterThanOrEqual) x y) x y))
      (umax ty x y))
(rule (simplify
       (vselect ty (icmp _ (IntCC.SignedLessThan) x y) x y))
      (smin ty x y))
(rule (simplify
       (vselect ty (icmp _ (IntCC.SignedLessThanOrEqual) x y) x y))
      (smin ty x y))
(rule (simplify
       (vselect ty (icmp _ (IntCC.UnsignedLessThan) x y) x y))
      (umin ty x y))
(rule (simplify
       (vselect ty (icmp _ (IntCC.UnsignedLessThanOrEqual) x y) x y))
      (umin ty x y))

;; These are the same rules as above, but when the operands for select are swapped
(rule (simplify
       (vselect ty (icmp _ (IntCC.SignedLessThan) x y) y x))
      (smax ty x y))
(rule (simplify
       (vselect ty (icmp _ (IntCC.SignedLessThanOrEqual) x y) y x))
      (smax ty x y))
(rule (simplify
       (vselect ty (icmp _ (IntCC.UnsignedLessThan) x y) y x))
      (umax ty x y))
(rule (simplify
       (vselect ty (icmp _ (IntCC.UnsignedLessThanOrEqual) x y) y x))
      (umax ty x y))
(rule (simplify
       (vselect ty (icmp _ (IntCC.SignedGreaterThan) x y) y x))
      (smin ty x y))
(rule (simplify
       (vselect ty (icmp _ (IntCC.SignedGreaterThanOrEqual) x y) y x))
      (smin ty x y))
(rule (simplify
       (vselect ty (icmp _ (IntCC.UnsignedGreaterThan) x y) y x))
      (umin ty x y))
(rule (simplify
       (vselect ty (icmp _ (IntCC.UnsignedGreaterThanOrEqual) x y) y x))
      (umin ty x y))

;; For floats convert fcmp lt into pseudo_min and gt into pseudo_max
;;
;; fmax_pseudo docs state:
;; The behaviour for this operations is defined as  fmax_pseudo(a, b) = (a < b) ? b : a, and the behaviour for zero
;; or NaN inputs follows from the behaviour of < with such inputs.
;;
;; That is exactly the operation that we match here!
(rule (simplify
       (select ty (fcmp _ (FloatCC.LessThan) x y) x y))
      (fmin_pseudo ty x y))
(rule (simplify
       (select ty (fcmp _ (FloatCC.GreaterThan) x y) x y))
      (fmax_pseudo ty x y))

;; Do the same for vectors
(rule (simplify
       (vselect ty (fcmp _ (FloatCC.LessThan) x y) x y))
      (fmin_pseudo ty x y))
(rule (simplify
       (vselect ty (fcmp _ (FloatCC.GreaterThan) x y) x y))
      (fmax_pseudo ty x y))

;; If both of the multiplied arguments to an `fma` are negated then remove
;; both of them since they cancel out.
(rule (simplify (fma ty (fneg ty x) (fneg ty y) z))
      (fma ty x y z))

;; If both of the multiplied arguments to an `fmul` are negated then remove
;; both of them since they cancel out.
(rule (simplify (fmul ty (fneg ty x) (fneg ty y)))
      (fmul ty x y))