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

module Morley.Michelson.Macro
  (
  -- * Macros types
    CadrStruct (..)
  , PairStruct (..)
  , UnpairStruct (..)
  , Macro (..)

  -- * Morley Parsed value types
  , ParsedValue

  -- * Morley Parsed instruction types
  , ParsedInstr
  , ParsedOp (..)
  , ParsedUExtInstr

    -- * For utilities
  , expandContract
  , expandValue

    -- * For parsing
  , mapPairLeaves

    -- * Internals exported for tests
  , expand
  , expandList
  , expandMacro
  , expandPapair
  , expandUnpapair
  , expandCadr
  , expandSetCadr
  , expandMapCadr
  ) where

import Data.Aeson.TH (deriveJSON)
import Data.Data (Data(..))
import Fmt (Buildable(build), GenericBuildable(..), (+|), (|+))

import Morley.Michelson.ErrorPos
import Morley.Michelson.Printer (RenderDoc(..))
import Morley.Michelson.Untyped
import Morley.Util.Aeson

data PairStruct
  = F FieldAnn
  | P PairStruct PairStruct
  deriving stock (Eq, Show, Data, Generic)
  deriving Buildable via GenericBuildable PairStruct

instance NFData PairStruct

data UnpairStruct
  = UF
  | UP UnpairStruct UnpairStruct
  deriving stock (Eq, Show, Data, Generic)
  deriving Buildable via GenericBuildable UnpairStruct

instance NFData UnpairStruct

data CadrStruct
  = A
  | D
  deriving stock (Eq, Show, Data, Generic)
  deriving Buildable via GenericBuildable CadrStruct

instance NFData CadrStruct

-- | Unexpanded instructions produced directly by the @ops@ parser, which
-- contains primitive Michelson Instructions, inline-able macros and sequences
data ParsedOp
  = Prim ParsedInstr SrcPos -- ^ Primitive Michelson instruction
  | Mac Macro        SrcPos -- ^ Built-in Michelson macro defined by the specification
  | Seq [ParsedOp]   SrcPos -- ^ A sequence of instructions
  deriving stock (Eq, Show, Data, Generic)

instance RenderDoc ParsedOp where
  renderDoc pn parsedOp = renderDoc pn $ expand parsedOp

instance Buildable ParsedOp where
  build = \case
    Prim parseInstr _ -> "<Prim: "+|parseInstr|+">"
    Mac macro _       -> "<Mac: "+|macro|+">"
    Seq parsedOps _   -> "<Seq: "+|parsedOps|+">"

instance NFData ParsedOp

-------------------------------------
-- Types produced by parser
-------------------------------------

type ParsedUExtInstr = ExtInstrAbstract ParsedOp

type ParsedInstr = InstrAbstract ParsedOp

type ParsedValue = Value' ParsedOp

-- | Built-in Michelson Macros defined by the specification
data Macro
  = CMP ParsedInstr VarAnn
  | IFX ParsedInstr [ParsedOp] [ParsedOp]
  | IFCMP ParsedInstr VarAnn [ParsedOp] [ParsedOp]
  | FAIL
  | PAPAIR PairStruct TypeAnn VarAnn
  | UNPAPAIR UnpairStruct
  | CADR [CadrStruct] VarAnn FieldAnn
  | CARN VarAnn Word
  | CDRN VarAnn Word
  | SET_CADR [CadrStruct] VarAnn FieldAnn
  | MAP_CADR [CadrStruct] VarAnn FieldAnn [ParsedOp]
  | DIIP Word [ParsedOp]
  | DUUP Word VarAnn
  | ASSERT
  | ASSERTX ParsedInstr
  | ASSERT_CMP ParsedInstr
  | ASSERT_NONE
  | ASSERT_SOME
  | ASSERT_LEFT
  | ASSERT_RIGHT
  | IF_SOME [ParsedOp] [ParsedOp]
  | IF_RIGHT [ParsedOp] [ParsedOp]
  deriving stock (Eq, Show, Data, Generic)

