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

-- NOTE this pragmas.
-- We disable some wargnings for the sake of speed up.
-- Write code with care.
{-# OPTIONS_GHC -Wno-incomplete-patterns #-}
{-# OPTIONS_GHC -Wno-overlapping-patterns #-}

-- | Typical usages of FAILWITH instruction.

module Morley.Michelson.FailPattern
  ( TypicalFailWith
  , typicalFailWithTag
  , isTypicalFailWith
  , modifyTypicalFailWith
  ) where

import Data.Singletons (withSingI)

import Morley.Michelson.Text (MText)
import Morley.Michelson.Typed

-- | This data type captures typical ways to use `FAILWITH` instruction.
-- Each constructor corresponds to a usage pattern.
data TypicalFailWith
  = FailWithString MText
  -- ^ Push a constant string and fail with it.
  | FailWithConstantPair MText SomeConstant
  -- ^ Push a constant pair where the first item is a string and the
  -- second one is an arbitrary constant. Fail afterwards.
  | FailWithStackValue MText
  -- ^ Push a constant string and apply the 'PAIR' instruction, then fail.

-- | Extract error tag out of 'TypicalFailWith'.
typicalFailWithTag :: TypicalFailWith -> MText
typicalFailWithTag = \case
  FailWithString str -> str
  FailWithConstantPair str _ -> str
  FailWithStackValue str -> str

-- | Check whether given instruction ends with a typical 'FAILWITH'
-- usage. It does not recursively check instructions that can be
-- passed to other instructions.
--
-- The instruction MUST be linearized to the left (see 'linearizeLeft').
isTypicalFailWith :: Instr inp out -> Maybe TypicalFailWith
isTypicalFailWith = \case
  Seq i1 FAILWITH -> isTypicalPreFailWith i1
  _ -> Nothing
  where
    isTypicalPreFailWith ::
      Instr inp (a ': out) -> Maybe TypicalFailWith
    isTypicalPreFailWith =
      \case
        PUSH v -> isTypicalErrorConstant v
        Seq _ (PUSH v) -> isTypicalErrorConstant v

        Seq (PUSH v) PAIR{} -> FailWithStackValue <$> isStringValue v
        Seq (Seq _ (PUSH v)) PAIR{} -> FailWithStackValue <$> isStringValue v

        _ -> Nothing

isTypicalErrorConstant ::
  forall t. ConstantScope t => Value t -> Maybe TypicalFailWith
isTypicalErrorConstant v
  | Just str <- isStringValue v = Just (FailWithString str)
  | VPair (VString str, secondItem) <- v =
      -- We need to pattern match to deduce `singI` for second item of the pair.
      case sing @t of
        STPair l r -> withSingI l $ withSingI r $ Just (FailWithConstantPair str (SomeConstant secondItem))
  | otherwise = Nothing

-- | If given instruction ends with a typical 'FAILWITH' usage, modify
-- the tag used there using given transformation function. It can
-- return any value, not necessarily a string.
modifyTypicalFailWith
  :: HasCallStack
  => (MText -> SomeConstant)
  -> Instr inp out
  -> Instr inp out
modifyTypicalFailWith f = modifyTypicalFailWith' . linearizeLeft
  where
    modifyTypicalFailWith' :: HasCallStack => Instr inp out -> Instr inp out
    modifyTypicalFailWith' =
      \case
        Seq i1 FAILWITH ->
          case i1 of
            PUSH v -> onPush v
            Seq i0 (PUSH v) -> Seq i0 (onPush v)

            -- It is quite hard to move the body of these very similar
            -- cases into a separate function because involved types
            -- are quite sophisticated.
            Seq (AnnPUSH ann@(Anns2' s _) v) (AnnPAIR ann')
              | _ :: Value a <- v
              , _ :: Instr (b ': s) ('TPair a b ': s) <- i1 ->
                case sing @('TPair a b) of
                  STPair l r -> withSingI l $ withSingI r $ case isStringValue v of
                    Just (f -> SomeConstant v') ->
                      AnnPUSH (Anns2' s starNotes) v' `Seq` AnnPAIR ann' `Seq` FAILWITH
                    Nothing ->
                      AnnPUSH ann v `Seq` AnnPAIR ann' `Seq` FAILWITH
            Seq (Seq i0 (AnnPUSH ann@(Anns2' s _) v)) (AnnPAIR ann')
              | _ :: Value a <- v
              , _ :: Instr s0 ('TPair a b ': s) <- i1 ->
                case sing @('TPair a b) of
                  STPair l r -> withSingI l $ withSingI r $ Seq i0 $ case isStringValue v of
                    Just (f -> SomeConstant v') ->
                      AnnPUSH (Anns2' s starNotes) v' `Seq` AnnPAIR ann' `Seq` FAILWITH
                    Nothing ->
                      AnnPUSH ann v `Seq` AnnPAIR ann' `Seq` FAILWITH

            _ -> Seq i1 FAILWITH

        i -> i

    onPush ::
      (HasCallStack, ConstantScope v) => Value v -> Instr inp out
    onPush v = case isTypicalErrorConstant v of
      Just (FailWithString (f -> SomeConstant v')) -> PUSH v' `Seq` FAILWITH
      Just (FailWithConstantPair (f -> SomeConstant v') (SomeConstant arg)) ->
        PUSH (VPair (v', arg)) `Seq` FAILWITH
      Just _ -> error "Unexpected TypicalFailWith"
      Nothing -> PUSH v `Seq` FAILWITH