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

module Morley.Michelson.TypeCheck.Types
    ( HST (..)
    , (-:&)
    , SomeHST (..)
    , SomeInstrOut (..)
    , SomeInstr (..)
    , 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(..), pretty, (+|), (|+))
import Prelude hiding (EQ, GT, LT)
import Text.PrettyPrint.Leijen.Text (Doc, (<+>))

import Morley.Michelson.Printer.Util
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 :: HST ts -> ()
rnf (HST ts
SNil) = ()
  rnf ((SingT x
d, Dict (WellTyped x)
b) ::& HST xs
hst) = (SingT x, Dict (WellTyped x), HST xs) -> ()
forall a. NFData a => a -> ()
rnf (SingT x
d, Dict (WellTyped x)
b, HST xs
hst)

instance Buildable (HST ts) where
  build :: HST ts -> Builder
build = HST ts -> Builder
forall a. RenderDoc a => a -> Builder
buildRenderDocExtended

instance RenderDoc (HST ts) where
  renderDoc :: RenderContext -> HST ts -> Doc
renderDoc RenderContext
_ HST ts
SNil = Doc
"[]"
  renderDoc RenderContext
context ((SingT x, Dict (WellTyped x))
r ::& HST xs
rs) = Doc
"[" Doc -> Doc -> Doc
<+> HST (x : xs) -> Doc
forall (t :: T) (ts_ :: [T]). HST (t : ts_) -> Doc
doRender ((SingT x, Dict (WellTyped x))
r (SingT x, Dict (WellTyped x)) -> HST xs -> HST (x : xs)
forall (out :: T) (xs :: [T]).
(SingI out, SingI xs) =>
(SingT out, Dict (WellTyped out)) -> HST xs -> HST (out : xs)
::& HST xs
rs) Doc -> Doc -> Doc
<+> Doc
"]"
    where
      doRender :: HST (t ': ts_) -> Doc
      doRender :: forall (t :: T) (ts_ :: [T]). HST (t : ts_) -> Doc
doRender ((SingT x
t, Dict (WellTyped x)
Dict) ::& ((SingT x, Dict (WellTyped x))
b ::& HST xs
c)) =
          RenderContext -> SingT x -> Doc
forall a. RenderDoc a => RenderContext -> a -> Doc
renderDoc RenderContext
context SingT x
t Doc -> Doc -> Doc
forall a. Semigroup a => a -> a -> a
<> Doc
"," Doc -> Doc -> Doc
<+> HST (x : xs) -> Doc
forall (t :: T) (ts_ :: [T]). HST (t : ts_) -> Doc
doRender ((SingT x, Dict (WellTyped x))
b (SingT x, Dict (WellTyped x)) -> HST xs -> HST (x : xs)
forall (out :: T) (xs :: [T]).
(SingI out, SingI xs) =>
(SingT out, Dict (WellTyped out)) -> HST xs -> HST (out : xs)
::& HST xs
c)
      doRender ((SingT x
t, Dict (WellTyped x)
Dict) ::& HST xs
SNil) = RenderContext -> SingT x -> Doc
forall a. RenderDoc a => RenderContext -> a -> Doc
renderDoc RenderContext
context SingT x
t
infixr 7 ::&

instance Eq (HST ts) where
  HST ts
SNil == :: HST ts -> HST ts -> Bool
== HST ts
SNil = Bool
True
  (SingT x
n1, Dict (WellTyped x)
Dict) ::& HST xs
h1 == (SingT x
n2, Dict (WellTyped x)
Dict) ::& HST xs
h2 =
    SingT x
n1 SingT x -> SingT x -> Bool
forall a. Eq a => a -> a -> Bool
== SingT x
SingT x
n2 Bool -> Bool -> Bool
forall a. Boolean a => a -> a -> a
&& HST xs
h1 HST xs -> HST xs -> Bool
forall a. Eq a => a -> a -> Bool
== HST xs
HST xs
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)
Sing x
s -:& :: forall (x :: T) (xs :: [T]).
(WellTyped x, SingI xs) =>
Sing x -> HST xs -> HST (x : xs)
-:& HST xs
hst = (Sing x
SingT x
s, Dict (WellTyped x)
forall (a :: Constraint). a => Dict a
Dict) (SingT x, Dict (WellTyped x)) -> HST xs -> HST (x : xs)
forall (out :: T) (xs :: [T]).
(SingI out, SingI xs) =>
(SingT out, Dict (WellTyped out)) -> HST xs -> HST (out : xs)
::& HST xs
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 -> ()
rnf (SomeHST HST ts
h) = HST ts -> ()
forall a. NFData a => a -> ()
rnf HST ts
h

instance Eq SomeHST where
  SomeHST HST ts
hst1 == :: SomeHST -> SomeHST -> Bool
== SomeHST HST ts
hst2 = HST ts
hst1 HST ts -> HST ts -> Bool
forall {k} (a1 :: k) (a2 :: k) (t :: k -> *).
(SingI a1, SingI a2, SDecide k, Eq (t a1)) =>
t a1 -> t a2 -> Bool
`eqParamSing` HST ts
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.
  (:::)
    :: SingI 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 :::

deriving stock instance Show (SomeInstrOut inp)

instance Buildable (SomeInstrOut inp) where
  build :: SomeInstrOut inp -> Builder
