{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE DerivingVia #-}
{-# LANGUAGE ImpredicativeTypes #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE LinearTypes #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE MultiWayIf #-}
{-# LANGUAGE OverloadedLabels #-}
{-# LANGUAGE OverloadedRecordDot #-}
{-# LANGUAGE QualifiedDo #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UnboxedTuples #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE NoImplicitPrelude #-}
{-# OPTIONS_GHC -Wno-name-shadowing #-}
{-# OPTIONS_GHC -Wno-partial-type-signatures #-}

module Control.Concurrent.DivideConquer.Utils.QueuePool (
  QueuePool,
  newQueuePool,
  pushWork,
  pushWorks,
  popWork,
  pushWorkMaster,
) where

import Control.Applicative (Alternative (..))
import Control.Applicative qualified as P
import Control.Concurrent (yield)
import Control.Concurrent.STM (STM, atomically, retry)
import Control.Concurrent.STM.TMDeque (TMDeque, closeTMDeque, isClosedTMDeque, newTMDequeIO, pushFrontTMDeque, sizeTMDeque, tryPopBackTMDeque, tryPopFrontTMDeque)
import Control.Monad qualified as NonLinear
import Control.Monad qualified as P
import Control.Monad.Borrow.Pure.BO
import Control.Monad.Borrow.Pure.BO.Unsafe (Alias (..), unsafeSystemIOToBO)
import Data.Coerce (coerce)
import Data.Foldable qualified as P
import Data.Function (fix)
import Data.List qualified as L
import Data.Monoid (Alt (..))
import Data.Ord (Down (..))
import Data.Ord qualified as P
import Data.V.Linear (V, theLength)
import Data.V.Linear.Internal (V (..))
import Data.Vector qualified as V
import Data.Vector.Algorithms.Intro qualified as AI
import Data.Vector.Hybrid.Mutable qualified as HMV
import Data.Vector.Mutable (RealWorld)
import GHC.Exts qualified as GHC
import GHC.IO qualified as GHC
import GHC.TypeLits (KnownNat)
import Prelude.Linear
import Unsafe.Linear qualified as Unsafe
import Prelude qualified as P

data QueuePool a = QueuePool
  { forall a. QueuePool a -> TMDeque a
mine :: !(TMDeque a)
  , forall a. QueuePool a -> MVector RealWorld (TMDeque a)
others :: !(V.MVector RealWorld (TMDeque a))
  , forall a. QueuePool a -> Int
num :: !Int
  }

newtype MasterQueuePool a = MasterQueuePool [TMDeque a]

instance Consumable (MasterQueuePool a) where
  consume :: MasterQueuePool a %1 -> ()
consume = [()] %1 -> ()
forall a. Consumable a => a %1 -> ()
consume ([()] %1 -> ())
-> (MasterQueuePool a %1 -> [()]) -> MasterQueuePool a %1 -> ()
forall b c a (q :: Multiplicity) (m :: Multiplicity)
       (n :: Multiplicity).
(b %1 -> c) %q -> (a %1 -> b) %m -> a %n -> c
. (TMDeque a %1 -> ()) -> [TMDeque a] %1 -> [()]
forall a b. (a %1 -> b) -> [a] %1 -> [b]
map TMDeque a %1 -> ()
forall a. TMDeque a %1 -> ()
consumeTMDQ ([TMDeque a] %1 -> [()])
-> (MasterQueuePool a %1 -> [TMDeque a])
-> MasterQueuePool a
%1 -> [()]
forall b c a (q :: Multiplicity) (m :: Multiplicity)
       (n :: Multiplicity).
(b %1 -> c) %q -> (a %1 -> b) %m -> a %n -> c
. forall a b. a %1 -> b
Unsafe.coerce @_ @[TMDeque a]

consumeTMDQ :: TMDeque a %1 -> ()
{-# NOINLINE consumeTMDQ #-}
consumeTMDQ :: forall a. TMDeque a %1 -> ()
consumeTMDQ = (TMDeque a %1 -> ()) -> TMDeque a %1 -> ()
forall a. a -> a
GHC.noinline ((TMDeque a %1 -> ()) -> TMDeque a %1 -> ())
-> (TMDeque a %1 -> ()) -> TMDeque a %1 -> ()
forall a b (p :: Multiplicity) (q :: Multiplicity).
(a %p -> b) %q -> a %p -> b
$ (TMDeque a -> ()) %1 -> TMDeque a %1 -> ()
forall a b (p :: Multiplicity) (x :: Multiplicity).
(a %p -> b) %1 -> a %x -> b
Unsafe.toLinear \TMDeque a
q -> IO () -> ()
forall a. IO a -> a
GHC.unsafePerformIO do
  !() <- STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b (p :: Multiplicity) (q :: Multiplicity).
(a %p -> b) %q -> a %p -> b
$ TMDeque a -> STM ()
forall a. TMDeque a -> STM ()
closeTMDeque TMDeque a
q
  P.pure ()

newQueuePool ::
  forall n a α.
  (KnownNat n) =>
  BO α (V n (Mut α (QueuePool a)), MasterQueuePool a)
newQueuePool :: forall (n :: Nat) a (α :: Lifetime).
KnownNat n =>
BO α (V n (Mut α (QueuePool a)), MasterQueuePool a)
newQueuePool = IO (V n (Mut α (QueuePool a)), MasterQueuePool a)
%1 -> BO α (V n (Mut α (QueuePool a)), MasterQueuePool a)
forall a (α :: Lifetime). IO a %1 -> BO α a
unsafeSystemIOToBO do
  let n :: Int
n = forall (n :: Nat). KnownNat n => Int
theLength @n

  qs <- Int -> IO (TMDeque a) -> IO [TMDeque a]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
NonLinear.replicateM Int
n IO (TMDeque a)
forall a. IO (TMDeque a)
newTMDequeIO
  pools <-
    P.mapM
      ( \(Int
num, [TMDeque a]
ini, TMDeque a
mine, [TMDeque a]
tl) -> do
          others <- Vector (TMDeque a) -> IO (MVector (PrimState IO) (TMDeque a))
forall (m :: * -> *) a.
PrimMonad m =>
Vector a -> m (MVector (PrimState m) a)
V.unsafeThaw (Vector (TMDeque a) -> IO (MVector (PrimState IO) (TMDeque a)))
-> Vector (TMDeque a) -> IO (MVector (PrimState IO) (TMDeque a))
forall a b (p :: Multiplicity) (q :: Multiplicity).
(a %p -> b) %q -> a %p -> b
$ [TMDeque a] -> Vector (TMDeque a)
forall a. [a] -> Vector a
V.fromList ([TMDeque a] -> Vector (TMDeque a))
-> [TMDeque a] -> Vector (TMDeque a)
forall a b (p :: Multiplicity) (q :: Multiplicity).
(a %p -> b) %q -> a %p -> b
$ [TMDeque a]
tl [TMDeque a] %1 -> [TMDeque a] %1 -> [TMDeque a]
forall a. Semigroup a => a %1 -> a %1 -> a
<> [TMDeque a]
ini
          P.pure P.$ QueuePool {others, ..}
      )
      P.$ L.zip4
        [0 ..]
        (L.inits qs)
        qs
        (P.drop 1 $ L.tails qs)
  let master = [TMDeque a] -> MasterQueuePool a
forall a. [TMDeque a] -> MasterQueuePool a
MasterQueuePool ([TMDeque a] -> MasterQueuePool a)
-> [TMDeque a] -> MasterQueuePool a
forall a b (p :: Multiplicity) (q :: Multiplicity).
(a %p -> b) %q -> a %p -> b
$ (QueuePool a -> TMDeque a) -> [QueuePool a] -> [TMDeque a]
forall a b. (a -> b) -> [a] -> [b]
P.map (QueuePool a -> TMDeque a
forall a. QueuePool a -> TMDeque a
mine (QueuePool a -> TMDeque a)
-> (QueuePool a -> QueuePool a) -> QueuePool a -> TMDeque a
forall b c a. (b -> c) -> (a -> b) -> a -> c
P.. QueuePool a -> QueuePool a
forall a b. Coercible a b => a -> b
coerce) [QueuePool a]
pools
  P.pure (V $ V.fromList $ map UnsafeAlias pools, master)

pushWorkMaster :: Mut α (MasterQueuePool a) %1 -> a %1 -> BO α (Mut α (MasterQueuePool a))
pushWorkMaster :: forall (α :: Lifetime) a.
Mut α (MasterQueuePool a)
%1 -> a %1 -> BO α (Mut α (MasterQueuePool a))
pushWorkMaster = (Mut α (MasterQueuePool a)
 -> a -> BO α (Mut α (MasterQueuePool a)))
%1 -> Mut α (MasterQueuePool a)
%1 -> a
%1 -> BO α (Mut α (MasterQueuePool a))
forall a b c (p :: Multiplicity) (q :: Multiplicity)
       (x :: Multiplicity) (y :: Multiplicity).
(a %p -> b %q -> c) %1 -> a %x -> b %y -> c
Unsafe.toLinear2 \(UnsafeAlias (MasterQueuePool [TMDeque a]
pools)) a
work ->
  case [TMDeque a]
pools of
    (TMDeque a
q : [TMDeque a]
qs) -> IO (Mut α (MasterQueuePool a))
%1 -> BO α (Mut α (MasterQueuePool a))
forall a (α :: Lifetime). IO a %1 -> BO α a
unsafeSystemIOToBO do
      STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b (p :: Multiplicity) (q :: Multiplicity).
(a %p -> b) %q -> a %p -> b
$ TMDeque a -> a -> STM ()
forall a. TMDeque a -> a -> STM ()
pushFrontTMDeque TMDeque a
q a
work
      Mut α (MasterQueuePool a) -> IO (Mut α (MasterQueuePool a))
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
P.pure (Mut α (MasterQueuePool a) -> IO (Mut α (MasterQueuePool a)))
-> Mut α (MasterQueuePool a) -> IO (Mut α (MasterQueuePool a))
forall a b (p :: Multiplicity) (q :: Multiplicity).
(a %p -> b) %q -> a %p -> b
$ MasterQueuePool a -> Mut α (MasterQueuePool a)
forall (ak :: AliasKind) (α :: Lifetime) a. a -> Alias ak α a
UnsafeAlias (MasterQueuePool a -> Mut α (MasterQueuePool a))
-> MasterQueuePool a -> Mut α (MasterQueuePool a)
forall a b (p :: Multiplicity) (q :: Multiplicity).
(a %p -> b) %q -> a %p -> b
$ [TMDeque a] -> MasterQueuePool a
forall a. [TMDeque a] -> MasterQueuePool a
MasterQueuePool (TMDeque a
q TMDeque a -> [TMDeque a] -> [TMDeque a]
forall a. a -> [a] -> [a]
: [TMDeque a]
qs)
    [] -> [Char] -> BO α (Mut α (MasterQueuePool a))
forall a. HasCallStack => [Char] -> a
error [Char]
"impossible: the length of pools is determined by the type-level nat n and cannot be zero"

pushWork :: Mut α (QueuePool a) %1 -> a %1 -> BO α (Mut α (QueuePool a))
pushWork :: forall (α :: Lifetime) a.
Mut α (QueuePool a) %1 -> a %1 -> BO α (Mut α (QueuePool a))
pushWork = (Mut α (QueuePool a) -> a -> BO α (Mut α (QueuePool a)))
%1 -> Mut α (QueuePool a) %1 -> a %1 -> BO α (Mut α (QueuePool a))
forall a b c (p :: Multiplicity) (q :: Multiplicity)
       (x :: Multiplicity) (y :: Multiplicity).
(a %p -> b %q -> c) %1 -> a %x -> b %y -> c
Unsafe.toLinear2 \(UnsafeAlias QueuePool {Int
MVector RealWorld (TMDeque a)
TMDeque a
mine :: forall a. QueuePool a -> TMDeque a
others :: forall a. QueuePool a -> MVector RealWorld (TMDeque a)
num :: forall a. QueuePool a -> Int
mine :: TMDeque a
others :: MVector RealWorld (TMDeque a)
num :: Int
..}) a
work ->
  IO (Mut α (QueuePool a)) %1 -> BO α (Mut α (QueuePool a))
forall a (α :: Lifetime). IO a %1 -> BO α a
unsafeSystemIOToBO do
    STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b (p :: Multiplicity) (q :: Multiplicity).
(a %p -> b) %q -> a %p -> b
$ TMDeque a -> a -> STM ()
forall a. TMDeque a -> a -> STM ()
pushFrontTMDeque TMDeque a
mine a
work
    Mut α (QueuePool a) -> IO (Mut α (QueuePool a))
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
P.pure (Mut α (QueuePool a) -> IO (Mut α (QueuePool a)))
-> Mut α (QueuePool a) -> IO (Mut α (QueuePool a))
forall a b (p :: Multiplicity) (q :: Multiplicity).
(a %p -> b) %q -> a %p -> b
$ QueuePool a -> Mut α (QueuePool a)
forall (ak :: AliasKind) (α :: Lifetime) a. a -> Alias ak α a
UnsafeAlias QueuePool {Int
MVector RealWorld (TMDeque a)
TMDeque a
mine :: TMDeque a
others :: MVector RealWorld (TMDeque a)
num :: Int
mine :: TMDeque a
others :: MVector RealWorld (TMDeque a)
num :: Int
..}

newtype Backwards f a = Backwards {forall {k} (f :: k -> *) (a :: k). Backwards f a -> f a
runBackwards :: f a}
  deriving newtype ((forall a b. (a -> b) -> Backwards f a -> Backwards f b)
-> (forall a b. a -> Backwards f b -> Backwards f a)
-> Functor (Backwards f)
forall a b. a -> Backwards f b -> Backwards f a
forall a b. (a -> b) -> Backwards f a -> Backwards f b
forall (f :: * -> *) a b.
Functor f =>
a -> Backwards f b -> Backwards f a
forall (f :: * -> *) a b.
Functor f =>
(a -> b) -> Backwards f a -> Backwards f b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall (f :: * -> *) a b.
Functor f =>
(a -> b) -> Backwards f a -> Backwards f b
fmap :: forall a b. (a -> b) -> Backwards f a -> Backwards f b
$c<$ :: forall (f :: * -> *) a b.
Functor f =>
a -> Backwards f b -> Backwards f a
<$ :: forall a b. a -> Backwards f b -> Backwards f a
P.Functor)

instance (P.Applicative f) => P.Applicative (Backwards f) where
  pure :: forall a. a -> Backwards f a
pure = f a -> Backwards f a
forall {k} (f :: k -> *) (a :: k). f a -> Backwards f a
Backwards (f a -> Backwards f a) -> (a -> f a) -> a -> Backwards f a
forall b c a. (b -> c) -> (a -> b) -> a -> c
P.. a -> f a
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
P.pure
  Backwards f (a -> b)
f <*> :: forall a b. Backwards f (a -> b) -> Backwards f a -> Backwards f b
<*> Backwards f a
x = f b -> Backwards f b
forall {k} (f :: k -> *) (a :: k). f a -> Backwards f a
Backwards (f a
x f a -> f (a -> b) -> f b
forall (f :: * -> *) a b. Applicative f => f a -> f (a -> b) -> f b
P.<**> f (a -> b)
f)

-- | Pushes works, the first element is on top.
pushWorks :: Mut α (QueuePool a) %1 -> [a] %1 -> BO α (Mut α (QueuePool a))
pushWorks :: forall (α :: Lifetime) a.
Mut α (QueuePool a) %1 -> [a] %1 -> BO α (Mut α (QueuePool a))
pushWorks = (Mut α (QueuePool a) -> [a] -> BO α (Mut α (QueuePool a)))
%1 -> Mut α (QueuePool a)
%1 -> [a]
%1 -> BO α (Mut α (QueuePool a))
forall a b c (p :: Multiplicity) (q :: Multiplicity)
       (x :: Multiplicity) (y :: Multiplicity).
(a %p -> b %q -> c) %1 -> a %x -> b %y -> c
Unsafe.toLinear2 \(UnsafeAlias QueuePool {Int
MVector RealWorld (TMDeque a)
TMDeque a
mine :: forall a. QueuePool a -> TMDeque a
others :: forall a. QueuePool a -> MVector RealWorld (TMDeque a)
num :: forall a. QueuePool a -> Int
mine :: TMDeque a
others :: MVector RealWorld (TMDeque a)
num :: Int
..}) [a]
work ->
  IO (Mut α (QueuePool a)) %1 -> BO α (Mut α (QueuePool a))
forall a (α :: Lifetime). IO a %1 -> BO α a
unsafeSystemIOToBO do
    STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b (p :: Multiplicity) (q :: Multiplicity).
(a %p -> b) %q -> a %p -> b
$ Backwards STM () -> STM ()
forall {k} (f :: k -> *) (a :: k). Backwards f a -> f a
runBackwards (Backwards STM () -> STM ()) -> Backwards STM () -> STM ()
forall a b. (a -> b) -> a -> b
P.$ (a -> Backwards STM ()) -> [a] -> Backwards STM ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
P.traverse_ (STM () -> Backwards STM ()
forall {k} (f :: k -> *) (a :: k). f a -> Backwards f a
Backwards (STM () -> Backwards STM ())
-> (a -> STM ()) -> a -> Backwards STM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
P.. TMDeque a -> a -> STM ()
forall a. TMDeque a -> a -> STM ()
pushFrontTMDeque TMDeque a
mine) [a]
work
    Mut α (QueuePool a) -> IO (Mut α (QueuePool a))
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
P.pure (Mut α (QueuePool a) -> IO (Mut α (QueuePool a)))
-> Mut α (QueuePool a) -> IO (Mut α (QueuePool a))
forall a b (p :: Multiplicity) (q :: Multiplicity).
(a %p -> b) %q -> a %p -> b
$ QueuePool a -> Mut α (QueuePool a)
forall (ak :: AliasKind) (α :: Lifetime) a. a -> Alias ak α a
UnsafeAlias QueuePool {Int
MVector RealWorld (TMDeque a)
TMDeque a
mine :: TMDeque a
others :: MVector RealWorld (TMDeque a)
num :: Int
mine :: TMDeque a
others :: MVector RealWorld (TMDeque a)
num :: Int
..}

popWork :: Mut α (QueuePool a) %1 -> BO α (Maybe (a, Mut α (QueuePool a)))
popWork :: forall (α :: Lifetime) a.
Mut α (QueuePool a) %1 -> BO α (Maybe (a, Mut α (QueuePool a)))
popWork = (Mut α (QueuePool a) -> BO α (Maybe (a, Mut α (QueuePool a))))
%1 -> Mut α (QueuePool a)
%1 -> BO α (Maybe (a, Mut α (QueuePool a)))
forall a b (p :: Multiplicity) (x :: Multiplicity).
(a %p -> b) %1 -> a %x -> b
Unsafe.toLinear \qs :: Mut α (QueuePool a)
qs@(UnsafeAlias QueuePool {Int
MVector RealWorld (TMDeque a)
TMDeque a
mine :: forall a. QueuePool a -> TMDeque a
others :: forall a. QueuePool a -> MVector RealWorld (TMDeque a)
num :: forall a. QueuePool a -> Int
mine :: TMDeque a
others :: MVector RealWorld (TMDeque a)
num :: Int
..}) ->
  IO (Maybe (a, Mut α (QueuePool a)))
%1 -> BO α (Maybe (a, Mut α (QueuePool a)))
forall a (α :: Lifetime). IO a %1 -> BO α a
unsafeSystemIOToBO do
    STM (Maybe (Maybe a)) -> IO (Maybe (Maybe a))
forall a. STM a -> IO a
atomically (TMDeque a -> STM (Maybe (Maybe a))
forall a. TMDeque a -> STM (Maybe (Maybe a))
tryPopFrontTMDeque TMDeque a
mine) IO (Maybe (Maybe a))
-> (Maybe (Maybe a) -> IO (Maybe (a, Mut α (QueuePool a))))
-> IO (Maybe (a, Mut α (QueuePool a)))
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
P.>>= \case
      Maybe (Maybe a)
Nothing -> Maybe (a, Mut α (QueuePool a))
-> IO (Maybe (a, Mut α (QueuePool a)))
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
P.pure Maybe (a, Mut α (QueuePool a))
forall a. Maybe a
Nothing
      Just (Just a
x) -> Maybe (a, Mut α (QueuePool a))
-> IO (Maybe (a, Mut α (QueuePool a)))
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
P.pure (Maybe (a, Mut α (QueuePool a))
 -> IO (Maybe (a, Mut α (QueuePool a))))
-> Maybe (a, Mut α (QueuePool a))
-> IO (Maybe (a, Mut α (QueuePool a)))
forall a b (p :: Multiplicity) (q :: Multiplicity).
(a %p -> b) %q -> a %p -> b
$ (a, Mut α (QueuePool a)) -> Maybe (a, Mut α (QueuePool a))
forall a. a -> Maybe a
Just (a
x, Mut α (QueuePool a)
qs)
      Just Maybe a
Nothing -> (IO (Maybe (a, Mut α (QueuePool a)))
 -> IO (Maybe (a, Mut α (QueuePool a))))
-> IO (Maybe (a, Mut α (QueuePool a)))
forall a. (a -> a) -> a
fix \IO (Maybe (a, Mut α (QueuePool a)))
self -> do
        !ranks <-
          Vector Int -> IO (MVector (PrimState IO) Int)
forall (m :: * -> *) a.
PrimMonad m =>
Vector a -> m (MVector (PrimState m) a)
V.unsafeThaw
            (Vector Int -> IO (MVector (PrimState IO) Int))
-> IO (Vector Int) -> IO (MVector (PrimState IO) Int)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
P.=<< STM (Vector Int) -> IO (Vector Int)
forall a. STM a -> IO a
atomically (STM (Vector Int) -> IO (Vector Int))
-> (Vector (TMDeque a) -> STM (Vector Int))
-> Vector (TMDeque a)
-> IO (Vector Int)
forall b c a. (b -> c) -> (a -> b) -> a -> c
P.. (\Vector (TMDeque a)
x -> do xs <- (TMDeque a -> STM Int) -> Vector (TMDeque a) -> STM (Vector Int)
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Vector a -> m (Vector b)
V.mapM TMDeque a -> STM Int
forall a. TMDeque a -> STM Int
sizeTMDeque Vector (TMDeque a)
x; xs P.<$ P.unless (V.any (P.> 0) xs) retry)
            (Vector (TMDeque a) -> IO (Vector Int))
-> IO (Vector (TMDeque a)) -> IO (Vector Int)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
P.=<< MVector (PrimState IO) (TMDeque a) -> IO (Vector (TMDeque a))
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> m (Vector a)
V.unsafeFreeze MVector RealWorld (TMDeque a)
MVector (PrimState IO) (TMDeque a)
others
        let ranked = MVector (PrimState IO) Int
-> MVector (PrimState IO) (TMDeque a)
-> MVector MVector MVector (PrimState IO) (Int, TMDeque a)
forall (u :: * -> * -> *) s a (v :: * -> * -> *) b.
u s a -> v s b -> MVector u v s (a, b)
HMV.unsafeZip MVector (PrimState IO) Int
ranks MVector RealWorld (TMDeque a)
MVector (PrimState IO) (TMDeque a)
others
        !() <- AI.sortBy (P.comparing P.$ Down P.. P.fst) ranked
        others' <- V.unsafeFreeze others

        progress <-
          atomically do
            ( isClosedTMDeque mine P.>>= \Bool
closed ->
                if Bool
closed then Maybe (Maybe a) -> STM (Maybe (Maybe a))
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
P.pure Maybe (Maybe a)
forall a. Maybe a
Nothing else STM (Maybe (Maybe a))
forall a. STM a
retry
              )
              <|> getAlt (P.foldMap' (Alt P.. (P.fmap Just P.. fromJustSTM P.<=< tryPopBackTMDeque)) P.$ others')
              <|> P.pure (Just Nothing)
        case progress of
          Maybe (Maybe a)
Nothing -> Maybe (a, Mut α (QueuePool a))
-> IO (Maybe (a, Mut α (QueuePool a)))
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
P.pure Maybe (a, Mut α (QueuePool a))
forall a. Maybe a
Nothing
          Just Maybe a
Nothing -> IO ()
yield IO ()
-> IO (Maybe (a, Mut α (QueuePool a)))
-> IO (Maybe (a, Mut α (QueuePool a)))
forall a b. IO a -> IO b -> IO b
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
P.*> IO (Maybe (a, Mut α (QueuePool a)))
self
          Just (Just a
x) -> Maybe (a, Mut α (QueuePool a))
-> IO (Maybe (a, Mut α (QueuePool a)))
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
P.pure (Maybe (a, Mut α (QueuePool a))
 -> IO (Maybe (a, Mut α (QueuePool a))))
-> Maybe (a, Mut α (QueuePool a))
-> IO (Maybe (a, Mut α (QueuePool a)))
forall a b (p :: Multiplicity) (q :: Multiplicity).
(a %p -> b) %q -> a %p -> b
$ (a, Mut α (QueuePool a)) -> Maybe (a, Mut α (QueuePool a))
forall a. a -> Maybe a
Just (a
x, Mut α (QueuePool a)
qs)

fromJustSTM :: Maybe (Maybe a) -> STM (Maybe a)
fromJustSTM :: forall a. Maybe (Maybe a) -> STM (Maybe a)
fromJustSTM = STM (Maybe a)
-> (Maybe a -> STM (Maybe a)) -> Maybe (Maybe a) -> STM (Maybe a)
forall b a. b -> (a -> b) -> Maybe a -> b
P.maybe (Maybe a -> STM (Maybe a)
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
P.pure Maybe a
forall a. Maybe a
Nothing) ((Maybe a -> STM (Maybe a)) -> Maybe (Maybe a) -> STM (Maybe a))
-> (Maybe a -> STM (Maybe a)) -> Maybe (Maybe a) -> STM (Maybe a)
forall a b (p :: Multiplicity) (q :: Multiplicity).
(a %p -> b) %q -> a %p -> b
$ STM (Maybe a) -> (a -> STM (Maybe a)) -> Maybe a -> STM (Maybe a)
forall b a. b -> (a -> b) -> Maybe a -> b
P.maybe STM (Maybe a)
forall a. STM a
retry (Maybe a -> STM (Maybe a)
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
P.pure (Maybe a -> STM (Maybe a)) -> (a -> Maybe a) -> a -> STM (Maybe a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
P.. a -> Maybe a
forall a. a -> Maybe a
Just)