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

{-# LANGUAGE FunctionalDependencies #-}

-- | Referenced-by-type versions of some instructions.
--
-- They allow to "dip" into stack or copy elements of stack referring them
-- by type. Their use is justified, because in most cases there is only
-- one element of each type of stack, and in cases when this does not hold
-- (e.g. entrypoint with multiple parameters of the same type), it might be
-- a good idea to wrap those types into a newtype or to use primitives from
-- @named@ package.
--
-- This module is experimental, i.e. everything here should work but may be
-- removed in favor of better development practices.
--
-- Each instruction is followed with usage example.
module Lorentz.Referenced
  ( DupT(..)
  , DipT(..)
  , dropT
  ) where

import Prelude hiding (drop, swap)

import Data.Eq.Singletons (DefaultEq, DefaultEqSym1)
import Data.List.Singletons qualified as LS
import Data.Type.Bool (If)
import Data.Vinyl.TypeLevel qualified as Peano
import GHC.TypeLits (ErrorMessage(..), TypeError)
import Type.Errors (ShowTypeQuoted)

import Lorentz.Base
import Lorentz.Constraints
import Lorentz.Instr
import Morley.Util.Peano
import Morley.Util.Type (IsElem)

-- $setup
-- >>> import Prelude ()
-- >>> import Lorentz
-- >>> import Lorentz.Zip (zipInstr)
-- >>> import Fmt (pretty, Buildable(..))

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

type StackElemNotFound :: [Type] -> Type -> ErrorMessage
type StackElemNotFound st a =
  'Text "Element of type " ':<>: ShowTypeQuoted a ':<>:
  'Text " is not present on stack" ':$$: 'ShowType st

type StackElemAmbiguous :: [Type] -> Type -> ErrorMessage
type StackElemAmbiguous st a =
  'Text "Ambiguous reference to element of type " ':<>: ShowTypeQuoted a ':<>:
  'Text " for stack" ':$$: 'ShowType st

-- Dup
----------------------------------------------------------------------------

{- | Allows duplicating stack elements referring them by type.

>>> :{
dupSample1 :: [Integer, MText, ()] :-> [MText, Integer, MText, ()]
dupSample1 = dupT @MText
:}

>>> pretty $ dupSample1 # zipInstr -$ 123 ::: [mt|Hello|] ::: ()
Hello : 123 : Hello : ()

>>> :{
dupSample2 :: [Integer, MText, ()] :-> [MText, Integer, MText, ()]
dupSample2 = dupT
:}

>>> pretty $ dupSample2 # zipInstr -$ 123 ::: [mt|Hello|] ::: ()
Hello : 123 : Hello : ()

>>> :{
dupSampleErr1 :: '[] :-> a
dupSampleErr1 = dupT @Bool
:}
...
... • Element of type 'Bool' is not present on stack
...   '[]
...

>>> :{
-- Should fully infer both wildcards
dupSampleErr2 :: _ :-> [Bool, Integer, _, ()]
dupSampleErr2 = dupT
:}
...
... • Found type wildcard ‘_’
...    standing for ‘'[Integer, Bool, ()] :: [*]’
...
... • Found type wildcard ‘_’ standing for ‘Bool’
...   To use the inferred type, enable PartialTypeSignatures
...

>>> :{
-- Should fully infer both wildcards
_dupSampleErr3 :: [Integer, _, ()] :-> (Bool ': _)
_dupSampleErr3 = dupT
:}
...
... • Found type wildcard ‘_’ standing for ‘Bool’
...
... • Found type wildcard ‘_’
...     standing for ‘'[Integer, Bool, ()] :: [*]’
...

-}
class st ~ (Head st ': Tail st) => DupT (a :: Type) (st :: [Type]) where
  -- | Duplicate an element of stack referring it by type.
  --
  -- If stack contains multiple entries of this type, compile error is raised.
  dupT :: st :-> a : st

instance ( EnsureElem a st
         , TheOnlyC (StackElemNotFound st a)
                    (StackElemAmbiguous st a)
                    (LS.FindIndices (DefaultEqSym1 a) st)
                    indexGHC
         , succ_index ~ 'Peano.S (ToPeano indexGHC)
         , ConstraintDUPNLorentz succ_index st (a ': st) a
         , Dupable a
         ) =>
  DupT a st where
  dupT = dupNPeano @succ_index

-- Dip
----------------------------------------------------------------------------

{- | Allows diving into stack referring expected new tip by type.

>>> :{
dipSample1
  :: [Natural, ()] :-> '[()]
  -> [Integer, MText, Natural, ()] :-> [Integer, MText, ()]
dipSample1 = dipT @Natural
:}

>>> pretty $ dipSample1 drop # zipInstr -$ 123 ::: [mt|Hello|] ::: 321 ::: ()
123 : Hello : ()

>>> :{
dipSample2
  :: [Natural, ()] :-> '[()]
  -> [Integer, MText, Natural, ()] :-> [Integer, MText, ()]
dipSample2 = dipT -- No type application needed
:}

>>> pretty $ dipSample2 drop # zipInstr -$ 123 ::: [mt|Hello|] ::: 321 ::: ()
123 : Hello : ()

>>> :{
-- An implementation of dropT that demands a bit more from inference.
dipSample3
  :: forall a inp dinp dout out.
     ( DipT a inp dinp dout out
     , dinp ~ (a ': dout)
     )
  => inp :-> out
dipSample3 = dipT (drop @a)
:}

>>> :{
pretty $ dipSample3 @Natural @'[Integer, MText, Natural, ()] # zipInstr
  -$ 123 ::: [mt|Hello|] ::: 321 ::: ()
:}
123 : Hello : ()

>>> :{
_dipSampleErr1
  :: [Natural, ()] :-> '[()]
  -> [Integer, MText, ()] :-> [Integer, MText, ()]
_dipSampleErr1 = dipT @Natural
:}
...
... • Element of type 'Natural' is not present on stack
...   '[Integer, MText, ()]
...

>>> :{
_dipSampleErr2
  :: [Natural, ()] :-> '[()]
  -> [Integer, MText, Natural, (), Natural] :-> [Integer, MText, ()]
_dipSampleErr2 = dipT @Natural
:}
...
... • Ambiguous reference to element of type 'Natural' for stack
...   '[Integer, MText, Natural, (), Natural]
...

>>> :{
_dipSampleErr3
  :: '[] :-> '[()]
  -> [Integer, MText, Natural, ()] :-> [Integer, MText, ()]
_dipSampleErr3 = dipT @Natural
:}
...
... • dipT requires a Lorentz instruction that takes input on the stack.
...

-}
type DipT :: Type -> [Type] -> [Type] -> [Type] -> [Type] -> Constraint
class dipInp ~ (a ': Tail dipInp) => DipT a inp dipInp dipOut out
  | inp a -> dipInp, dipOut inp a -> out, inp out a -> dipOut where
  -- | Dip down until an element of the given type is on top of the stack.
  --
  -- If the stack does not contain an element of this type, or contains more
  -- than one, then a compile-time error is raised.
  dipT :: (dipInp :-> dipOut) -> (inp :-> out)


{-
-- We'd like to be able to write something like this. Unfortunately, due to GHC
-- issue #22126, this causes numerous copies of the type error to be sprayed to
-- the screen.

instance ( dipInp ~ (a ': tlDI)
         , EnsureElem a inp
         , RequireNonEmpty
             ('Text "dipT requires a Lorentz instruction that takes input on the stack.")
             dipInp
         , index ~ ToPeano (TheOnly (StackElemNotFound inp a)
                                    (StackElemAmbiguous inp a)
                                    (LS.FindIndices (DefaultEqSym1 a) inp))
         , ConstraintDIPNLorentz index inp out dipInp dipOut
         ) =>
  DipT a inp dipInp dipOut out where
  dipT = dipNPeano @index

type family TheOnly (empty_err :: ErrorMessage) (many_err :: ErrorMessage) (xs :: [k]) :: k where
  TheOnly e_err _ '[] = TypeError e_err
  TheOnly _ _ '[x] = x
  TheOnly _ m_err _ = TypeError m_err
-}

instance ( dipInp ~ (a ': tail_dipInp)
         , EnsureElem a inp
         , RequireNonEmpty
             ('Text "dipT requires a Lorentz instruction that takes input on the stack.")
             dipInp
         , TheOnlyC (StackElemNotFound inp a)
                   (StackElemAmbiguous inp a)
                   (LS.FindIndices (DefaultEqSym1 a) inp)
                   indexGHC
         , index ~ ToPeano indexGHC
         , ConstraintDIPNLorentz index inp out dipInp dipOut
         ) =>
  DipT a inp dipInp dipOut out where
  dipT = dipNPeano @index

-- | @EnsureElem x xs@ constrains @x@ to be an element of @xs@. Unlike
-- @'IsElem' x xs ~ 'True@, @EnsureElem x xs@ can help infer @xs@. For
-- example, given @EnsureElem Int '[Char, b, Bool]@, where @b@ is otherwise
-- unknown, GHC will infer that @b ~ Int@. @EnsureElem@ can also increase
-- information about the length of @xs@. For example, given
-- @EnsureElem Int (Char ': more)@, GHC will infer that @more@ has at least
-- one element.
type EnsureElem :: forall k. k -> [k] -> Constraint
type EnsureElem x xs = (xs ~ (Head xs ': Tail xs), EnsureElem' x xs)
type family EnsureElem' x xs where
  EnsureElem' x (y ': ys) =
    ( If (x `DefaultEq` y) (() :: Constraint) (EnsureElem x ys)
    , If (IsElem x ys) (() :: Constraint) (x ~ y))

-- | @TheOnlyC empty_err many_err xs x@ constrains @x@ to be the only
-- element of @xs@. It produces the type error @empty_err@ if @xs@ is empty,
-- and the type error @many_err@ if @xs@ has more than one element.
type TheOnlyC :: ErrorMessage -> ErrorMessage -> [k] -> k -> Constraint
-- NB: This is not a type family because of GHC issue #22126, see above
class TheOnlyC empty_err many_err xs x | xs -> x
instance (TypeError e_err, y ~ Determined) => TheOnlyC e_err m_err '[] y
instance x ~ y => TheOnlyC e_err m_err '[x] y
instance (TypeError m_err, y ~ Determined) => TheOnlyC e_err m_err (x1 ': x2 ': xs) y

type RequireNonEmpty :: ErrorMessage -> [k] -> Constraint
type family RequireNonEmpty e xs where
  RequireNonEmpty e '[] = TypeError e
  RequireNonEmpty _ _ = ()

-- We just use this to satisfy fundeps in error cases.
type family Determined where {}

-- Drop
----------------------------------------------------------------------------

{- | Remove element with the given type from the stack.

>>> :{
dropSample1 :: [Integer, (), Natural] :-> [Integer, Natural]
dropSample1 = dropT @()
:}

>>> pretty $ dropSample1 # zipInstr -$ 123 ::: () ::: 321
123 : 321

>>> :{
dropSampleErr1 :: [Integer, Natural] :-> [Integer, Natural]
dropSampleErr1 = dropT @()
:}
...
... • Element of type '()' is not present on stack
...   '[Integer, Natural]
...

>>> :{
dropSampleErr1 :: [Integer, Integer] :-> '[Integer]
dropSampleErr1 = dropT @Integer
:}
...
... • Ambiguous reference to element of type 'Integer' for stack
...   '[Integer, Integer]
...
-}
dropT
  :: forall a inp dinp dout out.
     ( DipT a inp dinp dout out
     , dinp ~ (a ': dout)
     )
  => inp :-> out
dropT = dipT @a drop

-- Framing
----------------------------------------------------------------------------

{- Note that there instructions are only usable for concrete stacks.

When you know your stack only partially, and you try to refer to element of
type "X", then with the current approach compiler will require the unknown
part of stack to contain no elements of type "X", and this is annoying
at least because it ruins modularity.

This issue can be resolved with using 'framed' instruction wrapper and family.

-}