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

module Morley.Michelson.TypeCheck.Types
    ( HST (..)
    , (-:&)
    , SomeHST (..)
    , SomeTcInstrOut (..)
    , SomeTcInstr (..)
    , BoundVars (..)
    , withWTPm
    , unsafeWithWTP
    , mapSomeContract
    , mapSomeInstr
    , mapSomeInstrOut
    , noBoundVars
    ) where

import Data.Constraint (Dict(..))
import Data.Map.Lazy qualified as Map
import Data.Singletons (Sing, SingI(..))
import Fmt (Buildable(..), Doc, pretty, (+|), (<+>), (|+))
import Prelude hiding (EQ, GT, LT)

import Morley.Michelson.Typed (SomeContract(..), T(..))
import Morley.Michelson.Typed qualified as T
import Morley.Michelson.Typed.Contract
import Morley.Michelson.Typed.Instr
import Morley.Michelson.Typed.Scope (WellTyped, getWTP)
import Morley.Michelson.Untyped (Ty, Var)
import Morley.Util.Sing (eqParamSing)

-- | 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 '[]
  (::&) :: (T.SingI x, T.SingI xs)
        => (T.SingT x, Dict (WellTyped x))
        -> HST xs
        -> HST (x ': xs)

deriving stock instance Show (HST ts)

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

instance Buildable (HST ts) where
  build SNil = "[]"
  build (r ::& rs) = "[" <+> doRender (r ::& rs) <+> "]"
    where
      doRender :: HST (t ': ts_) -> Doc
      doRender ((t, Dict) ::& (b ::& c)) =
          build t <> "," <+> doRender (b ::& c)
      doRender ((t, Dict) ::& SNil) = build t
infixr 7 ::&

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

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

-- | No-argument type wrapper for @HST@ data type.
data SomeHST where
  SomeHST :: SingI 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 `eqParamSing` hst2

-- | This data type keeps part of type check result - instruction and
-- corresponding output stack.
data SomeTcInstrOut 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.
  (:::)
    :: SingI out
    => Instr inp out
    -> HST out
    -> SomeTcInstrOut 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)
    -> SomeTcInstrOut inp
infix 9 :::

deriving stock instance Show (SomeTcInstrOut inp)

instance Buildable (SomeTcInstrOut inp) where
  build (i ::: out) = i |+ " :: " +| build out
  build (AnyOutInstr i) = i |+ " :: *"

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

deriving stock instance Show (SomeTcInstr inp)

mapSomeInstrOut
  :: (forall out. Instr inp out -> Instr inp' out)
  -> SomeTcInstrOut inp
  -> SomeTcInstrOut inp'
mapSomeInstrOut f (i ::: out) = f i ::: out
mapSomeInstrOut f (AnyOutInstr i) = AnyOutInstr (f i)

mapSomeInstr
  :: (forall out. Instr inp out -> Instr inp out)
  -> SomeTcInstr inp
  -> SomeTcInstr inp
mapSomeInstr f (inp :/ instrAndOut) = inp :/ mapSomeInstrOut f instrAndOut

instance Buildable (SomeTcInstr inp) where
  build (inp :/ out) = inp |+ " -> " +| build out

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

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

noBoundVars :: BoundVars
noBoundVars = BoundVars Map.empty Nothing

-- | 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
  Right Dict -> a
  Left e -> fail (pretty e)

-- | Similar to @withWTPm@ but is meant to be used within tests.
unsafeWithWTP :: forall t a. SingI t => (WellTyped t => a) -> a
unsafeWithWTP fn = case getWTP @t of
  Right Dict -> fn
  Left e -> error $ pretty e