Expr.hs

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

-- | Tests for Indigo Expr

module Test.Expr
  ( test_SmallIndigoExpr
  ) where

import Prelude

import Data.Bits qualified as B
import Data.Map qualified as M
import Data.Set qualified as S
import Hedgehog (Gen)
import Hedgehog.Gen qualified as Gen
import Hedgehog.Range qualified as Range
import Test.Tasty (TestTree)

import Hedgehog.Gen.Michelson (genMText)
import Hedgehog.Gen.Michelson.Typed (genBigMap)
import Hedgehog.Gen.Tezos.Address (genAddress)
import Hedgehog.Gen.Tezos.Core (genChainId, genMutez)
import Hedgehog.Gen.Tezos.Crypto (genKeyHash, genPublicKey, genSignature)
import Indigo qualified as I
import Indigo.Lorentz
import Morley.Michelson.Interpret (runUnpack)
import Morley.Michelson.Interpret.Pack
import Morley.Michelson.Runtime.GState (genesisAddress)
import Morley.Michelson.Text
import Morley.Michelson.Typed qualified as T
import Morley.Michelson.Typed.Haskell.Value (BigMap(..))
import Morley.Tezos.Address
import Morley.Tezos.Core (dummyChainId)
import Morley.Tezos.Crypto qualified as C
import Test.Cleveland.Util (genTuple2)
import Test.Code.Expr
import Test.Util

genMyStore :: Gen MyStore
genMyStore = MyStore
  <$> Gen.map (Range.linear -100 100)
        ((, ()) <$> Gen.integral (Range.linearFrom 0 -1000 1000))
  <*> Gen.bool

genMySum :: Gen MySum
genMySum = Gen.choice [MySumA <$> Gen.bool, MySumB <$> Gen.integral (Range.linear 0 1000)]

