-- 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 = 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 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 :: [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 _ =
  -- generate dummy type variables
  case (Proxy, Proxy) of
    (_ :: Proxy s', _ :: Proxy (s0 :: [Kind.Type])) ->
      dipNPeano @n (dup :: var : s0 :-> var : var : s0) #
      digPeano @n
        \\ provideConstraintUnsafe @(ConstraintDIPNLorentz n s s' (var : s0) (var : var : s0))
        \\ 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 l = case varPosition @s @name @var of
  VarPosition sn -> dupLUnsafe sn # fromNamed 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 =
  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