-- SPDX-FileCopyrightText: 2021 Oxhead Alpha
-- SPDX-License-Identifier: LicenseRef-MIT-OA

-- | Module, containing some boilerplate for support of
-- arithmetic operations in Michelson language.

module Morley.Michelson.Typed.Arith
  ( ArithOp (..)
  , UnaryArithOp (..)
  , ToIntArithOp (..)
  , ToBytesArithOp (..)
  , evalToNatOp
  , ArithError (..)
  , ShiftArithErrorType (..)
  , MutezArithErrorType (..)
  , Add
  , Sub
  , SubMutez
  , Mul
  , EDiv
  , Abs
  , Neg
  , Or
  , And
  , Xor
  , Not
  , Lsl
  , Lsr
  , Compare
  , Eq'
  , Neq
  , Lt
  , Gt
  , Le
  , Ge
  , compareOp

  -- * Misc
  , Bls12381MulBadOrder
  ) where

import Crypto.Number.Basic (numBits)
import Crypto.Number.Serialize (i2osp, os2ip)
import Data.Bits (complement, setBit, shift, zeroBits, (.&.), (.|.))
import Data.ByteString qualified as BS
import Data.Constraint (Bottom(..), Dict(..))
import Fmt (Buildable(build), (+|), (|+))
import Type.Errors (DelayError)
import Unsafe qualified (fromIntegral)

