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

-- | Type families from "Morley.Michelson.Typed.Arith" lifted to Haskell types.
module Lorentz.Arith
  ( ArithOpHs (..)
  , DefArithOp (..)
  , UnaryArithOpHs (..)
  , DefUnaryArithOp (..)
  , ToIntegerArithOpHs (..)
  ) where

import Prelude hiding (natVal)

import Lorentz.Base
import Lorentz.Value
import Morley.Michelson.Typed.Arith
import Morley.Michelson.Typed.Instr qualified as M

-- | Lifted 'ArithOp'.
class ArithOpHs (aop :: Type) (n :: Type) (m :: Type) (r :: Type) where
  evalArithOpHs :: n : m : s :-> r : s
  default evalArithOpHs
    :: ( DefArithOp aop
       , ArithOp aop (ToT n) (ToT m)
       , ToT r ~ ArithRes aop (ToT n) (ToT m)
       )
    => n : m : s :-> r : s
  evalArithOpHs = I (defEvalOpHs @aop)

-- | Helper typeclass that provides default definition of 'evalArithOpHs'.
class DefArithOp aop where
  defEvalOpHs
    :: ( ArithOp aop n m
       , r ~ ArithRes aop n m
       )
    => M.Instr (n : m : s) (r : s)

-- | Lifted 'UnaryArithOp'.
class UnaryArithOpHs (aop :: Type) (n :: Type) where
  type UnaryArithResHs aop n :: Type
  evalUnaryArithOpHs :: n : s :-> UnaryArithResHs aop n : s
  default evalUnaryArithOpHs
    :: ( DefUnaryArithOp aop
       , UnaryArithOp aop (ToT n)
       , ToT (UnaryArithResHs aop n) ~ UnaryArithRes aop (ToT n)
       )
    => n : s :-> UnaryArithResHs aop n : s
  evalUnaryArithOpHs = I (defUnaryArithOpHs @aop)

-- | Helper typeclass that provides default definition of 'evalUnaryArithOpHs'.
class DefUnaryArithOp aop where
  defUnaryArithOpHs
    :: ( UnaryArithOp aop n
       , r ~ UnaryArithRes aop n
       )
    => M.Instr (n : s) (r : s)

class ToIntegerArithOpHs (n :: Type) where
  evalToIntOpHs :: n : s :-> Integer : s
  default evalToIntOpHs
    :: (ToIntArithOp (ToT n))
    => n : s :-> Integer : s
  evalToIntOpHs = I M.INT

instance DefArithOp Add where
  defEvalOpHs = M.ADD
instance DefArithOp Sub where
  defEvalOpHs = M.SUB
instance DefArithOp Mul where
  defEvalOpHs = M.MUL
instance DefArithOp And where
  defEvalOpHs = M.AND
instance DefArithOp Or where
  defEvalOpHs = M.OR
instance DefArithOp Xor where
  defEvalOpHs = M.XOR
instance DefArithOp Lsl where
  defEvalOpHs = M.LSL
instance DefArithOp Lsr where
  defEvalOpHs = M.LSR
instance DefArithOp EDiv where
  defEvalOpHs = M.EDIV
instance DefUnaryArithOp Not where
  defUnaryArithOpHs = M.NOT
instance DefUnaryArithOp Abs where
  defUnaryArithOpHs = M.ABS
instance DefUnaryArithOp Eq' where
  defUnaryArithOpHs = M.EQ
instance DefUnaryArithOp Neq where
  defUnaryArithOpHs = M.NEQ
instance DefUnaryArithOp Lt where
  defUnaryArithOpHs = M.LT
instance DefUnaryArithOp Le where
  defUnaryArithOpHs = M.LE
instance DefUnaryArithOp Gt where
  defUnaryArithOpHs = M.GT
instance DefUnaryArithOp Ge where
  defUnaryArithOpHs = M.GE
instance DefUnaryArithOp Neg where
  defUnaryArithOpHs = M.NEG

instance (r ~ Integer) => ArithOpHs Add Natural Integer r
instance (r ~ Integer) => ArithOpHs Add Integer Natural r
instance (r ~ Natural) => ArithOpHs Add Natural Natural r
instance (r ~ Integer) => ArithOpHs Add Integer Integer r
instance (r ~ Timestamp) => ArithOpHs Add Timestamp Integer r
instance (r ~ Timestamp) => ArithOpHs Add Integer Timestamp r
instance (r ~ Mutez) => ArithOpHs Add Mutez Mutez r
instance (r ~ Bls12381Fr) => ArithOpHs Add Bls12381Fr Bls12381Fr r
instance (r ~ Bls12381G1) => ArithOpHs Add Bls12381G1 Bls12381G1 r
instance (r ~ Bls12381G2) => ArithOpHs Add Bls12381G2 Bls12381G2 r


instance (r ~ Integer) => ArithOpHs Sub Natural Integer r
instance (r ~ Integer) => ArithOpHs Sub Integer Natural r
instance (r ~ Integer) => ArithOpHs Sub Natural Natural r
instance (r ~ Integer) => ArithOpHs Sub Integer Integer r
instance (r ~ Timestamp) => ArithOpHs Sub Timestamp Integer r
instance (r ~ Integer) => ArithOpHs Sub Timestamp Timestamp r

