-- 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 = error "impossible"

instance ( ElemHasNamedVar (ty : s) name var
             (VarNamePretty ty == 'VarNamed name)
         ) =>
         HasNamedVar (ty : s) name var where
  varPosition = 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 (SS SZ)

instance (HasNamedVar s name var) =>
         ElemHasNamedVar (ty : s) name var 'False where
  elemVarPosition = case varPosition @s @name @var of
    VarPosition n -> VarPosition (SS 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 _ =
  dupNPeano @n
    \\ provideConstraintUnsafe @(Take (Decrement n) s ++ (var : Drop n s) ~ s)
    \\ 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 _ =
  case varPosition @s @name @var of
    VarPosition sn -> dupLUnsafe 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 l = dupLNamed l # fromNamed l

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

_dupSample1
  :: [Integer, Natural, "a" :! (), "b" :! MText, "c" :! ByteString]
  :-> [MText, Integer, Natural, "a" :! (), "b" :! MText, "c" :! ByteString]
_dupSample1 =
  dupL #b

_dupSample2
  :: (HasNamedVar s "x" Natural)
  => (Integer : s)
     :-> (Natural : Integer : s)
_dupSample2 =
  dupL #x

_dupSample3
  :: (HasNamedVar s "x" Natural, VarIsUnnamed store)
  => (store : s)
     :-> (Natural : store : s)
_dupSample3 =
  dupL #x