import Morley.Michelson.Typed.Polymorphic
import Morley.Michelson.Typed.Scope (Comparable)
import Morley.Michelson.Typed.T (T(..))
import Morley.Michelson.Typed.Value (Value'(..))
import Morley.Tezos.Core (addMutez, mulMutez, subMutez, timestampFromSeconds, timestampToSeconds)
import Morley.Tezos.Crypto.BLS12381 qualified as BLS
import Morley.Util.TypeLits

-- | Class for binary arithmetic operation.
--
-- Takes binary operation marker as @op@ parameter,
-- types of left operand @n@ and right operand @m@.
--
-- 'Typeable' constraints in superclass are necessary for error messages.
class (Typeable n, Typeable m) =>
      ArithOp aop (n :: T) (m :: T) where

  -- | Type family @ArithRes@ denotes the type resulting from
  -- computing operation @op@ from operands of types @n@ and @m@.
  --
  -- For instance, adding integer to natural produces integer,
  -- which is reflected in following instance of type family:
  -- @ArithRes Add CNat CInt = CInt@.
  type ArithRes aop n m :: T

  -- | Evaluate arithmetic operation on given operands.
  evalOp
    :: proxy aop
    -> Value' instr n
    -> Value' instr m
    -> Either (ArithError (Value' instr n) (Value' instr m)) (Value' instr (ArithRes aop n m))

  -- | An operation can marked as commutative, it does not affect its
  -- runtime behavior, but enables certain optimization in the optimizer.
  -- We conservatively consider operations non-commutative by default.
  --
  -- Note that there is one unusual case: @AND@ works with @int : nat@
  -- but not with @nat : int@. That's how it's specified in Michelson.
  commutativityProof :: Maybe $ Dict (ArithRes aop n m ~ ArithRes aop m n, ArithOp aop m n)
  commutativityProof = Nothing

-- | Denotes the error type occurred in the arithmetic shift operation.
data ShiftArithErrorType
  = LslOverflow
  | LsrUnderflow
  deriving stock (Show, Eq, Ord, Generic)

instance NFData ShiftArithErrorType

-- | Denotes the error type occurred in the arithmetic operation involving mutez.
data MutezArithErrorType
  = AddOverflow
  | MulOverflow
  deriving stock (Show, Eq, Ord, Generic)

instance NFData MutezArithErrorType

-- | Represents an arithmetic error of the operation.
data ArithError n m
  = MutezArithError MutezArithErrorType n m
  | ShiftArithError ShiftArithErrorType n m
  deriving stock (Show, Eq, Ord, Generic)

instance (NFData n, NFData m) => NFData (ArithError n m)

-- | Class for unary arithmetic operation.
class UnaryArithOp aop (n :: T) where
  type UnaryArithRes aop n :: T
  evalUnaryArithOp :: proxy aop -> Value' instr n -> Value' instr (UnaryArithRes aop n)

-- | Class for conversions to an integer value.
class ToIntArithOp (n :: T) where
  evalToIntOp :: Value' instr n -> Value' instr 'TInt

-- | Class for conversions to bytes.
class ToBytesArithOp (n :: T) where
  evalToBytesOp :: Value' instr n -> Value' instr 'TBytes

data Add
data Sub
data SubMutez
data Mul
data EDiv
data Abs
data Neg

data Or
data And
data Xor
data Not
data Lsl
data Lsr

data Compare
data Eq'
data Neq
data Lt
data Gt
data Le
data Ge

-- For implementation hints see the reference implementation:
-- (note that you may need to change the branch)
-- https://gitlab.com/metastatedev/tezos/blob/master/src/proto_alpha/lib_protocol/script_ir_translator.ml#L4601

instance ArithOp Add 'TNat 'TInt where
  type ArithRes Add 'TNat 'TInt = 'TInt
  evalOp _ (VNat i) (VInt j) = Right $ VInt (toInteger i + j)
  commutativityProof = Just Dict
instance ArithOp Add 'TInt 'TNat where
  type ArithRes Add 'TInt 'TNat = 'TInt
  evalOp _ (VInt i) (VNat j) = Right $ VInt (i + toInteger j)
  commutativityProof = Just Dict
instance ArithOp Add 'TNat 'TNat where
  type ArithRes Add 'TNat 'TNat = 'TNat
  evalOp _ (VNat i) (VNat j) = Right $ VNat (i + j)
  commutativityProof = Just Dict
instance ArithOp Add 'TInt 'TInt where
  type ArithRes Add 'TInt 'TInt = 'TInt
  evalOp _ (VInt i) (VInt j) = Right $ VInt (i + j)
  commutativityProof = Just Dict
instance ArithOp Add 'TTimestamp 'TInt where
  type ArithRes Add 'TTimestamp 'TInt = 'TTimestamp
  evalOp _ (VTimestamp i) (VInt j) =
    Right $ VTimestamp $ timestampFromSeconds $ timestampToSeconds i + j
  commutativityProof = Just Dict
instance ArithOp Add 'TInt 'TTimestamp where
  type ArithRes Add 'TInt 'TTimestamp = 'TTimestamp
  evalOp _ (VInt i) (VTimestamp j) =
    Right $ VTimestamp $ timestampFromSeconds $ timestampToSeconds j + i
  commutativityProof = Just Dict
instance ArithOp Add 'TMutez 'TMutez where
  type ArithRes Add 'TMutez 'TMutez = 'TMutez
  evalOp _ n@(VMutez i) m@(VMutez j) = res
    where
      res = maybe (Left $ MutezArithError AddOverflow n m) (Right . VMutez) $ i `addMutez` j
  commutativityProof = Just Dict
instance ArithOp Add 'TBls12381Fr 'TBls12381Fr where
  type ArithRes Add 'TBls12381Fr 'TBls12381Fr = 'TBls12381Fr
  evalOp _ (VBls12381Fr i) (VBls12381Fr j) =
    Right $ VBls12381Fr (BLS.add i j)
  commutativityProof = Just Dict
instance ArithOp Add 'TBls12381G1 'TBls12381G1 where
  type ArithRes Add 'TBls12381G1 'TBls12381G1 = 'TBls12381G1
  evalOp _ (VBls12381G1 i) (VBls12381G1 j) =
    Right $ VBls12381G1 (BLS.add i j)
  commutativityProof = Just Dict
instance ArithOp Add 'TBls12381G2 'TBls12381G2 where
  type ArithRes Add 'TBls12381G2 'TBls12381G2 = 'TBls12381G2
  evalOp _ (VBls12381G2 i) (VBls12381G2 j) =
    Right $ VBls12381G2 (BLS.add i j)
  commutativityProof = Just Dict

instance ArithOp Sub 'TNat 'TInt where
  type ArithRes Sub 'TNat 'TInt = 'TInt
  evalOp _ (VNat i) (VInt j) = Right $ VInt (toInteger i - j)
instance ArithOp Sub 'TInt 'TNat where
  type ArithRes Sub 'TInt 'TNat = 'TInt
  evalOp _ (VInt i) (VNat j) = Right $ VInt (i - toInteger j)
instance ArithOp Sub 'TNat 'TNat where
  type ArithRes Sub 'TNat 'TNat = 'TInt
  evalOp _ (VNat i) (VNat j) = Right $ VInt (toInteger i - toInteger j)
instance ArithOp Sub 'TInt 'TInt where
  type ArithRes Sub 'TInt 'TInt = 'TInt
  evalOp _ (VInt i) (VInt j) = Right $ VInt (i - j)
instance ArithOp Sub 'TTimestamp 'TInt where
  type ArithRes Sub 'TTimestamp 'TInt = 'TTimestamp
  evalOp _ (VTimestamp i) (VInt j) =
    Right $ VTimestamp $ timestampFromSeconds $ timestampToSeconds i - j
instance ArithOp Sub 'TTimestamp 'TTimestamp where
  type ArithRes Sub 'TTimestamp 'TTimestamp = 'TInt
  evalOp _ (VTimestamp i) (VTimestamp j) =
    Right $ VInt $ timestampToSeconds i - timestampToSeconds j
instance ArithOp SubMutez 'TMutez 'TMutez where
  type ArithRes SubMutez 'TMutez 'TMutez = 'TOption 'TMutez
  evalOp _ (VMutez i) (VMutez j) =
    Right $ VOption $ fmap VMutez $ i `subMutez` j

instance ArithOp Mul 'TNat 'TInt where
  type ArithRes Mul 'TNat 'TInt = 'TInt
  evalOp _ (VNat i) (VInt j) = Right $ VInt (toInteger i * j)
  commutativityProof = Just Dict
instance ArithOp Mul 'TInt 'TNat where
  type ArithRes Mul 'TInt 'TNat = 'TInt
  evalOp _ (VInt i) (VNat j) = Right $ VInt (i * toInteger j)
  commutativityProof = Just Dict
instance ArithOp Mul 'TNat 'TNat where
  type ArithRes Mul 'TNat 'TNat = 'TNat
  evalOp _ (VNat i) (VNat j) = Right $ VNat (i * j)
  commutativityProof = Just Dict
instance ArithOp Mul 'TInt 'TInt where
  type ArithRes Mul 'TInt 'TInt = 'TInt
  evalOp _ (VInt i) (VInt j) = Right $ VInt (i * j)
  commutativityProof = Just Dict
instance ArithOp Mul 'TNat 'TMutez where
  type ArithRes Mul 'TNat 'TMutez = 'TMutez
  evalOp _ n@(VNat i) m@(VMutez j) = res
    where
      res = maybe (Left $ MutezArithError MulOverflow n m) (Right . VMutez) $ j `mulMutez` i
  commutativityProof = Just Dict
instance ArithOp Mul 'TMutez 'TNat where
  type ArithRes Mul 'TMutez 'TNat = 'TMutez
  evalOp _ n@(VMutez i) m@(VNat j) = res
    where
      res = maybe (Left $ MutezArithError MulOverflow n m) (Right . VMutez) $ i `mulMutez` j
  commutativityProof = Just Dict
instance ArithOp Mul 'TInt 'TBls12381Fr where
  type ArithRes Mul 'TInt 'TBls12381Fr = 'TBls12381Fr
  evalOp _ (VInt i) (VBls12381Fr j) = Right $ VBls12381Fr (fromIntegralOverflowing @Integer @BLS.Bls12381Fr i * j)
  commutativityProof = Just Dict
instance ArithOp Mul 'TNat 'TBls12381Fr where
  type ArithRes Mul 'TNat 'TBls12381Fr = 'TBls12381Fr
  evalOp _ (VNat i) (VBls12381Fr j) = Right $ VBls12381Fr (fromIntegralOverflowing @Natural @BLS.Bls12381Fr i * j)
  commutativityProof = Just Dict
instance ArithOp Mul 'TBls12381Fr 'TInt where
  type ArithRes Mul 'TBls12381Fr 'TInt = 'TBls12381Fr
  evalOp _ (VBls12381Fr i) (VInt j) = Right $ VBls12381Fr (i * fromIntegralOverflowing @Integer @BLS.Bls12381Fr j)
  commutativityProof = Just Dict
instance ArithOp Mul 'TBls12381Fr 'TNat where
  type ArithRes Mul 'TBls12381Fr 'TNat = 'TBls12381Fr
  evalOp _ (VBls12381Fr i) (VNat j) = Right $ VBls12381Fr (i * fromIntegralOverflowing @Natural @BLS.Bls12381Fr j)
  commutativityProof = Just Dict
instance ArithOp Mul 'TBls12381Fr 'TBls12381Fr where
  type ArithRes Mul 'TBls12381Fr 'TBls12381Fr = 'TBls12381Fr
  evalOp _ (VBls12381Fr i) (VBls12381Fr j) = Right $ VBls12381Fr (i * j)
  commutativityProof = Just Dict
instance ArithOp Mul 'TBls12381G1 'TBls12381Fr where
  type ArithRes Mul 'TBls12381G1 'TBls12381Fr = 'TBls12381G1
  evalOp _ (VBls12381G1 i) (VBls12381Fr j) = Right $ VBls12381G1 (BLS.multiply j i)
  commutativityProof = Nothing
instance ArithOp Mul 'TBls12381G2 'TBls12381Fr where
  type ArithRes Mul 'TBls12381G2 'TBls12381Fr = 'TBls12381G2
  evalOp _ (VBls12381G2 i) (VBls12381Fr j) = Right $ VBls12381G2 (BLS.multiply j i)
  commutativityProof = Nothing
instance (Bottom, Bls12381MulBadOrder BLS.Bls12381Fr BLS.Bls12381G1) =>
         ArithOp Mul 'TBls12381Fr 'TBls12381G1 where
  type ArithRes Mul 'TBls12381Fr 'TBls12381G1 = 'TBls12381G1
  evalOp = no
  commutativityProof = no
instance (Bottom, Bls12381MulBadOrder BLS.Bls12381Fr BLS.Bls12381G2) =>
         ArithOp Mul 'TBls12381Fr 'TBls12381G2 where
  type ArithRes Mul 'TBls12381Fr 'TBls12381G2 = 'TBls12381G2
  evalOp = no
  commutativityProof = no

instance ArithOp EDiv 'TNat 'TInt where
  type ArithRes EDiv 'TNat 'TInt = 'TOption ('TPair (EDivOpRes 'TNat 'TInt) (EModOpRes 'TNat 'TInt))
  evalOp _ i j = Right $ evalEDivOp i j
instance ArithOp EDiv 'TInt 'TNat where
  type ArithRes EDiv 'TInt 'TNat = 'TOption ('TPair (EDivOpRes 'TInt 'TNat) (EModOpRes 'TInt 'TNat))
  evalOp _ i j = Right $ evalEDivOp i j
