module Indexation.Vector
where
import Indexation.Prelude
import Data.Vector
import qualified Data.Vector.Mutable as MutableVector
import qualified Data.HashMap.Strict as HashMap
import qualified DeferredFolds.UnfoldM as UnfoldM
import qualified Data.HashTable.IO as HashtablesIO
import qualified Data.HashTable.Class as HashtablesClass
{-# NOINLINE populate #-}
populate :: Monad effect => Int -> effect (Int, element) -> effect (Vector element)
populate size effect =
do
mv <- return (unsafeDupablePerformIO (MutableVector.unsafeNew size))
let
loop stepsRemaining =
if stepsRemaining > 0
then do
(index, element) <- effect
() <- return (unsafeDupablePerformIO (MutableVector.write mv index element))
loop (pred stepsRemaining)
else do
!v <- return (unsafeDupablePerformIO (freeze mv))
return v
in loop size
{-# INLINE indexHashMapWithSize #-}
indexHashMapWithSize :: Int -> HashMap element Int -> Vector element
indexHashMapWithSize size hashMap =
runST $ do
mv <- MutableVector.new size
HashMap.foldrWithKey
(\ element index action -> MutableVector.write mv index element >> action)
(return ())
hashMap
freeze mv
{-# NOINLINE unfoldM #-}
unfoldM :: Monad m => Int -> UnfoldM m (Int, element) -> m (Vector element)
unfoldM size unfoldM =
let
step mv (index, element) = return (unsafeDupablePerformIO (MutableVector.write mv index element $> mv))
in do
!mv <- return (unsafeDupablePerformIO (MutableVector.unsafeNew size))
UnfoldM.foldlM' step mv unfoldM
!iv <- return (unsafeDupablePerformIO (unsafeFreeze mv))
return iv
hashTable :: HashtablesClass.HashTable t => Int -> HashtablesIO.IOHashTable t element Int -> IO (Vector element)
hashTable size hashTable = do
mv <- MutableVector.unsafeNew size
flip HashtablesIO.mapM_ hashTable $ \ (element, index) -> MutableVector.write mv index element
unsafeFreeze mv