-- | Tests on single Indigo `Expr`s or simple combinations of them.
-- Param and storage for these are generated randomly and their resulting stack
-- is validated against an Haskell function.
test_SmallIndigoExpr :: [TestTree]
test_SmallIndigoExpr =
  [ testIndigo "Cast" genInteger genInteger (validateStSuccess const) (exprUnary @Integer I.cast)
  , testIndigo "Size" genIntegerList genNatural (validateStSuccess (const . fromIntegralOverflowing . length)) exprSize
  , testIndigo "Add" genInteger genInteger (validateStSuccess (+)) (exprBinary @Integer (I.+))
  , testIndigo "Sub" genInteger genInteger (validateStSuccess (-)) (exprBinary @Integer (I.-))
  , testIndigo "Mul" genInteger genInteger (validateStSuccess (*)) (exprBinary @Integer (I.*))
  , testIndigo "Neg" genInteger genInteger (validateStSuccess (const . negate)) (exprUnary @Integer I.neg)
  , testIndigo "Abs" genInteger genNatural (validateStSuccess (\p _ -> fromIntegralOverflowing $ abs p)) exprAbs
  , testIndigo "DivEq" genInteger genInteger (validateStEither divEqCheck) exprDivEq
  , testIndigo "ModNeq" genInteger genInteger (validateStEither modNeqCheck) exprModNeq
  , testIndigo "Le3" genInteger Gen.bool (validateStSuccess (const . (<= 3))) exprLe3
  , testIndigo "Lt3OrGt10" genInteger Gen.bool (validateStSuccess (\p _ -> p < 3 || p > 10)) exprLt3OrGt10
  , testIndigo "Lt3OrGt10" genInteger Gen.bool (validateStSuccess (\p _ -> p >= 3 && p < 10)) exprGe3AndNotGe10
  , testIndigo "Xor" genNatural genNatural (validateStSuccess xor) (exprBinary @Natural (I.^))
  , testIndigo "Lsl" genNatural genShiftNatural (validateStSuccess (\p st -> B.shiftL p (fromIntegralOverflowing st))) (exprBinary @Natural (I.<<<))
  , testIndigo "Lsr" genNatural genShiftNatural (validateStSuccess (\p st -> B.shiftR p (fromIntegralOverflowing st))) (exprBinary @Natural (I.>>>))
  , testIndigo "Ge4OrNeq5AndEq6" genInteger Gen.bool (validateStSuccess (\p _ -> p >= 4 || p /= 5 && p == 6)) exprGe4OrNeq5AndEq6
  , testIndigo "Not" Gen.bool Gen.bool (validateStSuccess (\p _ -> not p)) exprNot
  , testIndigo "IsNat" genInteger (Gen.maybe genNatural) (validateStSuccess isNatCheck) exprIsNat
  , testIndigo "Fst" genIntegerPair genInteger (validateStSuccess (\(a,_) _ -> a)) exprFst
  , testIndigo "Snd" genIntegerPair genInteger (validateStSuccess (\(_,b) _ -> b)) exprSnd
  , testIndigo "Some" genInteger genIntegerMaybe (validateStSuccess (\p _ -> Just p)) (exprSome @Integer)
  , testIndigo "None" genInteger genIntegerMaybe (validateStSuccess (\_ _ -> Nothing)) (exprNone @Integer)
  , testIndigo "UStore" genInteger genMyStore (validateStack2 storeCheck) exprStore

  -- TODO: no `Arbitrary` instance for `Named`
  -- , ToField
  -- , SetField
  -- , Name
  -- , UnName
  -- , Construct
  -- , ConstructT

  , testIndigo "Set" genIntegerSet genInteger (validateStack2 setCheck) exprSet
  , testIndigo "EmptySet" genUnit genIntegerSet (validateStSuccess (\_ _ -> S.empty)) exprEmptySet
  , testIndigo "BigMapLookup" genBigMapInt genIntegerMaybe (validateStSuccess (\(bmMap -> p) _st -> M.lookup 2 p)) exprBigMapLookup
  , testIndigo "BigMapDelete" genInteger genBigMapInt (validateStSuccess (\p (BigMap i st) -> BigMap i $ M.delete p st)) exprBigMapDelete
  , testIndigo "BigMapInsert" genInteger genBigMapInt (validateStSuccess (\p (BigMap i st) -> BigMap i $ M.insert p p st)) exprBigMapInsert
  , testIndigo "Pack" genSignature genByteString (validateStSuccess (\p _ -> packValue' $ T.VSignature p)) exprPack
  , testIndigo "Unpack" genByteString (Gen.maybe genSignature) (validateStSuccess unpackCheck) exprUnpack
  , testIndigo "Cons" genInteger genIntegerList (validateStSuccess (\(p :: Integer) s -> p : s)) exprCons
  , testIndigo "Concat" genMText' genMText' (validateStSuccess @MText (\p s -> p <> s)) exprConcat
  , testIndigo "Slice" genNatural (Gen.maybe genMText') (validateStSuccess sliceCheck) exprSlice

  -- TODO: Our current testing framework uses storage type for
  -- validation, meaning that we cannot test contracts that way
  -- because we prohibit contract type from appearing in storage.

  -- , Contract
  -- , ConvertEpAddressToContract
  -- , ContractAddress
  -- , Self
  -- , ContractCallingUnsafe
  -- , RunFutureContract
  -- , ImplicitAccount

  , testIndigo "CheckSignature" Gen.bool Gen.bool (validateStSuccess checkSignatureCheck) exprCheckSignature
  , testIndigo "Crypto" genByteString genByteString (validateStack2 cryptoCheck) exprCrypto
  , testIndigo "HashKey" genPublicKey genKeyHash (validateStSuccess (\p _ -> C.hashKey p)) exprHashKey
  , testIndigo "ChainId" genUnit genChainId (validateStSuccess (\_ _ -> dummyChainId)) (exprNullary I.chainId)
  , testIndigo "Amount" genUnit genMutez' (validateStSuccess (\_ _ -> [tz|100u|])) (exprNullary I.amount)
  , testIndigo "Balance" genUnit genMutez' (validateStSuccess (\_ _ -> [tz|100u|])) (exprNullary I.balance)
  , testIndigo "Sender" genUnit genAddress (validateStSuccess (\_ _ -> MkAddress genesisAddress)) (exprNullary I.sender)

  -- TODO: ContractEnv needed
  -- , Now

  , testIndigo "NonZero" genInteger genIntegerMaybe (validateStSuccess nonZeroCheck) exprNonZero

  , testIndigo "Wrap" Gen.bool genMySum (validateStSuccess wrapCheck) exprWrap
  ]

  where
    genMutez' = genMutez def
    genMText' = genMText def
    genIntegerList = Gen.list (Range.linear 0 100) genInteger
    genIntegerSet = Gen.set (Range.linear 0 100) genInteger
    genIntegerPair = genTuple2 genInteger genInteger
    genIntegerMaybe = Gen.maybe genInteger
    genNatural = Gen.integral @_ @Natural (Range.linear 0 1000)
    genInteger = Gen.integral @_ @Integer (Range.linearFrom 0 -1000 1000)
    genByteString = Gen.bytes (Range.linear 0 100)
    genUnit = pure ()
    genBigMapInt = genBigMap def genInteger genInteger

    -- Cannot shift by more than 256 bits
    genShiftNatural = Gen.integral @_ @Natural (Range.linear 0 256)

----------------------------------------------------------------------------
-- Expected behavior
----------------------------------------------------------------------------

divEqCheck :: Integer -> Integer -> Either MichelsonFailureWithStack Integer
divEqCheck param st
  | param == 0 = Left zeroDivFail
  | otherwise = Right $ st `T.divMich` param

modNeqCheck :: Integer -> Integer -> Either MichelsonFailureWithStack Integer
modNeqCheck param st
  | param == 0 = Left zeroDivFail
  | st `T.modMich` param /= 0 = Right 0
  | otherwise = Right 1

isNatCheck :: Integer -> Maybe Natural -> Maybe Natural
isNatCheck param _st
  | param >= 0 = Just (fromIntegralOverflowing param)
  | otherwise = Nothing

unpackCheck :: ByteString -> Maybe Signature -> Maybe Signature
unpackCheck param _st =
  fmap unwrap . rightToMaybe . runUnpack $ param
  where
    unwrap :: Value 'T.TSignature -> Signature
    unwrap (T.VSignature signature) = signature

setCheck :: Set Integer -> Integer -> Either MichelsonFailureWithStack (Set Integer, Integer)
setCheck param _st = Right (newParam, newSt)
  where
    newParam
      | S.member 0 param = S.delete 0 param
      | otherwise = S.insert 1 param
    newSt
      | S.size newParam == 1 = 0
      | otherwise = 1

sliceCheck :: Natural -> Maybe MText -> Maybe MText
sliceCheck param (Just st) = Just . takeMText (fromIntegralOverflowing param) $ st
sliceCheck _param Nothing = Nothing

checkSignatureCheck :: Bool -> Bool -> Bool
checkSignatureCheck _param _st = check sampleSignature
  where
    check SignatureData{..} = C.checkSignature
      (partialParse C.parsePublicKey sdPublicKey)
      (partialParse C.parseSignature sdSignature)
      sdBytes

storeCheck
  :: Integer
  -> MyStore
  -> Either MichelsonFailureWithStack (Integer, MyStore)
storeCheck param st
  | param == 0 || M.member 0 stBigMap = Left notNewKeyFail
  | M.member -1 st1BigMap = Right (param, st)
  | otherwise = Right (param, st2)
  where
    stBigMap = ints st
    st1BigMap = M.insert param () stBigMap
    st2BigMap = M.insert 0 () st1BigMap
    -- st1 = mkUStore $ myTemplate {ints = UStoreSubMap st1BigMap}
    st2 = st {ints = st2BigMap}

cryptoCheck
  :: ByteString
  -> ByteString
  -> Either MichelsonFailureWithStack (ByteString, ByteString)
cryptoCheck param _st = Right (C.sha512 param, C.blake2b param)

nonZeroCheck :: Integer -> Maybe Integer -> Maybe Integer
nonZeroCheck param _st
  | param == 0 = Nothing
  | otherwise = Just param

wrapCheck :: Bool -> MySum -> MySum
wrapCheck param _st = MySumA param