-- 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 ((\\))
import Data.Singletons (Sing)

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 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, 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 :: Peano)
    -> 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.
unsafeDupL
  :: forall n s var.
     (ConstraintVarPosition s n, Dupable var)
  => Sing (n :: Peano) -> s :-> var : s
unsafeDupL _ =
  dupNPeano @n
    \\ unsafeProvideConstraint @(Take (Decrement n) s ++ (var : Drop n s) ~ s)
    \\ unsafeProvideConstraint @(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, Dupable var)
  => Label name -> s :-> (name :! var) : s
dupLNamed Label{} =
  case varPosition @s @name @var of
    VarPosition sn -> unsafeDupL 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 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