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

{-# LANGUAGE QuantifiedConstraints #-}

-- | Module, containing data types for Michelson value.

module Michelson.Typed.Instr
  ( Instr (..)
  , ExtInstr (..)
  , CommentType (..)
  , StackRef (..)
  , mkStackRef
  , PrintComment (..)
  , TestAssert (..)
  , ContractCode
  , Contract (..)
  , mapContractCode
  , mapEntriesOrdered
  , pattern CAR
  , pattern CDR
  , pattern PAIR
  , pattern UNPAIR
  , PackedNotes(..)
  , ConstraintDIPN
  , ConstraintDIPN'
  , ConstraintDIG
  , ConstraintDIG'
  , ConstraintDUG
  , ConstraintDUG'
  ) where

import Data.Singletons (Sing)
import Fmt (Buildable(..), (+||), (||+))
import qualified GHC.TypeNats as GHC (Nat)
import qualified Text.Show

import Michelson.Doc
import Michelson.ErrorPos
import Michelson.Printer.Util (RenderDoc(..), buildRenderDoc, needsParens, printDocS)
import Michelson.Typed.Annotation (Notes(..))
import Michelson.Typed.Arith
import Michelson.Typed.Entrypoints
import Michelson.Typed.Polymorphic
import Michelson.Typed.Scope
import Michelson.Typed.Sing (KnownT)
import Michelson.Typed.T (T(..))
import Michelson.Typed.Value (Comparable, ContractInp, ContractOut, Value'(..))
import Michelson.Untyped
  (Annotation(..), EntriesOrder(..), FieldAnn, TypeAnn, VarAnn, entriesOrderToInt)
import Util.Peano
import Util.TH
import Util.Type (type (++), KnownList)

-- | A wrapper to wrap annotations and corresponding singleton.
-- Apart from packing notes along with the corresponding Singleton,
-- this wrapper type, when included with `Instr` also helps to derive
-- the `Show` instance for `Instr` as `Sing a` does not have a `Show`
-- instance on its own.
data PackedNotes a where
  PackedNotes :: SingI a => Notes a -> PackedNotes (a ': s)

instance NFData (PackedNotes a) where
  rnf (PackedNotes n) = rnf n

instance Show (PackedNotes a) where
  show = printDocS True . renderDoc needsParens

instance Buildable (PackedNotes a) where
  build = buildRenderDoc

instance RenderDoc (PackedNotes a) where
  renderDoc pn (PackedNotes n) = renderDoc pn n

-- | Constraint that is used in DIPN, we want to share it with
-- typechecking code and eDSL code.
type ConstraintDIPN' kind (n :: Peano) (inp :: [kind])
  (out :: [kind]) (s :: [kind]) (s' :: [kind]) =
  ( SingI n, KnownPeano n, RequireLongerOrSameLength inp n
  , ((Take n inp) ++ s) ~ inp
  , ((Take n inp) ++ s') ~ out
  )

type ConstraintDIPN n inp out s s' = ConstraintDIPN' T n inp out s s'

type ConstraintDIG' kind (n :: Peano) (inp :: [kind])
  (out :: [kind]) (a :: kind) =
  ( SingI n, KnownPeano n, RequireLongerThan inp n
  , inp ~ (Take n inp ++ (a ': Drop ('S n) inp))
  , out ~ (a ': Take n inp ++ Drop ('S n) inp)
  )

type ConstraintDIG n inp out a = ConstraintDIG' T n inp out a

type ConstraintDUG' kind (n :: Peano) (inp :: [kind])
  (out :: [kind]) (a :: kind) =
  ( SingI n, KnownPeano n, RequireLongerThan out n
  , inp ~ (a ': Drop ('S 'Z) inp)
  , out ~ (Take n (Drop ('S 'Z) inp) ++ (a ': Drop ('S n) inp))
  )

type ConstraintDUG n inp out a = ConstraintDUG' T n inp out a

-- | Representation of Michelson instruction or sequence
-- of instructions.
--
-- Each Michelson instruction is represented by exactly one
-- constructor of this data type. Sequence of instructions
-- is represented with use of @Seq@ constructor in following
-- way: @SWAP; DROP ; DUP;@ -> @SWAP `Seq` DROP `Seq` DUP@.
-- Special case where there are no instructions is represented
-- by constructor @Nop@, e.g. @IF_NONE {} { SWAP; DROP; }@ ->
-- @IF_NONE Nop (SWAP `Seq` DROP)@.
--
-- Type parameter @inp@ states for input stack type. That is,
-- type of the stack that is required for operation to execute.
--
-- Type parameter @out@ states for output stack type or type
-- of stack that will be left after instruction's execution.
data Instr (inp :: [T]) (out :: [T]) where
  -- | A wrapper carrying original source location of the instruction.
  --
  -- TODO [#283]: replace this wrapper with something more clever and abstract.
  WithLoc :: InstrCallStack -> Instr a b -> Instr a b

  -- | A wrapper for instruction that also contain annotations for the
  -- top type on the result stack.
  --
  -- As of now, when converting from untyped representation,
  -- we only preserve field annotations and type annotations.
  -- Variable annotations are not preserved.
  --
  -- This can wrap only instructions with at least one non-failing execution
  -- branch.
  InstrWithNotes :: PackedNotes b -> Instr a b -> Instr a b

  -- | A wrapper for instruction with variable annotations.
  InstrWithVarNotes :: NonEmpty VarAnn -> Instr a b -> Instr a b

  -- | Execute given instruction on truncated stack.
  --
  -- This can wrap only instructions with at least one non-failing execution
  -- branch.
  --
  -- Morley has no such instruction, it is used solely in eDSLs.
  -- This instruction is sound because for all Michelson instructions
  -- the following property holds: if some code accepts stack @i@ and
  -- produces stack @o@, when it can also be run on stack @i + s@
  -- producing stack @o + s@; and also because Michelson never makes
  -- implicit assumptions on types, rather you have to express all
  -- "yet ambiguous" type information in code.
  -- We could make this not an instruction but rather a function
  -- which modifies an instruction (this would also automatically prove
  -- soundness of used transformation), but it occured to be tricky
  -- (in particular for TestAssert and DipN and family), so let's leave
  -- this for future work.
  FrameInstr
    :: forall a b s.
       (KnownList a, KnownList b)
    => Proxy s -> Instr a b -> Instr (a ++ s) (b ++ s)

  Seq :: Instr a b -> Instr b c -> Instr a c
  -- | Nop operation. Missing in Michelson spec, added to parse construction
  -- like  `IF {} { SWAP; DROP; }`.
  Nop :: Instr s s
  Ext :: ExtInstr s -> Instr s s
  -- | Nested wrapper is going to wrap a sequence of instructions with { }.
  -- It is crucial because serialisation of a contract
  -- depends on precise structure of its code.
  Nested :: Instr inp out -> Instr inp out
  -- | Places documentation generated for given instruction under some group.
  -- This is not part of 'ExtInstr' because it does not behave like 'Nop';
  -- instead, it inherits behaviour of instruction put within it.
  DocGroup :: DocGrouping -> Instr inp out -> Instr inp out

  -- | Variants of CAR/CDR to retain field annotations as they relate to the input
  -- stack, and hence won't be available from the annotation notes from
  -- the result stack we pack with the instructions during type check.
  AnnCAR :: FieldAnn -> Instr ('TPair a b ': s) (a ': s)
  AnnCDR :: FieldAnn -> Instr ('TPair a b ': s) (b ': s)

  -- Note that we can not merge DROP and DROPN into one instruction
  -- because they are packed differently.
  DROP :: Instr (a ': s) s
  DROPN
    :: forall (n :: Peano) s.
    (SingI n, KnownPeano n, RequireLongerOrSameLength s n, NFData (Sing n))
    => Sing n -> Instr s (Drop n s)
  DUP  :: Instr (a ': s) (a ': a ': s)
  SWAP :: Instr (a ': b ': s) (b ': a ': s)
  DIG
    :: forall (n :: Peano) inp out a. (ConstraintDIG n inp out a, NFData (Sing n))
    => Sing n -> Instr inp out
  DUG
    :: forall (n :: Peano) inp out a. (ConstraintDUG n inp out a, NFData (Sing n))
    => Sing n -> Instr inp out
  PUSH
    :: forall t s . ConstantScope t
    => Value' Instr t -> Instr s (t ': s)
  SOME :: Instr (a ': s) ('TOption a ': s)
  NONE :: forall a s . KnownT a => Instr s ('TOption a ': s)
  UNIT :: Instr s ('TUnit ': s)
  IF_NONE
    :: Instr s s'
    -> Instr (a ': s) s'
    -> Instr ('TOption a ': s) s'
  -- | Annotations for PAIR instructions can be different from notes presented on the stack
  -- in case of special field annotations, so we carry annotations for instruction
  -- separately from notes.
  AnnPAIR :: TypeAnn -> FieldAnn -> FieldAnn -> Instr (a ': b ': s) ('TPair a b ': s)
  LEFT :: forall b a s . KnownT b => Instr (a ': s) ('TOr a b ': s)
  RIGHT :: forall a b s . KnownT a => Instr (b ': s) ('TOr a b ': s)
  IF_LEFT
    :: Instr (a ': s) s'
    -> Instr (b ': s) s'
    -> Instr ('TOr a b ': s) s'
  NIL :: KnownT p => Instr s ('TList p ': s)
  CONS :: Instr (a ': 'TList a ': s) ('TList a ': s)
  IF_CONS
    :: Instr (a ': 'TList a ': s) s'
    -> Instr s s'
    -> Instr ('TList a ': s) s'
  SIZE :: SizeOp c => Instr (c ': s) ('TNat ': s)
  EMPTY_SET :: (KnownT e, Comparable e) => Instr s ('TSet e ': s)
  EMPTY_MAP :: (KnownT a, KnownT b, Comparable a) => Instr s ('TMap a b ': s)
  EMPTY_BIG_MAP :: (KnownT a, KnownT b, Comparable a) => Instr s ('TBigMap a b ': s)
  MAP :: (MapOp c, KnownT b)
      => Instr (MapOpInp c ': s) (b ': s)
      -> Instr (c ': s) (MapOpRes c b ': s)
  ITER :: IterOp c => Instr (IterOpEl c ': s) s -> Instr (c ': s) s
  MEM :: MemOp c => Instr (MemOpKey c ': c ': s) ('TBool ': s)
  GET
    :: (GetOp c, KnownT (GetOpVal c))
    => Instr (GetOpKey c ': c ': s) ('TOption (GetOpVal c) ': s)
  UPDATE
    :: UpdOp c
    => Instr (UpdOpKey c ': UpdOpParams c ': c ': s) (c ': s)
  IF :: Instr s s'
     -> Instr s s'
     -> Instr ('TBool ': s) s'
  LOOP :: Instr s ('TBool ': s)
       -> Instr ('TBool ': s) s
  LOOP_LEFT
    :: Instr (a ': s) ('TOr a b ': s)
    -> Instr ('TOr a b ': s) (b ': s)
  LAMBDA :: forall i o s . (KnownT i, KnownT o)
         => Value' Instr ('TLambda i o) -> Instr s ('TLambda i o ': s)
  EXEC :: Instr (t1 ': 'TLambda t1 t2 ': s) (t2 ': s)
  APPLY
    :: forall a b c s . (ConstantScope a, KnownT b)
    => Instr (a ': 'TLambda ('TPair a b) c ': s) ('TLambda b c ': s)
  DIP :: Instr a c -> Instr (b ': a) (b ': c)
  DIPN
    :: forall (n :: Peano) inp out s s'. (ConstraintDIPN n inp out s s', (NFData (Sing n)))
    => Sing n -> Instr s s' -> Instr inp out
  FAILWITH :: (KnownT a) => Instr (a ': s) t
  CAST :: forall a s . SingI a => Instr (a ': s) (a ': s)
  RENAME :: Instr (a ': s) (a ': s)
  PACK :: PackedValScope a => Instr (a ': s) ('TBytes ': s)
  UNPACK :: (UnpackedValScope a, KnownT a) => Instr ('TBytes ': s) ('TOption a ': s)
  CONCAT :: ConcatOp c => Instr (c ': c ': s) (c ': s)
  CONCAT' :: ConcatOp c => Instr ('TList c ': s) (c ': s)
  SLICE
    :: (SliceOp c, KnownT c)
    => Instr ('TNat ': 'TNat ': c ': s) ('TOption c ': s)
  ISNAT :: Instr ('TInt ': s) ('TOption ('TNat) ': s)
  ADD
    :: (ArithOp Add n m, Typeable n, Typeable m)
    => Instr (n ': m ': s) (ArithRes Add n m ': s)
  SUB
    :: (ArithOp Sub n m, Typeable n, Typeable m)
    => Instr (n ': m ': s) (ArithRes Sub n m ': s)
  MUL
    :: (ArithOp Mul n m, Typeable n, Typeable m)
    => Instr (n ':  m ': s) (ArithRes Mul n m ': s)
  EDIV
    :: EDivOp n m
    => Instr (n ':  m ': s)
                 (('TOption ('TPair (EDivOpRes n m)
                                      (EModOpRes n m))) ': s)
  ABS
    :: UnaryArithOp Abs n
    => Instr (n ': s) (UnaryArithRes Abs n ': s)
  NEG
    :: UnaryArithOp Neg n
    => Instr (n ': s) (UnaryArithRes Neg n ': s)
  LSL
    :: (ArithOp Lsl n m, Typeable n, Typeable m)
    => Instr (n ': m ': s) (ArithRes Lsl n m ': s)
  LSR
    :: (ArithOp Lsr n m, Typeable n, Typeable m)
    => Instr (n ':  m ': s) (ArithRes Lsr n m ': s)
  OR
    :: (ArithOp Or n m, Typeable n, Typeable m)
    => Instr (n ': m ': s) (ArithRes Or n m ': s)
  AND
    :: (ArithOp And n m, Typeable n, Typeable m)
    => Instr (n ': m ': s) (ArithRes And n m ': s)
  XOR
    :: (ArithOp Xor n m, Typeable n, Typeable m)
    => Instr (n ': m ': s) (ArithRes Xor n m ': s)
  NOT
    :: UnaryArithOp Not n
    => Instr (n ': s) (UnaryArithRes Not n ': s)
  COMPARE
    :: (Comparable n, KnownT n)
    => Instr (n ': n ': s) ('TInt ': s)
  EQ
    :: UnaryArithOp Eq' n
    => Instr (n ': s) (UnaryArithRes Eq' n ': s)
  NEQ
    :: UnaryArithOp Neq n
    => Instr (n ': s) (UnaryArithRes Neq n ': s)
  LT
    :: UnaryArithOp Lt n
    => Instr (n ': s) (UnaryArithRes Lt n ': s)
  GT
    :: UnaryArithOp Gt n
    => Instr (n ': s) (UnaryArithRes Gt n ': s)
  LE
    :: UnaryArithOp Le n
    => Instr (n ': s) (UnaryArithRes Le n ': s)
  GE
    :: UnaryArithOp Ge n
    => Instr (n ': s) (UnaryArithRes Ge n ': s)
  INT :: Instr ('TNat ': s) ('TInt ': s)
  SELF
    :: forall (arg :: T) s .
       (ParameterScope arg)
    => SomeEntrypointCallT arg
    -> Instr s ('TContract arg ': s)
  CONTRACT
    :: (ParameterScope p)
    => Notes p   -- Store Notes to be able to verify CONTRACT in typechecker
    -> EpName
    -> Instr ('TAddress ': s) ('TOption ('TContract p) ': s)
  TRANSFER_TOKENS
    :: (ParameterScope p) =>
       Instr (p ': 'TMutez ': 'TContract p ': s)
                   ('TOperation ': s)
  SET_DELEGATE
    :: Instr ('TOption 'TKeyHash ': s) ('TOperation ': s)

  CREATE_CONTRACT
    :: (ParameterScope p, StorageScope g)
    => Contract p g
    -> Instr ('TOption 'TKeyHash ':
              'TMutez ':
               g ': s)
             ('TOperation ': 'TAddress ': s)
  IMPLICIT_ACCOUNT
    :: Instr ('TKeyHash ': s) ('TContract 'TUnit ': s)
  NOW :: Instr s ('TTimestamp ': s)
  AMOUNT :: Instr s ('TMutez ': s)
  BALANCE :: Instr s ('TMutez ': s)
  CHECK_SIGNATURE
    :: Instr ('TKey ': 'TSignature ': 'TBytes ': s)
                   ('TBool ': s)
  SHA256 :: Instr ('TBytes ': s) ('TBytes ': s)
  SHA512 :: Instr ('TBytes ': s) ('TBytes ': s)
  BLAKE2B :: Instr ('TBytes ': s) ('TBytes ': s)
  HASH_KEY :: Instr ('TKey ': s) ('TKeyHash ': s)
  SOURCE :: Instr s ('TAddress ': s)
  SENDER :: Instr s ('TAddress ': s)
  ADDRESS :: Instr ('TContract a ': s) ('TAddress ': s)
  CHAIN_ID :: Instr s ('TChainId ': s)

deriving stock instance Show (Instr inp out)

instance Semigroup (Instr s s) where
  (<>) = Seq
instance Monoid (Instr s s) where
  mempty = Nop

-- We have to write down pattern like this because simple
-- @Instr (TPair a b : s) (a : s)@ signature would assume that we /expect/
-- the input stack to have pair on top, but we want to /provide/ this info
-- in scope of a pattern-match.
-- In pattern declaration we have to write down the two mentioned constraints
-- explicitly.
--
-- Note that internally GADT constructors are rewritten in the very same way.
pattern CAR :: () => (i ~ ('TPair a b : s), o ~ (a : s)) => Instr i o
pattern CAR = AnnCAR (AnnotationUnsafe "")

pattern CDR :: () => (i ~ ('TPair a b : s), o ~ (b : s)) => Instr i o
pattern CDR = AnnCDR (AnnotationUnsafe "")

pattern UNPAIR :: () => (i ~ ('TPair a b : s), o ~ (a : b : s)) => Instr i o
pattern UNPAIR = Seq DUP (Seq CAR (DIP CDR))

pattern PAIR :: () => (i ~ (a ': b ': s), o ~ ('TPair a b ': s)) => Instr i o
pattern PAIR = AnnPAIR (AnnotationUnsafe "") (AnnotationUnsafe "") (AnnotationUnsafe "")

data TestAssert (s :: [T]) where
  TestAssert
    :: (Typeable out)
    => Text
    -> PrintComment inp
    -> Instr inp ('TBool ': out)
    -> TestAssert inp

deriving stock instance Show (TestAssert s)
instance NFData (TestAssert s) where
  rnf (TestAssert a b c) = rnf (a, b, c)

-- | A reference into the stack of a given type.
data StackRef (st :: [T]) where
  -- | Keeps 0-based index to a stack element counting from the top.
  StackRef
    :: (KnownPeano idx, SingI idx, RequireLongerThan st idx)
    => Sing (idx :: Peano) -> StackRef st

instance NFData (StackRef st) where
  rnf (StackRef s) = rnf s

instance Eq (StackRef st) where
  StackRef snat1 == StackRef snat2 = peanoVal snat1 == peanoVal snat2

instance Show (StackRef st) where
  show (StackRef snat) = "StackRef {" +|| peanoVal snat ||+ "}"

-- | Create a stack reference, performing checks at compile time.
mkStackRef
  :: forall (gn :: GHC.Nat) st n.
      (n ~ ToPeano gn, SingI n, KnownPeano n, RequireLongerThan st n)
  => StackRef st
mkStackRef = StackRef $ sing @(ToPeano gn)

-- | A print format with references into the stack
newtype PrintComment (st :: [T]) = PrintComment
  { unPrintComment :: [Either Text (StackRef st)]
  } deriving stock (Eq, Show, Generic)
    deriving newtype (Semigroup, Monoid)

instance NFData (PrintComment st)

instance IsString (PrintComment st) where
  fromString = PrintComment . one . Left . fromString

data CommentType
 = FunctionStarts Text
 | FunctionEnds Text
 | StatementStarts Text
 | StatementEnds Text
 | JustComment Text
 | StackTypeComment (Maybe [T]) -- ^ 'Nothing' for any stack type
 deriving stock (Show, Generic)

instance NFData CommentType

data ExtInstr s
  = TEST_ASSERT (TestAssert s)
  | PRINT (PrintComment s)
  | DOC_ITEM SomeDocItem
  | COMMENT_ITEM CommentType
  deriving stock (Show, Generic)

instance NFData (ExtInstr s)

---------------------------------------------------

type ContractCode cp st = Instr (ContractInp cp st) (ContractOut st)

-- | Typed contract and information about annotations
-- which is not present in the contract code.
data Contract cp st = (ParameterScope cp, StorageScope st) => Contract
  { cCode       :: ContractCode cp st
  , cParamNotes :: ParamNotes cp
  , cStoreNotes :: Notes st
  , cEntriesOrder :: EntriesOrder
  }

deriving stock instance Show (Contract cp st)
deriving stock instance Eq (ContractCode cp st) => Eq (Contract cp st)
instance NFData (Contract cp st) where
  rnf (Contract a b c d) = rnf (a, b, c, d)

mapContractCode
  :: (ContractCode cp st -> ContractCode cp st)
  -> Contract cp st
  -> Contract cp st
mapContractCode f contract = contract { cCode = f $ cCode contract }

-- | Map each typed contract fields by the given function and sort the output
-- based on the 'EntriesOrder'.
mapEntriesOrdered
  :: Contract cp st
  -> (ParamNotes cp -> a)
  -> (Notes st -> a)
  -> (ContractCode cp st -> a)
  -> [a]
mapEntriesOrdered Contract{..} fParam fStorage fCode =
  fmap snd
    $ sortWith fst
        [ (paramPos, fParam cParamNotes)
        , (storagePos, fStorage cStoreNotes)
        , (codePos, fCode cCode)
        ]
  where
    (paramPos, storagePos, codePos) = entriesOrderToInt cEntriesOrder

$(deriveGADTNFData ''Instr)