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

{-# LANGUAGE FunctionalDependencies #-}

-- | Referenced-by-name versions of some instructions.
--
-- They allow to "dig" into stack or copy elements of stack referring them
-- by label.
module Lorentz.ReferencedByName
  ( -- * Constraints
    HasNamedVar
  , HasNamedVars
  , (:=)

    -- * Instructions
  , dupL
  , dupLNamed

    -- * Other
  , VarIsUnnamed
  ) where

import Data.Constraint ((\\))
import Data.Singletons (Sing)
import Named (NamedF(..))

import Lorentz.ADT
import Lorentz.Base
import Lorentz.Coercions
import Lorentz.Instr
import Michelson.Text
import Michelson.Typed
import Util.Label
import Util.Peano
import Util.Type
import Util.TypeLits

-- Errors
----------------------------------------------------------------------------

type family StackElemNotFound name :: ErrorMessage where
  StackElemNotFound name =
    'Text "Element with name `" ':<>: 'ShowType name ':<>:
    'Text "` is not present on stack"

data NamedVariableNotFound name

-- Helpers
----------------------------------------------------------------------------

-- | Name of a variable on stack.
data VarNamed = VarNamed Symbol | VarUnnamed | VarNameDummy

-- | Get variable name.
type family VarName (a :: Type) :: VarNamed where
  VarName (NamedF _ _ name) = 'VarNamed name
  VarName _ = 'VarUnnamed

type family AnyVN :: VarNamed

-- Attach an error message to given variable.
-- If its evaluation is stuck, compiler's attempt to display this type
-- will cause the given error being displayed.
type family Assert (vn :: VarNamed) (err :: Constraint) :: VarNamed where
  Assert 'VarNameDummy _ = AnyVN
  Assert vn _ = vn

type family VarNamePretty' (x :: Type) (vn :: VarNamed) :: VarNamed where
  VarNamePretty' x vn = Assert vn
    (TypeError
      ('Text "Not clear which name `" ':<>: 'ShowType x ':<>: 'Text "` variable has" ':$$:
       'Text "Consider adding `VarIsUnnamed " ':<>: 'ShowType x ':<>: 'Text "` constraint" ':$$:
       'Text "or carrying a named variable instead"
      )
    )

-- | 'VarName' with pretty error message.
type VarNamePretty x = VarNamePretty' x (VarName x)

-- | Requires type @x@ to be an unnamed variable.
--
-- When e.g. 'dupL' sees a polymorphic variable, it can't judge whether
-- is it a variable we are seeking for or not; @VarIsUnnamed@ helps to
-- assure the type system that given variable won't be named.
type VarIsUnnamed x = VarName x ~ 'VarUnnamed

----------------------------------------------------------------------------
-- Dup
----------------------------------------------------------------------------

-- | Indicates that stack @s@ contains a @name :! var@ or @name :? var@ value.

{- Implementation notes:
We intentially keep this typeclass as simple as possible so that if a user has

@
myFunc :: Integer : s :-> Integer : s
myFunc = something that requires @"globalVar" :! Integer@
@

then user gets clear @Missing `HasEnv s globalVar Integer` instance@ error message,
and can easily add this constraint to his methods.

-}
class HasNamedVar (s :: [Type]) (name :: Symbol) (var :: Type)
    | s name -> var where

  -- | 1-based position of the variable on stack.
  varPosition :: VarPosition s name var

type ConstraintVarPosition s n =
  ( SingI n, KnownPeano n, n > 'Z ~ 'True
  , RequireLongerOrSameLength s n, RequireLongerOrSameLength (ToTs s) n
  )

data VarPosition (s :: [Type]) (name :: Symbol) (var :: Type) where
  VarPosition
    :: (ConstraintVarPosition s n)
    => (Sing (n :: Nat))
    -> VarPosition s name var

instance ( TypeError (StackElemNotFound name)
         , var ~ NamedVariableNotFound name
         ) =>
         HasNamedVar '[] name var where
  varPosition :: VarPosition '[] name var
varPosition = Text -> VarPosition '[] name var
forall a. HasCallStack => Text -> a
error Text
"impossible"

instance ( ElemHasNamedVar (ty : s) name var
             (VarNamePretty ty == 'VarNamed name)
         ) =>
         HasNamedVar (ty : s) name var where
  varPosition :: VarPosition (ty : s) name var
varPosition = ElemHasNamedVar
  (ty : s) name var (VarNamePretty ty == 'VarNamed name) =>
VarPosition (ty : s) name var
forall (s :: [*]) (name :: Symbol) var (nameMatch :: Bool).
ElemHasNamedVar s name var nameMatch =>
VarPosition s name var
elemVarPosition @(ty : s) @name @var @(VarNamePretty ty == 'VarNamed name)

-- Helper for handling each separate variable on stack
class ElemHasNamedVar s name var (nameMatch :: Bool)
    | s name nameMatch -> var where
  elemVarPosition :: VarPosition s name var

instance (ty ~ NamedF f var name) =>
         ElemHasNamedVar (ty : s) name var 'True where
  elemVarPosition :: VarPosition (ty : s) name var
elemVarPosition = Sing ('S 'Z) -> VarPosition (ty : s) name var
forall (s :: [*]) (n :: Nat) (name :: Symbol) var.
ConstraintVarPosition s n =>
Sing n -> VarPosition s name var
VarPosition (SingNat 'Z -> SingNat ('S 'Z)
forall (n1 :: Nat).
(SingI n1, KnownPeano n1) =>
SingNat n1 -> SingNat ('S n1)
SS SingNat 'Z
SZ)

instance (HasNamedVar s name var) =>
         ElemHasNamedVar (ty : s) name var 'False where
  elemVarPosition :: VarPosition (ty : s) name var
elemVarPosition = case HasNamedVar s name var => VarPosition s name var
forall (s :: [*]) (name :: Symbol) var.
HasNamedVar s name var =>
VarPosition s name var
varPosition @s @name @var of
    VarPosition Sing n
n -> Sing ('S n) -> VarPosition (ty : s) name var
forall (s :: [*]) (n :: Nat) (name :: Symbol) var.
ConstraintVarPosition s n =>
Sing n -> VarPosition s name var
VarPosition (SingNat n -> SingNat ('S n)
forall (n1 :: Nat).
(SingI n1, KnownPeano n1) =>
SingNat n1 -> SingNat ('S n1)
SS Sing n
SingNat n
n)

-- | Version of 'HasNamedVar' for multiple variables.
--
-- >>> type HasContext = HasNamedVars s ["x" := Integer, "f" := Lambda MText MText]
type family HasNamedVars (s :: [Type]) (vs :: [NamedField]) :: Constraint where
  HasNamedVars _ '[] = ()
  HasNamedVars s ((n := ty) ': vs) = (HasNamedVar s n ty, HasNamedVars s vs)

-- | Get the variable at @n@-th position on stack, assuming that caller is sure
-- that stack is long enough.
--
-- @martoon: I'm not ready to fight the compiler regarding numerous
-- complex constraints, so just assuring it that those constraints will hold.
--
-- @heitor.toledo: GHC doesn't seem to enjoy it that we use 'ConstraintDUPNLorentz',
-- so I replaced it with the expanded definition.
dupLUnsafe
  :: forall n s var.
     (ConstraintVarPosition s n)
  => Sing (n :: Nat) -> s :-> var : s
dupLUnsafe :: Sing n -> s :-> (var : s)
dupLUnsafe Sing n
_ =
  forall (inp :: [*]) (out :: [*]) a.
ConstraintDUPNLorentz n inp out a =>
inp :-> out
forall (n :: Nat) (inp :: [*]) (out :: [*]) a.
ConstraintDUPNLorentz n inp out a =>
inp :-> out
dupNPeano @n
    (((Take (Decrement n) s ++ (var : Drop n s)) ~ s) =>
 s :-> (var : s))
-> Dict ((Take (Decrement n) s ++ (var : Drop n s)) ~ s)
-> s :-> (var : s)
forall (c :: Constraint) e r. HasDict c e => (c => r) -> e -> r
\\ MockableConstraint
  ((Take (Decrement n) s ++ (var : Drop n s)) ~ s) =>
Dict ((Take (Decrement n) s ++ (var : Drop n s)) ~ s)
forall (c :: Constraint). MockableConstraint c => Dict c
provideConstraintUnsafe @(Take (Decrement n) s ++ (var : Drop n s) ~ s)
    (((Take (Decrement n) (ToTs s) ++ (ToT var : Drop n (ToTs s)))
  ~ ToTs s) =>
 s :-> (var : s))
-> Dict
     ((Take (Decrement n) (ToTs s) ++ (ToT var : Drop n (ToTs s)))
      ~ ToTs s)
-> s :-> (var : s)
forall (c :: Constraint) e r. HasDict c e => (c => r) -> e -> r
\\ MockableConstraint
  ((Take (Decrement n) (ToTs s) ++ (ToT var : Drop n (ToTs s)))
   ~ ToTs s) =>
Dict
  ((Take (Decrement n) (ToTs s) ++ (ToT var : Drop n (ToTs s)))
   ~ ToTs s)
forall (c :: Constraint). MockableConstraint c => Dict c
provideConstraintUnsafe @(Take (Decrement n) (ToTs s) ++ (ToT var : Drop n (ToTs s)) ~ ToTs s)

-- | Version of 'dupL' that leaves a named variable on stack.
dupLNamed
  :: forall var name s.
     (HasNamedVar s name var)
  => Label name -> s :-> (name :! var) : s
dupLNamed :: Label name -> s :-> ((name :! var) : s)
dupLNamed Label name
_ =
  case HasNamedVar s name var => VarPosition s name var
forall (s :: [*]) (name :: Symbol) var.
HasNamedVar s name var =>
VarPosition s name var
varPosition @s @name @var of
    VarPosition Sing n
sn -> Sing n -> s :-> ((name :! var) : s)
forall (n :: Nat) (s :: [*]) var.
ConstraintVarPosition s n =>
Sing n -> s :-> (var : s)
dupLUnsafe Sing n
sn

-- | Take the element with given label on stack and copy it on top.
--
-- If there are multiple variables with given label, the one closest
-- to the top of the stack is picked.
dupL :: forall var name s.
        (HasNamedVar s name var)
     => Label name -> s :-> var : s
dupL :: Label name -> s :-> (var : s)
dupL Label name
l = Label name -> s :-> ((name :! var) : s)
forall var (name :: Symbol) (s :: [*]).
HasNamedVar s name var =>
Label name -> s :-> ((name :! var) : s)
dupLNamed Label name
l (s :-> ((name :! var) : s))
-> (((name :! var) : s) :-> (var : s)) -> s :-> (var : s)
forall (a :: [*]) (b :: [*]) (c :: [*]).
(a :-> b) -> (b :-> c) -> a :-> c
# Label name -> ((name :! var) : s) :-> (var : s)
forall (name :: Symbol) a (s :: [*]).
Label name -> (NamedF Identity a name : s) :-> (a : s)
fromNamed Label name
l

-- Samples
----------------------------------------------------------------------------

_dupSample1
  :: [Integer, Natural, "a" :! (), "b" :! MText, "c" :! ByteString]
  :-> [MText, Integer, Natural, "a" :! (), "b" :! MText, "c" :! ByteString]
_dupSample1 :: '[Integer, Natural, "a" :! (), "b" :! MText, "c" :! ByteString]
:-> '[MText, Integer, Natural, "a" :! (), "b" :! MText,
      "c" :! ByteString]
_dupSample1 =
  Label "b"
-> '[Integer, Natural, "a" :! (), "b" :! MText, "c" :! ByteString]
   :-> '[MText, Integer, Natural, "a" :! (), "b" :! MText,
         "c" :! ByteString]
forall var (name :: Symbol) (s :: [*]).
HasNamedVar s name var =>
Label name -> s :-> (var : s)
dupL IsLabel "b" (Label "b")
Label "b"
#b

_dupSample2
  :: (HasNamedVar s "x" Natural)
  => (Integer : s)
     :-> (Natural : Integer : s)
_dupSample2 :: (Integer : s) :-> (Natural : Integer : s)
_dupSample2 =
  Label "x" -> (Integer : s) :-> (Natural : Integer : s)
forall var (name :: Symbol) (s :: [*]).
HasNamedVar s name var =>
Label name -> s :-> (var : s)
dupL IsLabel "x" (Label "x")
Label "x"
#x

_dupSample3
  :: (HasNamedVar s "x" Natural, VarIsUnnamed store)
  => (store : s)
     :-> (Natural : store : s)
_dupSample3 :: (store : s) :-> (Natural : store : s)
_dupSample3 =
  Label "x" -> (store : s) :-> (Natural : store : s)
forall var (name :: Symbol) (s :: [*]).
HasNamedVar s name var =>
Label name -> s :-> (var : s)
dupL IsLabel "x" (Label "x")
Label "x"
#x