Gen.hs

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

module Util.Test.Gen
  ( genInternalByteString
  , genVar
  , genTyVar
  , genExpandedOp
  , genInstrCallStack
  , genSrcPos
  , genPos
  , genLetName
  , genExtInstrAbstract
  , genPrintComment
  , genStackRef
  , genTestAssert
  , genStackFn
  , genStackTypePattern
  , genInstrAbstract
  , genContract
  , genContract'
  , genValue
  , genValue'
  , genElt
  , genParameterType
  , genType
  , genEpName
  , genAnnotation
  , genT
  ) where

import Prelude hiding (EQ, GT, LT)

import qualified Data.Text as T
import Hedgehog (MonadGen(GenBase), Range)
import qualified Hedgehog.Gen as Gen
import qualified Hedgehog.Range as Range

import Michelson.ErrorPos (InstrCallStack(..), LetName(..), Pos(..), SrcPos(..))
import Michelson.Test.Util (genEither)
import Michelson.Text (genMText)
import Michelson.Untyped
  (Annotation, Contract, Contract'(..), Elt(..), EpName, ExpandedOp(..), ExtInstrAbstract(..),
  InstrAbstract(..), InternalByteString(..), ParameterType(..), PrintComment(..), StackFn(..),
  StackRef(..), StackTypePattern(..), T(..), TestAssert(..), TyVar(..), Type(..), Value,
  Value'(..), Var(..), epNameFromRefAnn, mkAnnotation, specialFieldAnn, specialVarAnns)

genInternalByteString :: MonadGen m => m InternalByteString
genInternalByteString = InternalByteString <$> Gen.bytes (Range.linear 0 100)

genVar :: MonadGen m => m Var
genVar = Var <$> genSmallText

genTyVar :: (MonadGen m, GenBase m ~ Identity) => m TyVar
genTyVar = Gen.choice [VarID <$> genVar, TyCon <$> genType]

genExpandedOp :: (MonadGen m, GenBase m ~ Identity) => m ExpandedOp
genExpandedOp = Gen.recursive Gen.choice
  -- non-recursive constructors
  [
    -- NB: When used together, genExpandedOp & genInstrAbstract are mutually recursive.
    -- So we use a generator of InstrAbstract that does not contain any ops as an escape hatch.
    PrimEx <$> Gen.choice instrAbstractWithoutOp
  ]
  -- recursive constructors
  [ PrimEx <$> genInstrAbstract genExpandedOp
  , SeqEx <$> genSmallList genExpandedOp
  , Gen.subtermM genExpandedOp $ \expandedOp -> WithSrcEx <$> genInstrCallStack <*> pure expandedOp
  ]

genInstrCallStack :: MonadGen m => m InstrCallStack
genInstrCallStack = InstrCallStack <$> genLetCallStack <*> genSrcPos
  where
    genLetCallStack = Gen.frequency
      [ (80, pure [])
      , (18, Gen.list (Range.singleton 1) genLetName)
      , (2, Gen.list (Range.singleton 2) genLetName)
      ]

genSrcPos :: MonadGen m => m SrcPos
genSrcPos = SrcPos <$> genPos <*> genPos

genPos :: MonadGen m => m Pos
genPos = Pos <$> Gen.word Range.linearBounded

genLetName :: MonadGen m => m LetName
genLetName = LetName <$> Gen.text (Range.linear 0 3) Gen.unicodeAll

genExtInstrAbstract :: (MonadGen m, GenBase m ~ Identity) => m op -> m (ExtInstrAbstract op)
genExtInstrAbstract genOp = Gen.choice
  [ STACKTYPE <$> genStackTypePattern
  , FN <$> genSmallText <*> genStackFn <*> genSmallList genOp
  , UTEST_ASSERT <$> genTestAssert genOp
  , UPRINT <$> genPrintComment
  , UCOMMENT <$> genSmallText
  ]

genPrintComment :: MonadGen m => m PrintComment
genPrintComment = PrintComment <$> Gen.list (Range.linear 0 5) (genEither genSmallText genStackRef)

genStackRef :: MonadGen m => m StackRef
genStackRef = StackRef <$> Gen.integral (Range.linear 0 (fromIntegral $ maxBound @Word64))

genTestAssert :: MonadGen m => m op -> m (TestAssert op)
genTestAssert genOp = TestAssert <$> genSmallText <*> genPrintComment <*> genSmallList genOp

genStackFn :: (MonadGen m, GenBase m ~ Identity) => m StackFn
genStackFn = StackFn
  <$> Gen.maybe (Gen.set smallCollectionRange genVar)
  <*> genStackTypePattern
  <*> genStackTypePattern

genStackTypePattern :: (MonadGen m, GenBase m ~ Identity) => m StackTypePattern
genStackTypePattern = Gen.recursive Gen.choice
  [ pure StkEmpty, pure StkRest ]
  [ Gen.subtermM genStackTypePattern $ \stp -> StkCons <$> genTyVar <*> pure stp ]

genInstrAbstract :: (MonadGen m, GenBase m ~ Identity) => m op -> m (InstrAbstract op)
genInstrAbstract genOp =
  Gen.choice $ instrAbstractWithOp genOp <> instrAbstractWithoutOp

instrAbstractWithOp :: (MonadGen m, GenBase m ~ Identity) => m op -> [m (InstrAbstract op)]
instrAbstractWithOp genOp =
  [ EXT <$> genExtInstrAbstract genOp
  , PUSH <$> genAnnotation <*> genType <*> genValue' genOp
  , IF_NONE <$> genSmallList genOp <*> genSmallList genOp
  , IF_LEFT <$> genSmallList genOp <*> genSmallList genOp
  , IF_CONS <$> genSmallList genOp <*> genSmallList genOp
  , MAP <$> genAnnotation <*> genSmallList genOp
  , ITER <$> genSmallList genOp
  , IF <$> genSmallList genOp <*> genSmallList genOp
  , LOOP <$> genSmallList genOp
  , LOOP_LEFT <$> genSmallList genOp
  , LAMBDA <$> genAnnotation <*> genType <*> genType <*> genSmallList genOp
  , DIP <$> genSmallList genOp
  , DIPN <$> Gen.word Range.linearBounded <*> genSmallList genOp
  , CREATE_CONTRACT <$> genAnnotation <*> genAnnotation <*> genContract' genOp
  ]

instrAbstractWithoutOp :: (MonadGen m, GenBase m ~ Identity) => [m (InstrAbstract op)]
instrAbstractWithoutOp =
  [ DROPN <$> Gen.word Range.linearBounded
  , pure DROP
  , DUP <$> genAnnotation
  , pure SWAP
  , DIG <$> Gen.word Range.linearBounded
  , DUG <$> Gen.word Range.linearBounded
  , SOME <$> genAnnotation <*> genAnnotation
  , NONE <$> genAnnotation <*> genAnnotation <*> genType
  , UNIT <$> genAnnotation <*> genAnnotation
  , PAIR <$> genAnnotation <*> genAnnotation <*> genAnnotation <*> genAnnotation
  , CAR <$> genAnnotation <*> genAnnotation
  , CDR <$> genAnnotation <*> genAnnotation
  , LEFT <$> genAnnotation <*> genAnnotation <*> genAnnotation <*> genAnnotation <*> genType
  , RIGHT <$> genAnnotation <*> genAnnotation <*> genAnnotation <*> genAnnotation <*> genType
  , NIL <$> genAnnotation <*> genAnnotation <*> genType
  , CONS <$> genAnnotation
  , SIZE <$> genAnnotation
  , EMPTY_SET <$> genAnnotation <*> genAnnotation <*> genType
  , EMPTY_MAP <$> genAnnotation <*> genAnnotation <*> genType <*> genType
  , EMPTY_BIG_MAP <$> genAnnotation <*> genAnnotation <*> genType <*> genType
  , MEM <$> genAnnotation
  , GET <$> genAnnotation
  , UPDATE <$> genAnnotation
  , EXEC <$> genAnnotation
  , APPLY <$> genAnnotation
  , pure FAILWITH
  , CAST <$> genAnnotation <*> genType
  , RENAME <$> genAnnotation
  , PACK <$> genAnnotation
  , UNPACK <$> genAnnotation <*> genAnnotation <*> genType
  , CONCAT <$> genAnnotation
  , SLICE <$> genAnnotation
  , ISNAT <$> genAnnotation
  , ADD <$> genAnnotation
  , SUB <$> genAnnotation
  , MUL <$> genAnnotation
  , EDIV <$> genAnnotation
  , ABS <$> genAnnotation
  , NEG <$> genAnnotation
  , LSL <$> genAnnotation
  , LSR <$> genAnnotation
  , OR <$> genAnnotation
  , AND <$> genAnnotation
  , XOR <$> genAnnotation
  , NOT <$> genAnnotation
  , COMPARE <$> genAnnotation
  , EQ <$> genAnnotation
  , NEQ <$> genAnnotation
  , LT <$> genAnnotation
  , GT <$> genAnnotation
  , LE <$> genAnnotation
  , GE <$> genAnnotation
  , INT <$> genAnnotation
  , SELF <$> genAnnotation <*> genAnnotation
  , CONTRACT <$> genAnnotation <*> genAnnotation <*> genType
  , TRANSFER_TOKENS <$> genAnnotation
  , SET_DELEGATE <$> genAnnotation
  , IMPLICIT_ACCOUNT <$> genAnnotation
  , NOW <$> genAnnotation
  , AMOUNT <$> genAnnotation
  , BALANCE <$> genAnnotation
  , CHECK_SIGNATURE <$> genAnnotation
  , SHA256 <$> genAnnotation
  , SHA512 <$> genAnnotation
  , BLAKE2B <$> genAnnotation
  , HASH_KEY <$> genAnnotation
  , SOURCE <$> genAnnotation
  , SENDER <$> genAnnotation
  , ADDRESS <$> genAnnotation
  , CHAIN_ID <$> genAnnotation
  ]

genContract :: (MonadGen m, GenBase m ~ Identity) => m Contract
genContract = genContract' genExpandedOp

genContract' :: (MonadGen m, GenBase m ~ Identity) => m op -> m (Contract' op)
genContract' genOp = Contract <$> genParameterType <*> genType <*> genSmallList genOp

