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

-- | Parsing of built-in Michelson macros.

module Morley.Michelson.Parser.Macro
  ( macro
  , allMacros
  , allSingleTokenMacros
  -- * These are handled separately to have better error messages
  , duupMac
  , diipMac
  , pairMac
  , ifCmpMac
  , mapCadrMac
  , cadrMac
  , carnMac
  , cdrnMac
  ) where

import Prelude hiding (note, try)

import Data.Default (Default(..))
import Text.Megaparsec (choice, label, skipMany)
import Text.Megaparsec.Char (string)
import Text.Megaparsec.Char.Lexer (decimal)

import Morley.Michelson.Macro
  (CadrStruct(..), Macro(..), PairStruct(..), ParsedOp(..), ParsedSeq, UnpairStruct(..))
import Morley.Michelson.Macro qualified as Macro
import Morley.Michelson.Parser.Annotations
import Morley.Michelson.Parser.Instr qualified as Instr
import Morley.Michelson.Parser.Lexer
import Morley.Michelson.Parser.Types (Parser)
import Morley.Michelson.Untyped (noAnn)

data AllowAnnotations = ForbidAnnotations | AllowAnnotations

annWithDef :: Default a => AllowAnnotations -> Parser a -> Parser a
annWithDef allow p = case allow of
  AllowAnnotations -> p
  ForbidAnnotations -> pure def

cmpOp :: AllowAnnotations -> Parser Macro.ParsedInstr
cmpOp = Instr.cmpOp . flip annWithDef noteDef

-- | Macros consisting of a single token are allowed in more contexts than other
-- macros.
singleTokenMacro
  :: AllowAnnotations
  -- ^ Whether we allow annotations. When used for parsing a single token, this
  -- is obviously disabled. Enabled when reused in the general context, e.g. in 'macro'.
  -> Parser Macro
singleTokenMacro allowAnns = choice
  [ setCadrMac allowAnns
  , string "CMP" *> (CMP <$> cmpOp allowAnns)
  , word "FAIL" FAIL
  , string "ASSERT_CMP" *> (ASSERT_CMP <$> cmpOp allowAnns)
  , word "ASSERT_NONE" ASSERT_NONE
  , word "ASSERT_SOME" ASSERT_SOME
  , word "ASSERT_LEFT" ASSERT_LEFT
  , word "ASSERT_RIGHT" ASSERT_RIGHT
  , string "ASSERT_" *> (ASSERTX <$> cmpOp allowAnns)
  , word "ASSERT" ASSERT
  , unpairMac allowAnns
  ]

-- | This includes all macros consisting of a single token, i.e. also the
-- specially-handled ones, like @PAIIR@, @CADR@, etc, but not @DUUP@, as it's
-- always represented as @DUPN@ and thus behaves like a primitive.
allSingleTokenMacros :: Parser Macro
allSingleTokenMacros = label "single-token macro" $ choice
  [ singleTokenMacro ForbidAnnotations
  , pairMac' ForbidAnnotations
  , cadrMac' ForbidAnnotations
  ]

-- | Includes all macros, also the specially-handled ones. Macros are allowed in
-- more contexts than instructions, so sometimes it's unambiguous.
--
-- Notably @DUUP@ and @DIIP@ are excluded, as they're always represented by a
-- single instruction and thus behave more like a primitive.
allMacros :: Parser (ParsedSeq ParsedOp) -> Parser Macro
allMacros ops = label "macro" $ choice
  [ macro ops
  , pairMac
  , carnMac
  , cdrnMac
  , cadrMac
  , ifCmpMac ops
  , mapCadrMac ops
  , ifX ops
  ]

macro :: Parser (ParsedSeq ParsedOp) -> Parser Macro
macro ops = label "macro"
   $  singleTokenMacro AllowAnnotations
  <|> word "IF_SOME" IF_SOME <*> ops <*> ops
  <|> word "IF_RIGHT" IF_RIGHT <*> ops <*> ops