instance ArithOp EDiv 'TNat 'TNat where
  type ArithRes EDiv 'TNat 'TNat = 'TOption ('TPair (EDivOpRes 'TNat 'TNat) (EModOpRes 'TNat 'TNat))
  evalOp _ i j = Right $ evalEDivOp i j
instance ArithOp EDiv 'TInt 'TInt where
  type ArithRes EDiv 'TInt 'TInt = 'TOption ('TPair (EDivOpRes 'TInt 'TInt) (EModOpRes 'TInt 'TInt))
  evalOp _ i j = Right $ evalEDivOp i j
instance ArithOp EDiv 'TMutez 'TMutez where
  type ArithRes EDiv 'TMutez 'TMutez = 'TOption ('TPair (EDivOpRes 'TMutez 'TMutez) (EModOpRes 'TMutez 'TMutez))
  evalOp _ i j = Right $ evalEDivOp i j
instance ArithOp EDiv 'TMutez 'TNat where
  type ArithRes EDiv 'TMutez 'TNat = 'TOption ('TPair (EDivOpRes 'TMutez 'TNat) (EModOpRes 'TMutez 'TNat))
  evalOp _ i j = Right $ evalEDivOp i j

type Bls12381MulBadOrder a1 a2 = DelayError
  ('Text "Multiplication of "
      ':<>: 'ShowType a2 ':<>: 'Text " and "
      ':<>: 'ShowType a1 ':<>: 'Text " works only other way around"
  )

