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

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

  -- * Morley Parsed value types
  , ParsedValue

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

    -- * For utilities
  , expandContract
  , expandValue

    -- * For parsing
  , mapPairLeaves
  , mapUnpairLeaves

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

import Data.Aeson.TH (deriveJSON)
import Data.Data (Data(..))
import qualified Data.Text as T
import Fmt (Buildable(build), genericF, (+|), (+||), (|+), (||+))
import qualified Text.PrettyPrint.Leijen.Text as PP (empty)

import Michelson.ErrorPos
import Michelson.Printer (RenderDoc(..))
import Michelson.Untyped
import Util.Aeson
import Util.Generic
import Util.Positive

-- | A programmer-defined macro
data LetMacro = LetMacro
  { lmName :: T.Text
  , lmSig :: StackFn
  , lmExpr :: [ParsedOp]
  } deriving stock (Eq, Show, Data, Generic)

instance Buildable LetMacro where
  build = genericF

instance NFData LetMacro

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

instance NFData PairStruct

instance Buildable PairStruct where
  build = genericF

data UnpairStruct
  = UF (VarAnn, FieldAnn)
  | UP UnpairStruct UnpairStruct
  deriving stock (Eq, Show, Data, Generic)

instance NFData UnpairStruct

instance Buildable UnpairStruct where
  build = genericF

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

instance NFData CadrStruct

instance Buildable CadrStruct where
  build = genericF

-- | 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
  | LMac LetMacro    SrcPos -- ^ User-defined macro with instructions to be inlined
  | Seq [ParsedOp]   SrcPos -- ^ A sequence of instructions
  deriving stock (Eq, Show, Data, Generic)

-- dummy value
instance RenderDoc ParsedOp where
  renderDoc _ _ = PP.empty

instance Buildable ParsedOp where
  build = \case
    Prim parseInstr _ -> "<Prim: "+|parseInstr|+">"
    Mac macro _       -> "<Mac: "+|macro|+">"
    LMac letMacro _   -> "<LMac: "+|letMacro|+">"
    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
  = CASE (NonEmpty [ParsedOp])
  | TAG Natural (NonEmpty Type)
  | ACCESS Natural Positive
  | SET Natural Positive
  | CONSTRUCT (NonEmpty [ParsedOp])
  | VIEW [ParsedOp]
  | VOID [ParsedOp]
  | CMP ParsedInstr VarAnn
  | IFX ParsedInstr [ParsedOp] [ParsedOp]
  | IFCMP ParsedInstr VarAnn [ParsedOp] [ParsedOp]
  | FAIL
  | PAPAIR PairStruct TypeAnn VarAnn
  | UNPAIR UnpairStruct
  | CADR [CadrStruct] VarAnn FieldAnn
  | 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
    TAG idx ty -> "<TAG: #"+||idx||+" from "+|toList ty|+""
    CASE parsedInstrs -> "<CASE: "+|toList parsedInstrs|+">"
    ACCESS idx size -> "<ACCESS: #"+||idx||+"/"+|size|+""
    SET idx size -> "<SET: #"+||idx||+"/"+|size|+""
    CONSTRUCT parsedInstrs -> "<CONSTRUCT: "+|toList parsedInstrs|+">"
    VIEW code -> "<VIEW: "+|code|+">"
    VOID code -> "<VOID: "+|code|+">"
    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|+">"
    UNPAIR pairStruct -> "<UNPAIR: "+|pairStruct|+">"
    CADR cadrStructs varAnn fieldAnn -> "<CADR: "+|cadrStructs|+", "+|varAnn|+", "+|fieldAnn|+">"
    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 [])

-- | Expand all macros in parsed contract.
expandContract :: Contract' ParsedOp -> Contract
expandContract contract =
  contract { contractCode = expandList (contractCode contract) }

-- Probably, some SYB can be used here
expandValue :: ParsedValue -> Value
expandValue = \case
  ValuePair l r -> ValuePair (expandValue l) (expandValue r)
  ValueLeft x -> ValueLeft (expandValue x)
  ValueRight x -> ValueRight (expandValue x)
  ValueSome x -> ValueSome (expandValue x)
  ValueNil -> ValueNil
  ValueSeq valueList -> ValueSeq (map expandValue valueList)
  ValueMap eltList -> ValueMap (map expandElt eltList)
  ValueLambda opList ->
    maybe ValueNil ValueLambda $
    nonEmpty (expandList $ toList opList)
  x -> fmap (expand []) x

