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

-- | Michelson instructions in untyped model.

module Morley.Michelson.Untyped.Instr
  ( InstrAbstract (..)
  , ExpandedOp (..)
  , ExpandedInstr
  , flattenExpandedOp
  ) where

import Prelude hiding (EQ, GT, LT)

import Data.Aeson.TH (deriveJSON)
import Data.Data (Data(..))
import Fmt (Buildable(build), (+|), (|+))
import Generics.SYB (everywhere, mkT)
import Text.PrettyPrint.Leijen.Text
  (Doc, align, braces, enclose, indent, integer, line, nest, space, text, (<$$>), (<+>))

import Morley.Michelson.ErrorPos (ErrorSrcPos)
import Morley.Michelson.Printer.Util
  (RenderDoc(..), buildRenderDoc, doesntNeedParens, needsParens, renderOpsList, spaces)
import Morley.Michelson.Untyped.Annotation
  (Annotation, AnyAnn, FieldAnn, KnownAnnTag, TypeAnn, VarAnn, fullAnnSet, renderAnyAnns,
  singleAnnSet)
import Morley.Michelson.Untyped.Contract (Contract'(..))
import Morley.Michelson.Untyped.Ext (ExtInstrAbstract)
import Morley.Michelson.Untyped.Type (Ty)
import Morley.Michelson.Untyped.Value (Value'(..))
import Morley.Michelson.Untyped.View
import Morley.Util.Aeson

-------------------------------------
-- Types after macroexpander
-------------------------------------

type ExpandedInstr = InstrAbstract ExpandedOp

data ExpandedOp
  = PrimEx ExpandedInstr
  | SeqEx [ExpandedOp]
  | WithSrcEx ErrorSrcPos ExpandedOp
  deriving stock (Show, Eq, Data, Generic)

instance NFData ExpandedOp

instance RenderDoc ExpandedOp where
  renderDoc pn (WithSrcEx _ op) = renderDoc pn op
  renderDoc pn (PrimEx i) = renderDoc pn i
  renderDoc _  (SeqEx i) = renderOpsList False i
  isRenderable =
    \case PrimEx i -> isRenderable i
          WithSrcEx _ op -> isRenderable op
          _ -> True

instance Buildable ExpandedOp where
  build (WithSrcEx _ op) = build op
  build (PrimEx expandedInstr) = "<PrimEx: "+|expandedInstr|+">"
  build (SeqEx expandedOps)    = "<SeqEx: "+|expandedOps|+">"

-- | Flatten all 'SeqEx' in 'ExpandedOp'. This function is mostly for
-- testing. It returns instructions with the same logic, but they are
-- not strictly equivalent, because they are serialized differently
-- (grouping instructions into sequences affects the way they are
-- PACK'ed).
flattenExpandedOp :: ExpandedOp -> [ExpandedInstr]
flattenExpandedOp =
  \case
    PrimEx i -> [flattenInstr i]
    SeqEx ops -> concatMap flattenExpandedOp ops
    WithSrcEx _ op -> flattenExpandedOp op
  where
    flattenInstr :: ExpandedInstr -> ExpandedInstr
    flattenInstr = everywhere (mkT flattenOps)

    flattenOps :: [ExpandedOp] -> [ExpandedOp]
    flattenOps [] = []
    flattenOps (SeqEx s : xs) = s ++ flattenOps xs
    flattenOps (x@(PrimEx _) : xs) = x : flattenOps xs
    flattenOps (WithSrcEx _ op : xs) = op : flattenOps xs

-------------------------------------
-- Abstract instruction
-------------------------------------

-- | Michelson instruction with abstract parameter @op@.  This
-- parameter is necessary, because at different stages of our pipeline
-- it will be different. Initially it can contain macros and
-- non-flattened instructions, but then it contains only vanilla
-- Michelson instructions.
data InstrAbstract op
  = EXT               (ExtInstrAbstract op)
  | DROPN              Word
  -- ^ "DROP n" instruction.
  -- Note: reference implementation permits int16 here.
  -- Negative numbers are parsed successfully there, but rejected later.
  -- Morley is more permissive, so we use 'Word' here,
  -- i. e. permit more positive numbers. We do not permit negative numbers
  -- at type level.
  -- In practice, probably nobody will ever have numbers greater than ≈1000
  -- here, at least due to gas limits.
  -- Same reasoning applies to other instructions which have a numeric
  -- parameter representing number of elements on stack.
  | DROP
  -- ^ 'DROP' is essentially as special case for 'DROPN', but we need
  -- both because they are packed differently.
  | DUP               VarAnn
  | DUPN              VarAnn Word
  | SWAP
  | DIG               Word
  | DUG               Word
  | PUSH              VarAnn Ty (Value' op)
  | SOME              TypeAnn VarAnn
  | NONE              TypeAnn VarAnn Ty
  | UNIT              TypeAnn VarAnn
  | IF_NONE           [op] [op]
  | PAIR              TypeAnn VarAnn FieldAnn FieldAnn
  | UNPAIR            VarAnn VarAnn FieldAnn FieldAnn
  | PAIRN             VarAnn Word
  | UNPAIRN           Word
  | CAR               VarAnn FieldAnn
  | CDR               VarAnn FieldAnn
  | LEFT              TypeAnn VarAnn FieldAnn FieldAnn Ty
  | RIGHT             TypeAnn VarAnn FieldAnn FieldAnn Ty
  | IF_LEFT           [op] [op]
  | NIL               TypeAnn VarAnn Ty
  | CONS              VarAnn
  | IF_CONS           [op] [op]
  | SIZE              VarAnn
  | EMPTY_SET         TypeAnn VarAnn Ty
  | EMPTY_MAP         TypeAnn VarAnn Ty Ty
  | EMPTY_BIG_MAP     TypeAnn VarAnn Ty Ty
  | MAP               VarAnn [op]
  | ITER              [op]
  | MEM               VarAnn
  | GET               VarAnn
  | GETN              VarAnn Word
  | UPDATE            VarAnn
  | UPDATEN           VarAnn Word
  | GET_AND_UPDATE    VarAnn
  | IF                [op] [op]
  | LOOP              [op]
  | LOOP_LEFT         [op]
  | LAMBDA            VarAnn Ty Ty [op]
  | EXEC              VarAnn
  | APPLY             VarAnn
  | DIP               [op]
  | DIPN              Word [op]
  | FAILWITH
  | CAST              VarAnn Ty
  | RENAME            VarAnn
  | PACK              VarAnn
  | UNPACK            TypeAnn VarAnn Ty
  | CONCAT            VarAnn
  | SLICE             VarAnn
  | ISNAT             VarAnn
  | ADD               VarAnn
  | SUB               VarAnn
  | SUB_MUTEZ         VarAnn
  | MUL               VarAnn
  | EDIV              VarAnn
  | ABS               VarAnn
  | NEG               VarAnn
  | LSL               VarAnn
  | LSR               VarAnn
  | OR                VarAnn
  | AND               VarAnn
  | XOR               VarAnn
  | NOT               VarAnn
  | COMPARE           VarAnn
  | EQ                VarAnn
  | NEQ               VarAnn
  | LT                VarAnn
  | GT                VarAnn
  | LE                VarAnn
  | GE                VarAnn
  | INT               VarAnn
  | VIEW              VarAnn ViewName Ty
  | SELF              VarAnn FieldAnn
  | CONTRACT          VarAnn FieldAnn Ty
  | TRANSFER_TOKENS   VarAnn
  | SET_DELEGATE      VarAnn
  | CREATE_CONTRACT   VarAnn VarAnn (Contract' op)
  | IMPLICIT_ACCOUNT  VarAnn
  | NOW               VarAnn
  | AMOUNT            VarAnn
  | BALANCE           VarAnn
  | VOTING_POWER      VarAnn
  | TOTAL_VOTING_POWER VarAnn
  | CHECK_SIGNATURE   VarAnn
  | SHA256            VarAnn
  | SHA512            VarAnn
  | BLAKE2B           VarAnn
  | SHA3              VarAnn
  | KECCAK            VarAnn
  | HASH_KEY          VarAnn
  | PAIRING_CHECK     VarAnn
  | SOURCE            VarAnn
  | SENDER            VarAnn
  | ADDRESS           VarAnn
  | CHAIN_ID          VarAnn
  | LEVEL             VarAnn
  | SELF_ADDRESS      VarAnn
  | NEVER
  | TICKET            VarAnn
  | READ_TICKET       VarAnn
  | SPLIT_TICKET      VarAnn
  | JOIN_TICKETS      VarAnn
  | OPEN_CHEST        VarAnn
  | SAPLING_EMPTY_STATE VarAnn Natural
  | SAPLING_VERIFY_UPDATE VarAnn
  | MIN_BLOCK_TIME    [AnyAnn]
  deriving stock (Eq, Functor, Data, Generic, Show)

instance NFData op => NFData (InstrAbstract op)

instance (RenderDoc op) => RenderDoc (InstrAbstract op) where
  renderDoc pn = \case
    EXT extInstr            -> renderDoc pn extInstr
    DROP                    -> "DROP"
    DROPN n                 -> "DROP" <+> text (show n)
    DUP va                  -> "DUP" <+> renderAnnot va
    DUPN va n               -> "DUP" <+> renderAnnot va <+> text (show n)
    SWAP                    -> "SWAP"
    DIG n                   -> "DIG" <+> text (show n)
    DUG n                   -> "DUG" <+> text (show n)
    PUSH va t v             ->
      let renderConsecutively =
            "PUSH" <+> renderAnnot va <+> renderTy t <+> renderDoc needsParens v
          renderAligned = "PUSH" <+> renderAnnot va <$$> (spaces 2 <> renderTy t)
                                <$$> spaces 2 <> nest 3 (renderDoc needsParens v)
      in case v of
        ValueNil      -> renderConsecutively
        ValueInt{}    -> renderConsecutively
        ValueString{} -> renderConsecutively
        ValueBytes{}  -> renderConsecutively
        ValueUnit     -> renderConsecutively
        ValueTrue     -> renderConsecutively
        ValueFalse    -> renderConsecutively
        ValueNone     -> renderConsecutively
        _             -> renderAligned
    SOME ta va              -> "SOME" <+> renderAnnots [ta] [] [va]
    NONE ta va t            -> "NONE" <+> renderAnnots [ta] [] [va] <+> renderTy t
    UNIT ta va              -> "UNIT" <+> renderAnnots [ta] [] [va]
    IF_NONE x y             -> "IF_NONE" <+> nest 9 (renderOps x) <$$> spaces 8 <> nest 9 (renderOps y)
    PAIR ta va fa1 fa2      -> "PAIR" <+> renderAnnots [ta] [fa1, fa2] [va]
    UNPAIR va1 va2 fa1 fa2  -> "UNPAIR" <+> renderAnnots [] [fa1, fa2] [va1, va2]
    PAIRN va n              -> "PAIR" <+> renderAnnots [] [] [va] <+> text (show n)
    UNPAIRN n               -> "UNPAIR" <+> text (show n)
    CAR va fa               -> "CAR" <+> renderAnnots [] [fa] [va]
    CDR va fa               -> "CDR" <+> renderAnnots [] [fa] [va]
    LEFT ta va fa1 fa2 t    -> "LEFT" <+> renderAnnots [ta] [fa1, fa2] [va] <+> renderTy t
    RIGHT ta va fa1 fa2 t   -> "RIGHT" <+> renderAnnots [ta] [fa1, fa2] [va] <+> renderTy t
    IF_LEFT x y             -> "IF_LEFT" <+> nest 9 (renderOps x) <$$> spaces 8 <> nest 9 (renderOps y)
    NIL ta va t             -> "NIL" <+> renderAnnots [ta] [] [va] <+> renderTy t
    CONS va                 -> "CONS" <+> renderAnnot va
    IF_CONS x y             -> "IF_CONS" <+> nest 9 (renderOps x) <$$> spaces 8 <> nest 9 (renderOps y)
    SIZE va                 -> "SIZE" <+> renderAnnot va
    EMPTY_SET ta va t       -> "EMPTY_SET" <+> renderAnnots [ta] [] [va] <+> renderComp t
    EMPTY_MAP ta va c t     -> "EMPTY_MAP" <+> renderAnnots [ta] [] [va] <+> renderComp c <+> renderTy t
    EMPTY_BIG_MAP ta va c t -> "EMPTY_BIG_MAP" <+> renderAnnots [ta] [] [va] <+> renderComp c <+> renderTy t
    MAP va s                -> "MAP" <+> renderAnnot va <$$> spaces 4 <> nest 5 (renderOps s)
    ITER s                  -> "ITER" <+> nest 6 (renderOps s)
    MEM va                  -> "MEM" <+> renderAnnot va
    GET va                  -> "GET" <+> renderAnnot va
    GETN va n               -> "GET" <+> renderAnnot va <+> text (show n)
    UPDATE va               -> "UPDATE" <+> renderAnnot va
    UPDATEN va n            -> "UPDATE" <+> renderAnnot va <+> text (show n)
    GET_AND_UPDATE va       -> "GET_AND_UPDATE" <+> renderAnnot va
    IF x y                  -> "IF" <+> nest 4 (renderOps x) <$$> spaces 3 <> nest 4 (renderOps y)
    LOOP s                  -> "LOOP" <+> nest 6 (renderOps s)
    LOOP_LEFT s             -> "LOOP_LEFT" <+> nest 11 (renderOps s)
    LAMBDA va t r s         -> "LAMBDA" <+> renderAnnot va <$$> (spaces 2 <> renderTy t) <$$> (spaces 2 <> renderTy r) <$$> spaces 2 <> nest 3 (renderOps s)
    EXEC va                 -> "EXEC" <+> renderAnnot va
    APPLY va                -> "APPLY" <+> renderAnnot va
    DIP s                   -> "DIP" <+> nest 5 (renderOps s)
    DIPN n s                -> "DIP" <+> text (show n) <> line <> indent 4 (renderOps s)
    FAILWITH                -> "FAILWITH"
    CAST va t               -> "CAST" <+> renderAnnot va <+> renderTy t
    RENAME va               -> "RENAME" <+> renderAnnot va
    PACK va                 -> "PACK" <+> renderAnnot va
    UNPACK ta va t          -> "UNPACK" <+> renderAnnots [ta] [] [va] <+> renderTy t
    CONCAT va               -> "CONCAT" <+> renderAnnot va
    SLICE va                -> "SLICE" <+> renderAnnot va
    ISNAT va                -> "ISNAT" <+> renderAnnot va
    ADD va                  -> "ADD" <+> renderAnnot va
    SUB va                  -> "SUB" <+> renderAnnot va
    SUB_MUTEZ va            -> "SUB_MUTEZ" <+> renderAnnot va
    MUL va                  -> "MUL" <+> renderAnnot va
    EDIV va                 -> "EDIV" <+> renderAnnot va
    ABS va                  -> "ABS" <+> renderAnnot va
    NEG va                  -> "NEG" <+> renderAnnot va
    LSL va                  -> "LSL" <+> renderAnnot va
    LSR va                  -> "LSR" <+> renderAnnot va
    OR  va                  -> "OR" <+> renderAnnot va
    AND va                  -> "AND" <+> renderAnnot va
    XOR va                  -> "XOR" <+> renderAnnot va
    NOT va                  -> "NOT" <+> renderAnnot va
    COMPARE va              -> "COMPARE" <+> renderAnnot va
    EQ va                   -> "EQ" <+> renderAnnot va
    NEQ va                  -> "NEQ" <+> renderAnnot va
    LT va                   -> "LT" <+> renderAnnot va
    GT va                   -> "GT" <+> renderAnnot va
    LE va                   -> "LE" <+> renderAnnot va
    GE va                   -> "GE" <+> renderAnnot va
    INT va                  -> "INT" <+> renderAnnot va
    VIEW va name ty         -> "VIEW" <+> renderAnnot va <+> renderViewName name <+> renderTy ty
    SELF va fa              -> "SELF" <+> renderAnnots [] [fa] [va]
    CONTRACT va fa t        -> "CONTRACT" <+> renderAnnots [] [fa] [va] <+> renderTy t
    TRANSFER_TOKENS va      -> "TRANSFER_TOKENS" <+> renderAnnot va
    SET_DELEGATE va         -> "SET_DELEGATE" <+> renderAnnot va
    CREATE_CONTRACT va1 va2 contract -> let
      body = enclose space space $ align $ (renderDoc doesntNeedParens contract)
      in "CREATE_CONTRACT" <+> renderAnnots [] [] [va1, va2] <$$> (indent 2 $ braces $ body)
    IMPLICIT_ACCOUNT va     -> "IMPLICIT_ACCOUNT" <+> renderAnnot va
    NOW va                  -> "NOW" <+> renderAnnot va
    AMOUNT va               -> "AMOUNT" <+> renderAnnot va
    BALANCE va              -> "BALANCE" <+> renderAnnot va
    VOTING_POWER va         -> "VOTING_POWER" <+> renderAnnot va
    TOTAL_VOTING_POWER va   -> "TOTAL_VOTING_POWER" <+> renderAnnot va
    CHECK_SIGNATURE va      -> "CHECK_SIGNATURE" <+> renderAnnot va
    SHA256 va               -> "SHA256" <+> renderAnnot va
    SHA512 va               -> "SHA512" <+> renderAnnot va
    BLAKE2B va              -> "BLAKE2B" <+> renderAnnot va
    SHA3 va                 -> "SHA3" <+> renderAnnot va
    KECCAK va               -> "KECCAK" <+> renderAnnot va
    HASH_KEY va             -> "HASH_KEY" <+> renderAnnot va
    PAIRING_CHECK va        -> "PAIRING_CHECK" <+> renderAnnot va
    SOURCE va               -> "SOURCE" <+> renderAnnot va
    SENDER va               -> "SENDER" <+> renderAnnot va
    ADDRESS va              -> "ADDRESS" <+> renderAnnot va
    CHAIN_ID va             -> "CHAIN_ID" <+> renderAnnot va
    LEVEL va                -> "LEVEL" <+> renderAnnot va
    SELF_ADDRESS va         -> "SELF_ADDRESS" <+> renderAnnot va
    NEVER                   -> "NEVER"
    TICKET va               -> "TICKET" <+> renderAnnot va
    READ_TICKET va          -> "READ_TICKET" <+> renderAnnot va
    SPLIT_TICKET va         -> "SPLIT_TICKET" <+> renderAnnot va
    JOIN_TICKETS va         -> "JOIN_TICKETS" <+> renderAnnot va
    OPEN_CHEST va           -> "OPEN_CHEST" <+> renderAnnot va
    SAPLING_EMPTY_STATE va n -> "SAPLING_EMPTY_STATE" <+> renderAnnot va <+> (integer $ toInteger n)
    SAPLING_VERIFY_UPDATE va -> "SAPLING_VERIFY_UPDATE" <+> renderAnnot va
    MIN_BLOCK_TIME anns      -> "MIN_BLOCK_TIME" <+> renderAnyAnns anns
    where
      renderTy = renderDoc @Ty needsParens
      renderComp = renderDoc @Ty needsParens
      renderOps = renderOpsList False

      renderAnnot :: KnownAnnTag tag => Annotation tag -> Doc
      renderAnnot = renderDoc doesntNeedParens . singleAnnSet

      renderAnnots :: [TypeAnn] -> [FieldAnn] -> [VarAnn] -> Doc
      renderAnnots ts fs vs = renderDoc doesntNeedParens $ fullAnnSet ts fs vs

  isRenderable = \case
    EXT extInstr -> isRenderable extInstr
    _ -> True

instance (RenderDoc op, Buildable op) => Buildable (InstrAbstract op) where
  build = \case
    EXT ext -> build ext
    mi -> buildRenderDoc mi

----------------------------------------------------------------------------
-- JSON serialization
----------------------------------------------------------------------------

deriveJSON morleyAesonOptions ''InstrAbstract
deriveJSON morleyAesonOptions ''ExpandedOp