instance UnaryArithOp Abs 'TInt where
  type UnaryArithRes Abs 'TInt = 'TNat
  evalUnaryArithOp _ (VInt i) = VNat (fromInteger $ abs i)

instance UnaryArithOp Neg 'TInt where
  type UnaryArithRes Neg 'TInt = 'TInt
  evalUnaryArithOp _ (VInt i) = VInt (-i)
instance UnaryArithOp Neg 'TNat where
  type UnaryArithRes Neg 'TNat = 'TInt
  evalUnaryArithOp _ (VNat i) = VInt (- fromIntegral i)
instance UnaryArithOp Neg 'TBls12381Fr where
  type UnaryArithRes Neg 'TBls12381Fr = 'TBls12381Fr
  evalUnaryArithOp _ (VBls12381Fr i) = VBls12381Fr (BLS.negate i)
instance UnaryArithOp Neg 'TBls12381G1 where
  type UnaryArithRes Neg 'TBls12381G1 = 'TBls12381G1
  evalUnaryArithOp _ (VBls12381G1 i) = VBls12381G1 (BLS.negate i)
instance UnaryArithOp Neg 'TBls12381G2 where
  type UnaryArithRes Neg 'TBls12381G2 = 'TBls12381G2
  evalUnaryArithOp _ (VBls12381G2 i) = VBls12381G2 (BLS.negate i)