expandElt :: Elt ParsedOp -> Elt ExpandedOp
expandElt (Elt l r) = Elt (expandValue l) (expandValue r)

expand :: LetCallStack -> ParsedOp -> ExpandedOp
expand cs = let ics pos = InstrCallStack cs pos in \case
  -- 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.
  (Mac (PAPAIR (P (F a) (F b)) t v) pos) ->
    WithSrcEx (ics pos) $ PrimEx (PAIR t v a b)
  -- DIIP is now always represented as a single instruction.
  -- `expandMacro` always returns a list which we wrap into `SeqEx`, so we
  -- can't use it.
  -- As the above comment says, we need to do something better eventually
  -- (e. g. to avoid `error` usage inside `expandMacro`).
  (Mac (DIIP n ops) pos) ->
    WithSrcEx (ics pos) $ PrimEx (DIPN n (expand cs <$> ops))
  (Mac m pos) -> WithSrcEx (ics pos) $ SeqEx $ expandMacro (ics pos) m
  (Prim i pos) -> WithSrcEx (ics pos) $ PrimEx $ expand cs <$> i
  (Seq s pos) -> WithSrcEx (ics pos) $ SeqEx $ expand cs <$> s
  (LMac l pos) -> expandLetMac l
    where
      expandLetMac :: LetMacro -> ExpandedOp
      expandLetMac LetMacro {..} =
        let newCS = LetName lmName : cs in
        let ics' = InstrCallStack newCS pos in
        WithSrcEx ics' $ PrimEx . EXT . FN lmName lmSig $ expand newCS <$> lmExpr

expandMacro :: InstrCallStack -> Macro -> [ExpandedOp]
expandMacro p@InstrCallStack{icsCallStack=cs,icsSrcPos=macroPos} = \case
  VIEW a             -> expandMacro p (UNPAIR $ UP (UF (noAnn, noAnn)) (UF (noAnn, noAnn))) ++
                        [ PrimEx (DIP $ expandMacro p $ DUUP 2 noAnn) ] ++
                        [ PrimEx $ PAIR noAnn noAnn noAnn noAnn ] ++
                        (expand cs <$> a) ++
                        [ PrimEx (DIP [PrimEx $ AMOUNT noAnn])
                        , PrimEx $ TRANSFER_TOKENS noAnn
                        , PrimEx $ NIL noAnn noAnn (Type TOperation noAnn)
                        , PrimEx $ SWAP
                        , PrimEx $ CONS noAnn
                        , PrimEx $ PAIR noAnn noAnn noAnn noAnn
                        ]
  VOID a             -> expandMacro p (UNPAIR (UP (UF (noAnn, noAnn)) (UF (noAnn, noAnn)))) ++
                        [ PrimEx SWAP
                        , PrimEx $ DIP $ expand cs <$> a
                        , PrimEx SWAP
                        , PrimEx $ EXEC noAnn
                        , PrimEx FAILWITH
                        ]
  CASE ops           -> expandCase (map (expand cs) <$> ops)
  TAG idx uty        -> expandTag idx uty
  ACCESS idx size    -> expandAccess idx size
  SET idx size       -> expandSet idx size
  CONSTRUCT ops      -> expandConstruct (map (expand cs) <$> ops)
  CMP i v            -> [PrimEx (COMPARE v), xo i]
  IFX i bt bf        -> [xo i, PrimEx $ IF (xp bt) (xp bf)]
  IFCMP i v bt bf    -> PrimEx <$> [COMPARE v, expand cs <$> i, IF (xp bt) (xp bf)]
  IF_SOME bt bf      -> [PrimEx (IF_NONE (xp bf) (xp bt))]
  IF_RIGHT bt bf     -> [PrimEx (IF_LEFT (xp bf) (xp bt))]
  FAIL               -> PrimEx <$> [UNIT noAnn noAnn, FAILWITH]
  ASSERT             -> oprimEx $ IF [] (expandMacro p FAIL)
  ASSERTX i          -> [expand cs $ mac $ IFX i [] [mac FAIL]]
  ASSERT_CMP i       -> [expand cs $ mac $ IFCMP i noAnn [] [mac FAIL]]
  ASSERT_NONE        -> oprimEx $ IF_NONE [] (expandMacro p FAIL)
  ASSERT_SOME        -> oprimEx $ IF_NONE (expandMacro p FAIL) []
  ASSERT_LEFT        -> oprimEx $ IF_LEFT [] (expandMacro p FAIL)
  ASSERT_RIGHT       -> oprimEx $ IF_LEFT (expandMacro p FAIL) []
  PAPAIR ps t v      -> expandPapair p ps t v
  UNPAIR ps          -> expandUnpapair p ps
  CADR c v f         -> expandCadr p c v f
  SET_CADR c v f     -> expandSetCadr p c v f
  MAP_CADR c v f ops -> expandMapCadr p c v f ops
  -- We handle DIIP outside.
  DIIP {}             -> error "expandMacro DIIP is unreachable"
  DUUP 1 v           -> oprimEx $ DUP v -- this case should be impossible in practice
  DUUP 2 v           -> PrimEx <$> [DIP [PrimEx $ DUP v], SWAP]
  DUUP n v           -> PrimEx <$> [DIPN (n - 1) [PrimEx $ DUP v], DIG n]
  where
    mac = flip Mac macroPos
    oprimEx = one . PrimEx
    xo = PrimEx . fmap (expand cs)
    xp = fmap (expand cs)

