-- SPDX-FileCopyrightText: 2020 Tocqueville Group
--
-- SPDX-License-Identifier: LicenseRef-MIT-TQ

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

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

import Data.Bits (complement, shift, (.&.), (.|.))
import Data.Constraint (Dict(..))
import Data.Singletons (Sing, SingI(..))
import Fmt (Buildable(build))

import Michelson.Typed.Annotation (AnnConvergeError, Notes(..), converge, convergeAnns, starNotes)
import Michelson.Typed.Sing (SingT(..))
import Michelson.Typed.T (T(..))
import Michelson.Typed.Value (Comparability(..), Comparable, Value'(..), checkComparability)
import Tezos.Core (addMutez, mulMutez, subMutez, timestampFromSeconds, timestampToSeconds)

-- | Class for binary arithmetic operation.
--
-- Takes binary operation marker as @op@ parameter,
-- types of left operand @n@ and right operand @m@.
class 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

  -- | Converge the notes of given operands.
  convergeArith
    :: proxy aop
    -> Notes n
    -> Notes m
    -> Either AnnConvergeError (Notes (ArithRes aop n m))

  -- | 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
  | SubUnderflow
  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)

-- | Marker data type for add operation.
class UnaryArithOp aop (n :: T) where
  type UnaryArithRes aop n :: T
  evalUnaryArithOp :: proxy aop -> Value' instr n -> Value' instr (UnaryArithRes aop n)

data Add
data Sub
data Mul
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

instance ArithOp Add 'TNat 'TInt where
  type ArithRes Add 'TNat 'TInt = 'TInt
  convergeArith _ _ n2 = Right n2
  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
  convergeArith _ n1 _ = Right n1
  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
  convergeArith _ n1 n2 = converge n1 n2
  evalOp _ (VNat i) (VNat j) = Right $ VNat (i + j)
  commutativityProof = Just Dict
instance ArithOp Add 'TInt 'TInt where
  type ArithRes Add 'TInt 'TInt = 'TInt
  convergeArith _ n1 n2 = converge n1 n2
  evalOp _ (VInt i) (VInt j) = Right $ VInt (i + j)
  commutativityProof = Just Dict
instance ArithOp Add 'TTimestamp 'TInt where
  type ArithRes Add 'TTimestamp 'TInt = 'TTimestamp
  convergeArith _ n1 _ = Right n1
  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
  convergeArith _ _ n2 = Right n2
  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
  convergeArith _ n1 n2 = converge n1 n2
  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 Sub 'TNat 'TInt where
  type ArithRes Sub 'TNat 'TInt = 'TInt
  convergeArith _ _ n2 = Right n2
  evalOp _ (VNat i) (VInt j) = Right $ VInt (toInteger i - j)
instance ArithOp Sub 'TInt 'TNat where
  type ArithRes Sub 'TInt 'TNat = 'TInt
  convergeArith _ n1 _ = Right n1
  evalOp _ (VInt i) (VNat j) = Right $ VInt (i - toInteger j)
instance ArithOp Sub 'TNat 'TNat where
  type ArithRes Sub 'TNat 'TNat = 'TInt
  -- | Subtraction between @Nat@ and @Nat@ does not retain annotation.
  convergeArith _ n1 n2 = (const starNotes) <$> converge n1 n2
  evalOp _ (VNat i) (VNat j) = Right $ VInt (toInteger i - toInteger j)
instance ArithOp Sub 'TInt 'TInt where
  type ArithRes Sub 'TInt 'TInt = 'TInt
  convergeArith _ n1 n2 = converge n1 n2
  evalOp _ (VInt i) (VInt j) = Right $ VInt (i - j)
instance ArithOp Sub 'TTimestamp 'TInt where
  type ArithRes Sub 'TTimestamp 'TInt = 'TTimestamp
  convergeArith _ n1 _ = Right n1
  evalOp _ (VTimestamp i) (VInt j) =
    Right $ VTimestamp $ timestampFromSeconds $ timestampToSeconds i - j
instance ArithOp Sub 'TTimestamp 'TTimestamp where
  type ArithRes Sub 'TTimestamp 'TTimestamp = 'TInt
  convergeArith _ (NTTimestamp a) (NTTimestamp b) = NTInt <$> (convergeAnns a b)
  evalOp _ (VTimestamp i) (VTimestamp j) =
    Right $ VInt $ timestampToSeconds i - timestampToSeconds j
instance ArithOp Sub 'TMutez 'TMutez where
  type ArithRes Sub 'TMutez 'TMutez = 'TMutez
  convergeArith _ n1 n2 = converge n1 n2
  evalOp _ n@(VMutez i) m@(VMutez j) = res
    where
      res = maybe (Left $ MutezArithError SubUnderflow n m) (Right . VMutez) $ i `subMutez` j

instance ArithOp Mul 'TNat 'TInt where
  type ArithRes Mul 'TNat 'TInt = 'TInt
  convergeArith _ _ n2 = Right n2
  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
  convergeArith _ n1 _ = Right n1
  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
  convergeArith _ n1 n2 = converge n1 n2
  evalOp _ (VNat i) (VNat j) = Right $ VNat (i * j)
  commutativityProof = Just Dict