-- | This might look strange, but at the time of writing, reference
-- implementation requires braced argument to @DIIP@, even if the argument is a
-- macro. Thus, @DIP 2 FAIL@ is fine, but @DIIP FAIL@ is a syntax error.
diipMac :: Parser [ParsedOp] -> Parser Macro
diipMac ops =
  (string "DI" *> (DIIP . succ <$> num "I") <* symbol1 "P") <*> (Macro.PSSequence <$> ops)

num :: Text -> Parser Word
num str = length <$> some (string str)

duupMac :: Parser Macro
duupMac = duupMac' AllowAnnotations

duupMac' :: AllowAnnotations -> Parser Macro
duupMac' allowAnns =
  (string "DU" *> (DUUP . succ <$> num "U") <* symbol1 "P") <*> annWithDef allowAnns noteDef

pairMacInner :: Parser PairStruct
pairMacInner = do
  string "P"
  l <- (string "A" $> F noAnn) <|> pairMacInner
  r <- (string "I" $> F noAnn) <|> pairMacInner
  return $ P l r

pairMac :: Parser Macro
pairMac = pairMac' AllowAnnotations

pairMac' :: AllowAnnotations -> Parser Macro
pairMac' allowAnns = do
  a <- pairMacInner
  symbol1 "R"
  (tn, vn, fns) <- annWithDef allowAnns $ permute3Def noteDef note (some note)
  let ps = Macro.mapPairLeaves fns a
  return $ PAPAIR ps tn vn

unpairMacInner :: Parser UnpairStruct
unpairMacInner = do
  string "P"
  l <- (string "A" $> UF) <|> unpairMacInner
  r <- (string "I" $> UF) <|> unpairMacInner
  return $ UP l r

unpairMac :: AllowAnnotations -> Parser Macro
unpairMac allowAnns =
  string "UN" *> (UNPAPAIR <$> unpairMacInner) <* symbol1 "R" <*
    annWithDef allowAnns (skipMany $ void noteF <|> void noteV)

cadrMac :: Parser Macro
cadrMac = cadrMac' AllowAnnotations

cadrMac' :: AllowAnnotations -> Parser Macro
cadrMac' allowAnns = do
  string "C"
  as <- (:) <$> cadrInner <*> some cadrInner
  symbol1 "R"
  (vn, fn) <- annWithDef allowAnns notesVF
  return $ CADR as vn fn

cadrInner :: Parser CadrStruct
cadrInner = (string "A" $> A) <|> (string "D" $> D)

carnMac :: Parser Macro
carnMac = symbol1 "CAR" *> (CARN <$> noteDef <*> lexeme decimal)

cdrnMac :: Parser Macro
cdrnMac = symbol1 "CDR" *> (CDRN <$> noteDef <*> lexeme decimal)

{-# ANN module ("HLint: ignore Reduce duplication" :: Text) #-}
setCadrMac :: AllowAnnotations -> Parser Macro
setCadrMac allowAnns = do
  string "SET_C"
  a <- some cadrInner
  symbol1 "R"
  (v, f) <- annWithDef allowAnns notesVF
  return $ SET_CADR a v f

mapCadrMac :: Parser (ParsedSeq ParsedOp) -> Parser Macro
mapCadrMac ops = do
  string "MAP_C"
  a <- some cadrInner
  symbol1 "R"
  (v, f) <- notesVF
  MAP_CADR a v f <$> ops

ifCmpMac :: Parser (ParsedSeq ParsedOp) -> Parser Macro
ifCmpMac ops = string "IFCMP" $> IFCMP <*> cmpOp ForbidAnnotations <*> ops <*> ops

ifX :: Parser (ParsedSeq ParsedOp) -> Parser Macro
ifX ops = symbol "IF" $> IFX <*> cmpOp ForbidAnnotations <*> ops <*> ops