build (Instr inp out
i ::: HST out
out) = Instr inp out
i Instr inp out -> Builder -> Builder
forall a b. (Buildable a, FromBuilder b) => a -> Builder -> b
|+ Builder
" :: " Builder -> Builder -> Builder
forall b. FromBuilder b => Builder -> Builder -> b
+| HST out -> Builder
forall p. Buildable p => p -> Builder
build HST out
out
  build (AnyOutInstr forall (out :: [T]). Instr inp out
i) = Instr inp Any
forall (out :: [T]). Instr inp out
i Instr inp Any -> Builder -> Builder
forall a b. (Buildable a, FromBuilder b) => a -> Builder -> b
|+ Builder
" :: *"

-- | 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 :/

deriving stock instance Show (SomeInstr inp)

mapSomeInstrOut
  :: (forall out. Instr inp out -> Instr inp' out)
  -> SomeInstrOut inp
  -> SomeInstrOut inp'
mapSomeInstrOut :: forall (inp :: [T]) (inp' :: [T]).
(forall (out :: [T]). Instr inp out -> Instr inp' out)
-> SomeInstrOut inp -> SomeInstrOut inp'
mapSomeInstrOut forall (out :: [T]). Instr inp out -> Instr inp' out
f (Instr inp out
i ::: HST out
out) = Instr inp out -> Instr inp' out
forall (out :: [T]). Instr inp out -> Instr inp' out
f Instr inp out
i Instr inp' out -> HST out -> SomeInstrOut inp'
forall (out :: [T]) (inp :: [T]).
SingI out =>
Instr inp out -> HST out -> SomeInstrOut inp
::: HST out
out
mapSomeInstrOut forall (out :: [T]). Instr inp out -> Instr inp' out
f (AnyOutInstr forall (out :: [T]). Instr inp out
i) = (forall (out :: [T]). Instr inp' out) -> SomeInstrOut inp'
forall (inp :: [T]).
(forall (out :: [T]). Instr inp out) -> SomeInstrOut inp
AnyOutInstr (Instr inp out -> Instr inp' out
forall (out :: [T]). Instr inp out -> Instr inp' out
f Instr inp out
forall (out :: [T]). Instr inp out
i)

mapSomeInstr
  :: (forall out. Instr inp out -> Instr inp out)
  -> SomeInstr inp
  -> SomeInstr inp
mapSomeInstr :: forall (inp :: [T]).
(forall (out :: [T]). Instr inp out -> Instr inp out)
-> SomeInstr inp -> SomeInstr inp
mapSomeInstr forall (out :: [T]). Instr inp out -> Instr inp out
f (HST inp
inp :/ SomeInstrOut inp
instrAndOut) = HST inp
inp HST inp -> SomeInstrOut inp -> SomeInstr inp
forall (inp :: [T]). HST inp -> SomeInstrOut inp -> SomeInstr inp
:/ (forall (out :: [T]). Instr inp out -> Instr inp out)
-> SomeInstrOut inp -> SomeInstrOut inp
forall (inp :: [T]) (inp' :: [T]).
(forall (out :: [T]). Instr inp out -> Instr inp' out)
-> SomeInstrOut inp -> SomeInstrOut inp'
mapSomeInstrOut forall (out :: [T]). Instr inp out -> Instr inp out
f SomeInstrOut inp
instrAndOut

instance Buildable (SomeInstr inp) where
  build :: SomeInstr inp -> Builder
build (HST inp
inp :/ SomeInstrOut inp
out) = HST inp
inp HST inp -> Builder -> Builder
forall a b. (Buildable a, FromBuilder b) => a -> Builder -> b
|+ Builder
" -> " Builder -> Builder -> Builder
forall b. FromBuilder b => Builder -> Builder -> b
+| SomeInstrOut inp -> Builder
forall p. Buildable p => p -> Builder
build SomeInstrOut inp
out

mapSomeContract ::
  (forall inp out. Instr inp out -> Instr inp out)
  -> SomeContract
  -> SomeContract
mapSomeContract :: (forall (inp :: [T]) (out :: [T]). Instr inp out -> Instr inp out)
-> SomeContract -> SomeContract
mapSomeContract forall (inp :: [T]) (out :: [T]). Instr inp out -> Instr inp out
f (SomeContract Contract cp st
fc) = Contract cp st -> SomeContract
forall (cp :: T) (st :: T). Contract cp st -> SomeContract
SomeContract (Contract cp st -> SomeContract) -> Contract cp st -> SomeContract
forall a b. (a -> b) -> a -> b
$ (forall (inp :: [T]) (out :: [T]). Instr inp out -> Instr inp out)
-> Contract cp st -> Contract cp st
forall (instr :: [T] -> [T] -> *) (cp :: T) (st :: T).
(forall (i :: [T]) (o :: [T]). instr i o -> instr i o)
-> Contract' instr cp st -> Contract' instr cp st
mapContractCode forall (inp :: [T]) (out :: [T]). Instr inp out -> Instr inp out
f Contract cp st
fc

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

noBoundVars :: BoundVars
noBoundVars :: BoundVars
noBoundVars = Map Var Ty -> Maybe SomeHST -> BoundVars
BoundVars Map Var Ty
forall k a. Map k a
Map.empty Maybe SomeHST
forall a. Maybe a
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 :: forall (t :: T) (m :: * -> *) a.
(MonadFail m, SingI t) =>
(WellTyped t => m a) -> m a
withWTPm WellTyped t => m a
a = case forall (t :: T).
SingI t =>
Either NotWellTyped (Dict (WellTyped t))
getWTP @t of
  Right Dict (WellTyped t)
Dict -> m a
WellTyped t => m a
a
  Left NotWellTyped
e -> String -> m a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (NotWellTyped -> String
forall a b. (Buildable a, FromBuilder b) => a -> b
pretty NotWellTyped
e)

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