{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DerivingVia #-}
{-# LANGUAGE ExplicitNamespaces #-}
{-# LANGUAGE ImpredicativeTypes #-}
{-# LANGUAGE PartialTypeSignatures #-}
{-# LANGUAGE QualifiedDo #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeAbstractions #-}
{-# LANGUAGE NoImplicitPrelude #-}
{-# OPTIONS_GHC -Wno-name-shadowing #-}

module Control.Monad.Borrow.Pure.Experimental.Reborrowable (
  Reborrowable (..),
  locally,
  locally_,
) where

import Control.Functor.Linear qualified as Control
import Control.Monad.Borrow.Pure.BO
import Prelude.Linear

class Reborrowable bor where
  {- |
  Executes an operation on a borrow in sub lifetime.
  You may need @-XImpredicativeTypes@ extension to use this function.

  Generalization of 'reborrowing'' and 'sharing'' that works for both 'Mut' and 'Share' borrows.
  -}
  locally' ::
    bor α a %1 ->
    (forall β. bor (β /\ α) a %1 -> BO (β /\ α') (After β r)) %1 ->
    BO α' (r, bor α a)

instance Reborrowable Mut where
  {-# SPECIALIZE instance Reborrowable Mut #-}
  locally' :: forall (α :: Lifetime) a (α' :: Lifetime) r.
Mut α a
%1 -> (forall (β :: Lifetime).
       Mut (β /\ α) a %1 -> BO (β /\ α') (After β r))
%1 -> BO α' (r, Mut α a)
locally' = Mut α a
%1 -> (forall (β :: Lifetime).
       Mut (β /\ α) a %1 -> BO (β /\ α') (After β r))
%1 -> BO α' (r, Mut α a)
forall (α :: Lifetime) a (α' :: Lifetime) r.
Mut α a
%1 -> (forall (β :: Lifetime).
       Mut (β /\ α) a %1 -> BO (β /\ α') (After β r))
%1 -> BO α' (r, Mut α a)
reborrowing'
  {-# INLINE locally' #-}

instance Reborrowable Share where
  {-# SPECIALIZE instance Reborrowable Share #-}
  locally' :: forall (α :: Lifetime) a (α' :: Lifetime) r.
Share α a
%1 -> (forall (β :: Lifetime).
       Share (β /\ α) a %1 -> BO (β /\ α') (After β r))
%1 -> BO α' (r, Share α a)
locally' Share α a
shr forall (β :: Lifetime).
Share (β /\ α) a %1 -> BO (β /\ α') (After β r)
k = Control.do
    let %1 !(Ur Share α a
sh) = Share α a %1 -> Ur (Share α a)
forall a. Movable a => a %1 -> Ur a
move Share α a
shr
    (,Share α a
sh) (r %1 -> (r, Share α a)) %1 -> BO α' r %1 -> BO α' (r, Share α a)
forall (f :: * -> *) a b.
Functor f =>
(a %1 -> b) %1 -> f a %1 -> f b
Control.<$> (forall (α :: Lifetime). BO (α /\ α') (After α r)) %1 -> BO α' r
forall (β :: Lifetime) a.
(forall (α :: Lifetime). BO (α /\ β) (After α a)) %1 -> BO β a
srunBO (Share (α /\ α) a %1 -> BO (α /\ α') (After α r)
forall (β :: Lifetime).
Share (β /\ α) a %1 -> BO (β /\ α') (After β r)
k (Share α a %1 -> Share (α /\ α) a
forall a b. (a <: b) => a %1 -> b
upcast Share α a
sh))
  {-# INLINE locally' #-}

locally ::
  (Reborrowable bor) =>
  bor α a %1 ->
  (forall β. bor (β /\ α) a %1 -> BO (β /\ α') r) %1 ->
  BO α' (r, bor α a)
{-# INLINE locally #-}
locally :: forall {k} (bor :: Lifetime -> k -> *) (α :: Lifetime) (a :: k)
       (α' :: Lifetime) r.
Reborrowable bor =>
bor α a
%1 -> (forall (β :: Lifetime). bor (β /\ α) a %1 -> BO (β /\ α') r)
%1 -> BO α' (r, bor α a)
locally bor α a
bor forall (β :: Lifetime). bor (β /\ α) a %1 -> BO (β /\ α') r
k = bor α a
%1 -> (forall (β :: Lifetime).
       bor (β /\ α) a %1 -> BO (β /\ α') (After β r))
%1 -> BO α' (r, bor α a)
forall {k} (bor :: Lifetime -> k -> *) (α :: Lifetime) (a :: k)
       (α' :: Lifetime) r.
Reborrowable bor =>
bor α a
%1 -> (forall (β :: Lifetime).
       bor (β /\ α) a %1 -> BO (β /\ α') (After β r))
%1 -> BO α' (r, bor α a)
forall (α :: Lifetime) (a :: k) (α' :: Lifetime) r.
bor α a
%1 -> (forall (β :: Lifetime).
       bor (β /\ α) a %1 -> BO (β /\ α') (After β r))
%1 -> BO α' (r, bor α a)
locally' bor α a
bor \bor (β /\ α) a
mut -> r %1 -> After β r
forall a. a %1 -> After β a
forall (f :: * -> *) a. Applicative f => a %1 -> f a
Control.pure (r %1 -> After β r)
%1 -> BO (β /\ α') r %1 -> BO (β /\ α') (After β r)
forall (f :: * -> *) a b.
Functor f =>
(a %1 -> b) %1 -> f a %1 -> f b
Control.<$> bor (β /\ α) a %1 -> BO (β /\ α') r
forall (β :: Lifetime). bor (β /\ α) a %1 -> BO (β /\ α') r
k bor (β /\ α) a
mut

locally_ ::
  (Reborrowable bor, Consumable r) =>
  bor α a %1 ->
  (forall β. bor (β /\ α) a %1 -> BO (β /\ α') r) %1 ->
  BO α' (bor α a)
{-# INLINE locally_ #-}
locally_ :: forall {k} (bor :: Lifetime -> k -> *) r (α :: Lifetime) (a :: k)
       (α' :: Lifetime).
(Reborrowable bor, Consumable r) =>
bor α a
%1 -> (forall (β :: Lifetime). bor (β /\ α) a %1 -> BO (β /\ α') r)
%1 -> BO α' (bor α a)
locally_ bor α a
bor forall (β :: Lifetime). bor (β /\ α) a %1 -> BO (β /\ α') r
k = (r %1 -> bor α a %1 -> bor α a) -> (r, bor α a) %1 -> bor α a
forall a (p :: Multiplicity) b c (q :: Multiplicity).
(a %p -> b %p -> c) %q -> (a, b) %p -> c
uncurry r %1 -> bor α a %1 -> bor α a
forall a b. Consumable a => a %1 -> b %1 -> b
lseq ((r, bor α a) %1 -> bor α a)
%1 -> BO α' (r, bor α a) %1 -> BO α' (bor α a)
forall (f :: * -> *) a b.
Functor f =>
(a %1 -> b) %1 -> f a %1 -> f b
Control.<$> bor α a
%1 -> (forall (β :: Lifetime). bor (β /\ α) a %1 -> BO (β /\ α') r)
%1 -> BO α' (r, bor α a)
forall {k} (bor :: Lifetime -> k -> *) (α :: Lifetime) (a :: k)
       (α' :: Lifetime) r.
Reborrowable bor =>
bor α a
%1 -> (forall (β :: Lifetime). bor (β /\ α) a %1 -> BO (β /\ α') r)
%1 -> BO α' (r, bor α a)
locally bor α a
bor bor (β /\ α) a %1 -> BO (β /\ α') r
forall (β :: Lifetime). bor (β /\ α) a %1 -> BO (β /\ α') r
k