Internal.hs

{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DefaultSignatures #-}
{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE DerivingVia #-}
{-# LANGUAGE ImpredicativeTypes #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE PartialTypeSignatures #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE QualifiedDo #-}
{-# LANGUAGE QuantifiedConstraints #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE RoleAnnotations #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneKindSignatures #-}
{-# LANGUAGE TypeData #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UnboxedTuples #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE UnliftedNewtypes #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE NoImplicitPrelude #-}
{-# OPTIONS_GHC -Wno-name-shadowing #-}
{-# OPTIONS_GHC -Wno-partial-type-signatures #-}
{-# OPTIONS_GHC -Wno-redundant-constraints #-}
{-# OPTIONS_HADDOCK hide #-}

module Control.Monad.Borrow.Pure.BO.Internal (
  module Control.Monad.Borrow.Pure.BO.Internal,
) where

import Control.Exception qualified as SystemIO
import Control.Functor.Linear qualified as Control
import Control.Monad qualified as NonLinear
import Control.Monad.Borrow.Pure.Affine.Internal
import Control.Monad.Borrow.Pure.Lifetime
import Control.Monad.Borrow.Pure.Lifetime.Token
import Control.Monad.Borrow.Pure.Lifetime.Token.Internal
import Control.Monad.Borrow.Pure.Utils (coerceLin)
import Control.Monad.ST.Strict (ST)
import Control.Syntax.DataFlow qualified as DataFlow
import Data.Coerce qualified
import Data.Coerce.Directed.Unsafe
import Data.Functor.Identity (Identity)
import Data.Functor.Linear qualified as Data
import Data.Kind (Type)
import Data.Monoid qualified as Mon
import Data.Ord qualified as Ord
import Data.Semigroup qualified as Sem
import Data.Tuple (Solo (..))
import Data.Type.Equality ((:~:) (Refl))
import GHC.Base (TYPE)
import GHC.Base qualified as GHC
import GHC.Exts (State#, runRW#)
import GHC.ST qualified as ST
import GHC.TypeError (ErrorMessage (..))
import Generics.Linear
import Prelude.Linear
import Prelude.Linear qualified as PL
import Prelude.Linear.Unsatisfiable (Unsatisfiable, unsatisfiable)
import System.IO.Linear qualified as L
import Unsafe.Coerce (unsafeCoerce#)
import Unsafe.Linear qualified as Unsafe

-- NOTE: NOINLINE here is REALLY important, otherwise GHC will inline 'UnsafeLinearly' and common subexpression elimination
-- causes severe soundness bug that the same expression reuses the same
-- linear resource and sometimes SEGV.
askLinearly :: BO α Linearly
{-# NOINLINE askLinearly #-}
askLinearly = GHC.noinline $ Control.pure UnsafeLinearly

asksLinearlyM :: (Linearly %1 -> BO α r) %1 -> BO α r
{-# INLINE asksLinearlyM #-}
asksLinearlyM k = Control.do
  lin <- askLinearly
  !a <- k lin
  Control.pure a

-- NOTE: We want to use @TypeData@ extension for 'ForBO', but it makes Haddock panic!

type ForBO :: Lifetime -> Type
data ForBO α

{- | Computation returning @a@ that can be performed only during the lifetime @α@.
     Internally it is a linear ST monad.
-}
newtype BO α a = BO (State# (ForBO α) %1 -> (# State# (ForBO α), a #))

instance (Semigroup w) => Semigroup (BO α w) where
  (<>) = Control.liftA2 (<>)
  {-# INLINE (<>) #-}

instance (Monoid w) => Monoid (BO α w) where
  mempty = Control.pure mempty
  {-# INLINE mempty #-}

unsafeUnBO :: BO α a %1 -> State# (ForBO α) %1 -> (# State# (ForBO α), a #)
{-# INLINE unsafeUnBO #-}
unsafeUnBO (BO f) = f

assocRBO :: BO ((α /\ β) /\ γ) a %1 -> BO (α /\ (β /\ γ)) a
{-# INLINE assocRBO #-}
assocRBO = unsafeCastBO

assocLBO :: BO (α /\ (β /\ γ)) a %1 -> BO ((α /\ β) /\ γ) a
{-# INLINE assocLBO #-}
assocLBO = unsafeCastBO

assocBOEq :: forall α β γ a. BO ((α /\ β) /\ γ) a :~: BO (α /\ (β /\ γ)) a
{-# INLINE assocBOEq #-}
assocBOEq = Unsafe.coerce $ Refl @(BO (α /\ β /\ γ) a)

instance Data.Functor (BO α) where
  fmap f (BO g) = BO \s -> case g s of
    (# s', a #) -> (# s', f a #)
  {-# INLINE fmap #-}

instance Control.Functor (BO α) where
  fmap f (BO g) = BO \s -> case g s of
    (# s', a #) -> (# s', f a #)
  {-# INLINE fmap #-}

instance Data.Applicative (BO α) where
  pure a = Control.pure a
  {-# INLINE pure #-}

  (<*>) = \f g -> f Control.<*> g
  {-# INLINE (<*>) #-}

instance Control.Applicative (BO α) where
  pure a = BO \s -> (# s, a #)
  {-# INLINE pure #-}

  BO f <*> BO g = BO \s -> case f s of
    (# s', h #) -> case g s' of
      (# s'', a #) -> (# s'', h a #)
  {-# INLINE (<*>) #-}

instance Control.Monad (BO α) where
  BO fa >>= f = BO \s -> case fa s of
    (# s', a #) -> (f a) PL.& \(BO g) -> g s'
  {-# INLINE (>>=) #-}

-- | Unsafely converts a 'BO' computation to linear 'L.IO'.
unsafeBOToLinIO :: BO α a %1 -> L.IO a
{-# INLINE unsafeBOToLinIO #-}
unsafeBOToLinIO (BO f) = L.IO (Unsafe.coerce f)

{- |
Unsafely performs a linear 'L.IO' computation in 'BO' monad.

This is really, really unsafe. If you don't know what you are doing,
you MUST NOT use this function, otherwise you can break purity in a hard way.
-}
unsafeLinIOToBO :: L.IO a %1 -> BO α a
{-# INLINE unsafeLinIOToBO #-}
unsafeLinIOToBO (L.IO f) = BO (Unsafe.coerce f)

runBO# :: forall {rep} α (o :: TYPE rep). (State# (ForBO α) %1 -> o) %1 -> o
{-# INLINE runBO# #-}
runBO# = Unsafe.toLinear \f -> runRW# \s ->
  f (unsafeCoerce# s)

execBO :: BO α a %1 -> Now α %1 -> (Now α, a)
{-# INLINE execBO #-}
execBO (BO f) !now =
  case runBO# f of
    (# s, !a #) -> dropState# s `PL.lseq` (now, a)

dropState# :: State# a %1 -> ()
{-# INLINE dropState# #-}
dropState# = Unsafe.toLinear \ !_ -> ()

-- | See also 'Control.Monad.Borrow.Pure.scope'.
sexecBO :: BO (α /\ β) a %1 -> Now α %1 -> BO β (Now α, a)
{-# INLINE sexecBO #-}
sexecBO f now = unsafeCastBO ((now,) PL.. Unsafe.toLinear (\ !a -> a) Control.<$> f)

{- |
Coerces lifetime in 'BO' computation usafely and brutally.

This is really, really unsafe. If you don't know what you are doing,
you MUST NOT use this function, otherwise you will break the soundness of the type system.
-}
unsafeCastBO :: BO α a %1 -> BO β a
{-# INLINE unsafeCastBO #-}
unsafeCastBO = Unsafe.coerce

-- | Unsafely peforms a 'ST' computation in 'BO' monad.
unsafeSTToBO :: ST s a %1 -> BO α a
{-# INLINE unsafeSTToBO #-}
unsafeSTToBO (ST.ST f) = BO (Unsafe.coerce f)

{- |
Unsafely peforms a 'BO' computation in 'ST' monad.

This is really unsafe. If you don't know what you are doing, you MUST NOT use this function, otherwise you can break purity in a hard way.
-}
unsafeBOToST :: BO α a %1 -> ST s a
{-# INLINE unsafeBOToST #-}
unsafeBOToST (BO f) = ST.ST (Unsafe.coerce f)

{- |
Unsafely performs a standard, non-linear 'IO' computation in 'BO' monad.

This is really, really unsafe. If you don't know what you are doing,
you MUST NOT use this function, otherwise you can break purity in a hard way.
-}
unsafeSystemIOToBO :: IO a %1 -> BO α a
{-# INLINE unsafeSystemIOToBO #-}
unsafeSystemIOToBO (GHC.IO a) = BO (Unsafe.coerce a)

-- | Unsafely performs a 'BO' in the standard, non-linear 'IO' monad.
unsafeBOToSystemIO :: BO α a %1 -> IO a
{-# INLINE unsafeBOToSystemIO #-}
unsafeBOToSystemIO (BO f) = GHC.IO (Unsafe.coerce f)

unsafePerformEvaluateUndupableBO :: BO α a %1 -> a
unsafePerformEvaluateUndupableBO (BO f) = runBO# \s ->
  case Unsafe.toLinear GHC.noDuplicate# s of
    s -> case f s of
      (# s, !a #) -> dropState# s `PL.lseq` a

-- | Run two computations in parallel, returning their results as a tuple.
parBO :: BO α a %1 -> BO α b %1 -> BO α (a, b)
parBO = Unsafe.toLinear2 \a b ->
  BO $
    Unsafe.toLinear \s ->
      case Unsafe.toLinear2 GHC.spark# (case unsafeUnBO a (GHC.noDuplicate# s) of (# _, a #) -> GHC.lazy a) s of
        (# _, a #) ->
          case Unsafe.toLinear2 GHC.spark# (case unsafeUnBO b (GHC.noDuplicate# s) of (# _, b #) -> GHC.lazy b) s of
            (# _, b #) ->
              case Unsafe.toLinear2 GHC.seq# a s of
                (# s, !a #) -> case Unsafe.toLinear2 GHC.seq# b s of
                  (# s, !b #) -> (# s, (a, b) #)

evaluateBO :: a %1 -> BO α a
{-# INLINE evaluateBO #-}
evaluateBO a = unsafeSystemIOToBO (Unsafe.toLinear SystemIO.evaluate a)

-- | Alias of kind 'ak' to a resource of type 'a'.
type Alias :: AliasKind -> Lifetime -> Type -> Type
newtype Alias ak α a = UnsafeAlias a

unsafeUnalias :: Alias ak α a %1 -> a
unsafeUnalias (UnsafeAlias x) = x

type role Alias nominal nominal representational

-- | Alias kind.
data AliasKind
  = -- | Borrower.
    Borrow BorrowKind
  | -- | Lender.
    Lend

-- | Borrower kind.
data BorrowKind
  = -- | Mutable.
    Mut
  | -- | Shared.
    Share

-- | Borrower of kind @bk@ that is active during the lifetime @α@.
type Borrow :: BorrowKind -> Lifetime -> Type -> Type
type Borrow bk = Alias ('Borrow bk)

-- | Mutable borrower, which is affine and can update the data.
type Mut :: Lifetime -> Type -> Type
type Mut = Borrow 'Mut

assocBorrowR ::
  Borrow bk ((α /\ β) /\ γ) a %1 ->
  Borrow bk (α /\ (β /\ γ)) a
{-# INLINE assocBorrowR #-}
assocBorrowR = coerceLin

assocBorrowL ::
  Borrow bk (α /\ (β /\ γ)) a %1 ->
  Borrow bk ((α /\ β) /\ γ) a
{-# INLINE assocBorrowL #-}
assocBorrowL = coerceLin

assocBorrowEq ::
  forall bk α β γ a.
  (Borrow bk ((α /\ β) /\ γ) a) :~: (Borrow bk (α /\ (β /\ γ)) a)
{-# INLINE assocBorrowEq #-}
assocBorrowEq = Unsafe.coerce $ Refl @(Borrow bk (α /\ β /\ γ) a)

assocLendR ::
  Lend ((α /\ β) /\ γ) a %1 ->
  Lend (α /\ (β /\ γ)) a
{-# INLINE assocLendR #-}
assocLendR = coerceLin

assocLendL ::
  Lend (α /\ (β /\ γ)) a %1 ->
  Lend ((α /\ β) /\ γ) a
{-# INLINE assocLendL #-}
assocLendL = coerceLin

assocLendEq :: forall α β γ a. (Lend ((α /\ β) /\ γ) a) :~: (Lend (α /\ (β /\ γ)) a)
{-# INLINE assocLendEq #-}
assocLendEq = Unsafe.coerce $ Refl @(Lend (α /\ β /\ γ) a)

instance (bk ~ 'Mut) => LinearOnly (Borrow bk α a) where
  linearOnly = UnsafeLinearOnly

deriving via AsAffine (Borrow bk α a) instance Consumable (Borrow bk α a)

-- | Shared borrower, which is unrestricted but usually can only read from the data.
type Share :: Lifetime -> Type -> Type
type Share = Borrow 'Share

instance Affine (Borrow bk α a) where
  aff = UnsafeAff
  {-# INLINE aff #-}

instance (k ~ 'Borrow 'Share) => Dupable (Alias k α a) where
  dup2 = Unsafe.toLinear $ NonLinear.join (,)
  {-# INLINE dup2 #-}

instance (k ~ 'Borrow 'Share) => Movable (Alias k α a) where
  move = Unsafe.toLinear Ur
  {-# INLINE move #-}

instance (α >= β, a <: b) => BO α a <: BO β b where
  subtype = UnsafeSubtype

instance (α >= β, a <: b, b <: a) => Mut α a <: Mut β b where
  subtype = UnsafeSubtype

instance (α >= β, a <: b) => Share α a <: Share β b where
  subtype = UnsafeSubtype

-- | Lender, which can retrieve the lifetime at the lifetime @α@.
type Lend :: Lifetime -> Type -> Type
type Lend = Alias 'Lend

instance (α <= β, a <: b) => Lend α a <: Lend β b where
  subtype = UnsafeSubtype

{- |
Borrow a resource linearly and obtain the mutable borrow to it and 'Lend' witness to 'reclaim the resource to lend at the 'End' of the lifetime.

For typical usage, you should use 'Control.Monad.Borrow.Pure.borrowM' to avoid type ambiguity.
-}
borrow :: forall α a. a %1 -> Linearly %1 -> (Mut α a, Lend α a)
borrow = Unsafe.toLinear2 \ !a !_ ->
  (UnsafeAlias a, UnsafeAlias a)

-- | Shares a mutable borrow, invalidating the original one.
share :: Borrow k α a %1 -> Ur (Share α a)
share = Unsafe.toLinear \(UnsafeAlias !a) -> Ur (UnsafeAlias a)

-- | Reclaims a 'borrow'ed resource at the 'End' of lifetime @α'.
reclaim' :: Lend α a %1 -> After α a
reclaim' l = After (reclaim l)

-- | Reclaims a 'borrow'ed resource at the 'End' of lifetime @α'.
reclaim :: (End α) => Lend α a %1 -> a
reclaim = \(UnsafeAlias !a) -> a

-- | Reborrow a mutable borrow into a sublifetime.
reborrow :: forall β α a. (α >= β) => Mut α a %1 -> (Mut β a, Lend β (Mut α a))
reborrow = Unsafe.toLinear \ !mutA ->
  (Data.Coerce.coerce mutA, Data.Coerce.coerce mutA)

-- | Collapse a borrower to a mutable borrower.
joinMut :: Borrow bk α (Mut β a) %1 -> Borrow bk (α /\ β) a
joinMut = coerceLin

joinLend :: Lend α (Lend α a) %1 -> Lend α a
joinLend = coerceLin

-- | Distribute an alias over a functor.
class DistributesAlias f where
  split_ :: Alias ak α (f x) %1 -> f (Alias ak α x)
  default split_ ::
    (GenericDistributesAlias f) =>
    Alias ak α (f x) %1 -> f (Alias ak α x)
  split_ = genericSplit

split ::
  forall f x ak α.
  (DistributesAlias f) =>
  Alias ak α (f x) %1 -> f (Alias ak α x)
{-# INLINE [1] split #-}
split = split_

deriving anyclass instance DistributesAlias Identity

deriving anyclass instance DistributesAlias []

deriving anyclass instance DistributesAlias Maybe

deriving anyclass instance DistributesAlias Solo

deriving anyclass instance DistributesAlias Ord.Down

deriving anyclass instance DistributesAlias Sem.Dual

deriving anyclass instance DistributesAlias Sem.Max

deriving anyclass instance DistributesAlias Sem.Min

deriving anyclass instance DistributesAlias Sem.First

deriving anyclass instance DistributesAlias Sem.Last

deriving anyclass instance DistributesAlias Mon.First

deriving anyclass instance DistributesAlias Mon.Last

splitPair :: Alias ak α (a, b) %1 -> (Alias ak α a, Alias ak α b)
{-# INLINE splitPair #-}
splitPair = coerceLin

splitEither :: Alias ak α (Either a b) %1 -> Either (Alias ak α a) (Alias ak α b)
{-# INLINE splitEither #-}
splitEither = coerceLin

instance (Unsatisfiable ('Text "Use splitEither directly!")) => DistributesAlias (Either e) where
  {-# INLINE split_ #-}
  split_ = unsatisfiable

instance (Unsatisfiable ('Text "Use splitPair instead!")) => DistributesAlias ((,) a) where
  {-# INLINE split_ #-}
  split_ = unsatisfiable

type GenericDistributesAlias f = (Generic1 f, GDistributeAlias (Rep1 f))

genericSplit ::
  forall f x ak α.
  (GenericDistributesAlias f) =>
  Alias ak α (f x) %1 -> f (Alias ak α x)
{-# INLINE genericSplit #-}
genericSplit =
  to1
    . gdistributeAlias @(Rep1 f)
    . unsafeMapAlias from1

unsafeMapAlias :: (a %1 -> b) %1 -> Alias ak α a %1 -> Alias ak α b
{-# INLINE unsafeMapAlias #-}
unsafeMapAlias f = coerceLin (\x -> let !y = f x in y)

instance (GenericDistributesAlias f) => DistributesAlias (Generically1 f) where
  {-# INLINE split_ #-}
  split_ = Generically1 . genericSplit . unsafeMapAlias \(Generically1 f) -> f

class GDistributeAlias f where
  gdistributeAlias :: Alias ak α (f x) %1 -> f (Alias ak α x)

instance
  ( GDistributeAlias f
  , GDistributeAlias g
  ) =>
  GDistributeAlias (f :*: g)
  where
  {-# INLINE gdistributeAlias #-}
  gdistributeAlias !(UnsafeAlias !(f :*: g)) =
    DataFlow.do
      !f <- gdistributeAlias $ UnsafeAlias f
      !g <- gdistributeAlias $ UnsafeAlias g
      f :*: g

instance
  ( GDistributeAlias f
  , GDistributeAlias g
  ) =>
  GDistributeAlias (f :+: g)
  where
  {-# INLINE gdistributeAlias #-}
  gdistributeAlias (UnsafeAlias x) = case x of
    L1 !l -> L1 (gdistributeAlias (UnsafeAlias l))
    R1 !r -> R1 (gdistributeAlias (UnsafeAlias r))

instance
  (Unsatisfiable (Text "Nonlinear fields cannot distribute borrows!")) =>
  GDistributeAlias (MP1 GHC.Many f)
  where
  {-# INLINE gdistributeAlias #-}
  gdistributeAlias = unsatisfiable

instance (GDistributeAlias f) => GDistributeAlias (MP1 GHC.One f) where
  {-# INLINE gdistributeAlias #-}
  gdistributeAlias =
    MP1 . gdistributeAlias . UnsafeAlias . unMP1 . unsafeUnalias

instance (GDistributeAlias f) => GDistributeAlias (M1 i c f) where
  {-# INLINE gdistributeAlias #-}
  gdistributeAlias (UnsafeAlias (M1 x)) =
    M1 $ gdistributeAlias $ UnsafeAlias x

instance DistributesAlias Par1 where
  {-# INLINE split_ #-}
  split_ (UnsafeAlias (Par1 a)) = Par1 (UnsafeAlias a)

instance
  ( DistributesAlias f
  , DistributesAlias g
  , Data.Functor f
  ) =>
  GDistributeAlias (f :.: g)
  where
  {-# INLINE gdistributeAlias #-}
  gdistributeAlias (UnsafeAlias (Comp1 !fg)) =
    Comp1 $ Data.fmap split_ $ split_ $ UnsafeAlias fg

instance GDistributeAlias Par1 where
  {-# INLINE gdistributeAlias #-}
  gdistributeAlias (UnsafeAlias (Par1 !a)) = Par1 (UnsafeAlias a)

instance
  (Unsatisfiable (Text "A type containing non-parametric field with type `" :<>: ShowType c :<>: Text "', which cannot be safely splitted!")) =>
  GDistributeAlias (K1 i c)
  where
  {-# INLINE gdistributeAlias #-}
  gdistributeAlias = unsatisfiable

instance GDistributeAlias U1 where
  gdistributeAlias = coerceLin
  {-# INLINE gdistributeAlias #-}