module VectorExtras.Immutable where

import qualified Data.HashMap.Strict as HashMap
import Data.Vector.Generic
import qualified Data.Vector.Generic.Mutable as Mutable
import qualified DeferredFolds.Unfoldr as Unfoldr
import VectorExtras.Prelude hiding (length)

-- |
-- Distribute the elements of a vector across the specified amount of chunks.
--
-- Depending on the size of the vector the first chunks may be larger in size then the others by one.
chunk :: (Vector v1 a, Vector v2 (v1 a)) => Int -> v1 a -> v2 (v1 a)
chunk :: Int -> v1 a -> v2 (v1 a)
chunk Int
chunksAmount v1 a
vector =
  let vectorLength :: Int
vectorLength = v1 a -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
length v1 a
vector
      smallerChunkSize :: Int
smallerChunkSize = Int -> Int -> Int
forall a. Integral a => a -> a -> a
div Int
vectorLength Int
chunksAmount
      largerChunkSize :: Int
largerChunkSize = Int -> Int
forall a. Enum a => a -> a
succ Int
smallerChunkSize
      largerChunksAmount :: Int
largerChunksAmount = Int
vectorLength Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
smallerChunkSize Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
chunksAmount
      largerChunksElementsAmount :: Int
largerChunksElementsAmount = Int
largerChunksAmount Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
largerChunkSize
   in Int -> (Int -> v1 a) -> v2 (v1 a)
forall (v :: * -> *) a. Vector v a => Int -> (Int -> a) -> v a
generate Int
chunksAmount ((Int -> v1 a) -> v2 (v1 a)) -> (Int -> v1 a) -> v2 (v1 a)
forall a b. (a -> b) -> a -> b
$ \Int
chunkIndex ->
        let chunkOriginalIndex :: Int
chunkOriginalIndex =
              if Int
largerChunksAmount Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
chunkIndex
                then Int
chunkIndex Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
largerChunkSize
                else Int
largerChunksElementsAmount Int -> Int -> Int
forall a. Num a => a -> a -> a
+ (Int
chunkIndex Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
largerChunksAmount) Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
smallerChunkSize
         in Int -> (Int -> a) -> v1 a
forall (v :: * -> *) a. Vector v a => Int -> (Int -> a) -> v a
generate (if Int
chunkIndex Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
largerChunksAmount then Int
largerChunkSize else Int
smallerChunkSize) ((Int -> a) -> v1 a) -> (Int -> a) -> v1 a
forall a b. (a -> b) -> a -> b
$ \Int
elemIndex ->
              v1 a -> Int -> a
forall (v :: * -> *) a. Vector v a => v a -> Int -> a
unsafeIndex v1 a
vector (Int -> a) -> Int -> a
forall a b. (a -> b) -> a -> b
$
                Int
chunkOriginalIndex Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
elemIndex

-- |
-- Construct from an unfoldr of the specified size.
--
-- It is your responsibility to ensure that the unfoldr is of the same size as the one specified.
{-# INLINE unfoldrWithSize #-}
unfoldrWithSize :: (Mutable.MVector (Mutable vector) a, Vector vector a) => Int -> Unfoldr a -> vector a
unfoldrWithSize :: Int -> Unfoldr a -> vector a
unfoldrWithSize Int
size Unfoldr a
unfoldr = Int -> Unfoldr (Int, a) -> vector a
forall (vector :: * -> *) a.
(MVector (Mutable vector) a, Vector vector a) =>
Int -> Unfoldr (Int, a) -> vector a
assocUnfoldrWithSize Int
size (Unfoldr a -> Unfoldr (Int, a)
forall a. Unfoldr a -> Unfoldr (Int, a)
Unfoldr.zipWithIndex Unfoldr a
unfoldr)

-- |
-- Construct from an unfoldr of associations of the specified size.
--
-- It is your responsibility to ensure that the indices in the unfoldr are within the specified size.
{-# INLINE assocUnfoldrWithSize #-}
assocUnfoldrWithSize :: (Mutable.MVector (Mutable vector) a, Vector vector a) => Int -> Unfoldr (Int, a) -> vector a
assocUnfoldrWithSize :: Int -> Unfoldr (Int, a) -> vector a
assocUnfoldrWithSize Int
size Unfoldr (Int, a)
unfoldr =
  (forall s. ST s (vector a)) -> vector a
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (vector a)) -> vector a)
-> (forall s. ST s (vector a)) -> vector a
forall a b. (a -> b) -> a -> b
$ do
    Mutable vector s a
mv <- Int -> ST s (Mutable vector (PrimState (ST s)) a)
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
Int -> m (v (PrimState m) a)
Mutable.new Int
size
    Unfoldr (Int, a) -> ((Int, a) -> ST s ()) -> ST s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
VectorExtras.Prelude.forM_ Unfoldr (Int, a)
unfoldr (((Int, a) -> ST s ()) -> ST s ())
-> ((Int, a) -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \(Int
index, a
element) -> Mutable vector (PrimState (ST s)) a -> Int -> a -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
Mutable.write Mutable vector s a
Mutable vector (PrimState (ST s)) a
mv Int
index a
element
    Mutable vector (PrimState (ST s)) a -> ST s (vector a)
forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
Mutable v (PrimState m) a -> m (v a)
freeze Mutable vector s a
Mutable vector (PrimState (ST s)) a
mv

-- |
-- Construct from a hash-map of the specified size.
--
-- It is your responsibility to ensure that the indices in the unfoldr are within the specified size.
{-# INLINE indexHashMapWithSize #-}
indexHashMapWithSize :: (Mutable.MVector (Mutable vector) a, Vector vector a) => Int -> HashMap a Int -> vector a
indexHashMapWithSize :: Int -> HashMap a Int -> vector a
indexHashMapWithSize Int
size = Int -> Unfoldr (Int, a) -> vector a
forall (vector :: * -> *) a.
(MVector (Mutable vector) a, Vector vector a) =>
Int -> Unfoldr (Int, a) -> vector a
assocUnfoldrWithSize Int
size (Unfoldr (Int, a) -> vector a)
-> (HashMap a Int -> Unfoldr (Int, a)) -> HashMap a Int -> vector a
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. ((a, Int) -> (Int, a)) -> Unfoldr (a, Int) -> Unfoldr (Int, a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (a, Int) -> (Int, a)
forall a b. (a, b) -> (b, a)
swap (Unfoldr (a, Int) -> Unfoldr (Int, a))
-> (HashMap a Int -> Unfoldr (a, Int))
-> HashMap a Int
-> Unfoldr (Int, a)
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. HashMap a Int -> Unfoldr (a, Int)
forall key value. HashMap key value -> Unfoldr (key, value)
Unfoldr.hashMapAssocs

-- |
-- Construct from a hash-map.
{-# INLINE indexHashMap #-}
indexHashMap :: (Mutable.MVector (Mutable vector) a, Vector vector a) => HashMap a Int -> vector a
indexHashMap :: HashMap a Int -> vector a
indexHashMap HashMap a Int
hashMap = Int -> HashMap a Int -> vector a
forall (vector :: * -> *) a.
(MVector (Mutable vector) a, Vector vector a) =>
Int -> HashMap a Int -> vector a
indexHashMapWithSize (HashMap a Int -> Int
forall k v. HashMap k v -> Int
HashMap.size HashMap a Int
hashMap) HashMap a Int
hashMap