instance Buildable Macro where
  build = \case
    CMP parsedInstr carAnn -> "<CMP: "+|parsedInstr|+", "+|carAnn|+">"
    IFX parsedInstr parsedOps1 parsedOps2 -> "<IFX: "+|parsedInstr|+", "+|parsedOps1|+", "+|parsedOps2|+">"
    IFCMP parsedInstr varAnn parsedOps1 parsedOps2 -> "<IFCMP: "+|parsedInstr|+", "+|varAnn|+", "+|parsedOps1|+", "+|parsedOps2|+">"
    FAIL -> "FAIL"
    PAPAIR pairStruct typeAnn varAnn -> "<PAPAIR: "+|pairStruct|+", "+|typeAnn|+", "+|varAnn|+">"
    UNPAPAIR pairStruct -> "<UNPAPAIR: "+|pairStruct|+">"
    CADR cadrStructs varAnn fieldAnn -> "<CADR: "+|cadrStructs|+", "+|varAnn|+", "+|fieldAnn|+">"
    CARN varAnn idx -> "<CAR: #"+|idx|+","+|varAnn|+">"
    CDRN varAnn idx -> "<CDR: #"+|idx|+","+|varAnn|+">"
    SET_CADR cadrStructs varAnn fieldAnn -> "<SET_CADR: "+|cadrStructs|+", "+|varAnn|+", "+|fieldAnn|+">"
    MAP_CADR cadrStructs varAnn fieldAnn parsedOps -> "<MAP_CADR: "+|cadrStructs|+", "+|varAnn|+", "+|fieldAnn|+", "+|parsedOps|+">"
    DIIP integer parsedOps -> "<DIIP: "+|integer|+", "+|parsedOps|+">"
    DUUP integer varAnn -> "<DUUP: "+|integer|+", "+|varAnn|+">"
    ASSERT -> "ASSERT"
    ASSERTX parsedInstr -> "<ASSERTX: "+|parsedInstr|+">"
    ASSERT_CMP parsedInstr -> "<ASSERT_CMP: "+|parsedInstr|+">"
    ASSERT_NONE  -> "ASSERT_NONE"
    ASSERT_SOME  -> "ASSERT_SOME"
    ASSERT_LEFT  -> "ASSERT_LEFT"
    ASSERT_RIGHT -> "ASSERT_RIGHT"
    IF_SOME parsedOps1 parsedOps2 -> "<IF_SOME: "+|parsedOps1|+", "+|parsedOps2|+">"
    IF_RIGHT parsedOps1 parsedOps2 -> "<IF_RIGHT: "+|parsedOps1|+", "+|parsedOps2|+">"

instance NFData Macro

