vector-pull (empty) → 0.1.0.0
raw patch · 3 files changed
+1516/−0 lines, 3 filesdep +atomic-counterdep +basedep +exceptions
Dependencies added: atomic-counter, base, exceptions, hedgehog, hspec, hspec-hedgehog, optics-core, streaming, vector, vector-pull
Files
- lib/Data/Vector/Pull.hs +1025/−0
- tests/spec.hs +345/−0
- vector-pull.cabal +146/−0
+ lib/Data/Vector/Pull.hs view
@@ -0,0 +1,1025 @@+{-# LANGUAGE AllowAmbiguousTypes #-}+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE MagicHash #-}+{-# LANGUAGE NamedFieldPuns #-}+{-# LANGUAGE RankNTypes #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE Strict #-}+{-# LANGUAGE TypeApplications #-}+{-# LANGUAGE UnboxedTuples #-}+{-# OPTIONS_GHC -Wno-dodgy-imports #-}+{-# OPTIONS_GHC -Wno-name-shadowing #-}+{-# OPTIONS_GHC -Wno-unrecognised-pragmas #-}+{-# HLINT ignore "Redundant lambda" #-}+{-# HLINT ignore "Avoid lambda" #-}+{-# OPTIONS_GHC -Wno-unused-do-bind #-}++-- |+-- Module : Data.Vector.Pull+-- Copyright : (c) Michael Ledger 2024-2026+-- License : MPL-2.0+-- Maintainer : Michael Ledger <mike@quasimal.com>+--+-- "Pull" vectors are arrays represented as a function from an index to a value.+-- This representation allows for vector operations to be completely fused+-- together with little effort -- other than inlining. It works well in+-- scenarios where your vector operations are all chained together around a+-- single source location, *and* consumption/materialisation also takes place+-- there.+--+-- But beware! If operations do not fuse properly, you can pay for it dearly.+-- Every operation adds another closure which much be evaluated for every single+-- element access.+module Data.Vector.Pull (+ Pull,+ optimise,++ -- * Construction+ fromVector,+ safeFromVector,+ fromMVector,+ mapFromVector,+ fromList,+ empty,+ singleton,+ append,+ enumFromTo,+ enumFromLen,+ replicate,+ generate,+ surround,+ surroundMaybes,+ intersperse,+ intersperseWith,+ intersperseMapWith,+ cons,+ snoc,++ -- ** Manipulation+ take,+ drop,++ -- * Use+ (!),+ (!?),+ set,+ modify,++ -- * Consuming 'Pull' vectors+ toVector,+ toVectorA,+ toVectorM,+ chunkPullM,+ toListM,+ toStreamM,+ toStream,+ writeToVector,+ writeToVectorM,+ mapToVector,+ uncons,+ head,+ last,++ -- ** Concurrent consumers+ concurrentToVectorM,+ concurrentChunkedToVectorM,++ -- * Size+ length,++ -- * Mapping+ map,+ imap,+ zipWith,+ mapWithNext,+ enumerate,+ Enumerated (Enumerated),++ -- * Folds+ foldr,+ foldr',+ foldl,+ foldl',+ folded,++ -- ** Indexed variants+ Data.Vector.Pull.ifoldr,+ Data.Vector.Pull.ifoldr',+ Data.Vector.Pull.ifoldl,+ Data.Vector.Pull.ifoldl',++ -- * Traversals+ traverse,+ mapM,+ mapM_,++ -- ** Indexed variants+ imapM,+ imapM_,++ -- * Traversals+ update,+)+where++import Control.Concurrent (newEmptyMVar, takeMVar, tryPutMVar)+import Control.Concurrent.Counter qualified as Counter+import Control.Monad (ap, liftM2, void, when, (>=>))+import Control.Monad.Catch (MonadCatch, MonadThrow (throwM), catchAll)+import Control.Monad.IO.Class (MonadIO (..))+import Data.Coerce+import Data.Foldable qualified as Foldable+import Data.Functor ((<&>))+import Data.Monoid+import Data.Vector (Vector)+import Data.Vector qualified as V+import Data.Vector.Generic qualified as G+import Data.Vector.Generic.Mutable qualified as GM+import Data.Vector.Unboxed qualified as UV+import GHC.Base (Int#, quotInt#, quotRemInt#, tagToEnum#, (*#), (+#), (-#), (<#), (<=#), (==#), (>#), (>=#))+import GHC.Exts (TYPE)+import GHC.Int (Int (..))+import Optics.AffineTraversal (atraversalVL)+import Optics.At.Core (Index, IxValue, Ixed (ix))+import Optics.Fold (Fold, foldring)+import Optics.Indexed.Core (FoldableWithIndex (..), ifor_)+import Streaming (Of, Stream)+import Streaming.Internal qualified as S+import Streaming.Prelude qualified as S+import Text.Printf (printf)+import Text.Show (Show (show))+import Prelude hiding (+ atomically,+ cons,+ drop,+ empty,+ enumFromTo,+ folded,+ foldl,+ foldl',+ foldr,+ fromList,+ head,+ ifoldl,+ ifoldl',+ ifoldr,+ ifoldr',+ imap,+ last,+ length,+ map,+ mapM,+ mapM_,+ modify,+ newTVarIO,+ replicate,+ set,+ show,+ snoc,+ take,+ toList,+ traverse,+ uncons,+ zipWith,+ (!?),+ )++-- | An index-based data structure. Conceptually it is just a @ Int -> a @. This+-- should have faster indexing, but slower everything-else. When your+-- construction of a 'Pull' is able to inline, the results can be extremely+-- efficient. When inlining is not available, modifications to the generator+-- layer closures on top of eachother.+data Pull (a :: TYPE r) = Pull+ { index# :: Int# -> a+ , length# :: Int#+ }++instance Functor Pull where+ {-# INLINE fmap #-}+ fmap = map++instance Applicative Pull where+ {-# INLINE pure #-}+ {-# INLINE (<*>) #-}+ {-# INLINE liftA2 #-}+ pure = singleton+ liftA2 = liftM2+ (<*>) = ap++instance Monad Pull where+ {-# INLINE (>>=) #-}+ m >>= k = foldr (append . k) empty m++instance Foldable Pull where+ {-# INLINE foldr #-}+ {-# INLINE foldr' #-}+ {-# INLINE foldl #-}+ {-# INLINE foldl' #-}+ {-# INLINE toList #-}+ foldr = Data.Vector.Pull.foldr+ foldr' = Data.Vector.Pull.foldr'+ foldl = Data.Vector.Pull.foldl+ foldl' = Data.Vector.Pull.foldl'+ length = Data.Vector.Pull.length+ toList = Data.Vector.Pull.toList++instance FoldableWithIndex Int Pull where+ ifoldMap = Data.Vector.Pull.ifoldMap+ ifoldMap' = Data.Vector.Pull.ifoldMap'+ ifoldr = Data.Vector.Pull.ifoldr+ ifoldr' = Data.Vector.Pull.ifoldr'+ ifoldl = Data.Vector.Pull.ifoldl+ ifoldl' = Data.Vector.Pull.ifoldl'++instance (Show a) => Show (Pull a) where+ show = show . toVector @V.Vector++instance Semigroup (Pull a) where+ (<>) = append++instance Monoid (Pull a) where+ mempty = empty++type instance Index (Pull a) = Int++type instance IxValue (Pull a) = a++instance Ixed (Pull a) where+ {-# INLINE ix #-}+ ix i = atraversalVL \point f x ->+ if 0 <= i && i < length x+ then f (x ! i) <&> set x i+ else point x++{-# INLINE optimise #-}++-- | /O(N * D)/+optimise :: forall v a. (G.Vector v a) => Pull a -> Pull a+optimise = fromVector @v . toVector++--------------------------------------------------------------------------------+-- Construction++{-# INLINE fromVector #-}+fromVector :: (G.Vector v a) => v a -> Pull a+fromVector !vec =+ Pull+ { length# = case G.length vec of I# a -> a+ , index# = G.unsafeIndex vec .# I#+ }++{-# INLINE safeFromVector #-}+safeFromVector :: (G.Vector v a) => v a -> Pull a+safeFromVector !vec =+ Pull+ { length# = case G.length vec of I# a -> a+ , index# = \i -> vec G.! I# i+ }++{-# INLINE fromMVector #-}+fromMVector :: (G.Vector v a, GM.PrimMonad m) => G.Mutable v (GM.PrimState m) a -> Pull (m a)+fromMVector !vec =+ Pull+ { length# = case GM.length vec of I# a -> a+ , index# = \i -> GM.unsafeRead vec (I# i)+ }++{-# INLINE mapFromVector #-}+mapFromVector :: forall v b a. (G.Vector v a) => (a -> b) -> v a -> Pull b+mapFromVector f = fmap f . fromVector++fromList :: [a] -> Pull a+fromList = fromVector . V.fromList++{-# INLINE empty #-}+empty :: Pull a+empty =+ Pull+ { index# = outOfBoundsError# "empty" 0#+ , length# = 0#+ }++{-# INLINE singleton #-}+singleton :: a -> Pull a+singleton a =+ Pull+ { index# = \i -> case i of+ 0# -> a+ _ -> outOfBoundsError# "singleton" 1# i+ , length# = 1#+ }++{-# INLINE cons #-}+cons :: a -> Pull a -> Pull a+cons x Pull {index#, length#} =+ Pull+ { index# = \i# -> case i# ># 0# of+ 1# -> index# (i# -# 1#)+ _ -> x+ , length# = length# +# 1#+ }++{-# INLINE snoc #-}+snoc :: Pull a -> a -> Pull a+snoc Pull {index#, length#} x =+ Pull+ { index# = \i# -> case i# <# length# of+ 1# -> index# i#+ _ -> x+ , length# = length# +# 1#+ }++{-# INLINE append #-}+append :: Pull a -> Pull a -> Pull a+append a b =+ Pull+ { length# = length# a +# length# b+ , index# =+ \i -> case i <# length# a of+ 1# -> index# a i+ _ -> index# b (i -# length# a)+ }++{-# INLINE surround #-}++-- | 'cons' and 'snoc' at the same time+surround :: a -> Pull a -> a -> Pull a+surround l Pull {length#, index#} r =+ Pull+ { length# = length# +# 2#+ , index# = \i# -> case i# <# 1# of+ 1# -> l+ _ -> case i# <# (length# +# 1#) of+ 1# -> index# (i# -# 1#)+ _ -> r+ }++{-# INLINE surroundMaybes #-}++-- | 'cons' and 'snoc' at the same time with optional elements+surroundMaybes :: Maybe a -> Pull a -> Maybe a -> Pull a+surroundMaybes (Just l) g (Just r) = surround l g r+surroundMaybes (Just l) g _ = cons l g+surroundMaybes _ g (Just r) = snoc g r+surroundMaybes _ g _ = g++{-# INLINE intersperseMapWith #-}+intersperseMapWith+ :: (b -> b -> b)+ -- ^ How to compute midpoints between elements+ -> (a -> b)+ -- ^ Mapping function+ -> Pull a+ -> Pull b+intersperseMapWith avg lift Pull {index#, length#} = case length# <=# 1# of+ 1# -> Pull {index# = \i# -> lift (index# i#), length#}+ _ ->+ Pull+ { length# = length# *# 2# -# 1#+ , index# = \j# ->+ case quotRemInt# j# 2# of+ (# i#, 0# #) -> lift (index# i#)+ (# i#, _ #) -> avg (lift (index# i#)) (lift (index# (i# +# 1#)))+ }++{-# INLINE intersperseWith #-}+intersperseWith :: (a -> a -> a) -> Pull a -> Pull a+intersperseWith fn = intersperseMapWith fn id++{-# INLINE intersperse #-}+intersperse :: a -> Pull a -> Pull a+intersperse x = intersperseMapWith (\_ _ -> x) id++{-# INLINE enumFromTo #-}+enumFromTo :: (Enum a) => a -> a -> Pull a+enumFromTo (fromEnum -> I# a) (fromEnum -> I# b) =+ Pull+ { length# = (b -# a) +# 1#+ , index# = \i -> toEnum (I# (a +# i))+ }++{-# INLINE enumFromLen #-}+enumFromLen :: (Enum a) => a -> Int -> Pull a+enumFromLen (fromEnum -> I# a) (I# length#) =+ Pull+ { length#+ , index# = \i -> toEnum (I# (a +# i))+ }++{-# INLINE replicate #-}+replicate :: Int -> a -> Pull a+replicate (I# l) a =+ Pull+ { length# = l+ , index# = const# a+ }++{-# INLINE generate #-}+generate :: Int -> (Int -> a) -> Pull a+generate (I# l) f =+ Pull+ { length# = l+ , index# = \i -> f (I# i)+ }++--------------------------------------------------------------------------------++{-# INLINE take #-}+take :: Int -> Pull a -> Pull a+take (I# newLength#) Pull {index#, length#} =+ Pull {index#, length# = max# 0# (min# length# newLength#)}++{-# INLINE drop #-}+drop :: Int -> Pull a -> Pull a+drop (I# dropAmount#) it@Pull {index#, length#}+ | tagToEnum# (dropAmount# <# 0#) = it+ | tagToEnum# (length# ># dropAmount#) =+ Pull+ { length# = length# -# dropAmount#+ , index# = index# .+# dropAmount#+ }+ | otherwise = empty++--------------------------------------------------------------------------------+-- Updates++{-# INLINE set #-}++-- | /O(N)/ Beware+set :: Pull a -> Int -> a -> Pull a+set Pull {index#, length#} (I# i) a =+ Pull+ { length# = length#+ , index# = \j -> case j ==# i of+ 1# -> a+ _ -> index# j+ }++{-# INLINE modify #-}++-- | /O(N)/ Beware+modify :: Pull a -> Int -> (a -> a) -> Pull a+modify Pull {index#, length#} (I# i) f =+ Pull+ { length# = length#+ , index# = \j -> case j ==# i of+ 1# -> f (index# j)+ _ -> index# j+ }++{-# INLINE update #-}+update :: (Coercible i Int, UV.Unbox i) => Vector a -> (a -> b) -> UV.Vector i -> (Int -> b) -> Vector b+update v fn indices _indexFn+ | UV.null indices = V.map fn v+update v fn indices indexFn =+ V.unfoldrExactN+ (V.length v)+ step+ UpdState {i = 0, j = coerce (UV.head indices), next = 1}+ where+ numUpdates = UV.length indices+ step UpdState {i, j, next}+ | i == j =+ ( indexFn i+ , if next < numUpdates+ then UpdState (i + 1) (coerce (UV.unsafeIndex indices next)) (next + 1)+ else UpdState (i + 1) (-1) maxBound+ )+ | otherwise = (fn (v `V.unsafeIndex` i), UpdState {i = i + 1, j, next})++data UpdState = UpdState+ { i, j, next :: !Int+ }++--------------------------------------------------------------------------------+-- Materialisation++{-# INLINE toVector #-}++-- | /O(N * D)/+toVector :: forall v a. (G.Vector v a) => Pull a -> v a+toVector = mapToVector id++{-# INLINE toVectorA #-}++-- | /O(N * D)/+toVectorA :: forall v a f. (G.Vector v a, Applicative f) => Pull (f a) -> f (v a)+toVectorA g =+ G.fromListN (length g) <$> foldr (liftA2 (:)) (pure []) g++{-# INLINE traverse #-}+traverse :: (Applicative f) => (a -> f b) -> Pull a -> f (Pull b)+traverse f x = fromVector @Vector <$> toVectorA (map f x)++{-# INLINE toVectorM #-}+toVectorM :: forall v a m. (GM.PrimMonad m, G.Vector v a) => Pull (m a) -> m (v a)+toVectorM = toMutVectorM >=> G.unsafeFreeze++{-# INLINE toMutVectorM #-}+toMutVectorM :: forall v a m. (GM.PrimMonad m, G.Vector v a) => Pull (m a) -> m (G.Mutable v (GM.PrimState m) a)+toMutVectorM Pull {length#, index#} = do+ result <- GM.unsafeNew (I# length#)+ let go i = case i <# length# of+ 1# -> do+ GM.unsafeWrite result (I# i) =<< index# i+ go (i +# 1#)+ _ -> pass+ go 0#+ pure result++data Enumerated a = Enumerated# Int# ~a++pattern Enumerated :: Int -> a -> Enumerated a+pattern Enumerated x a <- Enumerated# (I# -> x) a+ where+ Enumerated (I# x) a = Enumerated# x a++{-# COMPLETE Enumerated #-}++{-# INLINE enumerate #-}+enumerate :: Pull a -> Pull (Enumerated a)+enumerate Pull {length#, index#} =+ Pull+ { length#+ , index# = \i# -> Enumerated# i# (index# i#)+ }++{-# INLINE chunkPullM #-}+chunkPullM :: forall a m. (Monad m) => Int -> Pull (m a) -> Pull (Pull (m a))+chunkPullM chunkSize _gen | chunkSize <= 0 = error "chunkSize must be positive"+chunkPullM (I# chunkSize#) Pull {length#, index#} =+ Pull+ { length# =+ case r# of+ 0# -> chunks#+ _ -> chunks# +# 1#+ , index# = \chunkIndex# ->+ Pull+ { index# = \j# -> index# ((chunkIndex# *# chunkSize#) +# j#)+ , length# = case chunkIndex# <# chunks# of+ 1# -> chunkSize#+ _ -> r#+ }+ }+ where+ !(# chunks#, r# #) = length# `quotRemInt#` chunkSize#++{-# INLINE concurrentChunkedToVectorM #-}+concurrentChunkedToVectorM+ :: forall v a m+ . (G.Vector v a, MonadIO m, MonadCatch m, MonadThrow m, GM.PrimMonad m)+ => (m () -> m ())+ -- ^ Fork function, e.g. @ (void . forkIO) @+ -> Int+ -> Pull (m a)+ -> m (v a)+concurrentChunkedToVectorM _ _ Pull {length# = 0#} = pure G.empty+concurrentChunkedToVectorM forkIO caps gen@Pull {length#, index#} =+ if caps <= 1+ then toVectorM gen+ else do+ dest <- GM.unsafeNew (I# length#)+ signal <- liftIO newEmptyMVar++ let+ !(I# caps#) = caps+ !chunkSize# = max# 1# (length# `quotInt#` caps#)+ !(# chunks#, r# #) = length# `quotRemInt#` chunkSize#++ remaining <- liftIO case r# of+ 0# -> Counter.new (I# chunks#)+ _ -> Counter.new (I# (chunks# +# 1#))++ let+ go# i# = case i# <# chunks# of+ 1# -> do+ forkIO do+ catchAll+ (goChunk# (i# *# chunkSize#) chunkSize#)+ (liftIO . void . tryPutMVar signal . Just)+ go# (i# +# 1#)+ _ -> case r# of+ 0# -> pass+ _ -> forkIO (goChunk# (i# *# chunkSize#) r#)++ goChunk# j# len# =+ case len# of+ 0# -> do+ prevRemaining <- liftIO (Counter.sub remaining 1)+ when (prevRemaining <= 1) (void (liftIO (tryPutMVar signal Nothing {- don't block in case the atomic counter failed and somehow multiple threads wrote to the signal -})))+ _ -> do+ x <- index# j#+ GM.unsafeWrite dest (I# j#) x+ goChunk# (j# +# 1#) (len# -# 1#)++ go# 0#++ anyError <- liftIO (takeMVar signal)+ case anyError of+ Just err -> throwM err+ Nothing -> G.unsafeFreeze dest++{-# INLINE concurrentToVectorM #-}+concurrentToVectorM+ :: (G.Vector v a, MonadIO m, MonadCatch m, MonadThrow m, GM.PrimMonad m)+ => (m () -> m ())+ -- ^ Fork function e.g. @ (void . forkIO) @+ -> Pull (m a)+ -> m (v a)+concurrentToVectorM _ Pull {length# = 0#} = pure G.empty+concurrentToVectorM forkIO gen = do+ dest <- GM.unsafeNew (length gen)+ signal <- liftIO newEmptyMVar+ remaining <- liftIO (Counter.new (length gen))+ ifor_ gen \i x -> forkIO do+ catchAll+ ( do+ result <- x+ GM.unsafeWrite dest i result+ prevRemaining <- liftIO (Counter.sub remaining 1)+ when (prevRemaining <= 1) (void (liftIO (tryPutMVar signal Nothing {- don't block in case atomic counter fail -})))+ )+ (void . liftIO . tryPutMVar signal . Just)+ -- re-throwing won't do anything meaningful here++ anyError <- liftIO (takeMVar signal)+ case anyError of+ Just e -> throwM e+ Nothing -> G.unsafeFreeze dest++{-# INLINE toListM #-}+toListM :: (Monad m) => Pull (m a) -> m [a]+toListM Pull {length#, index#} =+ go 0#+ where+ go i = case i <# length# of+ 1# -> do+ x <- index# i+ xs <- go (i +# 1#)+ pure (x : xs)+ _ ->+ pure []++{-# INLINE toStream #-}+toStream :: (Monad m) => Pull a -> Stream (Of a) m ()+toStream Pull {length#, index#} = go 0#+ where+ go i = case i <# length# of+ 1# -> S.Step (index# i S.:> go (i +# 1#))+ _ -> S.Return ()++{-# INLINE toStreamM #-}+toStreamM :: (Monad m) => Pull (m a) -> Stream (Of a) m ()+toStreamM Pull {length#, index#} = go 0#+ where+ go i = case i <# length# of+ 1# ->+ S.Effect+ ( index# i <&> \val ->+ S.Step (val S.:> go (i +# 1#))+ )+ _ -> S.Return ()++{-# INLINE writeToVector #-}+writeToVector :: (GM.PrimMonad m, GM.PrimState m ~ s, G.Vector v a) => G.Mutable v s a -> Pull a -> m ()+writeToVector dest Pull {length#, index#} = do+ let+ !(I# finalLength#) = min (GM.length dest) (I# length#)+ go i = case i <# finalLength# of+ 1# -> do+ GM.unsafeWrite dest (I# i) (index# i)+ go (i +# 1#)+ _ -> pass+ go 0#++{-# INLINE writeToVectorM #-}+writeToVectorM :: (GM.PrimMonad m, GM.PrimState m ~ s, G.Vector v a) => G.Mutable v s a -> Pull (m a) -> m ()+writeToVectorM dest Pull {length#, index#} = do+ let+ !(I# finalLength#) = min (GM.length dest) (I# length#)+ go i = case i <# finalLength# of+ 1# -> do+ GM.unsafeWrite dest (I# i) =<< index# i+ go (i +# 1#)+ _ -> pass+ go 0#++{-# INLINE toList #-}++-- | /O(N * D)/+toList :: Pull a -> [a]+toList = foldr (:) []++{-# INLINE uncons #-}+uncons :: Pull a -> Maybe (a, Pull a)+uncons Pull {length#, index#} = case length# of+ 0# -> Nothing+ _ ->+ Just+ ( index# 0#+ , Pull+ { length# = length# -# 1#+ , index# = \j -> index# (j +# 1#)+ }+ )++--------------------------------------------------------------------------------+-- Size++{-# INLINE length #-}+length :: Pull a -> Int+length Pull {length#} = I# length#++--------------------------------------------------------------------------------+-- Internal utility++{-# INLINE slice# #-}+slice#+ :: Pull a+ -> (Int# -> a -> b -> b)+ -> b+ -> Int#+ -- ^ starting index (inclusive)+ -> Int#+ -- ^ end index (non-inclusive)+ -> b+slice# Pull {index#} f z0 start end = go start+ where+ go i = case i <# end of+ 1# -> f i (index# i) (go (i +# 1#))+ _ -> z0++{-# INLINE mapToVector #-}+mapToVector :: forall v b a. (G.Vector v b) => (a -> b) -> Pull a -> v b+mapToVector _ Pull {length# = 0#} = G.empty+mapToVector f Pull {length#, index#} =+ G.create do+ m <- GM.new (I# length#)+ let go i# = case i# <# length# of+ 1# -> do+ GM.unsafeWrite m (I# i#) $! f (index# i#)+ go (i# +# 1#)+ _ -> pass+ go 0#+ pure m++--------------------------------------------------------------------------------+-- Indexing++{-# INLINE (!) #-}++-- | /O(D)/+(!) :: Pull a -> Int -> a+(!) Pull {index#, length#} (I# i#)+ | tagToEnum# (i# <# length#) = index# i#+ | otherwise = outOfBoundsError# "(!)" i# length#++{-# INLINE (!?) #-}++-- | /O(D)/+(!?) :: Pull a -> Int -> Maybe a+(!?) Pull {length#, index#} (I# i#)+ | tagToEnum# (i# <# length#) && tagToEnum# (i# >=# 0#) = Just (index# i#)+ | otherwise = Nothing++--------------------------------------------------------------------------------+-- Maps++{-# INLINE map #-}++{-# INLINE imap #-}++{-# INLINE zipWith #-}++{-# INLINE mapWithNext #-}++-- | /O(1)/+map :: (a -> b) -> Pull a -> Pull b+map f Pull {length#, index#} =+ Pull+ { length#+ , index# = \i -> f (index# i)+ }++-- | /O(1)/+imap :: (Int -> a -> b) -> Pull a -> Pull b+imap f Pull {length#, index#} =+ Pull+ { length#+ , index# = \i -> f (I# i) (index# i)+ }++zipWith :: (a -> b -> c) -> Pull a -> Pull b -> Pull c+zipWith f Pull {length# = lenA#, index# = indexA#} Pull {length# = lenB#, index# = indexB#} =+ Pull+ { length# = min# lenA# lenB#+ , index# = \i# -> f (indexA# i#) (indexB# i#)+ }++mapWithNext :: (a -> a -> b) -> Pull a -> a -> Pull b+mapWithNext f Pull {length#, index#} endMarker =+ Pull+ { length#+ , index# = \i# ->+ f+ (index# i#)+ ( if tagToEnum# (i# <# lenMinus1#)+ then index# (i# +# 1#)+ else endMarker+ )+ }+ where+ lenMinus1# = length# -# 1#++--------------------------------------------------------------------------------+-- Folds++{-# INLINE foldr #-}++{-# INLINE ifoldr #-}++{-# INLINE foldl #-}++{-# INLINE ifoldl #-}++{-# INLINE foldl' #-}++{-# INLINE ifoldl' #-}++{-# INLINE ifoldMap #-}++{-# INLINE ifoldMap' #-}++-- | Right-associative fold+foldr :: (a -> b -> b) -> b -> Pull a -> b+foldr f z o@Pull {length#} = slice# o (\_ !x xs -> f x xs) z 0# length#++-- | Right-associative strict fold+foldr' :: (a -> b -> b) -> b -> Pull a -> b+foldr' f z o@Pull {length#} = slice# o (\_ !x !xs -> f x xs) z 0# length#++-- | Right-associative strict fold, with index+ifoldMap :: (Monoid m) => (Int -> a -> m) -> Pull a -> m+ifoldMap f o@Pull {length#} =+ slice#+ o+ (\i x xs -> f (I# i) x <> xs)+ mempty+ 0#+ length#++-- | Right-associative strict fold, with index+ifoldMap' :: (Monoid m) => (Int -> a -> m) -> Pull a -> m+ifoldMap' f o@Pull {length#} =+ slice#+ o+ (\i !x !xs -> f (I# i) x <> xs)+ mempty+ 0#+ length#++-- | Right-associative strict fold, with index+ifoldr :: (Int -> a -> b -> b) -> b -> Pull a -> b+ifoldr f z o@Pull {length#} = slice# o (\i !x xs -> f (I# i) x xs) z 0# length#++-- | Right-associative strict fold, with index+ifoldr' :: (Int -> a -> b -> b) -> b -> Pull a -> b+ifoldr' f z o@Pull {length#} = slice# o (\i !x !xs -> f (I# i) x xs) z 0# length#++-- | Left-associative fold+foldl :: (b -> a -> b) -> b -> Pull a -> b+foldl f z0 o@Pull {length#} = slice# o (\_ !x k z -> k (f z x)) id 0# length# z0++-- | Strict left-associative fold+foldl' :: (b -> a -> b) -> b -> Pull a -> b+foldl' f z0 o@Pull {length#} = slice# o (\_ !x k z -> k $! f z x) id 0# length# z0++-- | Left-associative fold, with index+ifoldl :: (Int -> b -> a -> b) -> b -> Pull a -> b+ifoldl f z0 o@Pull {length#} = slice# o (\i !x k z -> k (f (I# i) z x)) id 0# length# z0++-- | Strict left-associative fold, with index+ifoldl' :: (Int -> b -> a -> b) -> b -> Pull a -> b+ifoldl' f z0 o@Pull {length#} = slice# o (\i !x k z -> k $! f (I# i) z x) id 0# length# z0++folded :: Fold (Pull a) a+folded = foldring foldr++--------------------------------------------------------------------------------+-- Traversals++{-# INLINE mapM #-}++{-# INLINE imapM #-}++{-# INLINE mapM_ #-}++{-# INLINE imapM_ #-}++-- | /O(N * D)/ optimising+mapM :: (GM.PrimMonad m) => (a -> m b) -> Pull a -> m (Pull b)+mapM f = fmap (fromVector @Vector) . toVectorM . map f++-- | /O(N * D)/ optimising+imapM :: (Monad m) => (Int -> a -> m b) -> Pull a -> m (Pull b)+imapM f = fmap fromVector . V.imapM f . toVector++-- | /O(N * D)/+mapM_ :: (Monad m) => (a -> m ()) -> Pull a -> m ()+mapM_ f Pull {index#, length#} = go 0#+ where+ go i = case i <# length# of+ 1# -> f (index# i) >> go (i +# 1#)+ _ -> pass++-- | /O(N * D)/+imapM_ :: (Monad m) => (Int -> a -> m ()) -> Pull a -> m ()+imapM_ f Pull {index#, length#} = go 0#+ where+ go i = case i <# length# of+ 1# -> f (I# i) (index# i) >> go (i +# 1#)+ _ -> pass++{-# NOINLINE outOfBoundsError# #-}+outOfBoundsError# :: String -> Int# -> Int# -> a+outOfBoundsError# a len# i# =+ outOfBoundsError a (I# i#) (I# len#)+ where+ {-# INLINE outOfBoundsError #-}+ outOfBoundsError :: String -> Int -> Int -> a+ outOfBoundsError a i len = error (printf "Data.Vector.Pull.%s: out of bounds i=%d len=%d" a i len)++--+-- {-# INLINE isfoldr# #-}+-- isfoldr#+-- :: forall v a b+-- . (G.Vector v a)+-- => v a+-- -> (Int# -> a -> b -> b)+-- -> b+-- -> Int#+-- -> Int#+-- -> b+-- isfoldr# v = \f z start end ->+-- let go i = case i <# end of+-- 1# -> f i (G.unsafeIndex v (I# i)) (go (i +# 1#))+-- _ -> z+-- in go start+--+-- {-# INLINE isfoldr2# #-}+-- isfoldr2#+-- :: forall v a b+-- . (G.Vector v a)+-- => v a+-- -> (Int# -> a -> b -> b)+-- -> b+-- -> Int#+-- -> Int#+-- -> Int#+-- -> b+-- isfoldr2# v = \f z fake0 start end ->+-- let go fake i = case fake <# end of+-- 1# -> f fake (G.unsafeIndex v (I# i)) (go (fake +# 1#) (i +# 1#))+-- _ -> z+-- in go (start +# fake0) start+--+--+--+-- # at the end of a function name? That means it goes fast!!!1++--------------------------------------------------------------------------------+-- Utilities / internals++{-# INLINE max# #-}+max# :: Int# -> Int# -> Int#+max# a b = if tagToEnum# (a <# b) then b else a++{-# INLINE min# #-}+min# :: Int# -> Int# -> Int#+min# a b = if tagToEnum# (a <# b) then a else b++head :: Pull a -> Maybe a+head Pull {index#, length#} = case length# of+ 0# -> Nothing+ _ -> Just (index# 0#)++last :: Pull a -> Maybe a+last Pull {index#, length#} = case length# of+ 0# -> Nothing+ _ -> Just (index# (length# -# 1#))++{-# INLINE pass #-}+pass :: (Monad m) => m ()+pass = pure ()++{-# INLINE (.#) #-}+(.#) :: (b -> a) -> (Int# -> b) -> Int# -> a+(.#) f g a = f (g a)++{-# INLINE (.+#) #-}+(.+#) :: (Int# -> a) -> Int# -> Int# -> a+(.+#) f n a = f (n +# a)++{-# INLINE const# #-}+const# :: a -> Int# -> a+const# a _ = a
+ tests/spec.hs view
@@ -0,0 +1,345 @@+{-# LANGUAGE Strict #-}++-- |+-- Module : vector-pull/tests/spec.hs+-- Copyright : (c) Michael Ledger 2026+-- License : MPL-2.0+-- Maintainer : Michael Ledger <mike@quasimal.com>+--+-- Largely generated with the help of claude-code, hence the smelly quality+module Main where++import Data.Foldable+import Data.List qualified as List+import Data.Maybe (isNothing)+import Data.Vector (Vector)+import Data.Vector qualified as V+import Data.Vector.Pull qualified as P+import Hedgehog+import Hedgehog.Gen qualified as Gen+import Hedgehog.Range qualified as Range+import Test.Hspec+import Test.Hspec.Hedgehog++main :: IO ()+main = hspec spec++genInt :: Gen Int+genInt = Gen.int (Range.linear (-1000) 1000)++genIntList :: Gen [Int]+genIntList = Gen.list (Range.linear 0 100) genInt++genIntVector :: Gen (Vector Int)+genIntVector = V.fromList <$> genIntList++genNonEmptyIntList :: Gen [Int]+genNonEmptyIntList = Gen.list (Range.linear 1 100) genInt++genNonNegativeInt :: Gen Int+genNonNegativeInt = Gen.int (Range.linear 0 100)++spec :: Spec+spec = do+ describe "Data.Vector.Pull" do+ constructionSpec+ manipulationSpec+ indexingSpec+ updateSpec+ consumingSpec+ mappingSpec+ foldSpec+ instanceSpec++constructionSpec :: Spec+constructionSpec = describe "Construction" do+ describe "fromList" do+ it "roundtrips with toList" $ hedgehog do+ xs <- forAll genIntList+ toList (P.fromList xs) === xs++ describe "empty" do+ it "has length 0" do+ P.length (P.empty @Int) `shouldBe` 0++ it "converts to empty vector" do+ P.toVector @Vector (P.empty @Int) `shouldBe` V.empty++ describe "singleton" do+ it "creates a Pull with one element" $ hedgehog do+ x <- forAll genInt+ P.toVector @Vector (P.singleton x) === V.singleton x++ it "has length 1" $ hedgehog do+ x <- forAll genInt+ P.length (P.singleton x) === 1++ describe "append" do+ it "concatenates two Pulls" $ hedgehog do+ xs <- forAll genIntVector+ ys <- forAll genIntVector+ P.toVector @Vector (P.fromVector xs `P.append` P.fromVector ys) === (xs <> ys)++ describe "enumFromTo" do+ it "creates a range" do+ P.toVector @Vector (P.enumFromTo (1 :: Int) 5) `shouldBe` V.fromList [1, 2, 3, 4, 5]++ it "handles single element range" do+ P.toVector @Vector (P.enumFromTo (3 :: Int) 3) `shouldBe` V.singleton 3++ describe "enumFromLen" do+ it "creates a range with specific length" do+ P.toVector @Vector (P.enumFromLen (1 :: Int) 5) `shouldBe` V.fromList [1, 2, 3, 4, 5]++ it "handles zero length" do+ P.toVector @Vector (P.enumFromLen (1 :: Int) 0) `shouldBe` V.empty++ describe "replicate" do+ it "creates n copies of an element" $ hedgehog do+ n <- forAll genNonNegativeInt+ x <- forAll genInt+ P.toVector @Vector (P.replicate n x) === V.replicate n x++ describe "generate" do+ it "creates elements from a function" do+ P.toVector @Vector (P.generate 5 (* 2)) `shouldBe` V.fromList [0, 2, 4, 6, 8]++ describe "cons" do+ it "prepends an element" $ hedgehog do+ x <- forAll genInt+ xs <- forAll genIntVector+ P.toVector @Vector (P.cons x (P.fromVector xs)) === V.cons x xs++ describe "snoc" do+ it "appends an element" $ hedgehog do+ xs <- forAll genIntVector+ x <- forAll genInt+ P.toVector @Vector (P.snoc (P.fromVector xs) x) === V.snoc xs x++ describe "surround" do+ it "adds elements at both ends" $ hedgehog do+ l <- forAll genInt+ xs <- forAll genIntVector+ r <- forAll genInt+ P.toVector @Vector (P.surround l (P.fromVector xs) r) === V.cons l (V.snoc xs r)++ describe "intersperse" do+ it "intersperses an element" do+ P.toVector @Vector (P.intersperse 0 (P.fromList [1, 2, 3 :: Int])) `shouldBe` V.fromList [1, 0, 2, 0, 3]++ it "handles empty list" do+ P.toVector @Vector (P.intersperse 0 (P.empty @Int)) `shouldBe` V.empty++ it "handles singleton" do+ P.toVector @Vector (P.intersperse 0 (P.singleton (1 :: Int))) `shouldBe` V.singleton 1++manipulationSpec :: Spec+manipulationSpec = describe "Manipulation" do+ describe "take" do+ it "takes first n elements" $ hedgehog do+ n <- forAll genNonNegativeInt+ xs <- forAll genIntVector+ P.toVector @Vector (P.take n (P.fromVector xs)) === V.take n xs++ it "handles taking more than length" do+ P.toVector @Vector (P.take 10 (P.fromVector (V.fromList [1, 2, 3 :: Int]))) `shouldBe` V.fromList [1, 2, 3]++ it "handles negative take (should clamp to 0)" do+ P.toVector @Vector (P.take (-5) (P.fromVector (V.fromList [1, 2, 3 :: Int]))) `shouldBe` V.empty++ describe "drop" do+ it "drops first n elements" $ hedgehog do+ n <- forAll genNonNegativeInt+ xs <- forAll genIntVector+ (P.toVector @Vector $! P.drop n $! P.fromVector xs) === V.drop n xs++ it "handles dropping more than length" do+ (P.toVector @Vector $! P.drop 10 $! P.fromVector $! V.fromList [1, 2, 3 :: Int]) `shouldBe` V.empty++ it "handles negative drop (should be identity)" do+ (P.toVector @Vector $! P.drop (-5) $! P.fromVector $! V.fromList [1, 2, 3 :: Int]) `shouldBe` V.fromList [1, 2, 3]++ it "property: works with any Int (including negative)" $ hedgehog do+ n <- forAll genInt+ xs <- forAll genIntVector+ (P.toVector @Vector $! P.drop n $! P.fromVector xs) === V.drop n xs++indexingSpec :: Spec+indexingSpec = describe "Indexing" do+ describe "(!)" do+ it "indexes correctly" do+ let p = P.fromList [10, 20, 30 :: Int]+ p P.! 0 `shouldBe` 10+ p P.! 1 `shouldBe` 20+ p P.! 2 `shouldBe` 30++ it "throws on out of bounds" do+ let p = P.fromList [1 :: Int]+ (pure $! p P.! 5) `shouldThrow` anyErrorCall++ describe "(!?)" do+ it "returns Just for valid index" do+ let p = P.fromList [10, 20, 30 :: Int]+ (p P.!? 1) `shouldBe` Just 20++ it "returns Nothing for invalid index" do+ let p = P.fromList [1, 2, 3 :: Int]+ (p P.!? 10) `shouldBe` Nothing+ (p P.!? (-1)) `shouldBe` Nothing++updateSpec :: Spec+updateSpec = describe "Updates" do+ describe "set" do+ it "sets an element at index" do+ let p = P.fromList [1, 2, 3 :: Int]+ P.toVector @Vector (P.set p 1 99) `shouldBe` V.fromList [1, 99, 3]++ describe "modify" do+ it "modifies an element at index" do+ let p = P.fromList [1, 2, 3 :: Int]+ P.toVector @Vector (P.modify p 1 (* 10)) `shouldBe` V.fromList [1, 20, 3]++consumingSpec :: Spec+consumingSpec = describe "Consuming" do+ describe "toVector" do+ it "converts to Vector" $ hedgehog do+ xs <- forAll genIntList+ P.toVector (P.fromList xs) === V.fromList xs++ describe "toList" do+ it "converts to list" $ hedgehog do+ xs <- forAll genIntList+ toList (P.fromList xs) === xs++ describe "uncons" do+ it "returns Nothing for empty" do+ P.uncons (P.empty @Int) `shouldSatisfy` isNothing++ it "returns head and tail" do+ let Just (h, t) = P.uncons (P.fromList [1, 2, 3 :: Int])+ h `shouldBe` 1+ P.toVector @Vector t `shouldBe` V.fromList [2, 3]++ describe "head" do+ it "returns Nothing for empty" do+ P.head (P.empty @Int) `shouldBe` Nothing++ it "returns first element" $ hedgehog do+ x <- forAll genInt+ xs <- forAll genIntList+ P.head (P.fromList (x : xs)) === Just x++ describe "last" do+ it "returns Nothing for empty" do+ P.last (P.empty @Int) `shouldBe` Nothing++ it "returns last element" $ hedgehog do+ xs <- forAll genNonEmptyIntList+ P.last (P.fromList xs) === Just (List.last xs)++ describe "length" do+ it "returns correct length" $ hedgehog do+ xs <- forAll genIntVector+ P.length (P.fromVector xs) === length xs++mappingSpec :: Spec+mappingSpec = describe "Mapping" do+ describe "map" do+ it "maps a function over elements" $ hedgehog do+ xs <- forAll genIntVector+ P.toVector @Vector (fmap (* 2) (P.fromVector xs)) === fmap (* 2) xs++ describe "imap" do+ it "maps with index" do+ P.toVector @Vector (P.imap (+) (P.fromList [10, 20, 30 :: Int])) `shouldBe` V.fromList [10, 21, 32]++ describe "zipWith" do+ it "zips two Pulls" $ hedgehog do+ xs <- forAll genIntVector+ ys <- forAll genIntVector+ P.toVector @Vector (P.zipWith (+) (P.fromVector xs) (P.fromVector ys)) === V.zipWith (+) xs ys++ describe "enumerate" do+ it "pairs elements with indices" do+ let result = P.toVector @Vector (P.map (\(P.Enumerated i x) -> (i, x)) (P.enumerate (P.fromList "abc")))+ result `shouldBe` V.fromList [(0, 'a'), (1, 'b'), (2, 'c')]++foldSpec :: Spec+foldSpec = describe "Folds" do+ describe "foldr" do+ it "right folds" $ hedgehog do+ xs <- forAll genIntVector+ P.foldr (+) 0 (P.fromVector xs) === foldr (+) 0 xs++ it "preserves order" do+ P.foldr (:) [] (P.fromVector (V.fromList [1, 2, 3 :: Int])) `shouldBe` [1, 2, 3]++ describe "foldr'" do+ it "strict right folds" $ hedgehog do+ xs <- forAll genIntVector+ P.foldr' (+) 0 (P.fromVector xs) === foldr (+) 0 xs++ describe "foldl" do+ it "left folds" $ hedgehog do+ xs <- forAll genIntVector+ P.foldl (+) 0 (P.fromVector xs) === foldl (+) 0 xs++ describe "foldl'" do+ it "strict left folds" $ hedgehog do+ xs <- forAll genIntVector+ P.foldl' (+) 0 (P.fromVector xs) === foldl' (+) 0 xs++ describe "ifoldr" do+ it "right folds with index" do+ P.ifoldr (\i x acc -> (i, x) : acc) [] (P.fromList "ab")+ `shouldBe` [(0, 'a'), (1, 'b')]++ describe "ifoldl'" do+ it "strict left folds with index" do+ P.ifoldl' (\i acc x -> acc + i + x) 0 (P.fromList [10, 20, 30 :: Int])+ `shouldBe` 63 -- 0 + (0+10) + (1+20) + (2+30)++instanceSpec :: Spec+instanceSpec = describe "Instances" do+ describe "Functor" do+ it "fmap is map" $ hedgehog do+ xs <- forAll genIntVector+ P.toVector @Vector (fmap (* 2) (P.fromVector xs)) === fmap (* 2) xs++ describe "Applicative" do+ it "pure creates singleton" do+ P.toVector @Vector (pure (42 :: Int) :: P.Pull Int) `shouldBe` V.singleton 42++ it "liftA2 works" do+ let+ p1 = P.fromList [1, 2 :: Int]+ p2 = P.fromList [10, 20 :: Int]+ P.toVector @Vector (liftA2 (+) p1 p2) `shouldBe` V.fromList [11, 21, 12, 22]++ describe "Monad" do+ it "bind works" do+ let p = P.fromList [1, 2, 3 :: Int]+ P.toVector @Vector (p >>= \x -> P.fromList [x, x * 10]) `shouldBe` V.fromList [1, 10, 2, 20, 3, 30]++ describe "Foldable" do+ it "sum via Foldable" $ hedgehog do+ xs <- forAll genIntVector+ sum (P.fromVector xs) === sum xs++ it "length via Foldable" $ hedgehog do+ xs <- forAll genIntVector+ length (P.fromVector xs) === length xs++ describe "Semigroup" do+ it "(<>)" $ hedgehog do+ xs <- forAll genIntVector+ ys <- forAll genIntVector+ P.toVector @Vector (P.fromVector xs <> P.fromVector ys) === (xs <> ys)++ describe "Monoid" do+ it "mempty is empty" do+ P.toVector @Vector (mempty :: P.Pull Int) `shouldBe` V.empty++ describe "Show" do+ it "shows like a vector" do+ show (P.fromList [1, 2, 3 :: Int]) `shouldBe` "[1,2,3]"
+ vector-pull.cabal view
@@ -0,0 +1,146 @@+cabal-version: 2.4+name: vector-pull+version: 0.1.0.0+license: MPL-2.0+copyright: 2024 Michael Ledger+maintainer: mike@quasimal.com+author: Michael Ledger+category: Data+synopsis: Pull-array data structure+description:+ An implementation of `Int`-indexed arrays that internally are represented by+ an index function enclosed over whatever backend (e.g. `Data.Vector`) you+ like.++source-repository head+ type: git+ location: https://gitlab.com/combobulate.systems/vector-pull++library+ hs-source-dirs: lib+ exposed-modules:+ Data.Vector.Pull++ ghc-options:+ -Wall+ -Wincomplete-record-updates+ -Wincomplete-uni-patterns+ -Wmissing-deriving-strategies+ -Wunused-foralls+ -Wno-name-shadowing+ -Wno-partial-type-signatures+ -Wno-missing-home-modules+ -Wno-ambiguous-fields+ -fprint-explicit-foralls+ -fprint-explicit-kinds+ -fwrite-ide-info++ default-language: GHC2021+ default-extensions:+ AllowAmbiguousTypes+ BangPatterns+ BlockArguments+ ConstraintKinds+ DataKinds+ DeriveAnyClass+ DeriveDataTypeable+ DeriveFoldable+ DeriveFunctor+ DeriveGeneric+ DeriveLift+ DeriveTraversable+ DerivingStrategies+ DerivingVia+ DuplicateRecordFields+ EmptyCase+ EmptyDataDecls+ EmptyDataDeriving+ ExistentialQuantification+ ExplicitForAll+ FlexibleContexts+ FlexibleInstances+ FunctionalDependencies+ GADTSyntax+ GeneralisedNewtypeDeriving+ ImportQualifiedPost+ ImpredicativeTypes+ InstanceSigs+ KindSignatures+ LambdaCase+ MultiParamTypeClasses+ MultiWayIf+ NamedFieldPuns+ NoStarIsType+ NumericUnderscores+ OverloadedLabels+ OverloadedLists+ OverloadedRecordDot+ OverloadedStrings+ PackageImports+ PartialTypeSignatures+ PatternSynonyms+ PolyKinds+ PostfixOperators+ QualifiedDo+ QuasiQuotes+ RankNTypes+ ScopedTypeVariables+ StandaloneDeriving+ StandaloneKindSignatures+ StrictData+ TemplateHaskell+ TupleSections+ TypeAbstractions+ TypeApplications+ TypeFamilies+ TypeFamilyDependencies+ TypeOperators+ UndecidableInstances+ ViewPatterns++ build-depends:+ atomic-counter ^>=0.1,+ base >=4.13 && <4.22,+ exceptions ^>=0.10,+ optics-core ^>=0.4,+ streaming ^>=0.2,+ vector ^>=0.13,++test-suite vector-pull+ type: exitcode-stdio-1.0+ main-is: spec.hs+ hs-source-dirs: tests+ build-depends:+ base,+ hedgehog,+ hspec,+ hspec-hedgehog,+ vector,+ vector-pull,++ ghc-options:+ -Wall+ -Wno-partial-type-signatures++ default-language: GHC2021+ default-extensions:+ AllowAmbiguousTypes+ BangPatterns+ BlockArguments+ DataKinds+ DerivingStrategies+ DerivingVia+ DuplicateRecordFields+ ExplicitNamespaces+ FlexibleContexts+ FlexibleInstances+ ImportQualifiedPost+ OverloadedLists+ OverloadedStrings+ PartialTypeSignatures+ ScopedTypeVariables+ TypeApplications+ TypeFamilies+ TypeOperators+ TypeSynonymInstances+ UndecidableInstances