-- the correctness of type-annotation expansion is currently untested, as these
-- expansions are not explicitly documented in the Michelson Specification
expandPapair :: InstrCallStack -> PairStruct -> TypeAnn -> VarAnn -> [ExpandedOp]
expandPapair ics ps t v = case ps of
  P (F a) (F b) -> [PrimEx $ PAIR t v a b]
  P (F a) r     -> PrimEx <$> [ DIP $ expandMacro ics (PAPAIR r noAnn noAnn)
                              , PAIR t v a noAnn]
  P l     (F b) -> expandMacro ics (PAPAIR l noAnn noAnn) ++
                   [PrimEx $ PAIR t v noAnn b]
  P l     r     -> expandMacro ics (PAPAIR l noAnn noAnn) ++
                   [ PrimEx $ DIP $ expandMacro ics (PAPAIR r noAnn noAnn)
                   , PrimEx $ PAIR t v noAnn noAnn]
  F _           -> [] -- Do nothing in this case.
  -- It's impossible from the structure of PairStruct and considered cases above,
  -- but if it accidentally happened let's just do nothing.

expandUnpapair :: InstrCallStack -> UnpairStruct -> [ExpandedOp]
expandUnpapair ics = \case
  UP (UF (v,f)) (UF (w,g)) ->
    PrimEx <$> [ DUP noAnn
               , CAR v f
               , DIP [PrimEx $ CDR w g]
               ]
  UP (UF (v, f)) r ->
    PrimEx <$> [ DUP noAnn
               , CAR v f
               , DIP (PrimEx (CDR noAnn noAnn) : expandMacro ics (UNPAIR r))
               ]
  UP l (UF (v, f)) ->
    map PrimEx [ DUP noAnn
               , DIP [PrimEx $ CDR v f]
               , CAR noAnn noAnn
               ] ++
               expandMacro ics (UNPAIR l)
  UP l r ->
    expandMacro ics unpairOne ++
    [PrimEx $ DIP $ expandMacro ics $ UNPAIR r] ++
    expandMacro ics (UNPAIR l)
  UF _  -> [] -- Do nothing in this case.
  -- It's impossible from the structure of UnpairStruct and considered cases above,
  -- but if it accidentally happened let's just do nothing.
  where
    unpairOne = UNPAIR (UP fn fn)
    fn = UF (noAnn, noAnn)

