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

{-# OPTIONS_GHC -Wno-orphans #-}

{- | Evaluation of expressions.

Stack-based languages allow convenient expressions evaluation, for that
we just need binary instructions in infix notation, not in Polish postfix
notation that 'add' and other primitives provide. Compare:

> push 1; push 2; push 3; push 4; mul; rsub; add

vs

> push 1 |+| push 2 |-| push 3 |*| push 4

In these expressions each atom is some instruction providing a single value
on top of the stack, for example:

@
nthOdd :: Lambda Natural Natural
nthOdd = take |*| push @Natural 2 |+| push @Natural 1
@

For binary operations we provide the respective operators. Unary operations can
be lifted with 'unaryExpr':

@
implication :: [Bool, Bool] :-> '[Bool]
implication = unaryExpr not take |.|.| take
@

or with its alias in form of an operator:

@
implication :: [Bool, Bool] :-> '[Bool]
implication = not $: take |.|.| take
@

Stack changes are propagated from left to right. If an atom consumes an element
at the top of the stack, the next atom will accept only the remainder of the
stack.

In most cases you should prefer providing named variables to the formulas in
order to avoid messing up with the arguments:

@
f :: ("a" :! Natural) : ("b" :! Natural) : ("c" :! Natural) : s :-> Integer : s
f = fromNamed #a |+| fromNamed #b |-| fromNamed #c
@

Instead of putting all the elements on the stack upon applying the formula,
you may find it more convenient to evaluate most of the arguments right within
the formula:

@
withinRange
  :: Natural : a : b : c : ("config" :! Config) : s
  :-> Bool : a : b : c : ("config" :! Config) : s
withinRange =
  dup  |>=| do{ dupL #config; toField #minBound } |&|
  take |<=| do{ dupL #config; toField #maxBound }
@

-}
module Lorentz.Expr
  ( Expr
  , take
  , unaryExpr
  , ($:)
  , binaryExpr
  , (|+|)
  , (|-|)
  , (|*|)
  , (|==|)
  , (|/=|)
  , (|<|)
  , (|>|)
  , (|<=|)
  , (|>=|)
  , (|&|)
  , (|||)
  , (|.|.|)
  , (|^|)
  , (|<<|)
  , (|>>|)
  , (|:|)
  , pairE
  , (|@|)
  , (|%!|)
  , listE
  , transferTokensE
  , createContractE
  , viewE
  ) where

import Prelude (HasCallStack, foldr, uncurry, ($))

import Lorentz.Arith
import Lorentz.Base
import Lorentz.Constraints
import Lorentz.CustomArith
import Lorentz.Errors
import Lorentz.Instr
import Lorentz.Macro
import Lorentz.Rebinded
import Lorentz.Value
import Lorentz.ViewBase
import Morley.Michelson.Typed.Arith
import Morley.Util.Named
import Morley.Util.TypeLits

-- | Expression is just an instruction accepting stack @inp@ and producing
-- stack @out@ with evaluation result @res@ at the top.

-- Stack flows from left to right, so e.g. @ fromNamed #a |<| fromNamed #b @
-- consumes first the value with label @a@ from the stack's top, then the value
-- with label @b@.
type Expr inp out res = inp :-> res : out

{- Note on the type alias ↑

It was pretty difficult to decide whether 'Expr' should be a newtype or type
alias.

Pros of newtype:
* It is safer and clearer.
* It allows defining such 'IsLabel name (Expr ...)' instance that would allow
  writing @ #a |<| #b @ (which should consume values asserting their labels).

It would /not/ allow writing @5 :: Expr ... Natural@, because this way we lose
GHC checks on passing out-of-bounds numbers like @(-1) :: Natural@.

Cons:
* Any non-trivial operation like @push@, @dupL@ and others will require a
  wrapper.

To make all this as simple to use and understand as possible, we used a
type alias.
-}

-- | An expression producing 'Bool' can be placed as condition to 'if'.
instance (i ~ arg, o ~ argl, o ~ argr, r ~ Bool, outb ~ out) =>
         IsCondition (Expr i o r) arg argl argr outb out where
  ifThenElse e l r = e # if_ l r

-- | Consume an element at the top of stack. This is just an alias for @nop@.
take :: Expr (a : s) s a
take = nop

-- | Lift an instruction to an unary operation on expressions.
unaryExpr
  :: (forall s. a : s :-> r : s)
  -> Expr s0 s1 a -> Expr s0 s1 r
unaryExpr op a = a # op

-- | An alias for 'unaryExpr'.
infixr 9 $:
($:)
  :: (forall s. a : s :-> r : s)
  -> Expr s0 s1 a -> Expr s0 s1 r
($:) = unaryExpr

-- | Lift an instruction to a binary operation on expressions.
binaryExpr
  :: (forall s. a : b : s :-> r : s)
  -> Expr s0 s1 a -> Expr s1 s2 b -> Expr s0 s2 r
binaryExpr op l r = l # dip r # op

-- | Expressions addition.
infixl 6 |+|
(|+|)
  :: (ArithOpHs Add a b r)
  => Expr s0 s1 a -> Expr s1 s2 b -> Expr s0 s2 r
(|+|) = binaryExpr add

-- | Expressions subtraction.
infixl 6 |-|
(|-|)
  :: (ArithOpHs Sub a b r)
  => Expr s0 s1 a -> Expr s1 s2 b -> Expr s0 s2 r
(|-|) = binaryExpr sub

-- | Expressions multiplication.
infixl 7 |*|
(|*|)
  :: (ArithOpHs Mul a b r)
  => Expr s0 s1 a -> Expr s1 s2 b -> Expr s0 s2 r
(|*|) = binaryExpr mul

-- | Expressions comparison.
infix 4 |==|
infix 4 |/=|
infix 4 |<|
infix 4 |>|
infix 4 |<=|
infix 4 |>=|
(|==|), (|/=|), (|<|), (|>|), (|>=|), (|<=|)
  :: (NiceComparable a)
  => Expr s0 s1 a -> Expr s1 s2 a -> Expr s0 s2 Bool
(|==|) = binaryExpr eq
(|/=|) = binaryExpr neq
(|<|) = binaryExpr lt
(|>|) = binaryExpr gt
(|<=|) = binaryExpr le
(|>=|) = binaryExpr ge

-- | Bitwise/logical @AND@ on expressions.
infixl 2 |&|
(|&|)
  :: (ArithOpHs And a b r)
  => Expr s0 s1 a -> Expr s1 s2 b -> Expr s0 s2 r
(|&|) = binaryExpr and

-- | Bitwise/logical @OR@ on expressions.
--
-- In case you find this operator looking weird, see '|.|.|'
infixl 1 |||
(|||)
  :: (ArithOpHs Or a b r)
  => Expr s0 s1 a -> Expr s1 s2 b -> Expr s0 s2 r
(|||) = binaryExpr or

-- | An alias for '|||'.
infixl 1 |.|.|
(|.|.|)
  :: (ArithOpHs Or a b r)
  => Expr s0 s1 a -> Expr s1 s2 b -> Expr s0 s2 r
(|.|.|) = binaryExpr or

-- | Bitwise/logical @XOR@ on expressions.
infixl 3 |^|
(|^|)
  :: (ArithOpHs Xor a b r)
  => Expr s0 s1 a -> Expr s1 s2 b -> Expr s0 s2 r
(|^|) = binaryExpr xor

-- | Left shift on expressions.
infix 8 |<<|
(|<<|)
  :: (ArithOpHs Lsl a b r)
  => Expr s0 s1 a -> Expr s1 s2 b -> Expr s0 s2 r
(|<<|) = binaryExpr lsl

-- | Right shift on expressions.
infix 8 |>>|
(|>>|)
  :: (ArithOpHs Lsr a b r)
  => Expr s0 s1 a -> Expr s1 s2 b -> Expr s0 s2 r
(|>>|) = binaryExpr lsr

-- | 'cons' on expressions.
--
-- @
-- one :: a : s :-> [a] : s
-- one = take |:| nil
-- @
infixr 1 |:|
(|:|) :: Expr s0 s1 a -> Expr s1 s2 [a] -> Expr s0 s2 [a]
(|:|) = binaryExpr cons

-- | Construct a simple pair.
--
-- @
-- trivialContract :: ((), storage) :-> ([Operation], Storage)
-- trivialContract = nil |@| cdr
-- @
--
-- This is useful as pair appears even in simple contracts.
-- For more advanced types, use 'Lorentz.ADT.constructT'.
infixr 0 |@|
(|@|) :: Expr s0 s1 a -> Expr s1 s2 b -> Expr s0 s2 (a, b)
(|@|) = binaryExpr pair

-- | Construct a 'Rational' value from a given nominator and denominator.
-- In case denominator is 0, fails with custom exception.
infixl 7 |%!|
(|%!|) :: Expr s0 s1 Integer -> Expr s1 s2 Natural -> Expr s0 s2 Rational
(|%!|) = binaryExpr $ mkRational_ (failCustom #zero_denominator)

-- | An alias for '|@|'.
--
-- @
-- trivialContract :: ((), storage) :-> ([Operation], Storage)
-- trivialContract =
--   pairE
--     ( nil
--     , cdr
--     )
-- @
pairE :: (Expr s0 s1 a, Expr s1 s2 b) -> Expr s0 s2 (a, b)
pairE = uncurry (|@|)

-- | Construct a list given the constructor for each element.
listE :: KnownValue a => [Expr s s a] -> Expr s s [a]
listE = foldr (\pushElem pushRec -> pushElem # dip pushRec # cons) nil

-- | Version of 'transferTokens' instruction that accepts
-- all the arguments as expressions.
--
-- @
-- transferTokensE
--   ! #arg L.unit
--   ! #amount (push zeroMutez)
--   ! #contract take
--   |:| nil
-- @
--
-- You can provide arguments in arbitrary order, but direction of stack changes
-- flow is fixed: stack change in @arg@ expression affects stack available in
-- @amount@ expression, and stack changes in @amount@ expression affect
-- stack changes in @contract@ expression.
transferTokensE
  :: (NiceParameter p, IsNotInView)
  => ("arg" :! Expr s0 s1 p)
  -> ("amount" :! Expr s1 s2 Mutez)
  -> ("contract" :! Expr s2 s3 (ContractRef p))
  -> Expr s0 s3 Operation
transferTokensE (arg #arg -> paramE)
                (arg #amount -> amountE)
                (arg #contract -> contractE) =
  paramE # dip (amountE # dip contractE) # transferTokens

-- | Version of 'createContract' instruction that accepts
-- all the arguments as expressions.
--
-- @
-- createContractE
--   ! #delegate none
--   ! #balance (push zeroMutez)
--   ! #storage unit
--   ! #contract myContract
-- @
--
-- Note that this returns an operation, and pushes the address of the newly
-- created contract as a side-effect.
createContractE
  :: IsNotInView
  => ("delegate" :! Expr s0 s1 (Maybe KeyHash))
  -> ("balance" :! Expr s1 s2 Mutez)
  -> ("storage" :! Expr s2 s3 st)
  -> ("contract" :! Contract p st vd)
  -> Expr s0 (TAddress p vd : s3) Operation
createContractE (arg #delegate -> delegatorE)
                (arg #balance -> balanceE)
                (arg #storage -> storageE)
                (arg #contract -> contract') =
  delegatorE # dip (balanceE # dip storageE) # createContract contract'

-- | Version of 'view' instruction that accepts all the arguments as expressions.
--
-- @
-- viewE @"myview"
--   ! #arg (push zeroMutez)
--   ! #address (push addr)
-- @
viewE
  :: forall name arg ret p vd s0 s1 s2.
    ( HasCallStack, KnownSymbol name, KnownValue arg
    , NiceViewable ret, HasView vd name arg ret
    )
  => ("arg" :! Expr s0 s1 arg)
  -> ("address" :! Expr s1 s2 (TAddress p vd))
  -> Expr s0 s2 ret
viewE (arg #arg -> argE) (arg #address -> addrE) =
  argE # dip addrE # view @name