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

When operating on a polymorphic stack, you'll need 'HasNamedVar' constraint:

>>> :{
_foo :: s :-> ("bar" :! Natural) : s
_foo = dupLNamed #bar
--
_bar :: s :-> MText : s
_bar = dupL #baz
:}
...
... No instance for (HasNamedVar s "bar" Natural)
...
... No instance for (HasNamedVar s "baz" MText)
...

>>> :{
_foo :: HasNamedVar s "bar" Natural => s :-> ("bar" :! Natural) : s
_foo = dupLNamed #bar
--
_bar :: HasNamedVar s "bar" Natural => s :-> Natural : s
_bar = dupL #bar
:}

When the stack contains type variables, you may need 'VarIsUnnamed' constraint:

>>> :{
_bar
  :: HasNamedVar s "foo" MText
  => qux : s :-> ("foo" :! MText) : qux : s
_bar = dupLNamed #foo
--
_baz
  :: HasNamedVar s "bar" MText
  => corge : s :-> MText : corge : s
_baz = dupL #bar
:}
...
... Not clear which name `qux` variable has
... Consider adding `VarIsUnnamed qux` constraint
... or carrying a named variable instead
...
... Not clear which name `corge` variable has
...

>>> :{
_bar
  :: (HasNamedVar s "foo" MText, VarIsUnnamed qux)
  => qux : s :-> ("foo" :! MText) : qux : s
_bar = dupLNamed #foo
--
_baz
  :: (HasNamedVar s "bar" MText, VarIsUnnamed corge)
  => corge : s :-> MText : corge : s
_baz = dupL #bar
:}
-}
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.Util.Label
import Morley.Util.Named
import Morley.Util.Peano
import Morley.Util.Type
import Morley.Util.TypeLits

-- $setup
-- >>> import Prelude ()
-- >>> :m +Lorentz Lorentz.Zip Morley.Util.Named Fmt

-- 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@ value.

{- Implementation notes:
We intentionally 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
  -- | 0-based position of the variable on stack.
  varPosition :: VarPosition s name var

type VarPosition :: [Type] -> Symbol -> Type -> Type
data VarPosition s name var where
  VarPosition
    :: (x ~ name :! var, ConstraintDUPNLorentz ('S n) s (x : s) x)
    -- NB: 'ConstraintDUPNLorentz' is a slight overkill here, as it also asserts
    -- @x : s ~ x : s@, which is a triviality. However, introducing a new
    -- constraint synonym without that assertion doesn't seem justified at this
    -- point, and replacing 'ConstraintDUPNLorentz' with its definition is
    -- rather ugly in itself. If you're looking to implement dugNamed or
    -- something like this, you may want to do the legwork of introducing a new
    -- constraint synonym. -- @lierdakil
    => 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
forall a. 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 ~ (name :! var) => ElemHasNamedVar (ty : s) name var 'True where
  elemVarPosition :: VarPosition (ty : s) name var
elemVarPosition = Sing 'Z -> VarPosition (ty : s) name var
forall x (name :: Symbol) var (n :: Peano) (s :: [*]).
(x ~ (name :! var), ConstraintDUPNLorentz ('S n) s (x : s) x) =>
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 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 x (name :: Symbol) var (n :: Peano) (s :: [*]).
(x ~ (name :! var), ConstraintDUPNLorentz ('S n) s (x : s) x) =>
Sing n -> VarPosition s name var
VarPosition (Sing ('S n) -> VarPosition (ty : s) name var)
-> Sing ('S n) -> VarPosition (ty : s) name var
forall a b. (a -> b) -> a -> b
$ 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)

{- | Version of 'dupL' that leaves a named variable on stack.

>>> dupLNamed #foo # pair -$ (#foo :! (123 :: Integer))
(fromLabel @"foo" :! 123,fromLabel @"foo" :! 123)

>>> (dupLNamed #bar # ppaiir) -$ (#foo :! (123 :: Integer)) ::: (#bar :! (321 :: Integer))
(fromLabel @"bar" :! 321,(fromLabel @"foo" :! 123,fromLabel @"bar" :! 321))

>>> (dupLNamed #baz # ppaiir) -$ (#foo :! (123 :: Integer)) ::: (#bar :! (321 :: Integer))
...
... Element with name "baz" is not present 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
_ :: Sing n) -> forall (n :: Peano) a (inp :: [*]) (out :: [*]).
(ConstraintDUPNLorentz n inp out a, Dupable a) =>
inp :-> out
dupNPeano @('S n)

{- | 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 #foo # pair -$ (#foo :! (123 :: Integer))
(123,fromLabel @"foo" :! 123)

>>> (dupL #bar # ppaiir) -$ (#foo :! (123 :: Integer)) ::: (#bar :! (321 :: Integer))
(321,(fromLabel @"foo" :! 123,fromLabel @"bar" :! 321))

>>> (dupL #baz # ppaiir) -$ (#foo :! (123 :: Integer)) ::: (#bar :! (321 :: Integer))
...
... Element with name "baz" is not present on stack
...
-}
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 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 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 Label "x"
#x