expandList :: [ParsedOp] -> [ExpandedOp]
expandList = fmap expand
{-# DEPRECATED expandList "Use 'map expand' instead" #-}

-- | Expand all macros in parsed contract.
expandContract :: Contract' ParsedOp -> Contract
expandContract = fmap expand

-- | Expand all macros in parsed value.
expandValue :: ParsedValue -> Value
expandValue = fmap expand

expand :: ParsedOp -> ExpandedOp
expand = let ics pos = ErrorSrcPos pos in \case
  (Mac m pos) -> WithSrcEx (ics pos) $ either PrimEx SeqEx $ expandMacro' (ics pos) m
  (Prim i pos) -> WithSrcEx (ics pos) $ PrimEx $ expand <$> i
  (Seq s pos) -> WithSrcEx (ics pos) $ SeqEx $ expand <$> s

expandMacro' :: ErrorSrcPos -> Macro -> Either ExpandedInstr [ExpandedOp]
expandMacro' p@ErrorSrcPos{unErrorSrcPos=macroPos} = \case
  -- special cases
  -- We handle this case specially, because it's essentially just PAIR.
  -- It's needed because we have a hack in parser: we parse PAIR as PAPAIR.
  -- We need to do something better eventually.
  PAPAIR (P (F a) (F b)) t v -> Left $ PAIR t v a b
  -- DIIP is now always represented as a single instruction.
  DIIP n ops -> Left $ DIPN n (expand <$> ops)
  -- Similarly to above, DUUP is now always represented as a single instruction.
  DUUP n v -> Left $ DUPN v n

  -- regular cases
  CMP i v            -> Right $ [PrimEx (COMPARE v), xo i]
  IFX i bt bf        -> Right $ [xo i, PrimEx $ IF (xp bt) (xp bf)]
  IFCMP i v bt bf    -> Right $ PrimEx <$> [COMPARE v, expand <$> i, IF (xp bt) (xp bf)]
  IF_SOME bt bf      -> Right $ [PrimEx (IF_NONE (xp bf) (xp bt))]
  IF_RIGHT bt bf     -> Right $ [PrimEx (IF_LEFT (xp bf) (xp bt))]
  FAIL               -> Right $ PrimEx <$> [UNIT noAnn noAnn, FAILWITH]
  ASSERT             -> Right $ oprimEx $ IF [] fail'
  ASSERTX i          -> Right $ [expand $ mac $ IFX i [] [mac FAIL]]
  ASSERT_CMP i       -> Right $ [expand $ mac $ IFCMP i noAnn [] [mac FAIL]]
  ASSERT_NONE        -> Right $ oprimEx $ IF_NONE [] fail'
  ASSERT_SOME        -> Right $ oprimEx $ IF_NONE fail' []
  ASSERT_LEFT        -> Right $ oprimEx $ IF_LEFT [] fail'
  ASSERT_RIGHT       -> Right $ oprimEx $ IF_LEFT fail' []
  PAPAIR ps t v      -> Right $ expandPapair p ps t v
  UNPAPAIR ps        -> Right $ expandUnpapair p ps
  CADR c v f         -> Right $ expandCadr p c v f
  CARN v idx         -> Right $ [PrimEx (GETN v (2 * idx + 1))]
  CDRN v idx         -> Right $ [PrimEx (GETN v (2 * idx))]
  SET_CADR c v f     -> Right $ expandSetCadr p c v f
  MAP_CADR c v f ops -> Right $ expandMapCadr p c v f ops

  where
    fail' = expandMacro p FAIL
    mac = flip Mac macroPos
    oprimEx = one . PrimEx
    xo = PrimEx . fmap expand
    xp = fmap expand

expandMacro :: ErrorSrcPos -> Macro -> [ExpandedOp]
expandMacro = either (pure . PrimEx) id ... expandMacro'

-- | The macro expansion rules below were taken from: https://tezos.gitlab.io/active/michelson.html#syntactic-conveniences
--
-- The correctness of type-annotation expansion is currently untested, as these
-- expansions are not explicitly documented in the Michelson Specification.
expandPapair :: ErrorSrcPos -> PairStruct -> TypeAnn -> VarAnn -> [ExpandedOp]
expandPapair ics ps t v = case ps of
  P (F a) (F b) -> [PrimEx $ PAIR t v a b]

  -- > PA(\right)R / S => DIP ((\right)R) ; PAIR / S
  P (F a) r     -> PrimEx <$> [ DIP $ expandMacro ics (PAPAIR r noAnn noAnn)
                              , PAIR t v a noAnn]

  -- > P(\left)IR / S => (\left)R ; PAIR / S
  P l     (F b) -> expandMacro ics (PAPAIR l noAnn noAnn) ++
                   [PrimEx $ PAIR t v noAnn b]

  -- > P(\left)(\right)R =>  (\left)R ; DIP ((\right)R) ; PAIR / S
  P l     r     -> expandMacro ics (PAPAIR l noAnn noAnn) ++
                   [ PrimEx $ DIP $ expandMacro ics (PAPAIR r noAnn noAnn)
                   , PrimEx $ PAIR t v noAnn noAnn]

  -- It's impossible from the structure of PairStruct and considered cases above,
  -- but if it accidentally happened let's just do nothing.
  F _           -> []

-- | The macro expansion rules below were taken from: https://tezos.gitlab.io/active/michelson.html#syntactic-conveniences
expandUnpapair :: ErrorSrcPos -> UnpairStruct -> [ExpandedOp]
expandUnpapair ics = \case
  UP UF UF ->
    [ PrimEx (UNPAIR noAnn noAnn noAnn noAnn) ]

  -- > UNPA(\right)R / S => UNPAIR ; DIP (UN(\right)R) / S
  UP UF r ->
    PrimEx <$> [ UNPAIR noAnn noAnn noAnn noAnn
               , DIP (expandMacro ics (UNPAPAIR r))
               ]

  -- > UNP(\left)IR / S => UNPAIR ; UN(\left)R / S
  UP l UF ->
    PrimEx (UNPAIR noAnn noAnn noAnn noAnn)
      : expandMacro ics (UNPAPAIR l)

  -- > UNP(\left)(\right)R => UNPAIR ; DIP (UN(\right)R) ; UN(\left)R / S
  UP l r ->
    [ PrimEx (UNPAIR noAnn noAnn noAnn noAnn)
    , PrimEx $ DIP $ expandMacro ics $ UNPAPAIR r
    ] ++
    expandMacro ics (UNPAPAIR l)

  -- It's impossible from the structure of UnpairStruct and considered cases above,
  -- but if it accidentally happened let's just do nothing.
  UF -> []

expandCadr :: ErrorSrcPos -> [CadrStruct] -> VarAnn -> FieldAnn -> [ExpandedOp]
expandCadr ics cs v f = case cs of
  []    -> []
  [A]  -> [PrimEx $ CAR v f]
  [D]  -> [PrimEx $ CDR v f]
  A:css -> PrimEx (CAR noAnn noAnn) : expandMacro ics (CADR css v f)
  D:css -> PrimEx (CDR noAnn noAnn) : expandMacro ics (CADR css v f)

carNoAnn :: InstrAbstract op
carNoAnn = CAR [annQ|%%|] noAnn

cdrNoAnn :: InstrAbstract op
cdrNoAnn = CDR [annQ|%%|] noAnn

pairNoAnn :: VarAnn -> InstrAbstract op
pairNoAnn v = PAIR noAnn v [annQ|@|] [annQ|@|]

expandSetCadr :: ErrorSrcPos -> [CadrStruct] -> VarAnn -> FieldAnn -> [ExpandedOp]
expandSetCadr ics cs v f = PrimEx <$> case cs of
  []    -> []
  [A]   -> [DUP noAnn, CAR noAnn f, DROP,
           -- ↑ These operations just check that the left element of pair has %f
           cdrNoAnn, SWAP, PAIR noAnn v f [annQ|@|]]
  [D]   -> [DUP noAnn, CDR noAnn f, DROP,
           -- ↑ These operations just check that the right element of pair has %f
           carNoAnn, PAIR noAnn v [annQ|@|] f]
  A:css -> [DUP noAnn, DIP (PrimEx carNoAnn : expandMacro ics (SET_CADR css noAnn f)), cdrNoAnn, SWAP, pairNoAnn v]
  D:css -> [DUP noAnn, DIP (PrimEx cdrNoAnn : expandMacro ics (SET_CADR css noAnn f)), carNoAnn, pairNoAnn v]

expandMapCadr :: ErrorSrcPos -> [CadrStruct] -> VarAnn -> FieldAnn -> [ParsedOp] -> [ExpandedOp]
expandMapCadr ics cs v f ops = case cs of
  []    -> []
  [A]   -> PrimEx <$> [DUP noAnn, cdrNoAnn, DIP [PrimEx $ CAR noAnn f, SeqEx (expand <$> ops)], SWAP, PAIR noAnn v f [annQ|@|]]
  [D]   -> concat [PrimEx <$> [DUP noAnn, CDR noAnn f], [SeqEx (expand <$> ops)], PrimEx <$> [SWAP, carNoAnn, PAIR noAnn v [annQ|@|] f]]
  A:css -> PrimEx <$> [DUP noAnn, DIP (PrimEx carNoAnn : expandMacro ics (MAP_CADR css noAnn f ops)), cdrNoAnn, SWAP, pairNoAnn v]
  D:css -> PrimEx <$> [DUP noAnn, DIP (PrimEx cdrNoAnn : expandMacro ics (MAP_CADR css noAnn f ops)), carNoAnn, pairNoAnn v]

mapPairLeaves :: [FieldAnn] -> PairStruct -> PairStruct
mapPairLeaves fs p = evalState (pairLeavesST p) fs

pairLeavesST :: PairStruct -> State [FieldAnn] PairStruct
pairLeavesST = \case
  (P l r) -> do
    l' <- pairLeavesST l
    r' <- pairLeavesST r
    return $ P l' r'
  (F _) -> do
    f <- state getLeaf
    return $ F f
    where
      getLeaf (a:as) = (a, as)
      getLeaf _      = (noAnn, [])

deriveJSON morleyAesonOptions ''PairStruct
deriveJSON morleyAesonOptions ''UnpairStruct
deriveJSON morleyAesonOptions ''CadrStruct

$(mconcat
  [ deriveJSON morleyAesonOptions ''Macro
  , deriveJSON morleyAesonOptions ''ParsedOp
  ])