-- 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 (..)
  , ParsedSeq (..)
  , ParsedUExtInstr

    -- * For utilities
  , expandContract
  , expandValue

    -- * For parsing
  , mapPairLeaves

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

import Data.Aeson.TH (deriveJSON)
import Data.Data (Data(..))
import Data.Default (def)
import Fmt (Buildable(build), blockListF, (+|), (|+))

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

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

instance NFData PairStruct

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

instance NFData UnpairStruct

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

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
-------------------------------------

data ParsedSeq op
  = PSSingleMacro SrcPos Macro
  | PSSequence [op]
  deriving stock (Eq, Show, Data, Generic, Functor, Foldable)
  deriving anyclass NFData

type ParsedUExtInstr = ExtInstrAbstract ParsedSeq ParsedOp

type ParsedInstr = InstrAbstract ParsedSeq ParsedOp

type ParsedValue = Value' ParsedSeq ParsedOp

-- | Built-in Michelson Macros defined by the specification
data Macro
  = CMP ParsedInstr
  | IFX ParsedInstr (ParsedSeq ParsedOp) (ParsedSeq ParsedOp)
  | IFCMP ParsedInstr (ParsedSeq ParsedOp) (ParsedSeq 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 (ParsedSeq ParsedOp)
  | DIIP Word (ParsedSeq ParsedOp)
  | DUUP Word VarAnn
  | ASSERT
  | ASSERTX ParsedInstr
  | ASSERT_CMP ParsedInstr
  | ASSERT_NONE
  | ASSERT_SOME
  | ASSERT_LEFT
  | ASSERT_RIGHT
  | IF_SOME (ParsedSeq ParsedOp) (ParsedSeq ParsedOp)
  | IF_RIGHT (ParsedSeq ParsedOp) (ParsedSeq ParsedOp)
  deriving stock (Eq, Show, Data, Generic)

instance Buildable op => Buildable (ParsedSeq op) where
  build = \case
    PSSingleMacro _ mac -> build mac
    PSSequence xs -> blockListF xs

instance Buildable Macro where
  build = \case
    CMP parsedInstr -> "<CMP: " +| parsedInstr |+ ">"
    IFX parsedInstr parsedOps1 parsedOps2 -> "<IFX: " +| parsedInstr |+ ", " +| parsedOps1 |+ ", " +| parsedOps2 |+ ">"
    IFCMP parsedInstr parsedOps1 parsedOps2 -> "<IFCMP: " +| parsedInstr |+ ", " +| 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

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

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

expandSeq :: ParsedSeq ExpandedOp -> [ExpandedOp]
expandSeq = \case
  PSSingleMacro pos macro -> expandMacro (ErrorSrcPos pos) macro
  PSSequence xs -> xs

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 $ hoistInstr expandSeq $ 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
  -- DIIP is now always represented as a single instruction.
  DIIP n ops -> Left $ DIPN n $ expandSeq $ 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              -> Right $ [PrimEx (COMPARE def), xo i]
  IFX i bt bf        -> Right $ [xo i, PrimEx $ IF (xp bt) (xp bf)]
  IFCMP i bt bf      -> Right $ PrimEx <$> [COMPARE def, ex 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 (PSSequence []) (PSSequence [mac FAIL])]
  ASSERT_CMP i       -> Right $
    -- reference has weirdly inconsistent nesting between this and IFCMP
    [expand (Mac (CMP i) macroPos), PrimEx $ IF [] 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      -> 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' = one . expand $ mac FAIL
    mac = flip Mac macroPos
    oprimEx = one . PrimEx
    xo = PrimEx . ex
    ex = hoistInstr expandSeq . fmap expand
    xp = expandSeq . 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
  -> Either ExpandedInstr [ExpandedOp]
expandPapair ics@(ErrorSrcPos pos) ps t v = case ps of
  -- 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.
  P (F a) (F b) -> Left $ PAIR t v a b

  -- > PA(\right)R / S => DIP ((\right)R) ; PAIR / S
  -- docs lie a bit, reference actually does something weird, which we try to match
  -- to the best of our ability
  P (F a) r -> Right $
    concatMap papairDips (expandMacro ics (PAPAIR r noAnn noAnn))
    <> [PrimEx $ PAIR t v a noAnn]

  -- > P(\left)IR / S => (\left)R ; PAIR / S
  P l (F b) -> Right [ expand $ Mac (PAPAIR l noAnn noAnn) pos, PrimEx $ PAIR t v noAnn b ]

  -- > P(\left)(\right)R =>  (\left)R ; DIP ((\right)R) ; PAIR / S
  -- docs lie a bit, reference actually does something weird, which we try to match
  -- to the best of our ability
  P l r -> Right $
    (expand $ Mac (PAPAIR l noAnn noAnn) pos)
    : concatMap papairDips (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 _           -> Right []

-- | Mimic the weirdness that is the reference implementation
papairDips :: ExpandedOp -> [ExpandedOp]
papairDips = \case
  PrimEx (DIP xs) -> [PrimEx $ DIPN 2 xs]
  PrimEx (DIPN n xs) -> [PrimEx $ DIPN (succ n) xs]
  SeqEx xs -> concatMap papairDips xs
  PrimEx x -> [PrimEx $ DIP [PrimEx x]]
  WithSrcEx s x -> WithSrcEx s <$> papairDips x

-- | 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
  -- docs lie a bit, reference actually does something weird, which we try to match
  -- to the best of our ability
  UP UF r ->
    PrimEx (UNPAIR noAnn noAnn noAnn noAnn)
    : concatMap papairDips (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
  -- docs lie a bit, reference actually does something weird, which we try to match
  -- to the best of our ability
  UP l r ->
    PrimEx (UNPAIR noAnn noAnn noAnn noAnn)
    : concatMap papairDips (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 f op
carNoAnn = CAR [annQ|%%|] noAnn

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

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

expandSetCadr :: ErrorSrcPos -> [CadrStruct] -> VarAnn -> FieldAnn -> [ExpandedOp]
expandSetCadr (ErrorSrcPos ps) cs v f = PrimEx <$> case cs of
  []    -> []
  [A]   -> [cdrNoAnn, SWAP, PAIR noAnn v f [annQ|@|]]
  [D]   -> [carNoAnn, PAIR noAnn v [annQ|@|] f]
  A:css -> [DUP noAnn, DIP [PrimEx carNoAnn, expand $ Mac (SET_CADR css noAnn f) ps], cdrNoAnn, SWAP, pairNoAnn v]
  D:css -> [DUP noAnn, DIP [PrimEx cdrNoAnn, expand $ Mac (SET_CADR css noAnn f) ps], carNoAnn, pairNoAnn v]

expandMapCadr :: ErrorSrcPos -> [CadrStruct] -> VarAnn -> FieldAnn -> ParsedSeq ParsedOp -> [ExpandedOp]
expandMapCadr (ErrorSrcPos pos) cs v f ops = case cs of
  []    -> []
  -- > MAP_CAR code  =>  DUP ; CDR ; DIP { CAR ; code } ; SWAP ; PAIR
  [A]   -> PrimEx <$>
    [ DUP noAnn
    , cdrNoAnn
    , DIP [PrimEx (CAR noAnn f), SeqEx $ expandSeq $ expand <$> ops]
    , SWAP
    , PAIR noAnn v f [annQ|@|]]
  -- > MAP_CDR code  =>  DUP ; CDR ; code ; SWAP ; CAR ; PAIR
  [D]   -> (PrimEx <$> [DUP noAnn, CDR noAnn f]) <>
    (SeqEx (expandSeq $ expand <$> ops) : (PrimEx <$> [SWAP, carNoAnn, PAIR noAnn v [annQ|@|] f]))
  -- > MAP_CA(\rest=[AD]+)R code / S   =>
  --     { DUP ; DIP { CAR ; MAP_C(\rest)R code } ; CDR ; SWAP ; PAIR } / S
  A:css -> PrimEx <$> [DUP noAnn, DIP [PrimEx carNoAnn, expand $ Mac (MAP_CADR css noAnn f ops) pos], cdrNoAnn, SWAP, pairNoAnn v]
  D:css -> PrimEx <$> [DUP noAnn, DIP [PrimEx cdrNoAnn, expand $ Mac (MAP_CADR css noAnn f ops) pos], 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
  , deriveJSON morleyAesonOptions ''ParsedSeq
  ])