diff options
-rw-r--r-- | examples/vector-map.hs | 33 | ||||
-rw-r--r-- | lib/Data/Packed/Internal/Vector.hs | 46 | ||||
-rw-r--r-- | lib/Data/Packed/Vector.hs | 4 |
3 files changed, 79 insertions, 4 deletions
diff --git a/examples/vector-map.hs b/examples/vector-map.hs index f116946..7796cc0 100644 --- a/examples/vector-map.hs +++ b/examples/vector-map.hs | |||
@@ -30,12 +30,45 @@ addInitialM = mapVectorM_ (\x -> do | |||
30 | put $ x + i | 30 | put $ x + i |
31 | ) | 31 | ) |
32 | 32 | ||
33 | -- sum the values of the even indiced elements | ||
34 | sumEvens :: Vector Int -> Int | ||
35 | sumEvens = foldVectorWithIndex (\x a b -> if x `mod` 2 == 0 then a + b else b) 0 | ||
36 | |||
37 | -- sum and print running total of evens | ||
38 | sumEvensAndPrint :: Vector Int -> VectorMonadT () | ||
39 | sumEvensAndPrint = mapVectorWithIndexM_ (\ i x -> do | ||
40 | when (i `mod` 2 == 0) (do | ||
41 | v <- get | ||
42 | put $ v + x | ||
43 | v' <- get | ||
44 | liftIO $ putStr $ (show v') ++ " " | ||
45 | return ()) | ||
46 | return () | ||
47 | ) | ||
48 | |||
49 | indexPlusSum :: Vector Int -> VectorMonadT () | ||
50 | indexPlusSum v' = do | ||
51 | v <- mapVectorWithIndexM (\i x -> do | ||
52 | s <- get | ||
53 | let inc = x+s | ||
54 | liftIO $ putStr $ show (i,inc) ++ " " | ||
55 | put inc | ||
56 | return inc) v' | ||
57 | liftIO $ do | ||
58 | putStrLn "" | ||
59 | putStrLn $ show v | ||
60 | |||
33 | ------------------------------------------- | 61 | ------------------------------------------- |
62 | |||
34 | main = do | 63 | main = do |
35 | v' <- test1 v | 64 | v' <- test1 v |
36 | putStrLn "" | 65 | putStrLn "" |
37 | putStrLn $ show v' | 66 | putStrLn $ show v' |
38 | evalStateT (addInitialM v) 0 | 67 | evalStateT (addInitialM v) 0 |
39 | putStrLn "" | 68 | putStrLn "" |
69 | putStrLn $ show (sumEvens v) | ||
70 | evalStateT (sumEvensAndPrint v) 0 | ||
71 | putStrLn "" | ||
72 | evalStateT (indexPlusSum v) 0 | ||
40 | return () | 73 | return () |
41 | 74 | ||
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 | ||
368 | foldVectorWithIndex :: Storable a => (Int -> a -> b -> b) -> b -> Vector a -> b | ||
369 | foldVectorWithIndex 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 | |||
367 | foldLoop f s0 d = go (d - 1) s0 | 377 | foldLoop 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 | ||
422 | mapVectorWithIndexM :: (Storable a, Storable b, Monad m) => (Int -> a -> m b) -> Vector a -> m (Vector b) | ||
423 | mapVectorWithIndexM 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 | ||
440 | mapVectorWithIndexM_ :: (Storable a, Monad m) => (Int -> a -> m ()) -> Vector a -> m () | ||
441 | mapVectorWithIndexM_ 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 | ||
diff --git a/lib/Data/Packed/Vector.hs b/lib/Data/Packed/Vector.hs index a526caa..ad690f9 100644 --- a/lib/Data/Packed/Vector.hs +++ b/lib/Data/Packed/Vector.hs | |||
@@ -27,9 +27,9 @@ module Data.Packed.Vector ( | |||
27 | -- vectorMax, vectorMin, | 27 | -- vectorMax, vectorMin, |
28 | vectorMaxIndex, vectorMinIndex, | 28 | vectorMaxIndex, vectorMinIndex, |
29 | mapVector, zipVector, zipVectorWith, unzipVector, unzipVectorWith, | 29 | mapVector, zipVector, zipVectorWith, unzipVector, unzipVectorWith, |
30 | mapVectorM, mapVectorM_, | 30 | mapVectorM, mapVectorM_, mapVectorWithIndexM, mapVectorWithIndexM_, |
31 | fscanfVector, fprintfVector, freadVector, fwriteVector, | 31 | fscanfVector, fprintfVector, freadVector, fwriteVector, |
32 | foldLoop, foldVector, foldVectorG | 32 | foldLoop, foldVector, foldVectorG, foldVectorWithIndex |
33 | ) where | 33 | ) where |
34 | 34 | ||
35 | import Data.Packed.Internal.Vector | 35 | import Data.Packed.Internal.Vector |