instance ArithOp Mul 'TInt 'TInt where
  type ArithRes Mul 'TInt 'TInt = 'TInt
  convergeArith _ n1 n2 = converge n1 n2
  evalOp _ (VInt i) (VInt j) = Right $ VInt (i * j)
  commutativityProof = Just Dict
instance ArithOp Mul 'TNat 'TMutez where
  type ArithRes Mul 'TNat 'TMutez = 'TMutez
  convergeArith _ _ n2 = Right n2
  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
  convergeArith _ n1 _ = Right n1
  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 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 ArithOp Or 'TNat 'TNat where
  type ArithRes Or 'TNat 'TNat = 'TNat
  convergeArith _ n1 n2 = converge n1 n2
  evalOp _ (VNat i) (VNat j) = Right $ VNat (i .|. j)
  commutativityProof = Just Dict
instance ArithOp Or 'TBool 'TBool where
  type ArithRes Or 'TBool 'TBool = 'TBool
  convergeArith _ n1 n2 = converge n1 n2
  evalOp _ (VBool i) (VBool j) = Right $ VBool (i .|. j)
  commutativityProof = Just Dict

instance ArithOp And 'TInt 'TNat where
  type ArithRes And 'TInt 'TNat = 'TNat
  convergeArith _ _ n2 = Right n2
  evalOp _ (VInt i) (VNat j) = Right $ VNat (fromInteger (i .&. toInteger j))
instance ArithOp And 'TNat 'TNat where
  type ArithRes And 'TNat 'TNat = 'TNat
  convergeArith _ n1 n2 = converge n1 n2
  evalOp _ (VNat i) (VNat j) = Right $ VNat (i .&. j)
  commutativityProof = Just Dict
instance ArithOp And 'TBool 'TBool where
  type ArithRes And 'TBool 'TBool = 'TBool
  convergeArith _ n1 n2 = converge n1 n2
  evalOp _ (VBool i) (VBool j) = Right $ VBool (i .&. j)
  commutativityProof = Just Dict

instance ArithOp Xor 'TNat 'TNat where
  type ArithRes Xor 'TNat 'TNat = 'TNat
  convergeArith _ n1 n2 = converge n1 n2
  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
  convergeArith _ n1 n2 = converge n1 n2
  evalOp _ (VBool i) (VBool j) = Right $ VBool (i `xor` j)
  commutativityProof = Just Dict

instance ArithOp Lsl 'TNat 'TNat where
  type ArithRes Lsl 'TNat 'TNat = 'TNat
  convergeArith _ n1 n2 = converge n1 n2
  evalOp _ n@(VNat i) m@(VNat j) =
    if j > 256
    then Left $ ShiftArithError LslOverflow n m
    else Right $ VNat (fromInteger $ shift (toInteger i) (fromIntegral j))

instance ArithOp Lsr 'TNat 'TNat where
  type ArithRes Lsr 'TNat 'TNat = 'TNat
  convergeArith _ n1 n2 = converge n1 n2
  evalOp _ n@(VNat i) m@(VNat j) =
    if j > 256
    then Left $ ShiftArithError LsrUnderflow n m
    else Right $ VNat (fromInteger $ shift (toInteger i) (-(fromIntegral j)))

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)

compareOp :: forall t i. (Comparable t, SingI t) => Value' i t -> Value' i t -> Integer
compareOp a' b' = case (sing :: Sing t, a', b') of
  (STInt, i, j) -> toInteger $ fromEnum (compare i j) - 1
  (STNat, i, j) -> toInteger $ fromEnum (compare i j) - 1
  (STString, i, j) -> toInteger $ fromEnum (compare i j) - 1
  (STBytes, i, j) -> toInteger $ fromEnum (compare i j) - 1
  (STMutez, i, j) -> toInteger $ fromEnum (compare i j) - 1
  (STBool, i, j) -> toInteger $ fromEnum (compare i j) - 1
  (STKeyHash, i, j) -> toInteger $ fromEnum (compare i j) - 1
  (STTimestamp, i, j) -> toInteger $ fromEnum (compare i j) - 1
  (STAddress, i, j) -> toInteger $ fromEnum (compare i j) - 1
  (STPair l m, VPair (a, b), VPair (c, d)) ->
    case checkComparability l of
      CanBeCompared ->
        case compareOp a c of
          0  -> case checkComparability m of
            CanBeCompared -> compareOp b d
          r' -> r'

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 Buildable ShiftArithErrorType where
  build = \case
    LslOverflow -> "lsl overflow"
    LsrUnderflow -> "lsr underflow"

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

instance (Show n, Show m) => Buildable (ArithError n m) where
  build (MutezArithError errType n m) = "Mutez "
    <> build errType <> " with " <> show n <> ", " <> show m
  build (ShiftArithError errType n m) =
    build errType <> " with " <> show n <> ", " <> show m