{-# LANGUAGE QuantifiedConstraints #-}

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

module Michelson.Typed.Instr
  ( Instr (..)
  , ExtInstr (..)
  , StackRef (..)
  , mkStackRef
  , PrintComment (..)
  , TestAssert (..)
  , Contract
  , FullContract (..)
  , fcParamNotes
  , mapFullContractCode
  , pattern CAR
  , pattern CDR
  , PackedNotes(..)
  , ConstraintDIPN
  , ConstraintDIPN'
  , ConstraintDIG
  , ConstraintDIG'
  , ConstraintDUG
  , ConstraintDUG'
  ) where

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

import Michelson.Typed.Annotation (Notes(..))
import Michelson.Typed.Arith
import Michelson.Doc
import Michelson.Typed.EntryPoints
import Michelson.Typed.Polymorphic
import Michelson.Typed.Scope
import Michelson.Typed.T (CT(..), T(..))
import Michelson.Typed.Value (ContractInp, ContractOut, Value'(..))
import Michelson.Untyped (Annotation(..), FieldAnn, VarAnn)
import Util.Peano
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 :: Notes a -> Sing a -> PackedNotes (a ': s)

instance Show (PackedNotes a) where
  show (PackedNotes n _) = show 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 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)
    => 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
    => Sing n -> Instr inp out
  DUG
    :: forall (n :: Peano) inp out a. ConstraintDUG n inp out a
    => 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 . SingI 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'
  PAIR :: Instr (a ': b ': s) ('TPair a b ': s)
  LEFT :: forall b a s . SingI b => Instr (a ': s) ('TOr a b ': s)
  RIGHT :: forall a b s . SingI 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 :: SingI 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) ('Tc 'CNat ': s)
  EMPTY_SET :: (Typeable e, SingI e) => Instr s ('TSet e ': s)
  EMPTY_MAP :: (Typeable a, Typeable b, SingI a, SingI b) => Instr s ('TMap a b ': s)
  EMPTY_BIG_MAP :: (Typeable a, Typeable b, SingI a, SingI b) => Instr s ('TBigMap a b ': s)
  MAP :: MapOp c
      => 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 ('Tc (MemOpKey c) ': c ': s) ('Tc 'CBool ': s)
  GET
    :: GetOp c
    => Instr ('Tc (GetOpKey c) ': c ': s) ('TOption (GetOpVal c) ': s)
  UPDATE
    :: UpdOp c
    => Instr ('Tc (UpdOpKey c) ': UpdOpParams c ': c ': s) (c ': s)
  IF :: Instr s s'
     -> Instr s s'
     -> Instr ('Tc 'CBool ': s) s'
  LOOP :: Instr s ('Tc 'CBool ': s)
       -> Instr ('Tc 'CBool ': s) s
  LOOP_LEFT
    :: Instr (a ': s) ('TOr a b ': s)
    -> Instr ('TOr a b ': s) (b ': s)
  LAMBDA :: forall i o s . (Each [Typeable, SingI] [i, 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
    => 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'
    => Sing n -> Instr s s' -> Instr inp out
  FAILWITH :: (Typeable a, SingI 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) ('Tc 'CBytes ': s)
  UNPACK :: UnpackedValScope a => Instr ('Tc 'CBytes ': 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
    => Instr ('Tc 'CNat ': 'Tc 'CNat ': c ': s) ('TOption c ': s)
  ISNAT :: Instr ('Tc 'CInt ': s) ('TOption ('Tc 'CNat) ': s)
  ADD
    :: (ArithOp Add n m, Typeable n, Typeable m)
    => Instr ('Tc n ': 'Tc m ': s) ('Tc (ArithRes Add n m) ': s)
  SUB
    :: (ArithOp Sub n m, Typeable n, Typeable m)
    => Instr ('Tc n ': 'Tc m ': s) ('Tc (ArithRes Sub n m) ': s)
  MUL
    :: (ArithOp Mul n m, Typeable n, Typeable m)
    => Instr ('Tc n ': 'Tc m ': s) ('Tc (ArithRes Mul n m) ': s)
  EDIV
    :: EDivOp n m
    => Instr ('Tc n ': 'Tc m ': s)
                 (('TOption ('TPair ('Tc (EDivOpRes n m))
                                      ('Tc (EModOpRes n m)))) ': s)
  ABS
    :: UnaryArithOp Abs n
    => Instr ('Tc n ': s) ('Tc (UnaryArithRes Abs n) ': s)
  NEG
    :: UnaryArithOp Neg n
    => Instr ('Tc n ': s) ('Tc (UnaryArithRes Neg n) ': s)
  LSL
    :: (ArithOp Lsl n m, Typeable n, Typeable m)
    => Instr ('Tc n ': 'Tc m ': s) ('Tc (ArithRes Lsl n m) ': s)
  LSR
    :: (ArithOp Lsr n m, Typeable n, Typeable m)
    => Instr ('Tc n ': 'Tc m ': s) ('Tc (ArithRes Lsr n m) ': s)
  OR
    :: (ArithOp Or n m, Typeable n, Typeable m)
    => Instr ('Tc n ': 'Tc m ': s) ('Tc (ArithRes Or n m) ': s)
  AND
    :: (ArithOp And n m, Typeable n, Typeable m)
    => Instr ('Tc n ': 'Tc m ': s) ('Tc (ArithRes And n m) ': s)
  XOR
    :: (ArithOp Xor n m, Typeable n, Typeable m)
    => Instr ('Tc n ': 'Tc m ': s) ('Tc (ArithRes Xor n m) ': s)
  NOT
    :: UnaryArithOp Not n
    => Instr ('Tc n ': s) ('Tc (UnaryArithRes Not n) ': s)
  COMPARE
    :: (Comparable n, Typeable n, SingI n)
    => Instr (n ': n ': s) ('Tc 'CInt ': s)
  EQ
    :: UnaryArithOp Eq' n
    => Instr ('Tc n ': s) ('Tc (UnaryArithRes Eq' n) ': s)
  NEQ
    :: UnaryArithOp Neq n
    => Instr ('Tc n ': s) ('Tc (UnaryArithRes Neq n) ': s)
  LT
    :: UnaryArithOp Lt n
    => Instr ('Tc n ': s) ('Tc (UnaryArithRes Lt n) ': s)
  GT
    :: UnaryArithOp Gt n
    => Instr ('Tc n ': s) ('Tc (UnaryArithRes Gt n) ': s)
  LE
    :: UnaryArithOp Le n
    => Instr ('Tc n ': s) ('Tc (UnaryArithRes Le n) ': s)
  GE
    :: UnaryArithOp Ge n
    => Instr ('Tc n ': s) ('Tc (UnaryArithRes Ge n) ': s)
  INT :: Instr ('Tc 'CNat ': s) ('Tc 'CInt ': 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 ('Tc 'CAddress ': s) ('TOption ('TContract p) ': s)
  TRANSFER_TOKENS
    :: (ParameterScope p) =>
       Instr (p ': 'Tc 'CMutez ': 'TContract p ': s)
                   ('TOperation ': s)
  SET_DELEGATE
    :: Instr ('TOption ('Tc 'CKeyHash) ': s) ('TOperation ': s)

  CREATE_CONTRACT
    :: (ParameterScope p, StorageScope g)
    => FullContract p g
    -> Instr ('TOption ('Tc 'CKeyHash) ':
              'Tc 'CMutez ':
               g ': s)
             ('TOperation ': 'Tc 'CAddress ': s)
  IMPLICIT_ACCOUNT
    :: Instr ('Tc 'CKeyHash ': s) ('TContract 'TUnit ': s)
  NOW :: Instr s ('Tc 'CTimestamp ': s)
  AMOUNT :: Instr s ('Tc 'CMutez ': s)
  BALANCE :: Instr s ('Tc 'CMutez ': s)
  CHECK_SIGNATURE
    :: Instr ('TKey ': 'TSignature ': 'Tc 'CBytes ': s)
                   ('Tc 'CBool ': s)
  SHA256 :: Instr ('Tc 'CBytes ': s) ('Tc 'CBytes ': s)
  SHA512 :: Instr ('Tc 'CBytes ': s) ('Tc 'CBytes ': s)
  BLAKE2B :: Instr ('Tc 'CBytes ': s) ('Tc 'CBytes ': s)
  HASH_KEY :: Instr ('TKey ': s) ('Tc 'CKeyHash ': s)
  STEPS_TO_QUOTA :: Instr s ('Tc 'CNat ': s)
  SOURCE :: Instr s ('Tc 'CAddress ': s)
  SENDER :: Instr s ('Tc 'CAddress ': s)
  ADDRESS :: Instr ('TContract a ': s) ('Tc 'CAddress ': 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

instance NFData (Instr inp out) where
  -- This is the simplest way, and we don't care much about performance in tests.
  -- May be unsatisfactory for benchmarks though.
  -- If @instance Show@ is not automatically derived anymore, or you have other
  -- problems with this instance, then please reimplement it manually.
  rnf = rnf @String . show

pattern CAR :: Instr ('TPair a b : s) (a : s)
pattern CAR = AnnCAR (AnnotationUnsafe "")

pattern CDR :: Instr ('TPair a b : s) (b: s)
pattern CDR = AnnCDR (AnnotationUnsafe "")

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

deriving stock instance Show (TestAssert s)

-- | 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 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 IsString (PrintComment st) where
  fromString = PrintComment . one . Left . fromString

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

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

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

-- | Typed contract and information about annotations
-- which is not present in the contract code.
-- TODO [#12]: rename this to 'Contract' and the current 'Contract' to 'ContractCode'?
data FullContract cp st = (ParameterScope cp, StorageScope st) => FullContract
  { fcCode :: Contract cp st
  , fcParamNotesSafe :: ParamNotes cp
  , fcStoreNotes :: Notes st
  }

deriving stock instance Show (FullContract cp st)

fcParamNotes :: FullContract cp st -> Notes cp
fcParamNotes = unParamNotes . fcParamNotesSafe

mapFullContractCode
  :: (Contract cp st -> Contract cp st)
  -> FullContract cp st
  -> FullContract cp st
mapFullContractCode f contract = contract { fcCode = f $ fcCode contract }