bitwiseWithByteStringsEnd
  :: (Int -> Int -> Int) -- ^ Length comparator
  -> (Word8 -> Word8 -> Word8) -- ^ Zipper
  -> ByteString -> ByteString -> ByteString
bitwiseWithByteStringsEnd comp op i j = BS.packZipWith op i' j'
  where
    resultLen = comp lenI lenJ
    lenI = length i
    lenJ = length j
    normalize len = -- see Note [resultLen and normalize]
      (BS.replicate (resultLen - len) 0 <>) . BS.takeEnd resultLen
    i' = normalize lenI i
    j' = normalize lenJ j

{-
Note [resultLen and normalize]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

The idea with 'normalize' is to take exactly 'resultLen' bytes and left-pad with
zero bytes to the same length, so that 'packZipWith' is aligned correctly.

Note that depending on whether 'comp' is 'min' or 'max', either 'replicate' or
'takeEnd' is a no-op, respectively.

'takeEnd' is O(1), and replicate is O(1) for arguments <= 0, hence we do both
instead of making two functions or checking some condition.

-- @lierdakil
-}

instance ArithOp Or 'TNat 'TNat where
  type ArithRes Or 'TNat 'TNat = 'TNat
  evalOp _ (VNat i) (VNat j) = Right $ VNat (i .|. j)
  commutativityProof = Just Dict
instance ArithOp Or 'TBool 'TBool where
  type ArithRes Or 'TBool 'TBool = 'TBool
  evalOp _ (VBool i) (VBool j) = Right $ VBool (i .|. j)
  commutativityProof = Just Dict
instance ArithOp Or 'TBytes 'TBytes where
  type ArithRes Or 'TBytes 'TBytes = 'TBytes
  evalOp _ (VBytes i) (VBytes j) = Right $ VBytes $ bitwiseWithByteStringsEnd max (.|.) i j
  commutativityProof = Just Dict

instance ArithOp And 'TInt 'TNat where
  type ArithRes And 'TInt 'TNat = 'TNat
  evalOp _ (VInt i) (VNat j) = Right $ VNat (fromInteger (i .&. toInteger j))
instance ArithOp And 'TNat 'TNat where
  type ArithRes And 'TNat 'TNat = 'TNat
  evalOp _ (VNat i) (VNat j) = Right $ VNat (i .&. j)
  commutativityProof = Just Dict
instance ArithOp And 'TBool 'TBool where
  type ArithRes And 'TBool 'TBool = 'TBool
  evalOp _ (VBool i) (VBool j) = Right $ VBool (i .&. j)
  commutativityProof = Just Dict
