Expression.hs

-- SPDX-FileCopyrightText: 2020 Tocqueville Group
-- SPDX-FileCopyrightText: 2018 obsidian.systems
--
-- SPDX-License-Identifier: LicenseRef-MIT-TQ
-- SPDX-License-Identifier: LicenseRef-MIT-obsidian-systems

-- | Module that defines Expression type, its related types
-- and its JSON instance.
module Morley.Micheline.Expression
  ( Annotation(..)
  , Expression(..)
  , MichelinePrimAp(..)
  , MichelinePrimitive(..)
  , michelsonPrimitive
  , annotToText
  , annotFromText
  , genExpression
  , genExpressionSeq
  , genMichelinePrimAp
  , genExprAnnotation
  ) where

import Data.Aeson (FromJSON, ToJSON, parseJSON, toEncoding, toJSON, withObject, withText, (.!=), (.:), (.:?))
import Data.Aeson.Casing (aesonPrefix, snakeCase)
import qualified Data.Aeson.Encoding.Internal as Aeson
import Data.Aeson.TH (deriveToJSON)
import qualified Data.Aeson.Types as Aeson
import qualified Data.HashMap.Strict as HashMap
import qualified Data.Sequence as Seq
import qualified Data.Text as T (uncons)
import Fmt (pretty)
import Hedgehog (MonadGen(GenBase))
import qualified Hedgehog.Gen as Gen
import qualified Hedgehog.Range as Range

import qualified Michelson.Untyped.Annotation as MUA (Annotation)
import Michelson.Untyped.Annotation
  (KnownAnnTag(..), FieldAnn, TypeAnn, VarAnn, TypeTag, FieldTag, VarTag, ann, annPrefix)
import Morley.Micheline.Json
import Util.ByteString (HexJSONByteString(..))
import Util.Test.Gen (genAnnotation)

newtype MichelinePrimitive = MichelinePrimitive Text
  deriving newtype (Eq, Ord, ToJSON, FromJSON)
  deriving stock (Show)

michelsonPrimitive :: Seq Text
michelsonPrimitive = Seq.fromList [
  "parameter", "storage", "code", "False", "Elt", "Left", "None", "Pair",
  "Right", "Some", "True", "Unit", "PACK", "UNPACK", "BLAKE2B", "SHA256",
  "SHA512", "ABS", "ADD", "AMOUNT", "AND", "BALANCE", "CAR", "CDR",
  "CHECK_SIGNATURE", "COMPARE", "CONCAT", "CONS", "CREATE_ACCOUNT",
  "CREATE_CONTRACT", "IMPLICIT_ACCOUNT", "DIP", "DROP", "DUP", "EDIV",
  "EMPTY_MAP", "EMPTY_SET", "EQ", "EXEC", "FAILWITH", "GE", "GET", "GT",
  "HASH_KEY", "IF", "IF_CONS", "IF_LEFT", "IF_NONE", "INT", "LAMBDA",
  "LE", "LEFT", "LOOP", "LSL", "LSR", "LT", "MAP", "MEM", "MUL", "NEG",
  "NEQ", "NIL", "NONE", "NOT", "NOW", "OR", "PAIR", "PUSH", "RIGHT",
  "SIZE", "SOME", "SOURCE", "SENDER", "SELF", "STEPS_TO_QUOTA", "SUB",
  "SWAP", "TRANSFER_TOKENS", "SET_DELEGATE", "UNIT", "UPDATE", "XOR",
  "ITER", "LOOP_LEFT", "ADDRESS", "CONTRACT", "ISNAT", "CAST", "RENAME",
  "bool", "contract", "int", "key", "key_hash", "lambda", "list", "map",
  "big_map", "nat", "option", "or", "pair", "set", "signature", "string",
  "bytes", "mutez", "timestamp", "unit", "operation", "address", "SLICE",
  "DIG", "DUG", "EMPTY_BIG_MAP", "APPLY", "chain_id", "CHAIN_ID"
  ]

-- | Type for Micheline Expression
data Expression
  = ExpressionInt Integer
    -- ^ Micheline represents both nats and ints using the same decimal format.
    -- The Haskell Integer type spans all possible values that the final
    -- (Michelson) type could end up being, and then some, so we use
    -- (StringEncode Integer) to represent all integral values here for easy
    -- JSON encoding compatibility.
  | ExpressionString Text
  | ExpressionBytes ByteString
  | ExpressionSeq (Seq Expression)
  | ExpressionPrim MichelinePrimAp
  deriving stock (Eq, Show)

data Annotation
  = AnnotationType TypeAnn
  | AnnotationVariable VarAnn
  | AnnotationField FieldAnn
  deriving stock (Eq, Show)

data MichelinePrimAp = MichelinePrimAp
  { mpaPrim :: MichelinePrimitive
  , mpaArgs :: Seq Expression
  , mpaAnnots :: Seq Annotation
  } deriving stock (Eq, Show)