instance (r ~ Integer) => ArithOpHs Mul Natural Integer r
instance (r ~ Integer) => ArithOpHs Mul Integer Natural r
instance (r ~ Natural) => ArithOpHs Mul Natural Natural r
instance (r ~ Integer) => ArithOpHs Mul Integer Integer r
instance (r ~ Mutez) => ArithOpHs Mul Natural Mutez r
instance (r ~ Mutez) => ArithOpHs Mul Mutez Natural r
instance (r ~ Bls12381Fr) => ArithOpHs Mul Integer Bls12381Fr r
instance (r ~ Bls12381Fr) => ArithOpHs Mul Natural Bls12381Fr r
instance (r ~ Bls12381Fr) => ArithOpHs Mul Bls12381Fr Integer r
instance (r ~ Bls12381Fr) => ArithOpHs Mul Bls12381Fr Natural r
instance (r ~ Bls12381Fr) => ArithOpHs Mul Bls12381Fr Bls12381Fr r
instance (r ~ Bls12381G1) => ArithOpHs Mul Bls12381G1 Bls12381Fr r
instance (r ~ Bls12381G2) => ArithOpHs Mul Bls12381G2 Bls12381Fr r
instance (r ~ Bls12381G1) => ArithOpHs Mul Bls12381Fr Bls12381G1 r where
  evalArithOpHs = I (M.SWAP `M.Seq` M.MUL)
instance (r ~ Bls12381G2) => ArithOpHs Mul Bls12381Fr Bls12381G2 r where
  evalArithOpHs = I (M.SWAP `M.Seq` M.MUL)

instance (r ~ Maybe (Integer, Natural)) => ArithOpHs EDiv Natural Integer r
instance (r ~ Maybe (Integer, Natural)) => ArithOpHs EDiv Integer Natural r
instance (r ~ Maybe (Natural, Natural)) => ArithOpHs EDiv Natural Natural r
instance (r ~ Maybe (Integer, Natural)) => ArithOpHs EDiv Integer Integer r
instance (r ~ Maybe (Natural, Mutez)) => ArithOpHs EDiv Mutez Mutez r
instance (r ~ Maybe (Mutez, Mutez)) => ArithOpHs EDiv Mutez Natural r

instance UnaryArithOpHs Neg Integer where
  type UnaryArithResHs Neg Integer = Integer
instance UnaryArithOpHs Neg Natural where
  type UnaryArithResHs Neg Natural = Integer
instance UnaryArithOpHs Neg Bls12381Fr where
  type UnaryArithResHs Neg Bls12381Fr = Bls12381Fr
instance UnaryArithOpHs Neg Bls12381G1 where
  type UnaryArithResHs Neg Bls12381G1 = Bls12381G1
instance UnaryArithOpHs Neg Bls12381G2 where
  type UnaryArithResHs Neg Bls12381G2 = Bls12381G2


instance (r ~ Natural) => ArithOpHs Or Natural Natural r
instance (r ~ Bool) => ArithOpHs Or Bool Bool r

instance (r ~ Natural) => ArithOpHs And Integer Natural r
instance (r ~ Natural) => ArithOpHs And Natural Natural r
instance (r ~ Bool) => ArithOpHs And Bool Bool r

instance (r ~ Natural) => ArithOpHs Xor Natural Natural r
instance (r ~ Bool) => ArithOpHs Xor Bool Bool r

instance (r ~ Natural) => ArithOpHs Lsl Natural Natural r where

instance (r ~ Natural) => ArithOpHs Lsr Natural Natural r

instance UnaryArithOpHs Abs Integer where
  type UnaryArithResHs Abs Integer = Natural

instance UnaryArithOpHs Not Integer where
  type UnaryArithResHs Not Integer = Integer
instance UnaryArithOpHs Not Natural where
  type UnaryArithResHs Not Natural = Integer
instance UnaryArithOpHs Not Bool where
  type UnaryArithResHs Not Bool = Bool

instance UnaryArithOpHs Eq' Integer where
  type UnaryArithResHs Eq' Integer = Bool
instance UnaryArithOpHs Eq' Natural where
  type UnaryArithResHs Eq' Natural = Bool
  evalUnaryArithOpHs = evalToIntOpHs # evalUnaryArithOpHs @Eq'

instance UnaryArithOpHs Neq Integer where
  type UnaryArithResHs Neq Integer = Bool
instance UnaryArithOpHs Neq Natural where
  type UnaryArithResHs Neq Natural = Bool
  evalUnaryArithOpHs = evalToIntOpHs # evalUnaryArithOpHs @Neq

instance UnaryArithOpHs Lt Integer where
  type UnaryArithResHs Lt Integer = Bool
instance UnaryArithOpHs Lt Natural where
  type UnaryArithResHs Lt Natural = Bool
  evalUnaryArithOpHs = evalToIntOpHs # evalUnaryArithOpHs @Lt

instance UnaryArithOpHs Gt Integer where
  type UnaryArithResHs Gt Integer = Bool
instance UnaryArithOpHs Gt Natural where
  type UnaryArithResHs Gt Natural = Bool
  evalUnaryArithOpHs = evalToIntOpHs # evalUnaryArithOpHs @Gt

instance UnaryArithOpHs Le Integer where
  type UnaryArithResHs Le Integer = Bool
instance UnaryArithOpHs Le Natural where
  type UnaryArithResHs Le Natural = Bool
  evalUnaryArithOpHs = evalToIntOpHs # evalUnaryArithOpHs @Le

instance UnaryArithOpHs Ge Integer where
  type UnaryArithResHs Ge Integer = Bool
instance UnaryArithOpHs Ge Natural where
  type UnaryArithResHs Ge Natural = Bool
  evalUnaryArithOpHs = evalToIntOpHs # evalUnaryArithOpHs @Ge

instance ToIntegerArithOpHs Natural
instance ToIntegerArithOpHs Bls12381Fr