instance ArithOp And 'TBytes 'TBytes where
  type ArithRes And 'TBytes 'TBytes = 'TBytes
  evalOp _ (VBytes i) (VBytes j) = Right $ VBytes $ bitwiseWithByteStringsEnd min (.&.) i j
  commutativityProof = Just Dict

instance ArithOp Xor 'TNat 'TNat where
  type ArithRes Xor 'TNat 'TNat = 'TNat
  evalOp _ (VNat i) (VNat j) = Right $ VNat (i `xor` j)
  commutativityProof = Just Dict
instance ArithOp Xor 'TBool 'TBool where
  type ArithRes Xor 'TBool 'TBool = 'TBool
  evalOp _ (VBool i) (VBool j) = Right $ VBool (i `xor` j)
  commutativityProof = Just Dict
instance ArithOp Xor 'TBytes 'TBytes where
  type ArithRes Xor 'TBytes 'TBytes = 'TBytes
  evalOp _ (VBytes i) (VBytes j) = Right $ VBytes $ bitwiseWithByteStringsEnd max xor i j
  commutativityProof = Just Dict

instance ArithOp Lsl 'TNat 'TNat where
  type ArithRes Lsl 'TNat 'TNat = 'TNat
  evalOp _ n@(VNat i) m@(VNat j) =
    if j > 256
    then Left $ ShiftArithError LslOverflow n m
    else Right $ VNat (fromInteger $ shift (toInteger i) (Unsafe.fromIntegral @Natural @Int j))
instance ArithOp Lsl 'TBytes 'TNat where
  type ArithRes Lsl 'TBytes 'TNat = 'TBytes
  evalOp _ n@(VBytes i) m@(VNat j) =
    if j > 64000
    then Left $ ShiftArithError LslOverflow n m
    else Right $ VBytes $
      let -- according to the docs, the result is left-padded to this number of bytes
          len = length i + Unsafe.fromIntegral @Natural @Int ((j + 7) `div` 8)
          bs = i2osp $ shift (os2ip i) (Unsafe.fromIntegral @Natural @Int j)
      in BS.replicate (len - length bs) 0 <> bs

instance ArithOp Lsr 'TNat 'TNat where
  type ArithRes Lsr 'TNat 'TNat = 'TNat
  evalOp _ n@(VNat i) m@(VNat j) =
    if j > 256
    then Left $ ShiftArithError LsrUnderflow n m
    else Right $ VNat (fromInteger $ shift (toInteger i) (-(Unsafe.fromIntegral @Natural @Int j)))
instance ArithOp Lsr 'TBytes 'TNat where
  type ArithRes Lsr 'TBytes 'TNat = 'TBytes
  evalOp _ (VBytes i) (VNat j)
    -- Michelson docs claim there's a limit of 256, but actually there isn't.
    -- https://gitlab.com/tezos/tezos/-/issues/5018
    -- Very large shifts just return empty byte string
    -- https://gitlab.com/tezos/tezos/-/blob/b2a49caf2e6d4b34c4fd226996d3637e19397401/src/proto_alpha/lib_protocol/script_bytes.ml#L48
    | j > fromIntegralOverflowing @Int @Natural maxBound = Right $ VBytes mempty
    | otherwise = Right $ VBytes $
        let k = Unsafe.fromIntegral @Natural @Int j -- safe, we just checked bounds above
            -- empirical number of bytes to match octez impl
            -- this is inverse to LSL in some sense
            len = length i - (k `div` 8)
            bs = i2osp $ shift (os2ip i) (-k)
            finalLenDiff = len - length bs
        in BS.replicate finalLenDiff 0 <> BS.drop (-finalLenDiff) bs

instance UnaryArithOp Not 'TInt where
  type UnaryArithRes Not 'TInt = 'TInt
  evalUnaryArithOp _ (VInt i) = VInt (complement i)
instance UnaryArithOp Not 'TNat where
  type UnaryArithRes Not 'TNat = 'TInt
  evalUnaryArithOp _ (VNat i) = VInt (complement $ toInteger i)
instance UnaryArithOp Not 'TBool where
  type UnaryArithRes Not 'TBool = 'TBool
  evalUnaryArithOp _ (VBool i) = VBool (not i)
