{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE ExplicitNamespaces #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}

module DataFrame.Internal.AggKernelDirect (
    directThreshold,
    directReduce,
) where

import Control.Concurrent (forkIO, getNumCapabilities)
import Control.Concurrent.MVar (newEmptyMVar, putMVar, takeMVar)
import Control.Exception (SomeException, throwIO, try)
import Data.Type.Equality (TestEquality (..), type (:~:) (Refl))
import qualified Data.Vector.Unboxed as VU
import qualified Data.Vector.Unboxed.Mutable as VUM
import System.IO.Unsafe (unsafePerformIO)
import Type.Reflection (typeRep)

import DataFrame.Internal.AggKernel (Reduction (..))
import DataFrame.Internal.Column (
    Column (..),
    fromUnboxedVector,
    materializePacked,
 )

{- | Group-domain size at or below which the direct-indexed accumulator path is
taken; wider domains keep the group-range kernel. The admitted reductions are
order-independent, so the per-worker accumulator merge is exact.
-}
directThreshold :: Int
directThreshold :: Int
directThreshold = Int
262144

capabilities :: Int
capabilities :: Int
capabilities = IO Int -> Int
forall a. IO a -> a
unsafePerformIO IO Int
getNumCapabilities
{-# NOINLINE capabilities #-}

{- | Below this many rows the parallel fan-out is not worth it; a single
sequential direct pass runs instead (tiny accumulator, one tight loop). Matches
the grouping/scatter parallel threshold.
-}
parThreshold :: Int
parThreshold :: Int
parThreshold = Int
200000

{- | Run a recognised reduction through the direct-indexed path. 'Nothing' (so
the caller falls back to the order-preserving kernel) unless the reduction is
order-independent at this element type AND the column is a clean unboxed Int/Double.
-}
directReduce :: Reduction -> VU.Vector Int -> Int -> Column -> Maybe Column
directReduce :: Reduction -> Vector Int -> Int -> Column -> Maybe Column
directReduce Reduction
red Vector Int
g Int
nGroups Column
col = case Column
col of
    UnboxedColumn Maybe Bitmap
Nothing (Vector a
v :: VU.Vector a) ->
        case TypeRep a -> TypeRep Int -> Maybe (a :~: Int)
forall a b. TypeRep a -> TypeRep b -> Maybe (a :~: b)
forall {k} (f :: k -> *) (a :: k) (b :: k).
TestEquality f =>
f a -> f b -> Maybe (a :~: b)
testEquality (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @a) (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @Int) of
            Just a :~: Int
Refl -> Reduction -> Vector Int -> Int -> Vector Int -> Maybe Column
directInt Reduction
red Vector Int
g Int
nGroups Vector a
Vector Int
v
            Maybe (a :~: Int)
Nothing -> case TypeRep a -> TypeRep Double -> Maybe (a :~: Double)
forall a b. TypeRep a -> TypeRep b -> Maybe (a :~: b)
forall {k} (f :: k -> *) (a :: k) (b :: k).
TestEquality f =>
f a -> f b -> Maybe (a :~: b)
testEquality (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @a) (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @Double) of
                Just a :~: Double
Refl -> Reduction -> Vector Int -> Int -> Vector Double -> Maybe Column
directDouble Reduction
red Vector Int
g Int
nGroups Vector a
Vector Double
v
                Maybe (a :~: Double)
Nothing -> Maybe Column
forall a. Maybe a
Nothing
    p :: Column
p@(PackedText Maybe Bitmap
_ PackedTextData
_) -> Reduction -> Vector Int -> Int -> Column -> Maybe Column
directReduce Reduction
red Vector Int
g Int
nGroups (Column -> Column
materializePacked Column
p)
    Column
_ -> Maybe Column
forall a. Maybe a
Nothing
{-# INLINEABLE directReduce #-}

-- | The order-independent reductions over an Int column.
directInt :: Reduction -> VU.Vector Int -> Int -> VU.Vector Int -> Maybe Column
directInt :: Reduction -> Vector Int -> Int -> Vector Int -> Maybe Column
directInt Reduction
red Vector Int
g Int
nGroups Vector Int
v = case Reduction
red of
    Reduction
RCount -> Column -> Maybe Column
forall a. a -> Maybe a
Just (Vector Int -> Column
forall a. (Columnable a, Unbox a) => Vector a -> Column
fromUnboxedVector (Vector Int -> Int -> Int -> Vector Int
countDirect Vector Int
g Int
nGroups (Vector Int -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Int
v)))
    Reduction
RSum -> Column -> Maybe Column
forall a. a -> Maybe a
Just (Vector Int -> Column
forall a. (Columnable a, Unbox a) => Vector a -> Column
fromUnboxedVector (Vector Int -> Int -> Vector Int -> Vector Int
sumIntDirect Vector Int
g Int
nGroups Vector Int
v))
    Reduction
RMin -> Column -> Maybe Column
forall a. a -> Maybe a
Just (Vector Int -> Column
forall a. (Columnable a, Unbox a) => Vector a -> Column
fromUnboxedVector (Bool -> Vector Int -> Int -> Vector Int -> Vector Int
extremaIntDirect Bool
True Vector Int
g Int
nGroups Vector Int
v))
    Reduction
RMax -> Column -> Maybe Column
forall a. a -> Maybe a
Just (Vector Int -> Column
forall a. (Columnable a, Unbox a) => Vector a -> Column
fromUnboxedVector (Bool -> Vector Int -> Int -> Vector Int -> Vector Int
extremaIntDirect Bool
False Vector Int
g Int
nGroups Vector Int
v))
    Reduction
RMean -> Column -> Maybe Column
forall a. a -> Maybe a
Just (Vector Double -> Column
forall a. (Columnable a, Unbox a) => Vector a -> Column
fromUnboxedVector (Vector Int -> Int -> Vector Int -> Vector Double
meanIntDirect Vector Int
g Int
nGroups Vector Int
v))
    Reduction
_ -> Maybe Column
forall a. Maybe a
Nothing

{- | Over a Double column only @count@ is order-independent; the float
sum/mean/variance reductions must keep the order-preserving kernel.
-}
directDouble ::
    Reduction -> VU.Vector Int -> Int -> VU.Vector Double -> Maybe Column
directDouble :: Reduction -> Vector Int -> Int -> Vector Double -> Maybe Column
directDouble Reduction
red Vector Int
g Int
nGroups Vector Double
v = case Reduction
red of
    Reduction
RCount -> Column -> Maybe Column
forall a. a -> Maybe a
Just (Vector Int -> Column
forall a. (Columnable a, Unbox a) => Vector a -> Column
fromUnboxedVector (Vector Int -> Int -> Int -> Vector Int
countDirect Vector Int
g Int
nGroups (Vector Double -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Double
v)))
    Reduction
_ -> Maybe Column
forall a. Maybe a
Nothing

-- | Whether to fan out at this row count.
shouldPar :: Int -> Bool
shouldPar :: Int -> Bool
shouldPar Int
n = Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
parThreshold Bool -> Bool -> Bool
&& Int
capabilities Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
1

{- | Fork @caps@ workers over disjoint contiguous row ranges of @[0, n)@, each
producing its own private accumulator (no shared array, no sync). Returns the
partials in worker order for the caller's merge; rethrows the first failure.
-}
runPartialsOver ::
    Int -> Int -> (Int -> Int -> IO (VUM.IOVector Int)) -> IO [VUM.IOVector Int]
runPartialsOver :: Int
-> Int -> (Int -> Int -> IO (IOVector Int)) -> IO [IOVector Int]
runPartialsOver Int
n Int
caps Int -> Int -> IO (IOVector Int)
fill = do
    let !per :: Int
per = (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
caps Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
caps
        spawn :: Int -> IO (MVar (Either SomeException (IOVector Int)))
spawn Int
w = do
            MVar (Either SomeException (IOVector Int))
var <- IO (MVar (Either SomeException (IOVector Int)))
forall a. IO (MVar a)
newEmptyMVar
            let !lo :: Int
lo = Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
n (Int
w Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
per)
                !hi :: Int
hi = Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
n (Int
lo Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
per)
            ThreadId
_ <- IO () -> IO ThreadId
forkIO (IO (IOVector Int) -> IO (Either SomeException (IOVector Int))
forall e a. Exception e => IO a -> IO (Either e a)
try (Int -> Int -> IO (IOVector Int)
fill Int
lo Int
hi) IO (Either SomeException (IOVector Int))
-> (Either SomeException (IOVector Int) -> IO ()) -> IO ()
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= MVar (Either SomeException (IOVector Int))
-> Either SomeException (IOVector Int) -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar (Either SomeException (IOVector Int))
var)
            MVar (Either SomeException (IOVector Int))
-> IO (MVar (Either SomeException (IOVector Int)))
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure MVar (Either SomeException (IOVector Int))
var
    [MVar (Either SomeException (IOVector Int))]
vars <- (Int -> IO (MVar (Either SomeException (IOVector Int))))
-> [Int] -> IO [MVar (Either SomeException (IOVector Int))]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM Int -> IO (MVar (Either SomeException (IOVector Int)))
spawn [Int
0 .. Int
caps Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
    [Either SomeException (IOVector Int)]
results <- (MVar (Either SomeException (IOVector Int))
 -> IO (Either SomeException (IOVector Int)))
-> [MVar (Either SomeException (IOVector Int))]
-> IO [Either SomeException (IOVector Int)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM MVar (Either SomeException (IOVector Int))
-> IO (Either SomeException (IOVector Int))
forall a. MVar a -> IO a
takeMVar [MVar (Either SomeException (IOVector Int))]
vars
    (Either SomeException (IOVector Int) -> IO (IOVector Int))
-> [Either SomeException (IOVector Int)] -> IO [IOVector Int]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ((SomeException -> IO (IOVector Int))
-> (IOVector Int -> IO (IOVector Int))
-> Either SomeException (IOVector Int)
-> IO (IOVector Int)
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (forall e a. Exception e => e -> IO a
throwIO @SomeException) IOVector Int -> IO (IOVector Int)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure) [Either SomeException (IOVector Int)]
results

{- | As 'runPartialsOver' but each worker produces a PAIR of accumulators (e.g.
sum and count for the fused integer mean).
-}
runPartialsPairOver ::
    Int ->
    Int ->
    (Int -> Int -> IO (VUM.IOVector Int, VUM.IOVector Int)) ->
    IO [(VUM.IOVector Int, VUM.IOVector Int)]
runPartialsPairOver :: Int
-> Int
-> (Int -> Int -> IO (IOVector Int, IOVector Int))
-> IO [(IOVector Int, IOVector Int)]
runPartialsPairOver Int
n Int
caps Int -> Int -> IO (IOVector Int, IOVector Int)
fill = do
    let !per :: Int
per = (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
caps Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
caps
        spawn :: Int
-> IO (MVar (Either SomeException (IOVector Int, IOVector Int)))
spawn Int
w = do
            MVar (Either SomeException (IOVector Int, IOVector Int))
var <- IO (MVar (Either SomeException (IOVector Int, IOVector Int)))
forall a. IO (MVar a)
newEmptyMVar
            let !lo :: Int
lo = Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
n (Int
w Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
per)
                !hi :: Int
hi = Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
n (Int
lo Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
per)
            ThreadId
_ <- IO () -> IO ThreadId
forkIO (IO (IOVector Int, IOVector Int)
-> IO (Either SomeException (IOVector Int, IOVector Int))
forall e a. Exception e => IO a -> IO (Either e a)
try (Int -> Int -> IO (IOVector Int, IOVector Int)
fill Int
lo Int
hi) IO (Either SomeException (IOVector Int, IOVector Int))
-> (Either SomeException (IOVector Int, IOVector Int) -> IO ())
-> IO ()
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= MVar (Either SomeException (IOVector Int, IOVector Int))
-> Either SomeException (IOVector Int, IOVector Int) -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar (Either SomeException (IOVector Int, IOVector Int))
var)
            MVar (Either SomeException (IOVector Int, IOVector Int))
-> IO (MVar (Either SomeException (IOVector Int, IOVector Int)))
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure MVar (Either SomeException (IOVector Int, IOVector Int))
var
    [MVar (Either SomeException (IOVector Int, IOVector Int))]
vars <- (Int
 -> IO (MVar (Either SomeException (IOVector Int, IOVector Int))))
-> [Int]
-> IO [MVar (Either SomeException (IOVector Int, IOVector Int))]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM Int
-> IO (MVar (Either SomeException (IOVector Int, IOVector Int)))
spawn [Int
0 .. Int
caps Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
    [Either SomeException (IOVector Int, IOVector Int)]
results <- (MVar (Either SomeException (IOVector Int, IOVector Int))
 -> IO (Either SomeException (IOVector Int, IOVector Int)))
-> [MVar (Either SomeException (IOVector Int, IOVector Int))]
-> IO [Either SomeException (IOVector Int, IOVector Int)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM MVar (Either SomeException (IOVector Int, IOVector Int))
-> IO (Either SomeException (IOVector Int, IOVector Int))
forall a. MVar a -> IO a
takeMVar [MVar (Either SomeException (IOVector Int, IOVector Int))]
vars
    (Either SomeException (IOVector Int, IOVector Int)
 -> IO (IOVector Int, IOVector Int))
-> [Either SomeException (IOVector Int, IOVector Int)]
-> IO [(IOVector Int, IOVector Int)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ((SomeException -> IO (IOVector Int, IOVector Int))
-> ((IOVector Int, IOVector Int)
    -> IO (IOVector Int, IOVector Int))
-> Either SomeException (IOVector Int, IOVector Int)
-> IO (IOVector Int, IOVector Int)
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (forall e a. Exception e => e -> IO a
throwIO @SomeException) (IOVector Int, IOVector Int) -> IO (IOVector Int, IOVector Int)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure) [Either SomeException (IOVector Int, IOVector Int)]
results

-------------------------------------------------------------------------------
-- Count (order-independent: per-group row count)
-------------------------------------------------------------------------------

countDirect :: VU.Vector Int -> Int -> Int -> VU.Vector Int
countDirect :: Vector Int -> Int -> Int -> Vector Int
countDirect Vector Int
g Int
nGroups Int
n
    | Bool -> Bool
not (Int -> Bool
shouldPar Int
n) =
        IO (Vector Int) -> Vector Int
forall a. IO a -> a
unsafePerformIO (Vector Int -> Int -> Int -> Int -> IO (IOVector Int)
countChunk Vector Int
g Int
nGroups Int
0 Int
n IO (IOVector Int)
-> (IOVector Int -> IO (Vector Int)) -> IO (Vector Int)
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= IOVector Int -> IO (Vector Int)
MVector (PrimState IO) Int -> IO (Vector Int)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze)
    | Bool
otherwise = IO (Vector Int) -> Vector Int
forall a. IO a -> a
unsafePerformIO (IO (Vector Int) -> Vector Int) -> IO (Vector Int) -> Vector Int
forall a b. (a -> b) -> a -> b
$ do
        [IOVector Int]
parts <- Int
-> Int -> (Int -> Int -> IO (IOVector Int)) -> IO [IOVector Int]
runPartialsOver Int
n Int
capabilities (Vector Int -> Int -> Int -> Int -> IO (IOVector Int)
countChunk Vector Int
g Int
nGroups)
        Int -> [IOVector Int] -> IO (Vector Int)
mergeIntSum Int
nGroups [IOVector Int]
parts
{-# NOINLINE countDirect #-}

countChunk :: VU.Vector Int -> Int -> Int -> Int -> IO (VUM.IOVector Int)
countChunk :: Vector Int -> Int -> Int -> Int -> IO (IOVector Int)
countChunk Vector Int
g Int
nGroups Int
lo Int
hi = do
    IOVector Int
acc <- Int -> Int -> IO (MVector (PrimState IO) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate Int
nGroups (Int
0 :: Int)
    let go :: Int -> IO ()
go !Int
i
            | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
hi = () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
            | Bool
otherwise = do
                let !k :: Int
k = Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector Int
g Int
i
                Int
c <- MVector (PrimState IO) Int -> Int -> IO Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VUM.unsafeRead IOVector Int
MVector (PrimState IO) Int
acc Int
k
                MVector (PrimState IO) Int -> Int -> Int -> IO ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VUM.unsafeWrite IOVector Int
MVector (PrimState IO) Int
acc Int
k (Int
c Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
                Int -> IO ()
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
    Int -> IO ()
go Int
lo
    IOVector Int -> IO (IOVector Int)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure IOVector Int
acc

-------------------------------------------------------------------------------
-- Integer sum (exact: merge order irrelevant)
-------------------------------------------------------------------------------

sumIntDirect :: VU.Vector Int -> Int -> VU.Vector Int -> VU.Vector Int
sumIntDirect :: Vector Int -> Int -> Vector Int -> Vector Int
sumIntDirect Vector Int
g Int
nGroups Vector Int
v
    | Bool -> Bool
not (Int -> Bool
shouldPar Int
n) =
        IO (Vector Int) -> Vector Int
forall a. IO a -> a
unsafePerformIO (Vector Int -> Vector Int -> Int -> Int -> Int -> IO (IOVector Int)
sumIntChunk Vector Int
g Vector Int
v Int
nGroups Int
0 Int
n IO (IOVector Int)
-> (IOVector Int -> IO (Vector Int)) -> IO (Vector Int)
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= IOVector Int -> IO (Vector Int)
MVector (PrimState IO) Int -> IO (Vector Int)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze)
    | Bool
otherwise = IO (Vector Int) -> Vector Int
forall a. IO a -> a
unsafePerformIO (IO (Vector Int) -> Vector Int) -> IO (Vector Int) -> Vector Int
forall a b. (a -> b) -> a -> b
$ do
        [IOVector Int]
parts <- Int
-> Int -> (Int -> Int -> IO (IOVector Int)) -> IO [IOVector Int]
runPartialsOver Int
n Int
capabilities (Vector Int -> Vector Int -> Int -> Int -> Int -> IO (IOVector Int)
sumIntChunk Vector Int
g Vector Int
v Int
nGroups)
        Int -> [IOVector Int] -> IO (Vector Int)
mergeIntSum Int
nGroups [IOVector Int]
parts
  where
    !n :: Int
n = Vector Int -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Int
v
{-# NOINLINE sumIntDirect #-}

sumIntChunk ::
    VU.Vector Int -> VU.Vector Int -> Int -> Int -> Int -> IO (VUM.IOVector Int)
sumIntChunk :: Vector Int -> Vector Int -> Int -> Int -> Int -> IO (IOVector Int)
sumIntChunk Vector Int
g Vector Int
v Int
nGroups Int
lo Int
hi = do
    IOVector Int
acc <- Int -> Int -> IO (MVector (PrimState IO) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate Int
nGroups (Int
0 :: Int)
    let go :: Int -> IO ()
go !Int
i
            | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
hi = () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
            | Bool
otherwise = do
                let !k :: Int
k = Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector Int
g Int
i
                Int
c <- MVector (PrimState IO) Int -> Int -> IO Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VUM.unsafeRead IOVector Int
MVector (PrimState IO) Int
acc Int
k
                MVector (PrimState IO) Int -> Int -> Int -> IO ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VUM.unsafeWrite IOVector Int
MVector (PrimState IO) Int
acc Int
k (Int
c Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector Int
v Int
i)
                Int -> IO ()
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
    Int -> IO ()
go Int
lo
    IOVector Int -> IO (IOVector Int)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure IOVector Int
acc

-------------------------------------------------------------------------------
-- Integer min / max (order-independent)
-------------------------------------------------------------------------------

extremaIntDirect ::
    Bool -> VU.Vector Int -> Int -> VU.Vector Int -> VU.Vector Int
extremaIntDirect :: Bool -> Vector Int -> Int -> Vector Int -> Vector Int
extremaIntDirect Bool
isMin Vector Int
g Int
nGroups Vector Int
v
    | Bool -> Bool
not (Int -> Bool
shouldPar Int
n) =
        IO (Vector Int) -> Vector Int
forall a. IO a -> a
unsafePerformIO (Bool
-> Vector Int
-> Vector Int
-> Int
-> Int
-> Int
-> IO (IOVector Int)
extremaIntChunk Bool
isMin Vector Int
g Vector Int
v Int
nGroups Int
0 Int
n IO (IOVector Int)
-> (IOVector Int -> IO (Vector Int)) -> IO (Vector Int)
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= IOVector Int -> IO (Vector Int)
MVector (PrimState IO) Int -> IO (Vector Int)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze)
    | Bool
otherwise = IO (Vector Int) -> Vector Int
forall a. IO a -> a
unsafePerformIO (IO (Vector Int) -> Vector Int) -> IO (Vector Int) -> Vector Int
forall a b. (a -> b) -> a -> b
$ do
        [IOVector Int]
parts <- Int
-> Int -> (Int -> Int -> IO (IOVector Int)) -> IO [IOVector Int]
runPartialsOver Int
n Int
capabilities (Bool
-> Vector Int
-> Vector Int
-> Int
-> Int
-> Int
-> IO (IOVector Int)
extremaIntChunk Bool
isMin Vector Int
g Vector Int
v Int
nGroups)
        Bool -> Int -> [IOVector Int] -> IO (Vector Int)
mergeExtremaInt Bool
isMin Int
nGroups [IOVector Int]
parts
  where
    !n :: Int
n = Vector Int -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Int
v
{-# NOINLINE extremaIntDirect #-}

extremaIntChunk ::
    Bool ->
    VU.Vector Int ->
    VU.Vector Int ->
    Int ->
    Int ->
    Int ->
    IO (VUM.IOVector Int)
extremaIntChunk :: Bool
-> Vector Int
-> Vector Int
-> Int
-> Int
-> Int
-> IO (IOVector Int)
extremaIntChunk Bool
isMin Vector Int
g Vector Int
v Int
nGroups Int
lo Int
hi = do
    let !seed :: Int
seed = if Bool
isMin then Int
forall a. Bounded a => a
maxBound else Int
forall a. Bounded a => a
minBound
        combine :: Int -> Int -> Int
combine Int
a Int
b = if Bool
isMin then Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
a Int
b else Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
a Int
b
    IOVector Int
acc <- Int -> Int -> IO (MVector (PrimState IO) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate Int
nGroups Int
seed
    let go :: Int -> IO ()
go !Int
i
            | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
hi = () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
            | Bool
otherwise = do
                let !k :: Int
k = Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector Int
g Int
i
                Int
c <- MVector (PrimState IO) Int -> Int -> IO Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VUM.unsafeRead IOVector Int
MVector (PrimState IO) Int
acc Int
k
                MVector (PrimState IO) Int -> Int -> Int -> IO ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VUM.unsafeWrite IOVector Int
MVector (PrimState IO) Int
acc Int
k (Int -> Int -> Int
combine Int
c (Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector Int
v Int
i))
                Int -> IO ()
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
    Int -> IO ()
go Int
lo
    IOVector Int -> IO (IOVector Int)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure IOVector Int
acc

-------------------------------------------------------------------------------
-- Integer mean (exact integer sum + count, divided once -> order-independent)
-------------------------------------------------------------------------------

{- | Integer mean in ONE fused pass: a running integer sum and count per group,
divided once at finalize. The integer sum is exact, so the parallel partial
merge is byte-identical to the sequential single pass at any @-N@.
-}
meanIntDirect :: VU.Vector Int -> Int -> VU.Vector Int -> VU.Vector Double
meanIntDirect :: Vector Int -> Int -> Vector Int -> Vector Double
meanIntDirect Vector Int
g Int
nGroups Vector Int
v
    | Bool -> Bool
not (Int -> Bool
shouldPar Int
n) = IO (Vector Double) -> Vector Double
forall a. IO a -> a
unsafePerformIO (IO (Vector Double) -> Vector Double)
-> IO (Vector Double) -> Vector Double
forall a b. (a -> b) -> a -> b
$ do
        (IOVector Int
s, IOVector Int
c) <- Vector Int
-> Vector Int
-> Int
-> Int
-> Int
-> IO (IOVector Int, IOVector Int)
meanIntChunk Vector Int
g Vector Int
v Int
nGroups Int
0 Int
n
        Int -> IOVector Int -> IOVector Int -> IO (Vector Double)
finalizeMeanInt Int
nGroups IOVector Int
s IOVector Int
c
    | Bool
otherwise = IO (Vector Double) -> Vector Double
forall a. IO a -> a
unsafePerformIO (IO (Vector Double) -> Vector Double)
-> IO (Vector Double) -> Vector Double
forall a b. (a -> b) -> a -> b
$ do
        [(IOVector Int, IOVector Int)]
parts <- Int
-> Int
-> (Int -> Int -> IO (IOVector Int, IOVector Int))
-> IO [(IOVector Int, IOVector Int)]
runPartialsPairOver Int
n Int
capabilities (Vector Int
-> Vector Int
-> Int
-> Int
-> Int
-> IO (IOVector Int, IOVector Int)
meanIntChunk Vector Int
g Vector Int
v Int
nGroups)
        (IOVector Int
s, IOVector Int
c) <- Int
-> [(IOVector Int, IOVector Int)]
-> IO (IOVector Int, IOVector Int)
mergePair Int
nGroups [(IOVector Int, IOVector Int)]
parts
        Int -> IOVector Int -> IOVector Int -> IO (Vector Double)
finalizeMeanInt Int
nGroups IOVector Int
s IOVector Int
c
  where
    !n :: Int
n = Vector Int -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Int
v
{-# NOINLINE meanIntDirect #-}

meanIntChunk ::
    VU.Vector Int ->
    VU.Vector Int ->
    Int ->
    Int ->
    Int ->
    IO (VUM.IOVector Int, VUM.IOVector Int)
meanIntChunk :: Vector Int
-> Vector Int
-> Int
-> Int
-> Int
-> IO (IOVector Int, IOVector Int)
meanIntChunk Vector Int
g Vector Int
v Int
nGroups Int
lo Int
hi = do
    IOVector Int
s <- Int -> Int -> IO (MVector (PrimState IO) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate Int
nGroups (Int
0 :: Int)
    IOVector Int
c <- Int -> Int -> IO (MVector (PrimState IO) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate Int
nGroups (Int
0 :: Int)
    let go :: Int -> IO ()
go !Int
i
            | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
hi = () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
            | Bool
otherwise = do
                let !k :: Int
k = Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector Int
g Int
i
                Int
sv <- MVector (PrimState IO) Int -> Int -> IO Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VUM.unsafeRead IOVector Int
MVector (PrimState IO) Int
s Int
k
                MVector (PrimState IO) Int -> Int -> Int -> IO ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VUM.unsafeWrite IOVector Int
MVector (PrimState IO) Int
s Int
k (Int
sv Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
VU.unsafeIndex Vector Int
v Int
i)
                Int
cv <- MVector (PrimState IO) Int -> Int -> IO Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VUM.unsafeRead IOVector Int
MVector (PrimState IO) Int
c Int
k
                MVector (PrimState IO) Int -> Int -> Int -> IO ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VUM.unsafeWrite IOVector Int
MVector (PrimState IO) Int
c Int
k (Int
cv Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
                Int -> IO ()
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
    Int -> IO ()
go Int
lo
    (IOVector Int, IOVector Int) -> IO (IOVector Int, IOVector Int)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (IOVector Int
s, IOVector Int
c)

finalizeMeanInt ::
    Int -> VUM.IOVector Int -> VUM.IOVector Int -> IO (VU.Vector Double)
finalizeMeanInt :: Int -> IOVector Int -> IOVector Int -> IO (Vector Double)
finalizeMeanInt Int
nGroups IOVector Int
s IOVector Int
c = do
    MVector RealWorld Double
out <- Int -> IO (MVector (PrimState IO) Double)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.new Int
nGroups
    let go :: Int -> IO ()
go !Int
k
            | Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
nGroups = () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
            | Bool
otherwise = do
                Int
sv <- MVector (PrimState IO) Int -> Int -> IO Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VUM.unsafeRead IOVector Int
MVector (PrimState IO) Int
s Int
k
                Int
cv <- MVector (PrimState IO) Int -> Int -> IO Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VUM.unsafeRead IOVector Int
MVector (PrimState IO) Int
c Int
k
                MVector (PrimState IO) Double -> Int -> Double -> IO ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VUM.unsafeWrite
                    MVector RealWorld Double
MVector (PrimState IO) Double
out
                    Int
k
                    (if Int
cv Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 then Double
0 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
0 else Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
sv Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
cv)
                Int -> IO ()
go (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
    Int -> IO ()
go Int
0
    MVector (PrimState IO) Double -> IO (Vector Double)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector RealWorld Double
MVector (PrimState IO) Double
out

-------------------------------------------------------------------------------
-- Partial accumulation + merge
-------------------------------------------------------------------------------

mergeIntSum :: Int -> [VUM.IOVector Int] -> IO (VU.Vector Int)
mergeIntSum :: Int -> [IOVector Int] -> IO (Vector Int)
mergeIntSum Int
nGroups [IOVector Int]
parts = case [IOVector Int]
parts of
    [] -> IOVector Int -> IO (Vector Int)
MVector (PrimState IO) Int -> IO (Vector Int)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze (IOVector Int -> IO (Vector Int))
-> IO (IOVector Int) -> IO (Vector Int)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Int -> Int -> IO (MVector (PrimState IO) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate Int
nGroups Int
0
    (IOVector Int
p0 : [IOVector Int]
rest) -> do
        let add :: IOVector Int -> IO ()
add !IOVector Int
p = do
                let go :: Int -> IO ()
go !Int
k
                        | Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
nGroups = () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
                        | Bool
otherwise = do
                            Int
a <- MVector (PrimState IO) Int -> Int -> IO Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VUM.unsafeRead IOVector Int
MVector (PrimState IO) Int
p0 Int
k
                            Int
b <- MVector (PrimState IO) Int -> Int -> IO Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VUM.unsafeRead IOVector Int
MVector (PrimState IO) Int
p Int
k
                            MVector (PrimState IO) Int -> Int -> Int -> IO ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VUM.unsafeWrite IOVector Int
MVector (PrimState IO) Int
p0 Int
k (Int
a Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
b)
                            Int -> IO ()
go (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
                Int -> IO ()
go Int
0
        (IOVector Int -> IO ()) -> [IOVector Int] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ IOVector Int -> IO ()
add [IOVector Int]
rest
        MVector (PrimState IO) Int -> IO (Vector Int)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze IOVector Int
MVector (PrimState IO) Int
p0

{- | Merge per-worker (sum, count) partials into the first worker's pair by
exact integer addition; returns the accumulated pair for finalize.
-}
mergePair ::
    Int ->
    [(VUM.IOVector Int, VUM.IOVector Int)] ->
    IO (VUM.IOVector Int, VUM.IOVector Int)
mergePair :: Int
-> [(IOVector Int, IOVector Int)]
-> IO (IOVector Int, IOVector Int)
mergePair Int
nGroups [(IOVector Int, IOVector Int)]
parts = case [(IOVector Int, IOVector Int)]
parts of
    [] -> (,) (IOVector Int -> IOVector Int -> (IOVector Int, IOVector Int))
-> IO (IOVector Int)
-> IO (IOVector Int -> (IOVector Int, IOVector Int))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Int -> IO (MVector (PrimState IO) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate Int
nGroups Int
0 IO (IOVector Int -> (IOVector Int, IOVector Int))
-> IO (IOVector Int) -> IO (IOVector Int, IOVector Int)
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Int -> Int -> IO (MVector (PrimState IO) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate Int
nGroups Int
0
    ((IOVector Int
s0, IOVector Int
c0) : [(IOVector Int, IOVector Int)]
rest) -> do
        let add :: (IOVector Int, IOVector Int) -> IO ()
add (IOVector Int
s, IOVector Int
c) = do
                let go :: Int -> IO ()
go !Int
k
                        | Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
nGroups = () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
                        | Bool
otherwise = do
                            Int
sa <- MVector (PrimState IO) Int -> Int -> IO Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VUM.unsafeRead IOVector Int
MVector (PrimState IO) Int
s0 Int
k
                            Int
sb <- MVector (PrimState IO) Int -> Int -> IO Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VUM.unsafeRead IOVector Int
MVector (PrimState IO) Int
s Int
k
                            MVector (PrimState IO) Int -> Int -> Int -> IO ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VUM.unsafeWrite IOVector Int
MVector (PrimState IO) Int
s0 Int
k (Int
sa Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
sb)
                            Int
ca <- MVector (PrimState IO) Int -> Int -> IO Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VUM.unsafeRead IOVector Int
MVector (PrimState IO) Int
c0 Int
k
                            Int
cb <- MVector (PrimState IO) Int -> Int -> IO Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VUM.unsafeRead IOVector Int
MVector (PrimState IO) Int
c Int
k
                            MVector (PrimState IO) Int -> Int -> Int -> IO ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VUM.unsafeWrite IOVector Int
MVector (PrimState IO) Int
c0 Int
k (Int
ca Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
cb)
                            Int -> IO ()
go (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
                Int -> IO ()
go Int
0
        ((IOVector Int, IOVector Int) -> IO ())
-> [(IOVector Int, IOVector Int)] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (IOVector Int, IOVector Int) -> IO ()
add [(IOVector Int, IOVector Int)]
rest
        (IOVector Int, IOVector Int) -> IO (IOVector Int, IOVector Int)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (IOVector Int
s0, IOVector Int
c0)

mergeExtremaInt :: Bool -> Int -> [VUM.IOVector Int] -> IO (VU.Vector Int)
mergeExtremaInt :: Bool -> Int -> [IOVector Int] -> IO (Vector Int)
mergeExtremaInt Bool
isMin Int
nGroups [IOVector Int]
parts = case [IOVector Int]
parts of
    [] ->
        IOVector Int -> IO (Vector Int)
MVector (PrimState IO) Int -> IO (Vector Int)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze (IOVector Int -> IO (Vector Int))
-> IO (IOVector Int) -> IO (Vector Int)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Int -> Int -> IO (MVector (PrimState IO) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate Int
nGroups (if Bool
isMin then Int
forall a. Bounded a => a
maxBound else Int
forall a. Bounded a => a
minBound)
    (IOVector Int
p0 : [IOVector Int]
rest) -> do
        let combine :: Int -> Int -> Int
combine Int
a Int
b = if Bool
isMin then Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
a Int
b else Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
a Int
b
            add :: IOVector Int -> IO ()
add !IOVector Int
p = do
                let go :: Int -> IO ()
go !Int
k
                        | Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
nGroups = () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
                        | Bool
otherwise = do
                            Int
a <- MVector (PrimState IO) Int -> Int -> IO Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VUM.unsafeRead IOVector Int
MVector (PrimState IO) Int
p0 Int
k
                            Int
b <- MVector (PrimState IO) Int -> Int -> IO Int
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
VUM.unsafeRead IOVector Int
MVector (PrimState IO) Int
p Int
k
                            MVector (PrimState IO) Int -> Int -> Int -> IO ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
VUM.unsafeWrite IOVector Int
MVector (PrimState IO) Int
p0 Int
k (Int -> Int -> Int
combine Int
a Int
b)
                            Int -> IO ()
go (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
                Int -> IO ()
go Int
0
        (IOVector Int -> IO ()) -> [IOVector Int] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ IOVector Int -> IO ()
add [IOVector Int]
rest
        MVector (PrimState IO) Int -> IO (Vector Int)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze IOVector Int
MVector (PrimState IO) Int
p0