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

    -- * Other
  , VarIsUnnamed
  ) where

import Data.Constraint ((\\))
import qualified Data.Kind as Kind
import Data.Singletons (Sing, SingI)
import Named (NamedF(..))

import Lorentz.ADT
import Lorentz.Base
import Lorentz.Coercions
import Lorentz.Instr
import Michelson.Text
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 :: Kind.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 :: Kind.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 :: [Kind.Type]) (name :: Symbol) (var :: Kind.Type)
    | s name -> var where

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

data VarPosition (s :: [Kind.Type]) (name :: Symbol) (var :: Kind.Type) where
  VarPosition
    :: (SingI n, KnownPeano 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 "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 'Z -> VarPosition (ty : s) name var
forall (n :: Nat) (s :: [*]) (name :: Symbol) var.
(SingI n, KnownPeano n) =>
Sing n -> VarPosition s name var
VarPosition Sing 'Z
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 n :: Sing n
n -> Sing ('S n) -> VarPosition (ty : s) name var
forall (n :: Nat) (s :: [*]) (name :: Symbol) var.
(SingI n, KnownPeano 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 :: [Kind.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.
dupLUnsafe
  :: forall n s var.
     (SingI n, KnownPeano n)
  => Sing (n :: Nat) -> s :-> var : s
dupLUnsafe :: Sing n -> s :-> (var : s)
dupLUnsafe _ =
  -- generate dummy type variables
  case (Proxy Any
forall k (t :: k). Proxy t
Proxy, Proxy Any
forall k (t :: k). Proxy t
Proxy) of
    (Proxy Any
_ :: Proxy s', Proxy Any
_ :: Proxy (s0 :: [Kind.Type])) ->
      ((var : Any) :-> (var : var : Any)) -> s :-> Any
forall (n :: Nat) (inp :: [*]) (out :: [*]) (s :: [*]) (s' :: [*]).
ConstraintDIPNLorentz n inp out s s' =>
(s :-> s') -> inp :-> out
dipNPeano @n ((var : Any) :-> (var : var : Any)
forall a (s :: [*]). (a & s) :-> (a & (a & s))
dup :: var : s0 :-> var : var : s0) (s :-> Any) -> (Any :-> (var : s)) -> s :-> (var : s)
forall (a :: [*]) (b :: [*]) (c :: [*]).
(a :-> b) -> (b :-> c) -> a :-> c
#
      forall (inp :: [*]) (out :: [*]) a.
ConstraintDIGLorentz n inp out a =>
inp :-> out
forall (n :: Nat) (inp :: [*]) (out :: [*]) a.
ConstraintDIGLorentz n inp out a =>
inp :-> out
digPeano @n
        (((SingI n, KnownPeano n,
   RequireLongerOrSameLength
     (Take n (ToTs Any) ++ Drop ('S n) (ToTs Any)) n,
   (Take n (Take n (ToTs Any) ++ Drop ('S n) (ToTs Any))
    ++ (ToT var : ToTs Any))
   ~ (Take n (ToTs Any) ++ Drop ('S n) (ToTs Any)),
   (Take n (Take n (ToTs Any) ++ Drop ('S n) (ToTs Any))
    ++ (ToT var : ToT var : ToTs Any))
   ~ (Take n (ToTs Any) ++ (ToT var : Drop ('S n) (ToTs Any)))),
  (SingI n, KnownPeano n, RequireLongerOrSameLength s n,
   (Take n s ++ (var : Any)) ~ s,
   (Take n s ++ (var : var : Any)) ~ Any)) =>
 s :-> (var : s))
-> Dict
     ((SingI n, KnownPeano n,
       RequireLongerOrSameLength
         (Take n (ToTs Any) ++ Drop ('S n) (ToTs Any)) n,
       (Take n (Take n (ToTs Any) ++ Drop ('S n) (ToTs Any))
        ++ (ToT var : ToTs Any))
       ~ (Take n (ToTs Any) ++ Drop ('S n) (ToTs Any)),
       (Take n (Take n (ToTs Any) ++ Drop ('S n) (ToTs Any))
        ++ (ToT var : ToT var : ToTs Any))
       ~ (Take n (ToTs Any) ++ (ToT var : Drop ('S n) (ToTs Any)))),
      (SingI n, KnownPeano n, RequireLongerOrSameLength s n,
       (Take n s ++ (var : Any)) ~ s,
       (Take n s ++ (var : var : Any)) ~ Any))
-> s :-> (var : s)
forall (c :: Constraint) e r. HasDict c e => (c => r) -> e -> r
\\ MockableConstraint
  (ConstraintDIPNLorentz n s Any (var : Any) (var : var : Any)) =>
Dict (ConstraintDIPNLorentz n s Any (var : Any) (var : var : Any))
forall (c :: Constraint). MockableConstraint c => Dict c
provideConstraintUnsafe @(ConstraintDIPNLorentz n s s' (var : s0) (var : var : s0))
        (((SingI n, KnownPeano n, RequireLongerThan (ToTs Any) n,
   ToTs Any
   ~ (Take n (ToTs Any) ++ (ToT var : Drop ('S n) (ToTs Any))),
   (ToT var : ToTs s)
   ~ (ToT var : (Take n (ToTs Any) ++ Drop ('S n) (ToTs Any)))),
  (SingI n, KnownPeano n, RequireLongerThan Any n,
   Any ~ (Take n Any ++ (var : Drop ('S n) Any)),
   (var : s) ~ (var : (Take n Any ++ Drop ('S n) Any)))) =>
 s :-> (var : s))
-> Dict
     ((SingI n, KnownPeano n, RequireLongerThan (ToTs Any) n,
       ToTs Any
       ~ (Take n (ToTs Any) ++ (ToT var : Drop ('S n) (ToTs Any))),
       (ToT var : ToTs s)
       ~ (ToT var : (Take n (ToTs Any) ++ Drop ('S n) (ToTs Any)))),
      (SingI n, KnownPeano n, RequireLongerThan Any n,
       Any ~ (Take n Any ++ (var : Drop ('S n) Any)),
       (var : s) ~ (var : (Take n Any ++ Drop ('S n) Any))))
-> s :-> (var : s)
forall (c :: Constraint) e r. HasDict c e => (c => r) -> e -> r
\\ MockableConstraint (ConstraintDIGLorentz n Any (var : s) var) =>
Dict (ConstraintDIGLorentz n Any (var : s) var)
forall (c :: Constraint). MockableConstraint c => Dict c
provideConstraintUnsafe @(ConstraintDIGLorentz n s' (var : s) var)

-- | 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.
--
-- TODO: maybe call it duupL for consistency? Or better not? :thinking:
dupL :: forall var name s.
        (HasNamedVar s name var)
     => Label name -> s :-> var : s
dupL :: Label name -> s :-> (var : s)
dupL l :: Label name
l = 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 sn :: Sing n
sn -> Sing n -> s :-> (NamedF Identity var name : s)
forall (n :: Nat) (s :: [*]) var.
(SingI n, KnownPeano n) =>
Sing n -> s :-> (var : s)
dupLUnsafe Sing n
sn (s :-> (NamedF Identity var name : s))
-> ((NamedF Identity var name : s) :-> (var : s))
-> s :-> (var : s)
forall (a :: [*]) (b :: [*]) (c :: [*]).
(a :-> b) -> (b :-> c) -> a :-> c
# Label name -> (NamedF Identity var name : s) :-> (var : s)
forall (name :: Symbol) a (s :: [*]).
Label name -> (NamedF Identity a name : s) :-> (a : s)
fromNamed Label name
l

{- Note about optimizations:

In @duupX@ we have several separate instances in order to optimize
"duupX @1" and "duupX @2" cases, but I'm not sure whether can we do the same
here and at the same time preserve sane error messages
(like "missing `HasNamedVar s name Integer` instance", without mentioning any internal
details).

Since we are already delving into severely deceving the type system here,
I think we'd better optimize those cases via adding respective rules in
our optimizer (that would be difficult to do in an honest way).

-}

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