instance FromJSON MichelinePrimAp where
  parseJSON = withObject "Prim" $ \v -> MichelinePrimAp
    <$> v .: "prim"
    <*> v .:? "args" .!= mempty
    <*> v .:? "annots" .!= mempty

deriveToJSON (aesonPrefix snakeCase) ''MichelinePrimAp

annotFromText :: MonadFail m => Text -> m Annotation
annotFromText txt = case result of
    Just a -> pure a
    Nothing -> fail "Unknown annotation type"
  where
    result = (AnnotationType <$> stripPrefix @TypeTag txt)
         <|> (AnnotationVariable <$> stripPrefix @VarTag txt)
         <|> (AnnotationField <$> stripPrefix @FieldTag txt)

stripPrefix :: forall tag . KnownAnnTag tag => Text -> Maybe (MUA.Annotation tag)
stripPrefix txt = do
  (n, t) <- T.uncons txt
  guard (toText [n] == prefix)
  Just $ ann t
  where
    prefix = annPrefix @tag

annotToText :: Annotation -> Text
annotToText = \case
  AnnotationType n -> pretty n
  AnnotationVariable n -> pretty n
  AnnotationField n -> pretty n

instance FromJSON Annotation where
  parseJSON = withText "Annotation" annotFromText

instance ToJSON Annotation where
  toJSON = toJSON . annotToText
  toEncoding = toEncoding . annotToText

instance FromJSON Expression where
  parseJSON v = ExpressionSeq <$> parseJSON v
            <|> ExpressionPrim <$> parseJSON v
            <|> ExpressionString <$> withObject "ExpressionString" (.: "string") v
            <|> ExpressionInt . unStringEncode <$> withObject "ExpressionInt" (.: "int") v
            <|> ExpressionBytes . unHexJSONByteString <$> withObject "ExpressionBytes" (.: "bytes") v

instance ToJSON Expression where
  toJSON (ExpressionSeq xs) = toJSON xs
  toJSON (ExpressionPrim xs) = toJSON xs
  toJSON (ExpressionString x) = Aeson.Object (HashMap.singleton "string" $ toJSON x)
  toJSON (ExpressionInt x) = Aeson.Object (HashMap.singleton "int" $ toJSON $ StringEncode x)
  toJSON (ExpressionBytes x) = Aeson.Object (HashMap.singleton "bytes" $ toJSON $ HexJSONByteString x)

  toEncoding (ExpressionSeq xs) = toEncoding xs
  toEncoding (ExpressionPrim xs) = toEncoding xs
  toEncoding (ExpressionString x) = Aeson.pairs (Aeson.pair "string" (toEncoding x))
  toEncoding (ExpressionInt x) = Aeson.pairs (Aeson.pair "int" (toEncoding $ StringEncode x))
  toEncoding (ExpressionBytes x) = Aeson.pairs (Aeson.pair "bytes" (toEncoding $ HexJSONByteString x))

-----------------------------------------------------
-- Gen Function for testing
-----------------------------------------------------

genExpression :: forall m. (MonadGen m, GenBase m ~ Identity) => m Expression
genExpression = Gen.recursive Gen.choice
  [genExpressionInt, genExpressionString, genExpressionBytes]
  [genSeq, genExpressionPrim]
  where
    genExpressionInt = ExpressionInt <$> (Gen.integral (Range.linearFrom 0 -1000 1000))
    genExpressionString = ExpressionString <$> (Gen.text (Range.linear 0 10) Gen.unicodeAll)
    genExpressionBytes = ExpressionBytes <$> (Gen.bytes (Range.linear 0 100))
    genSeq = ExpressionSeq <$> genExpressionSeq
    genExpressionPrim = ExpressionPrim <$> genMichelinePrimAp

genExpressionSeq :: forall m. (MonadGen m, GenBase m ~ Identity) => m (Seq Expression)
genExpressionSeq = (Gen.seq (Range.linear 0 10) genExpression)

genMichelinePrimAp :: forall m. (MonadGen m, GenBase m ~ Identity) => m MichelinePrimAp
genMichelinePrimAp = MichelinePrimAp <$> genMichelinePrimitive <*> genExpressionSeq <*> genAnnots
  where
    genMichelinePrimitive = MichelinePrimitive <$> (Gen.element $ toList michelsonPrimitive)

    genAnnots :: m (Seq Annotation)
    genAnnots = Gen.seq (Range.linear 0 10) genExprAnnotation

genExprAnnotation :: forall m. (MonadGen m, GenBase m ~ Identity) => m Annotation
genExprAnnotation =  Gen.choice
  [genAnnotationType, genAnnotationVariable, genAnnotationField]
  where
    genAnnotationType = AnnotationType <$> genAnnotation
    genAnnotationVariable = AnnotationVariable <$> genAnnotation
    genAnnotationField = AnnotationField <$> genAnnotation