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

module Michelson.TypeCheck.Types
    ( HST (..)
    , (-:&)
    , pattern (::&+)
    , SomeHST (..)
    , SomeInstrOut (..)
    , SomeInstr (..)
    , SomeContract (..)
    , BoundVars (..)
    , TcExtFrames
    , getWTP
    , withWTPm
    , mapSomeContract
    , noBoundVars
    ) where

import Data.Constraint (Dict(..))
import qualified Data.Map.Lazy as Map
import Data.Singletons (Sing, SingI(..))
import Fmt (Buildable(..), Builder, (+|), (|+))
import Prelude hiding (EQ, GT, LT)
import qualified Text.Show

import Michelson.Typed (Notes(..), T(..), SingT(..), notesT, starNotes)
import Michelson.Typed.Haskell.Value (WellTyped)
import qualified Michelson.Typed as T
import Michelson.Typed.Instr
import Michelson.Untyped (Type, Var, noAnn)
import Michelson.Untyped.Annotation (VarAnn)
import Util.Typeable

-- | Data type holding type information for stack (Heterogeneous Stack Type).
--
-- This data type is used along with instruction data type @Instr@
-- to carry information about its input and output stack types.
--
-- That is, if there is value @instr :: Instr inp out@, along with this
-- @instr@ one may carry @inpHST :: HST inp@ and @outHST :: HST out@ which will
-- contain whole information about input and output stack types for @instr@.
--
-- Data type @HST@ is very similar to @Data.Vinyl.Rec@,
-- but is specialized for a particular purpose.
-- In particular, definition of @HST (t1 ': t2 ': ... tn ': '[])@ requires
-- constraints @(Typeable t1, Typeable t2, ..., Typeable tn)@ as well as
-- constraints @(Typeable '[ t1 ], Typeable '[ t1, t2 ], ...)@.
-- These applications of @Typeable@ class are required for convenient usage
-- of type encoded by @HST ts@ with some functions from @Data.Typeable@.
--
-- Data type @HST@ (Heterogeneous Stack Type) is a heterogenuous list of tuples.
-- First element of tuple is a structure, holding field and type annotations
-- for a given type.
-- Second element of tuple is an optional variable annotation for the stack
-- element.
-- Additionally constructor keeps 'SingI' constraint for the current type.
data HST (ts :: [T])  where
  SNil :: HST '[]
  (::&) :: (Typeable xs, T.KnownT x)
        => (Notes x, Dict (WellTyped x), VarAnn)
        -> HST xs
        -> HST (x ': xs)

instance NFData (HST ts) where
  rnf (SNil) = ()
  rnf ((a, d, b) ::& hst) = rnf (a, d, b, hst)

instance Show (HST ts) where
  show SNil = "[]"
  show (r ::& rs) = "[ " <> showDo (r ::& rs) <> " ]"
    where
      showDo :: HST (t ': ts_) -> String
      showDo ((notesT -> t, Dict, _vn) ::& (b ::& c)) =
          show t <> ", " <> showDo (b ::& c)
      showDo ((notesT -> t, Dict, _vn) ::& SNil) = show t

instance Buildable (HST ts) where
  build SNil = "[]"
  build (r ::& rs) = "[ " +| showDo (r ::& rs) |+ " ]"
    where
      showDo :: HST (t ': ts_) -> Builder
      showDo ((notesT -> t, Dict, _vn) ::& (b ::& c)) =
          build t |+ ", " +| showDo (b ::& c)
      showDo ((notesT -> t, Dict, _vn) ::& SNil) = build t

infixr 7 ::&

instance Eq (HST ts) where
  SNil == SNil = True
  (n1, Dict, a1) ::& h1 == (n2, Dict, a2) ::& h2 =
    n1 == n2 && a1 == a2 && h1 == h2

-- | Append a type to 'HST', assuming that notes and annotations
-- for this type are unknown.
(-:&)
  :: (Typeable xs, WellTyped x)
  => Sing x
  -> HST xs
  -> HST (x ': xs)
_ -:& hst = (starNotes, Dict, noAnn) ::& hst
infixr 7 -:&

-- | Extended pattern-match - adds @Sing x@ argument.
infixr 7 ::&+
pattern (::&+)
  :: ()
  => ( ys ~ (x ': xs)
     , T.KnownT x, Typeable xs
     )
  => (Sing x, Notes x, Dict (WellTyped x), VarAnn)
  -> HST xs
  -> HST ys
pattern x ::&+ hst <- ((\(n, d, v) -> (T.notesSing n, n, d, v)) -> x) ::& hst
  where (_, n, d, v) ::&+ hst = (n, d, v) ::& hst

-- | No-argument type wrapper for @HST@ data type.
data SomeHST where
  SomeHST :: Typeable ts => HST ts -> SomeHST

deriving stock instance Show SomeHST
instance NFData SomeHST where
  rnf (SomeHST h) = rnf h

instance Eq SomeHST where
  SomeHST hst1 == SomeHST hst2 = hst1 `eqParam1` hst2

-- | This data type keeps part of type check result - instruction and
-- corresponding output stack.
data SomeInstrOut inp where
  -- | Type-check result with concrete output stack, most common case.
  --
  -- Output stack type is wrapped inside the type and @Typeable@
  -- constraint is provided to allow convenient unwrapping.
  (:::)
    :: (Typeable out)
    => Instr inp out
    -> HST out
    -> SomeInstrOut inp

  -- | Type-check result which matches against arbitrary output stack.
  -- Information about annotations in the output stack is absent.
  --
  -- This case is only possible when the corresponding code terminates
  -- with @FAILWITH@ instruction in all possible executions.
  -- The opposite may be not true though (example: you push always-failing
  -- lambda and immediatelly execute it - stack type is known).
  AnyOutInstr
    :: (forall out. Instr inp out)
    -> SomeInstrOut inp
infix 9 :::

instance Show (ExtInstr inp) => Show (SomeInstrOut inp) where
  show (i ::: out) = show i <> " :: " <> show out
  show (AnyOutInstr i) = show i <> " :: *"

-- | Data type keeping the whole type check result: instruction and
-- type representations of instruction's input and output.
data SomeInstr inp where
  (:/) :: HST inp -> SomeInstrOut inp -> SomeInstr inp
infix 8 :/

instance Show (ExtInstr inp) => Show (SomeInstr inp) where
  show (inp :/ out) = show inp <> " -> " <> show out

data SomeContract where
  SomeContract :: Contract cp st -> SomeContract

instance NFData SomeContract where
  rnf (SomeContract c) = rnf c

mapSomeContract ::
  (forall inp out. Instr inp out -> Instr inp out)
  -> SomeContract
  -> SomeContract
mapSomeContract f (SomeContract fc) = SomeContract $ mapContractCode f fc

deriving stock instance Show SomeContract

-- | Set of variables defined in a let-block.
data BoundVars = BoundVars (Map Var Type) (Maybe SomeHST)

noBoundVars :: BoundVars
noBoundVars = BoundVars Map.empty Nothing

-- | State for type checking @nop@
type TcExtFrames = [BoundVars]

fromMDict :: Maybe (Dict a) -> (a => Maybe (Dict b)) -> Maybe (Dict b)
fromMDict ma b = ma >>= (\Dict -> b)

-- | Given a type, provide evidence that it is well typed w.r.t to the
--  Michelson rules regarding where comparable types are required.
getWTP :: forall t. SingI t => Maybe (Dict (WellTyped t))
getWTP = case sing @t of
  STKey -> Just Dict
  STUnit -> Just Dict
  STSignature -> Just Dict
  STChainId -> Just Dict
  STOption s -> fromMDict (getWTP_ s) (Just Dict)
  STList s -> fromMDict (getWTP_ s) (Just Dict)
  STSet s -> fromMDict (getWTP_ s) $ fromMDict (T.getComparableProofS s) (Just Dict)
  STOperation  -> Just Dict
  STContract s -> fromMDict (getWTP_ s) (Just Dict)
  STPair s1 s2 -> fromMDict (getWTP_ s1) $ fromMDict (getWTP_ s2) $ Just Dict
  STOr s1 s2 -> fromMDict (getWTP_ s1) $ fromMDict (getWTP_ s2) $ Just Dict
  STLambda s1 s2 -> fromMDict (getWTP_ s1) $ fromMDict (getWTP_ s2) $ Just Dict
  STMap s1 s2 ->
    fromMDict (getWTP_ s1) $
    fromMDict (getWTP_ s2) $
    fromMDict (T.getComparableProofS s1) $ Just Dict
  STBigMap s1 s2 ->
    fromMDict (getWTP_ s1) $
    fromMDict (getWTP_ s2) $
    fromMDict (T.getComparableProofS s1) $ Just Dict
  STInt -> Just Dict
  STNat -> Just Dict
  STString -> Just Dict
  STBytes -> Just Dict
  STMutez -> Just Dict
  STBool -> Just Dict
  STKeyHash -> Just Dict
  STTimestamp -> Just Dict
  STAddress -> Just Dict
  where
    getWTP_ :: forall t1. SingI t1 => Sing t1 -> Maybe (Dict (WellTyped t1))
    getWTP_ _ = getWTP @t1

-- | Given a type and an action that requires evidence that the type is well typed,
--  generate the evidence and execute the action, or else fail with an error.
withWTPm :: forall t m a. (MonadFail m, SingI t) => (WellTyped t => m a) -> m a
withWTPm a = case getWTP @t of
  Just Dict -> a
  Nothing ->
    fail ("This type is not well typed because it has an non-comparable type in it," <>
          "where a comparable type is required")