{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DerivingVia #-}
{-# LANGUAGE ExplicitNamespaces #-}
{-# LANGUAGE ImpredicativeTypes #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE RoleAnnotations #-}
{-# LANGUAGE UnboxedTuples #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE UnliftedNewtypes #-}
{-# OPTIONS_GHC -Wno-redundant-constraints #-}
{-# OPTIONS_HADDOCK hide #-}

module Control.Monad.Borrow.Pure.Lifetime.Token.Internal (
  module Control.Monad.Borrow.Pure.Lifetime.Token.Internal,
) where

import Control.Functor.Linear qualified as Control
import Control.Monad.Borrow.Pure.Affine.Internal
import Control.Monad.Borrow.Pure.Lifetime.Internal
import Data.Coerce.Directed.Unsafe
import Data.Functor.Linear qualified as Data
import Data.Kind (Constraint)
import Data.Unrestricted.Linear
import GHC.Base (TYPE, UnliftedType, noinline, withDict)
import GHC.Exts qualified as GHC
import GHC.Stack (HasCallStack)
import Unsafe.Linear qualified as Unsafe

type role Now nominal

-- | Witness that the lifetime @α@ is ongoing.
data Now (α :: Lifetime) = UnsafeNow

data SomeNow where
  MkSomeNow :: Now (Al i) %1 -> SomeNow

newLifetime :: Linearly %1 -> SomeNow
newLifetime :: Linearly %1 -> SomeNow
newLifetime Linearly
UnsafeLinearly = Now (Al (ZonkAny 0)) -> SomeNow
forall (i :: Nat). Now (Al i) -> SomeNow
MkSomeNow Now (Al (ZonkAny 0))
forall (α :: Lifetime). Now α
UnsafeNow

newLifetime' :: Linearly %1 -> (forall ι. Now (Al ι) %1 -> a) %1 -> a
newLifetime' :: forall a.
Linearly %1 -> (forall (ι :: Nat). Now (Al ι) %1 -> a) %1 -> a
newLifetime' Linearly
lin forall (ι :: Nat). Now (Al ι) %1 -> a
k =
  case Linearly %1 -> SomeNow
newLifetime Linearly
lin of
    MkSomeNow Now (Al i)
now -> Now (Al i) %1 -> a
forall (ι :: Nat). Now (Al ι) %1 -> a
k Now (Al i)
now

-- | Static Lifetime is always available.
nowStatic :: Now Static
nowStatic :: Now Static
nowStatic = Now Static
forall (α :: Lifetime). Now α
UnsafeNow

instance Affine (Now α) where
  aff :: Now α %1 -> Aff (Now α)
aff Now α
UnsafeNow = Now α -> Aff (Now α)
forall a. a -> Aff a
UnsafeAff Now α
forall (α :: Lifetime). Now α
UnsafeNow
  {-# INLINE aff #-}

instance LinearOnly (Now α) where
  linearOnly :: LinearOnlyWitness (Now α)
linearOnly = LinearOnlyWitness (Now α)
forall {k} (a :: k). LinearOnlyWitness a
UnsafeLinearOnly
  {-# INLINE linearOnly #-}

type role EndToken nominal

-- | Witness that the lifetime @α@ has ended.
data EndToken (α :: Lifetime) = UnsafeEnd

instance (α >= β) => EndToken α <: EndToken β where
  subtype :: SubtypeWitness (EndToken α) (EndToken β)
subtype = SubtypeWitness (EndToken α) (EndToken β)
forall {k} {k1} (a :: k) (b :: k1). SubtypeWitness a b
UnsafeSubtype

endLifetime :: Now (Al i) %1 -> (Ur (EndToken (Al i)))
endLifetime :: forall (i :: Nat). Now (Al i) %1 -> Ur (EndToken (Al i))
endLifetime Now (Al i)
UnsafeNow = EndToken (Al i) -> Ur (EndToken (Al i))
forall a. a -> Ur a
Ur EndToken (Al i)
forall (α :: Lifetime). EndToken α
UnsafeEnd

-- | Witness that the lifetime @α@ has ended.
class End (α :: Lifetime) where
  endToken :: EndToken α

-- | Static lifetime lasts forever.
neverEnds :: (HasCallStack, End Static) => a
neverEnds :: forall a. (HasCallStack, End Static) => a
neverEnds = [Char] -> a
forall a. HasCallStack => [Char] -> a
error [Char]
"Unreachable: if you see this, you created an End Static in the internal code!"

{- |
Utility type to represent an object available after the lifetime @α@.

You can use 'Control.Applicative' and 'Control.Monad' instances to write 'After' conveniently.
-}
newtype After α a = After ((End α) => a)

instance (α <= β, a <: b) => After α a <: After β b where
  subtype :: SubtypeWitness (After α a) (After β b)
subtype = SubtypeWitness (After α a) (After β b)
forall {k} {k1} (a :: k) (b :: k1). SubtypeWitness a b
UnsafeSubtype

unAfter :: (End α) => After α a %1 -> a
{-# INLINE unAfter #-}
unAfter :: forall (α :: Lifetime) a. End α => After α a %1 -> a
unAfter (After End α => a
r) = a
End α => a
r

withEnd :: forall α r. EndToken α -> After α r %1 -> r
{-# INLINE withEnd #-}
withEnd :: forall (α :: Lifetime) r. EndToken α -> After α r %1 -> r
withEnd EndToken α
end (After End α => r
a) = ((End α => r) -> r) %1 -> (End α => r) %1 -> r
forall a b (p :: Multiplicity) (x :: Multiplicity).
(a %p -> b) %1 -> a %x -> b
Unsafe.toLinear (forall (cls :: Constraint) meth r.
WithDict cls meth =>
meth -> (cls => r) -> r
withDict @(End α) EndToken α
end) r
End α => r
a

instance Data.Functor (After α) where
  fmap :: forall a b. (a %1 -> b) -> After α a %1 -> After α b
fmap a %1 -> b
f (After End α => a
r) = (End α => b) -> After α b
forall (α :: Lifetime) a. (End α => a) -> After α a
After (a %1 -> b
f a
End α => a
r)
  {-# INLINE fmap #-}

instance Control.Functor (After α) where
  fmap :: forall a b. (a %1 -> b) %1 -> After α a %1 -> After α b
fmap a %1 -> b
f (After End α => a
r) = (End α => b) -> After α b
forall (α :: Lifetime) a. (End α => a) -> After α a
After (a %1 -> b
f a
End α => a
r)
  {-# INLINE fmap #-}

instance Data.Applicative (After α) where
  pure :: forall a. a -> After α a
pure a
a = (End α => a) -> After α a
forall (α :: Lifetime) a. (End α => a) -> After α a
After a
End α => a
a
  {-# INLINE pure #-}
  After End α => a %1 -> b
f <*> :: forall a b. After α (a %1 -> b) %1 -> After α a %1 -> After α b
<*> After End α => a
r = (End α => b) -> After α b
forall (α :: Lifetime) a. (End α => a) -> After α a
After (a %1 -> b
End α => a %1 -> b
f a
End α => a
r)
  {-# INLINE (<*>) #-}

instance Control.Applicative (After α) where
  pure :: forall a. a %1 -> After α a
pure a
a = (End α => a) -> After α a
forall (α :: Lifetime) a. (End α => a) -> After α a
After a
End α => a
a
  {-# INLINE pure #-}
  After End α => a %1 -> b
f <*> :: forall a b. After α (a %1 -> b) %1 -> After α a %1 -> After α b
<*> After End α => a
r = (End α => b) -> After α b
forall (α :: Lifetime) a. (End α => a) -> After α a
After (a %1 -> b
End α => a %1 -> b
f a
End α => a
r)
  {-# INLINE (<*>) #-}

instance Control.Monad (After α) where
  After End α => a
r >>= :: forall a b. After α a %1 -> (a %1 -> After α b) %1 -> After α b
>>= a %1 -> After α b
k = (End α => b) -> After α b
forall (α :: Lifetime) a. (End α => a) -> After α a
After (After α b %1 -> b
forall (α :: Lifetime) a. End α => After α a %1 -> a
unAfter (a %1 -> After α b
k a
End α => a
r))
  {-# INLINE (>>=) #-}

-- | Witness that the current computation is in a linear context.
data Linearly = UnsafeLinearly

linearly :: (Movable a) => (Linearly %1 -> a) %1 -> a
{-# NOINLINE linearly #-}
linearly :: forall a. Movable a => (Linearly %1 -> a) %1 -> a
linearly = ((Linearly %1 -> a) %1 -> a) -> (Linearly %1 -> a) %1 -> a
forall a. a -> a
GHC.noinline \Linearly %1 -> a
f ->
  case a %1 -> Ur a
forall a. Movable a => a %1 -> Ur a
move (Linearly %1 -> a
f Linearly
UnsafeLinearly) of
    Ur !a
x -> a
x

data LinearOnlyWitness a = UnsafeLinearOnly

-- | A (non-bottom) value of the type @a@ can only live in a linear context.
type LinearOnly :: forall rep. TYPE rep -> Constraint
class LinearOnly a where
  linearOnly :: LinearOnlyWitness a

withLinearly :: (LinearOnly a) => a %1 -> (Linearly, a)
{-# NOINLINE withLinearly #-}
withLinearly :: forall a. LinearOnly a => a %1 -> (Linearly, a)
withLinearly = (a %1 -> (Linearly, a)) -> a %1 -> (Linearly, a)
forall a. a -> a
noinline \ !a
a -> (Linearly
UnsafeLinearly, a
a)

withLinearly# :: forall (a :: UnliftedType). (LinearOnly a) => a %1 -> (# Linearly, a #)
withLinearly# :: forall (a :: UnliftedType).
LinearOnly a =>
a %1 -> (# Linearly, a #)
withLinearly# = (a %1 -> (# Linearly, a #)) -> a %1 -> (# Linearly, a #)
forall a. a -> a
noinline \ !a
a -> (# Linearly
UnsafeLinearly, a
a #)

instance LinearOnly Linearly where
  linearOnly :: LinearOnlyWitness Linearly
linearOnly = LinearOnlyWitness Linearly
forall {k} (a :: k). LinearOnlyWitness a
UnsafeLinearOnly
  {-# INLINE linearOnly #-}

instance Consumable Linearly where
  consume :: Linearly %1 -> ()
consume = \Linearly
UnsafeLinearly -> ()
  {-# INLINE consume #-}

instance Dupable Linearly where
  -- NOTE: without inlining, GHC optimizer (especially, full-laziness and demand analysis)
  -- can eliminate duplicated 'Linearly's too eagerly, ruining the state-threading,
  -- and result in resource corruption in some cases.
  -- Such optimization can manifest when, for example, one duplicates 'Linearly'
  -- tokens multiple times and feed them to different allocation functions.
  -- Although we are not able to detect the exact situation, but we believe that
  -- GHC optimizer then eliminates every invocation on bulk allocation functions
  -- into a single one, which introduces unintended reuse of linear resources.
  -- Hence, we must instruct GHC not to inline this function and force
  dup2 :: Linearly %1 -> (Linearly, Linearly)
dup2 = (Linearly %1 -> (Linearly, Linearly))
-> Linearly %1 -> (Linearly, Linearly)
forall a. a -> a
GHC.noinline \Linearly
UnsafeLinearly -> (Linearly
UnsafeLinearly, Linearly
UnsafeLinearly)
  {-# NOINLINE dup2 #-}