instance UnaryArithOp Not 'TBytes where
  type UnaryArithRes Not 'TBytes = 'TBytes
  evalUnaryArithOp _ (VBytes i) = VBytes (BS.map complement i)

-- | Implementation for @COMPARE@ instruction.
compareOp :: Comparable t => Value' i t -> Value' i t -> Integer
compareOp a b =
  -- If at some point we need to return a number outside of [-1; 1] range,
  -- let's extend 'tcompare' respectively and use it here.
  -- Such situation seems unlikely to happen though, our previous communication
  -- with the Tezos developers shows that even if in Tezos 'COMPARE' returns
  -- something unusual, that is probably a bug.
  toInteger $ fromEnum (compare a b) - 1

instance UnaryArithOp Eq' 'TInt where
  type UnaryArithRes Eq' 'TInt = 'TBool
  evalUnaryArithOp _ (VInt i) = VBool (i == 0)

instance UnaryArithOp Neq 'TInt where
  type UnaryArithRes Neq 'TInt = 'TBool
  evalUnaryArithOp _ (VInt i) = VBool (i /= 0)


instance UnaryArithOp Lt 'TInt where
  type UnaryArithRes Lt 'TInt = 'TBool
  evalUnaryArithOp _ (VInt i) = VBool (i < 0)

instance UnaryArithOp Gt 'TInt where
  type UnaryArithRes Gt 'TInt = 'TBool
  evalUnaryArithOp _ (VInt i) = VBool (i > 0)

instance UnaryArithOp Le 'TInt where
  type UnaryArithRes Le 'TInt = 'TBool
  evalUnaryArithOp _ (VInt i) = VBool (i <= 0)

instance UnaryArithOp Ge 'TInt where
  type UnaryArithRes Ge 'TInt = 'TBool
  evalUnaryArithOp _ (VInt i) = VBool (i >= 0)

instance ToIntArithOp 'TNat where
  evalToIntOp (VNat i) = VInt (toInteger i)

instance ToIntArithOp 'TBls12381Fr where
  evalToIntOp (VBls12381Fr i) = VInt (toInteger i)

instance ToIntArithOp 'TBytes where
  evalToIntOp (VBytes i) = VInt $ case BS.uncons i of
    Nothing -> 0
    Just (msb, _)
      -- positive, msbit not set
      | msb .&. 0x80 == 0 -> x
      -- negative
      | otherwise -> x - comp
      where
        x = os2ip i
        nbytes = length i
        comp = setBit zeroBits (nbytes * 8)

instance ToBytesArithOp 'TNat where
  evalToBytesOp (VNat i)
    | 0 <- i = VBytes mempty
    | otherwise = VBytes $ i2osp $ toInteger i

instance ToBytesArithOp 'TInt where
  evalToBytesOp (VInt i)
    | 0 <- i = VBytes mempty
    | otherwise = VBytes $ BS.replicate (nbytes - length bs) 0 <> bs
    where
      x = abs i
      nbits = case numBits x of
        n | i < 0, setBit zeroBits (n - 1) == x
          -- exactly 0b10…0, by convention it's a negative, so doesn't need an extra bit
          -> n
          | otherwise -> n + 1 -- extra bit for sign.
      nbytes = (nbits + 7) `div` 8
      bs
        | i >= 0 = i2osp i
        | otherwise = i2osp $ setBit zeroBits (nbytes * 8) - x -- two's complement

evalToNatOp :: Value' instr 'TBytes -> Value' instr 'TNat
evalToNatOp (VBytes a)
  | null a = VNat 0
  | otherwise = VNat . Unsafe.fromIntegral @Integer @Natural $ os2ip a

instance Buildable ShiftArithErrorType where
  build = \case
    LslOverflow -> "lsl overflow"
    LsrUnderflow -> "lsr underflow"

instance Buildable MutezArithErrorType where
  build = \case
    AddOverflow -> "add overflow"
    MulOverflow -> "mul overflow"

instance (Buildable n, Buildable m) => Buildable (ArithError n m) where
  build (MutezArithError errType n m) = "Mutez "
    +| errType |+ " with " +| n |+ ", " +| m |+ ""
  build (ShiftArithError errType n m) = ""
    +| errType |+ " with " +| n |+ ", " +| m |+ ""