genValue :: (MonadGen m, GenBase m ~ Identity) => m Value
genValue = genValue' genExpandedOp

genValue' :: MonadGen m => m op -> m (Value' op)
genValue' genOp = Gen.recursive Gen.choice
  -- non-recursive constructors
  [ ValueInt <$> Gen.integral (Range.linearFrom 0 (fromIntegral $ minBound @Int64) (fromIntegral $ maxBound @Word64))
  , ValueString <$> genMText
  , ValueBytes <$> genInternalByteString
  , pure ValueUnit
  , pure ValueTrue
  , pure ValueFalse
  , pure ValueNone
  , pure ValueNil
  ]
  -- recursive constructors
  [ Gen.subterm2 (genValue' genOp) (genValue' genOp) ValuePair
  , Gen.subterm (genValue' genOp) ValueLeft
  , Gen.subterm (genValue' genOp) ValueRight
  , Gen.subterm (genValue' genOp) ValueSome
  , ValueSeq <$> genSmallNonEmpty (genValue' genOp)
  , ValueMap <$> genSmallNonEmpty (genElt genOp)
  , ValueLambda <$> genSmallNonEmpty genOp
  ]

genElt :: MonadGen m => m op -> m (Elt op)
genElt genOp = Elt <$> genValue' genOp <*> genValue' genOp

genParameterType :: (MonadGen m, GenBase m ~ Identity) => m ParameterType
genParameterType = ParameterType <$> genType <*> genAnnotation


genType :: (MonadGen m, GenBase m ~ Identity) => m Type
genType = Type <$> genT <*> genAnnotation

genEpName :: (MonadGen m, GenBase m ~ Identity) => m EpName
genEpName = Gen.mapMaybe (rightToMaybe . epNameFromRefAnn) genAnnotation

genAnnotation :: forall m a. (MonadGen m, GenBase m ~ Identity) => m (Annotation a)
genAnnotation = Gen.mapMaybe (rightToMaybe . mkAnnotation) genAnnotationText
  where
    genAnnotationText :: m Text
    genAnnotationText = Gen.frequency
      [ (1, pure "")
      , (17, T.cons <$> annStart <*> Gen.text (Range.linear 0 100) annBodyChar)
      , (1, pure specialFieldAnn)
      , (1, Gen.element specialVarAnns)
      ]

    annStart :: m Char
    annStart = Gen.choice [ pure '_', Gen.alphaNum ]

    annBodyChar :: m Char
    annBodyChar = Gen.choice [ annStart, Gen.element (".%@" :: String) ]

genT :: (MonadGen m, GenBase m ~ Identity) => m T
genT =
  Gen.recursive Gen.choice
    -- non-recursive constructors
    [ pure TKey
    , pure TUnit
    , pure TSignature
    , pure TChainId
    , pure TOperation
    , pure TInt
    , pure TNat
    , pure TString
    , pure TBytes
    , pure TMutez
    , pure TBool
    , pure TKeyHash
    , pure TTimestamp
    , pure TAddress
    ]
    -- recursive constructors
    [ TOption <$> genType
    , TList <$> genType
    , TSet <$> genType
    , TContract <$> genType
    , TPair <$> genAnnotation <*> genAnnotation <*> genType <*> genType
    , TOr <$> genAnnotation <*> genAnnotation <*> genType <*> genType
    , TLambda <$> genType <*> genType
    , TMap <$> genType <*> genType
    , TBigMap <$> genType <*> genType
    ]

smallCollectionRange :: Range Int
smallCollectionRange = Range.linear 0 3

genSmallList :: MonadGen m => m a -> m [a]
genSmallList = Gen.list smallCollectionRange

genSmallNonEmpty :: MonadGen m => m a -> m (NonEmpty a)
genSmallNonEmpty = Gen.nonEmpty smallCollectionRange

genSmallText :: MonadGen m => m Text
genSmallText = Gen.text (Range.linear 0 10) Gen.unicodeAll