summaryrefslogtreecommitdiff
path: root/lib/Data/Packed/Internal/Vector.hs
diff options
context:
space:
mode:
Diffstat (limited to 'lib/Data/Packed/Internal/Vector.hs')
-rw-r--r--lib/Data/Packed/Internal/Vector.hs46
1 files changed, 44 insertions, 2 deletions
diff --git a/lib/Data/Packed/Internal/Vector.hs b/lib/Data/Packed/Internal/Vector.hs
index be2fcbb..a47c376 100644
--- a/lib/Data/Packed/Internal/Vector.hs
+++ b/lib/Data/Packed/Internal/Vector.hs
@@ -18,8 +18,8 @@ module Data.Packed.Internal.Vector (
18 fromList, toList, (|>), 18 fromList, toList, (|>),
19 join, (@>), safe, at, at', subVector, takesV, 19 join, (@>), safe, at, at', subVector, takesV,
20 mapVector, zipVectorWith, unzipVectorWith, 20 mapVector, zipVectorWith, unzipVectorWith,
21 mapVectorM, mapVectorM_, 21 mapVectorM, mapVectorM_, mapVectorWithIndexM, mapVectorWithIndexM_,
22 foldVector, foldVectorG, foldLoop, 22 foldVector, foldVectorG, foldLoop, foldVectorWithIndex,
23 createVector, vec, 23 createVector, vec,
24 asComplex, asReal, float2DoubleV, double2FloatV, 24 asComplex, asReal, float2DoubleV, double2FloatV,
25 fwriteVector, freadVector, fprintfVector, fscanfVector, 25 fwriteVector, freadVector, fprintfVector, fscanfVector,
@@ -364,6 +364,16 @@ foldVector f x v = unsafePerformIO $
364 go (dim v -1) x 364 go (dim v -1) x
365{-# INLINE foldVector #-} 365{-# INLINE foldVector #-}
366 366
367-- the zero-indexed index is passed to the folding function
368foldVectorWithIndex :: Storable a => (Int -> a -> b -> b) -> b -> Vector a -> b
369foldVectorWithIndex f x v = unsafePerformIO $
370 unsafeWith v $ \p -> do
371 let go (-1) s = return s
372 go !k !s = do y <- peekElemOff p k
373 go (k-1::Int) (f k y s)
374 go (dim v -1) x
375{-# INLINE foldVectorWithIndex #-}
376
367foldLoop f s0 d = go (d - 1) s0 377foldLoop f s0 d = go (d - 1) s0
368 where 378 where
369 go 0 s = f (0::Int) s 379 go 0 s = f (0::Int) s
@@ -408,6 +418,38 @@ mapVectorM_ f v = do
408 mapVectorM' f' v' (k+1) t 418 mapVectorM' f' v' (k+1) t
409{-# INLINE mapVectorM_ #-} 419{-# INLINE mapVectorM_ #-}
410 420
421-- | monadic map over Vectors with the zero-indexed index passed to the mapping function
422mapVectorWithIndexM :: (Storable a, Storable b, Monad m) => (Int -> a -> m b) -> Vector a -> m (Vector b)
423mapVectorWithIndexM f v = do
424 w <- return $! unsafePerformIO $! createVector (dim v)
425 mapVectorM' f v w 0 (dim v -1)
426 return w
427 where mapVectorM' f' v' w' !k !t
428 | k == t = do
429 x <- return $! inlinePerformIO $! unsafeWith v' $! \p -> peekElemOff p k
430 y <- f' k x
431 return $! inlinePerformIO $! unsafeWith w' $! \q -> pokeElemOff q k y
432 | otherwise = do
433 x <- return $! inlinePerformIO $! unsafeWith v' $! \p -> peekElemOff p k
434 y <- f' k x
435 _ <- return $! inlinePerformIO $! unsafeWith w' $! \q -> pokeElemOff q k y
436 mapVectorM' f' v' w' (k+1) t
437{-# INLINE mapVectorWithIndexM #-}
438
439-- | monadic map over Vectors with the zero-indexed index passed to the mapping function
440mapVectorWithIndexM_ :: (Storable a, Monad m) => (Int -> a -> m ()) -> Vector a -> m ()
441mapVectorWithIndexM_ f v = do
442 mapVectorM' f v 0 (dim v -1)
443 where mapVectorM' f' v' !k !t
444 | k == t = do
445 x <- return $! inlinePerformIO $! unsafeWith v' $! \p -> peekElemOff p k
446 f' k x
447 | otherwise = do
448 x <- return $! inlinePerformIO $! unsafeWith v' $! \p -> peekElemOff p k
449 _ <- f' k x
450 mapVectorM' f' v' (k+1) t
451{-# INLINE mapVectorWithIndexM_ #-}
452
411------------------------------------------------------------------- 453-------------------------------------------------------------------
412 454
413 455