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

{-# 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 (Bottom(..), (\\))
import Data.Singletons (Sing)
import Fcf qualified
import Type.Errors (DelayError, IfStuck)

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

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

type StackElemNotFound :: Symbol -> k
type StackElemNotFound name = DelayError
  ('Text "Element with name " ':<>: 'ShowType name ':<>: 'Text " is not present on stack")
  -- NB: as @name@ is a Symbol, it'll be double-quoted in ShowType.

data NamedVariableNotFound name

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

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

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

-- | 'VarName' with pretty error message.
type VarNamePretty :: Type -> VarNamed
type VarNamePretty x =
  IfStuck (VarName x)
    (DelayError
      ('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"
      )
    )
    (Fcf.Pure (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.

-}
type HasNamedVar :: [Type] -> Symbol -> Type -> Constraint
class HasNamedVar s name var | s name -> var where
  -- | 1-based position of the variable on stack.
  varPosition :: VarPosition s name var

type ConstraintVarPosition :: [Type] -> Peano -> Constraint
type ConstraintVarPosition s n =
  ( SingI n, n > 'Z ~ 'True
  , RequireLongerOrSameLength s n, RequireLongerOrSameLength (ToTs s) n
  )

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

instance ( Bottom
         , StackElemNotFound name
         , var ~ NamedVariableNotFound name
         ) =>
         HasNamedVar '[] name var where
  varPosition :: VarPosition '[] name var
varPosition = VarPosition '[] name var
forall a. Bottom => a
no

instance ( ElemHasNamedVar (ty : s) name var
             (VarNamePretty ty == 'VarNamed name)
         ) =>
         HasNamedVar (ty : s) name var where
  varPosition :: VarPosition (ty : s) name var
varPosition = 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
type ElemHasNamedVar :: [Type] -> Symbol -> Type -> Bool -> Constraint
class ElemHasNamedVar s name var nameMatch | 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 :: Peano) (name :: Symbol) var.
ConstraintVarPosition s n =>
Sing n -> VarPosition s name var
VarPosition (Sing 'Z -> SingNat ('S 'Z)
forall (n :: Peano). Sing n -> SingNat ('S n)
SS 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 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 :: Peano) (name :: Symbol) var.
ConstraintVarPosition s n =>
Sing n -> VarPosition s name var
VarPosition (Sing n -> SingNat ('S n)
forall (n :: Peano). Sing n -> SingNat ('S n)
SS Sing n
n)

-- | Version of 'HasNamedVar' for multiple variables.
--
-- > type HasContext = HasNamedVars s ["x" := Integer, "f" := Lambda MText MText]
type HasNamedVars :: [Type] -> [NamedField] -> Constraint
type family HasNamedVars s vs 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.
unsafeDupL
  :: forall n s var.
     (ConstraintVarPosition s n, Dupable var)
  => Sing (n :: Peano) -> s :-> var : s
unsafeDupL :: forall (n :: Peano) (s :: [*]) var.
(ConstraintVarPosition s n, Dupable var) =>
Sing n -> s :-> (var : s)
unsafeDupL Sing n
_ =
  forall (n :: Peano) a (inp :: [*]) (out :: [*]).
(ConstraintDUPNLorentz n inp out a, Dupable a) =>
inp :-> out
dupNPeano @n
    (((LazyTake (Decrement n) s ++ (var : Drop n s)) ~ s) =>
 s :-> (var : s))
-> Dict ((LazyTake (Decrement n) s ++ (var : Drop n s)) ~ s)
-> s :-> (var : s)
forall (c :: Constraint) e r. HasDict c e => (c => r) -> e -> r
\\ forall (c :: Constraint). MockableConstraint c => Dict c
unsafeProvideConstraint @(LazyTake (Decrement n) s ++ (var : Drop n s) ~ s)
    (((LazyTake (Decrement n) (ToTs s) ++ (ToT var : Drop n (ToTs s)))
  ~ ToTs s) =>
 s :-> (var : s))
-> Dict
     ((LazyTake (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
\\ forall (c :: Constraint). MockableConstraint c => Dict c
unsafeProvideConstraint @(LazyTake (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, Dupable var)
  => Label name -> s :-> (name :! var) : s
dupLNamed :: forall var (name :: Symbol) (s :: [*]).
(HasNamedVar s name var, Dupable var) =>
Label name -> s :-> ((name :! var) : s)
dupLNamed Label{} =
  case 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 :: Peano) (s :: [*]) var.
(ConstraintVarPosition s n, Dupable var) =>
Sing n -> s :-> (var : s)
unsafeDupL 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, Dupable var)
     => Label name -> s :-> var : s
dupL :: forall var (name :: Symbol) (s :: [*]).
(HasNamedVar s name var, Dupable var) =>
Label name -> s :-> (var : s)
dupL Label name
l = Label name -> s :-> ((name :! var) : s)
forall var (name :: Symbol) (s :: [*]).
(HasNamedVar s name var, Dupable 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 -> ((name :! a) : 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, Dupable 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 :: forall (s :: [*]).
HasNamedVar s "x" Natural =>
(Integer : s) :-> (Natural : Integer : s)
_dupSample2 =
  Label "x" -> (Integer : s) :-> (Natural : Integer : s)
forall var (name :: Symbol) (s :: [*]).
(HasNamedVar s name var, Dupable 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 :: forall (s :: [*]) store.
(HasNamedVar s "x" Natural, VarIsUnnamed store) =>
(store : s) :-> (Natural : store : s)
_dupSample3 =
  Label "x" -> (store : s) :-> (Natural : store : s)
forall var (name :: Symbol) (s :: [*]).
(HasNamedVar s name var, Dupable var) =>
Label name -> s :-> (var : s)
dupL IsLabel "x" (Label "x")
Label "x"
#x