expandCadr :: InstrCallStack -> [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 noAnn noAnn

cdrNoAnn :: InstrAbstract op
cdrNoAnn = CDR noAnn noAnn

pairNoAnn :: VarAnn -> InstrAbstract op
pairNoAnn v = PAIR noAnn v noAnn noAnn

expandSetCadr :: InstrCallStack -> [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
           CDR (ann "%%") noAnn, SWAP, PAIR noAnn v f (ann "@")]
  [D]   -> [DUP noAnn, CDR noAnn f, DROP,
           -- ↑ These operations just check that the right element of pair has %f
           CAR (ann "%%") noAnn, PAIR noAnn v (ann "@") 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 :: InstrCallStack -> [CadrStruct] -> VarAnn -> FieldAnn -> [ParsedOp] -> [ExpandedOp]
expandMapCadr ics@InstrCallStack{icsCallStack=cls} cs v f ops = case cs of
  []    -> []
  [A]   -> PrimEx <$> [DUP noAnn, cdrNoAnn, DIP [PrimEx $ CAR noAnn f, SeqEx (expand cls <$> ops)], SWAP, pairNoAnn v]
  [D]   -> concat [PrimEx <$> [DUP noAnn, CDR noAnn f], [SeqEx (expand cls <$> ops)], PrimEx <$> [SWAP, carNoAnn, pairNoAnn v]]
  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]

expandCase :: NonEmpty [ExpandedOp] -> [ExpandedOp]
expandCase = mkGenericTree (\_ l r -> one . PrimEx $ IF_LEFT l r)

expandTag :: Natural -> NonEmpty Type -> [ExpandedOp]
expandTag idx unionTy =
  reverse . fst $ mkGenericTree merge (([], ) <$> unionTy)
  where
    merge i (li, lt) (ri, rt) =
      let ty = Type (TOr noAnn noAnn lt rt) noAnn
      in if idx < i
          then (PrimEx (LEFT noAnn noAnn noAnn noAnn rt) : li, ty)
          else (PrimEx (RIGHT noAnn noAnn noAnn noAnn lt) : ri, ty)

expandAccess :: Natural -> Positive -> [ExpandedOp]
expandAccess idx size =
  mkGenericTree merge (replicateNE size [])
  where
    merge i li ri =
      if idx < i
        then PrimEx (CAR noAnn noAnn) : li
        else PrimEx (CDR noAnn noAnn) : ri

expandSet :: Natural -> Positive -> [ExpandedOp]
expandSet idx size =
  PrimEx <$>
  appEndo (mkGenericTree merge (replicateNE size base)) []
  where
    base = pre $ DIP [PrimEx DROP]
    merge i li ri = mconcat $
      if idx < i
        then [ pre $ DIP
                 (map PrimEx [DUP n, DIP [PrimEx $ CDR n n], CAR n n])
             , li
             , pre $ PAIR n n n n
             ]
        else [ pre $ DIP
                 (map PrimEx [DUP n, DIP [PrimEx $ CAR n n], CDR n n])
             , ri
             , pre $ SWAP
             , pre $ PAIR n n n n
             ]
    pre e = Endo (e :)
    n = noAnn

expandConstruct :: NonEmpty [ExpandedOp] -> [ExpandedOp]
expandConstruct ctors =
  appEndo (mkGenericTree merge $ map toBase ctors) []
  where
    toBase ops = Endo (ops ++)
    merge _ li ri =
      mconcat
        [ li
        , pre . PrimEx $ DIP (appEndo ri [])
        , pre . PrimEx $ PAIR noAnn noAnn noAnn noAnn
        ]
    pre e = Endo (e :)

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, [])

mapUnpairLeaves :: [(VarAnn, FieldAnn)] -> UnpairStruct -> UnpairStruct
mapUnpairLeaves fs p = evalState (unpairLeavesST p) fs

unpairLeavesST :: UnpairStruct -> State [(VarAnn, FieldAnn)] UnpairStruct
unpairLeavesST = \case
  (UP l r) -> do
    l' <- unpairLeavesST l
    r' <- unpairLeavesST r
    return $ UP l' r'
  (UF _) -> do
    f <- state getLeaf
    return $ UF f
    where
      getLeaf (a:as) = (a, as)
      getLeaf _      = ((noAnn, noAnn), [])

deriveJSON morleyAesonOptions ''ParsedOp
deriveJSON morleyAesonOptions ''LetMacro
deriveJSON morleyAesonOptions ''PairStruct
deriveJSON morleyAesonOptions ''UnpairStruct
deriveJSON morleyAesonOptions ''CadrStruct
deriveJSON morleyAesonOptions ''Macro