{-# LANGUAGE MagicHash, UnboxedTuples, BangPatterns, FlexibleContexts #-}
{-# LANGUAGE TypeSynonymInstances #-}

{-# OPTIONS_GHC -fno-warn-orphans #-}

-- |
-- Module      :  Internal.Vector
-- Copyright   :  (c) Alberto Ruiz 2007-15
-- License     :  BSD3
-- Maintainer  :  Alberto Ruiz
-- Stability   :  provisional
--

module Internal.Vector(
    I,Z,R,C,
    fi,ti,
    Vector, fromList, unsafeToForeignPtr, unsafeFromForeignPtr, unsafeWith,
    createVector, avec, inlinePerformIO,
    toList, dim, (@>), at', (|>),
    vjoin, subVector, takesV, idxs,
    buildVector,
    asReal, asComplex,
    toByteString,fromByteString,
    zipVector, unzipVector, zipVectorWith, unzipVectorWith,
    foldVector, foldVectorG, foldVectorWithIndex, foldLoop,
    mapVector, mapVectorM, mapVectorM_,
    mapVectorWithIndex, mapVectorWithIndexM, mapVectorWithIndexM_
) where

import Foreign.Marshal.Array
import Foreign.ForeignPtr
import Foreign.Ptr
import Foreign.Storable
import Foreign.C.Types(CInt)
import Data.Int(Int64)
import Data.Complex
import System.IO.Unsafe(unsafePerformIO)
import GHC.ForeignPtr(mallocPlainForeignPtrBytes)
import GHC.Base(realWorld#, IO(IO), when)
import qualified Data.Vector.Storable as Vector
import Data.Vector.Storable(Vector, fromList, unsafeToForeignPtr, unsafeFromForeignPtr, unsafeWith)

import Data.Binary
import Data.Binary.Put
import Control.Monad(replicateM)
import qualified Data.ByteString.Internal as BS
import Data.Vector.Storable.Internal(updPtr)

type I = CInt
type Z = Int64
type R = Double
type C = Complex Double


-- | specialized fromIntegral
fi :: Int -> CInt
fi :: Int -> I
fi = forall a b. (Integral a, Num b) => a -> b
fromIntegral

-- | specialized fromIntegral
ti :: CInt -> Int
ti :: I -> Int
ti = forall a b. (Integral a, Num b) => a -> b
fromIntegral


-- | Number of elements
dim :: (Storable t) => Vector t -> Int
dim :: forall t. Storable t => Vector t -> Int
dim = forall t. Storable t => Vector t -> Int
Vector.length
{-# INLINE dim #-}


-- C-Haskell vector adapter
{-# INLINE avec #-}
avec :: Storable a => Vector a -> (f -> IO r) -> ((CInt -> Ptr a -> f) -> IO r)
avec :: forall a f r.
Storable a =>
Vector a -> (f -> IO r) -> (I -> Ptr a -> f) -> IO r
avec Vector a
v f -> IO r
f I -> Ptr a -> f
g = forall a b. Storable a => Vector a -> (Ptr a -> IO b) -> IO b
unsafeWith Vector a
v forall a b. (a -> b) -> a -> b
$ \Ptr a
ptr -> f -> IO r
f (I -> Ptr a -> f
g (forall a b. (Integral a, Num b) => a -> b
fromIntegral (forall t. Storable t => Vector t -> Int
Vector.length Vector a
v)) Ptr a
ptr)

-- allocates memory for a new vector
createVector :: Storable a => Int -> IO (Vector a)
createVector :: forall a. Storable a => Int -> IO (Vector a)
createVector Int
n = do
    forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
n forall a. Ord a => a -> a -> Bool
< Int
0) forall a b. (a -> b) -> a -> b
$ forall a. HasCallStack => [Char] -> a
error ([Char]
"trying to createVector of negative dim: "forall a. [a] -> [a] -> [a]
++forall a. Show a => a -> [Char]
show Int
n)
    ForeignPtr a
fp <- forall b. Storable b => b -> IO (ForeignPtr b)
doMalloc forall a. HasCallStack => a
undefined
    forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. Storable a => ForeignPtr a -> Int -> Int -> Vector a
unsafeFromForeignPtr ForeignPtr a
fp Int
0 Int
n
  where
    --
    -- Use the much cheaper Haskell heap allocated storage
    -- for foreign pointer space we control
    --
    doMalloc :: Storable b => b -> IO (ForeignPtr b)
    doMalloc :: forall b. Storable b => b -> IO (ForeignPtr b)
doMalloc b
dummy = do
        forall a. Int -> IO (ForeignPtr a)
mallocPlainForeignPtrBytes (Int
n forall a. Num a => a -> a -> a
* forall a. Storable a => a -> Int
sizeOf b
dummy)

{- | creates a Vector from a list:

@> fromList [2,3,5,7]
4 |> [2.0,3.0,5.0,7.0]@

-}

safeRead :: Storable a => Vector a -> (Ptr a -> IO c) -> c
safeRead :: forall a c. Storable a => Vector a -> (Ptr a -> IO c) -> c
safeRead Vector a
v = forall a. IO a -> a
inlinePerformIO forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. Storable a => Vector a -> (Ptr a -> IO b) -> IO b
unsafeWith Vector a
v
{-# INLINE safeRead #-}

inlinePerformIO :: IO a -> a
inlinePerformIO :: forall a. IO a -> a
inlinePerformIO (IO State# RealWorld -> (# State# RealWorld, a #)
m) = case State# RealWorld -> (# State# RealWorld, a #)
m State# RealWorld
realWorld# of (# State# RealWorld
_, a
r #) -> a
r
{-# INLINE inlinePerformIO #-}

{- extracts the Vector elements to a list

>>> toList (linspace 5 (1,10))
[1.0,3.25,5.5,7.75,10.0]

-}
toList :: Storable a => Vector a -> [a]
toList :: forall a. Storable a => Vector a -> [a]
toList Vector a
v = forall a c. Storable a => Vector a -> (Ptr a -> IO c) -> c
safeRead Vector a
v forall a b. (a -> b) -> a -> b
$ forall a. Storable a => Int -> Ptr a -> IO [a]
peekArray (forall t. Storable t => Vector t -> Int
dim Vector a
v)

{- | Create a vector from a list of elements and explicit dimension. The input
     list is truncated if it is too long, so it may safely
     be used, for instance, with infinite lists.

>>> 5 |> [1..]
[1.0,2.0,3.0,4.0,5.0]
it :: (Enum a, Num a, Foreign.Storable.Storable a) => Vector a

-}
(|>) :: (Storable a) => Int -> [a] -> Vector a
infixl 9 |>
Int
n |> :: forall a. Storable a => Int -> [a] -> Vector a
|> [a]
l
    | forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
l' forall a. Eq a => a -> a -> Bool
== Int
n = forall a. Storable a => [a] -> Vector a
fromList [a]
l'
    | Bool
otherwise      = forall a. HasCallStack => [Char] -> a
error [Char]
"list too short for |>"
  where
    l' :: [a]
l' = forall a. Int -> [a] -> [a]
take Int
n [a]
l


-- | Create a vector of indexes, useful for matrix extraction using '(??)'
idxs :: [Int] -> Vector I
idxs :: [Int] -> Vector I
idxs [Int]
js = forall a. Storable a => [a] -> Vector a
fromList (forall a b. (a -> b) -> [a] -> [b]
map forall a b. (Integral a, Num b) => a -> b
fromIntegral [Int]
js) :: Vector I

{- | takes a number of consecutive elements from a Vector

>>> subVector 2 3 (fromList [1..10])
[3.0,4.0,5.0]
it :: (Enum t, Num t, Foreign.Storable.Storable t) => Vector t

-}
subVector :: Storable t => Int       -- ^ index of the starting element
                        -> Int       -- ^ number of elements to extract
                        -> Vector t  -- ^ source
                        -> Vector t  -- ^ result
subVector :: forall t. Storable t => Int -> Int -> Vector t -> Vector t
subVector = forall t. Storable t => Int -> Int -> Vector t -> Vector t
Vector.slice
{-# INLINE subVector #-}




{- | Reads a vector position:

>>> fromList [0..9] @> 7
7.0

-}
(@>) :: Storable t => Vector t -> Int -> t
infixl 9 @>
Vector t
v @> :: forall t. Storable t => Vector t -> Int -> t
@> Int
n
    | Int
n forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
n forall a. Ord a => a -> a -> Bool
< forall t. Storable t => Vector t -> Int
dim Vector t
v = forall t. Storable t => Vector t -> Int -> t
at' Vector t
v Int
n
    | Bool
otherwise = forall a. HasCallStack => [Char] -> a
error [Char]
"vector index out of range"
{-# INLINE (@>) #-}

-- | access to Vector elements without range checking
at' :: Storable a => Vector a -> Int -> a
at' :: forall t. Storable t => Vector t -> Int -> t
at' Vector a
v Int
n = forall a c. Storable a => Vector a -> (Ptr a -> IO c) -> c
safeRead Vector a
v forall a b. (a -> b) -> a -> b
$ forall a b c. (a -> b -> c) -> b -> a -> c
flip forall a. Storable a => Ptr a -> Int -> IO a
peekElemOff Int
n
{-# INLINE at' #-}

{- | concatenate a list of vectors

>>> vjoin [fromList [1..5::Double], konst 1 3]
[1.0,2.0,3.0,4.0,5.0,1.0,1.0,1.0]
it :: Vector Double

-}
vjoin :: Storable t => [Vector t] -> Vector t
vjoin :: forall t. Storable t => [Vector t] -> Vector t
vjoin [] = forall a. Storable a => [a] -> Vector a
fromList []
vjoin [Vector t
v] = Vector t
v
vjoin [Vector t]
as = forall a. IO a -> a
unsafePerformIO forall a b. (a -> b) -> a -> b
$ do
    let tot :: Int
tot = forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum (forall a b. (a -> b) -> [a] -> [b]
map forall t. Storable t => Vector t -> Int
dim [Vector t]
as)
    Vector t
r <- forall a. Storable a => Int -> IO (Vector a)
createVector Int
tot
    forall a b. Storable a => Vector a -> (Ptr a -> IO b) -> IO b
unsafeWith Vector t
r forall a b. (a -> b) -> a -> b
$ \Ptr t
ptr ->
        forall {a} {t}.
(Num t, Storable a) =>
[Vector a] -> t -> Ptr a -> IO ()
joiner [Vector t]
as Int
tot Ptr t
ptr
    forall (m :: * -> *) a. Monad m => a -> m a
return Vector t
r
  where joiner :: [Vector a] -> t -> Ptr a -> IO ()
joiner [] t
_ Ptr a
_ = forall (m :: * -> *) a. Monad m => a -> m a
return ()
        joiner (Vector a
v:[Vector a]
cs) t
_ Ptr a
p = do
            let n :: Int
n = forall t. Storable t => Vector t -> Int
dim Vector a
v
            forall a b. Storable a => Vector a -> (Ptr a -> IO b) -> IO b
unsafeWith Vector a
v forall a b. (a -> b) -> a -> b
$ \Ptr a
pb -> forall a. Storable a => Ptr a -> Ptr a -> Int -> IO ()
copyArray Ptr a
p Ptr a
pb Int
n
            [Vector a] -> t -> Ptr a -> IO ()
joiner [Vector a]
cs t
0 (forall a. Storable a => Ptr a -> Int -> Ptr a
advancePtr Ptr a
p Int
n)


{- | Extract consecutive subvectors of the given sizes.

>>> takesV [3,4] (linspace 10 (1,10::Double))
[[1.0,2.0,3.0],[4.0,5.0,6.0,7.0]]
it :: [Vector Double]

-}
takesV :: Storable t => [Int] -> Vector t -> [Vector t]
takesV :: forall t. Storable t => [Int] -> Vector t -> [Vector t]
takesV [Int]
ms Vector t
w | forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Int]
ms forall a. Ord a => a -> a -> Bool
> forall t. Storable t => Vector t -> Int
dim Vector t
w = forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"takesV " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> [Char]
show [Int]
ms forall a. [a] -> [a] -> [a]
++ [Char]
" on dim = " forall a. [a] -> [a] -> [a]
++ (forall a. Show a => a -> [Char]
show forall a b. (a -> b) -> a -> b
$ forall t. Storable t => Vector t -> Int
dim Vector t
w)
            | Bool
otherwise = forall t. Storable t => [Int] -> Vector t -> [Vector t]
go [Int]
ms Vector t
w
    where go :: [Int] -> Vector t -> [Vector t]
go [] Vector t
_ = []
          go (Int
n:[Int]
ns) Vector t
v = forall t. Storable t => Int -> Int -> Vector t -> Vector t
subVector Int
0 Int
n Vector t
v
                      forall a. a -> [a] -> [a]
: [Int] -> Vector t -> [Vector t]
go [Int]
ns (forall t. Storable t => Int -> Int -> Vector t -> Vector t
subVector Int
n (forall t. Storable t => Vector t -> Int
dim Vector t
v forall a. Num a => a -> a -> a
- Int
n) Vector t
v)

---------------------------------------------------------------

-- | transforms a complex vector into a real vector with alternating real and imaginary parts
asReal :: (RealFloat a, Storable a) => Vector (Complex a) -> Vector a
asReal :: forall a.
(RealFloat a, Storable a) =>
Vector (Complex a) -> Vector a
asReal Vector (Complex a)
v = forall a. Storable a => ForeignPtr a -> Int -> Int -> Vector a
unsafeFromForeignPtr (forall a b. ForeignPtr a -> ForeignPtr b
castForeignPtr ForeignPtr (Complex a)
fp) (Int
2forall a. Num a => a -> a -> a
*Int
i) (Int
2forall a. Num a => a -> a -> a
*Int
n)
    where (ForeignPtr (Complex a)
fp,Int
i,Int
n) = forall a. Storable a => Vector a -> (ForeignPtr a, Int, Int)
unsafeToForeignPtr Vector (Complex a)
v

-- | transforms a real vector into a complex vector with alternating real and imaginary parts
asComplex :: (RealFloat a, Storable a) => Vector a -> Vector (Complex a)
asComplex :: forall a.
(RealFloat a, Storable a) =>
Vector a -> Vector (Complex a)
asComplex Vector a
v = forall a. Storable a => ForeignPtr a -> Int -> Int -> Vector a
unsafeFromForeignPtr (forall a b. ForeignPtr a -> ForeignPtr b
castForeignPtr ForeignPtr a
fp) (Int
i forall a. Integral a => a -> a -> a
`div` Int
2) (Int
n forall a. Integral a => a -> a -> a
`div` Int
2)
    where (ForeignPtr a
fp,Int
i,Int
n) = forall a. Storable a => Vector a -> (ForeignPtr a, Int, Int)
unsafeToForeignPtr Vector a
v

--------------------------------------------------------------------------------


-- | map on Vectors
mapVector :: (Storable a, Storable b) => (a-> b) -> Vector a -> Vector b
mapVector :: forall a b.
(Storable a, Storable b) =>
(a -> b) -> Vector a -> Vector b
mapVector a -> b
f Vector a
v = forall a. IO a -> a
unsafePerformIO forall a b. (a -> b) -> a -> b
$ do
    Vector b
w <- forall a. Storable a => Int -> IO (Vector a)
createVector (forall t. Storable t => Vector t -> Int
dim Vector a
v)
    forall a b. Storable a => Vector a -> (Ptr a -> IO b) -> IO b
unsafeWith Vector a
v forall a b. (a -> b) -> a -> b
$ \Ptr a
p ->
        forall a b. Storable a => Vector a -> (Ptr a -> IO b) -> IO b
unsafeWith Vector b
w forall a b. (a -> b) -> a -> b
$ \Ptr b
q -> do
            let go :: Int -> IO ()
go (-1) = forall (m :: * -> *) a. Monad m => a -> m a
return ()
                go !Int
k = do a
x <- forall a. Storable a => Ptr a -> Int -> IO a
peekElemOff Ptr a
p Int
k
                           forall a. Storable a => Ptr a -> Int -> a -> IO ()
pokeElemOff      Ptr b
q Int
k (a -> b
f a
x)
                           Int -> IO ()
go (Int
kforall a. Num a => a -> a -> a
-Int
1)
            Int -> IO ()
go (forall t. Storable t => Vector t -> Int
dim Vector a
v forall a. Num a => a -> a -> a
-Int
1)
    forall (m :: * -> *) a. Monad m => a -> m a
return Vector b
w
{-# INLINE mapVector #-}

-- | zipWith for Vectors
zipVectorWith :: (Storable a, Storable b, Storable c) => (a-> b -> c) -> Vector a -> Vector b -> Vector c
zipVectorWith :: forall a b c.
(Storable a, Storable b, Storable c) =>
(a -> b -> c) -> Vector a -> Vector b -> Vector c
zipVectorWith a -> b -> c
f Vector a
u Vector b
v = forall a. IO a -> a
unsafePerformIO forall a b. (a -> b) -> a -> b
$ do
    let n :: Int
n = forall a. Ord a => a -> a -> a
min (forall t. Storable t => Vector t -> Int
dim Vector a
u) (forall t. Storable t => Vector t -> Int
dim Vector b
v)
    Vector c
w <- forall a. Storable a => Int -> IO (Vector a)
createVector Int
n
    forall a b. Storable a => Vector a -> (Ptr a -> IO b) -> IO b
unsafeWith Vector a
u forall a b. (a -> b) -> a -> b
$ \Ptr a
pu ->
        forall a b. Storable a => Vector a -> (Ptr a -> IO b) -> IO b
unsafeWith Vector b
v forall a b. (a -> b) -> a -> b
$ \Ptr b
pv ->
            forall a b. Storable a => Vector a -> (Ptr a -> IO b) -> IO b
unsafeWith Vector c
w forall a b. (a -> b) -> a -> b
$ \Ptr c
pw -> do
                let go :: Int -> IO ()
go (-1) = forall (m :: * -> *) a. Monad m => a -> m a
return ()
                    go !Int
k = do a
x <- forall a. Storable a => Ptr a -> Int -> IO a
peekElemOff Ptr a
pu Int
k
                               b
y <- forall a. Storable a => Ptr a -> Int -> IO a
peekElemOff Ptr b
pv Int
k
                               forall a. Storable a => Ptr a -> Int -> a -> IO ()
pokeElemOff      Ptr c
pw Int
k (a -> b -> c
f a
x b
y)
                               Int -> IO ()
go (Int
kforall a. Num a => a -> a -> a
-Int
1)
                Int -> IO ()
go (Int
n forall a. Num a => a -> a -> a
-Int
1)
    forall (m :: * -> *) a. Monad m => a -> m a
return Vector c
w
{-# INLINE zipVectorWith #-}

-- | unzipWith for Vectors
unzipVectorWith :: (Storable (a,b), Storable c, Storable d)
                   => ((a,b) -> (c,d)) -> Vector (a,b) -> (Vector c,Vector d)
unzipVectorWith :: forall a b c d.
(Storable (a, b), Storable c, Storable d) =>
((a, b) -> (c, d)) -> Vector (a, b) -> (Vector c, Vector d)
unzipVectorWith (a, b) -> (c, d)
f Vector (a, b)
u = forall a. IO a -> a
unsafePerformIO forall a b. (a -> b) -> a -> b
$ do
      let n :: Int
n = forall t. Storable t => Vector t -> Int
dim Vector (a, b)
u
      Vector c
v <- forall a. Storable a => Int -> IO (Vector a)
createVector Int
n
      Vector d
w <- forall a. Storable a => Int -> IO (Vector a)
createVector Int
n
      forall a b. Storable a => Vector a -> (Ptr a -> IO b) -> IO b
unsafeWith Vector (a, b)
u forall a b. (a -> b) -> a -> b
$ \Ptr (a, b)
pu ->
          forall a b. Storable a => Vector a -> (Ptr a -> IO b) -> IO b
unsafeWith Vector c
v forall a b. (a -> b) -> a -> b
$ \Ptr c
pv ->
              forall a b. Storable a => Vector a -> (Ptr a -> IO b) -> IO b
unsafeWith Vector d
w forall a b. (a -> b) -> a -> b
$ \Ptr d
pw -> do
                  let go :: Int -> IO ()
go (-1) = forall (m :: * -> *) a. Monad m => a -> m a
return ()
                      go !Int
k   = do (a, b)
z <- forall a. Storable a => Ptr a -> Int -> IO a
peekElemOff Ptr (a, b)
pu Int
k
                                   let (c
x,d
y) = (a, b) -> (c, d)
f (a, b)
z
                                   forall a. Storable a => Ptr a -> Int -> a -> IO ()
pokeElemOff      Ptr c
pv Int
k c
x
                                   forall a. Storable a => Ptr a -> Int -> a -> IO ()
pokeElemOff      Ptr d
pw Int
k d
y
                                   Int -> IO ()
go (Int
kforall a. Num a => a -> a -> a
-Int
1)
                  Int -> IO ()
go (Int
nforall a. Num a => a -> a -> a
-Int
1)
      forall (m :: * -> *) a. Monad m => a -> m a
return (Vector c
v,Vector d
w)
{-# INLINE unzipVectorWith #-}

foldVector :: Storable a => (a -> b -> b) -> b -> Vector a -> b
foldVector :: forall a b. Storable a => (a -> b -> b) -> b -> Vector a -> b
foldVector a -> b -> b
f b
x Vector a
v = forall a. IO a -> a
unsafePerformIO forall a b. (a -> b) -> a -> b
$
    forall a b. Storable a => Vector a -> (Ptr a -> IO b) -> IO b
unsafeWith Vector a
v forall a b. (a -> b) -> a -> b
$ \Ptr a
p -> do
        let go :: Int -> b -> IO b
go (-1) b
s = forall (m :: * -> *) a. Monad m => a -> m a
return b
s
            go !Int
k !b
s = do a
y <- forall a. Storable a => Ptr a -> Int -> IO a
peekElemOff Ptr a
p Int
k
                          Int -> b -> IO b
go (Int
kforall a. Num a => a -> a -> a
-Int
1::Int) (a -> b -> b
f a
y b
s)
        Int -> b -> IO b
go (forall t. Storable t => Vector t -> Int
dim Vector a
v forall a. Num a => a -> a -> a
-Int
1) b
x
{-# INLINE foldVector #-}

-- the zero-indexed index is passed to the folding function
foldVectorWithIndex :: Storable a => (Int -> a -> b -> b) -> b -> Vector a -> b
foldVectorWithIndex :: forall a b.
Storable a =>
(Int -> a -> b -> b) -> b -> Vector a -> b
foldVectorWithIndex Int -> a -> b -> b
f b
x Vector a
v = forall a. IO a -> a
unsafePerformIO forall a b. (a -> b) -> a -> b
$
    forall a b. Storable a => Vector a -> (Ptr a -> IO b) -> IO b
unsafeWith Vector a
v forall a b. (a -> b) -> a -> b
$ \Ptr a
p -> do
        let go :: Int -> b -> IO b
go (-1) b
s = forall (m :: * -> *) a. Monad m => a -> m a
return b
s
            go !Int
k !b
s = do a
y <- forall a. Storable a => Ptr a -> Int -> IO a
peekElemOff Ptr a
p Int
k
                          Int -> b -> IO b
go (Int
kforall a. Num a => a -> a -> a
-Int
1::Int) (Int -> a -> b -> b
f Int
k a
y b
s)
        Int -> b -> IO b
go (forall t. Storable t => Vector t -> Int
dim Vector a
v forall a. Num a => a -> a -> a
-Int
1) b
x
{-# INLINE foldVectorWithIndex #-}

foldLoop :: (Int -> t -> t) -> t -> Int -> t
foldLoop :: forall t. (Int -> t -> t) -> t -> Int -> t
foldLoop Int -> t -> t
f t
s0 Int
d = Int -> t -> t
go (Int
d forall a. Num a => a -> a -> a
- Int
1) t
s0
     where
       go :: Int -> t -> t
go Int
0 t
s = Int -> t -> t
f (Int
0::Int) t
s
       go !Int
j !t
s = Int -> t -> t
go (Int
j forall a. Num a => a -> a -> a
- Int
1) (Int -> t -> t
f Int
j t
s)

foldVectorG :: Storable t1 => (Int -> (Int -> t1) -> t -> t) -> t -> Vector t1 -> t
foldVectorG :: forall t1 t.
Storable t1 =>
(Int -> (Int -> t1) -> t -> t) -> t -> Vector t1 -> t
foldVectorG Int -> (Int -> t1) -> t -> t
f t
s0 Vector t1
v = forall t. (Int -> t -> t) -> t -> Int -> t
foldLoop Int -> t -> t
g t
s0 (forall t. Storable t => Vector t -> Int
dim Vector t1
v)
    where g :: Int -> t -> t
g !Int
k !t
s = Int -> (Int -> t1) -> t -> t
f Int
k (forall a c. Storable a => Vector a -> (Ptr a -> IO c) -> c
safeRead Vector t1
v forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b c. (a -> b -> c) -> b -> a -> c
flip forall a. Storable a => Ptr a -> Int -> IO a
peekElemOff) t
s
          {-# INLINE g #-} -- Thanks to Ryan Ingram (http://permalink.gmane.org/gmane.comp.lang.haskell.cafe/46479)
{-# INLINE foldVectorG #-}

-------------------------------------------------------------------

-- | monadic map over Vectors
--    the monad @m@ must be strict
mapVectorM :: (Storable a, Storable b, Monad m) => (a -> m b) -> Vector a -> m (Vector b)
mapVectorM :: forall a b (m :: * -> *).
(Storable a, Storable b, Monad m) =>
(a -> m b) -> Vector a -> m (Vector b)
mapVectorM a -> m b
f Vector a
v = do
    Vector b
w <- forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$! forall a. IO a -> a
unsafePerformIO forall a b. (a -> b) -> a -> b
$! forall a. Storable a => Int -> IO (Vector a)
createVector (forall t. Storable t => Vector t -> Int
dim Vector a
v)
    Vector b -> Int -> Int -> m ()
mapVectorM' Vector b
w Int
0 (forall t. Storable t => Vector t -> Int
dim Vector a
v forall a. Num a => a -> a -> a
-Int
1)
    forall (m :: * -> *) a. Monad m => a -> m a
return Vector b
w
    where mapVectorM' :: Vector b -> Int -> Int -> m ()
mapVectorM' Vector b
w' !Int
k !Int
t
              | Int
k forall a. Eq a => a -> a -> Bool
== Int
t               = do
                                       a
x <- forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$! forall a. IO a -> a
inlinePerformIO forall a b. (a -> b) -> a -> b
$! forall a b. Storable a => Vector a -> (Ptr a -> IO b) -> IO b
unsafeWith Vector a
v forall a b. (a -> b) -> a -> b
$! \Ptr a
p -> forall a. Storable a => Ptr a -> Int -> IO a
peekElemOff Ptr a
p Int
k
                                       b
y <- a -> m b
f a
x
                                       forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$! forall a. IO a -> a
inlinePerformIO forall a b. (a -> b) -> a -> b
$! forall a b. Storable a => Vector a -> (Ptr a -> IO b) -> IO b
unsafeWith Vector b
w' forall a b. (a -> b) -> a -> b
$! \Ptr b
q -> forall a. Storable a => Ptr a -> Int -> a -> IO ()
pokeElemOff Ptr b
q Int
k b
y
              | Bool
otherwise            = do
                                       a
x <- forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$! forall a. IO a -> a
inlinePerformIO forall a b. (a -> b) -> a -> b
$! forall a b. Storable a => Vector a -> (Ptr a -> IO b) -> IO b
unsafeWith Vector a
v forall a b. (a -> b) -> a -> b
$! \Ptr a
p -> forall a. Storable a => Ptr a -> Int -> IO a
peekElemOff Ptr a
p Int
k
                                       b
y <- a -> m b
f a
x
                                       ()
_ <- forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$! forall a. IO a -> a
inlinePerformIO forall a b. (a -> b) -> a -> b
$! forall a b. Storable a => Vector a -> (Ptr a -> IO b) -> IO b
unsafeWith Vector b
w' forall a b. (a -> b) -> a -> b
$! \Ptr b
q -> forall a. Storable a => Ptr a -> Int -> a -> IO ()
pokeElemOff Ptr b
q Int
k b
y
                                       Vector b -> Int -> Int -> m ()
mapVectorM' Vector b
w' (Int
kforall a. Num a => a -> a -> a
+Int
1) Int
t
{-# INLINE mapVectorM #-}

-- | monadic map over Vectors
mapVectorM_ :: (Storable a, Monad m) => (a -> m ()) -> Vector a -> m ()
mapVectorM_ :: forall a (m :: * -> *).
(Storable a, Monad m) =>
(a -> m ()) -> Vector a -> m ()
mapVectorM_ a -> m ()
f Vector a
v = do
    Int -> Int -> m ()
mapVectorM' Int
0 (forall t. Storable t => Vector t -> Int
dim Vector a
v forall a. Num a => a -> a -> a
-Int
1)
    where mapVectorM' :: Int -> Int -> m ()
mapVectorM' !Int
k !Int
t
              | Int
k forall a. Eq a => a -> a -> Bool
== Int
t            = do
                                    a
x <- forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$! forall a. IO a -> a
inlinePerformIO forall a b. (a -> b) -> a -> b
$! forall a b. Storable a => Vector a -> (Ptr a -> IO b) -> IO b
unsafeWith Vector a
v forall a b. (a -> b) -> a -> b
$! \Ptr a
p -> forall a. Storable a => Ptr a -> Int -> IO a
peekElemOff Ptr a
p Int
k
                                    a -> m ()
f a
x
              | Bool
otherwise         = do
                                    a
x <- forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$! forall a. IO a -> a
inlinePerformIO forall a b. (a -> b) -> a -> b
$! forall a b. Storable a => Vector a -> (Ptr a -> IO b) -> IO b
unsafeWith Vector a
v forall a b. (a -> b) -> a -> b
$! \Ptr a
p -> forall a. Storable a => Ptr a -> Int -> IO a
peekElemOff Ptr a
p Int
k
                                    ()
_ <- a -> m ()
f a
x
                                    Int -> Int -> m ()
mapVectorM' (Int
kforall a. Num a => a -> a -> a
+Int
1) Int
t
{-# INLINE mapVectorM_ #-}

-- | monadic map over Vectors with the zero-indexed index passed to the mapping function
--    the monad @m@ must be strict
mapVectorWithIndexM :: (Storable a, Storable b, Monad m) => (Int -> a -> m b) -> Vector a -> m (Vector b)
mapVectorWithIndexM :: forall a b (m :: * -> *).
(Storable a, Storable b, Monad m) =>
(Int -> a -> m b) -> Vector a -> m (Vector b)
mapVectorWithIndexM Int -> a -> m b
f Vector a
v = do
    Vector b
w <- forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$! forall a. IO a -> a
unsafePerformIO forall a b. (a -> b) -> a -> b
$! forall a. Storable a => Int -> IO (Vector a)
createVector (forall t. Storable t => Vector t -> Int
dim Vector a
v)
    Vector b -> Int -> Int -> m ()
mapVectorM' Vector b
w Int
0 (forall t. Storable t => Vector t -> Int
dim Vector a
v forall a. Num a => a -> a -> a
-Int
1)
    forall (m :: * -> *) a. Monad m => a -> m a
return Vector b
w
    where mapVectorM' :: Vector b -> Int -> Int -> m ()
mapVectorM' Vector b
w' !Int
k !Int
t
              | Int
k forall a. Eq a => a -> a -> Bool
== Int
t               = do
                                       a
x <- forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$! forall a. IO a -> a
inlinePerformIO forall a b. (a -> b) -> a -> b
$! forall a b. Storable a => Vector a -> (Ptr a -> IO b) -> IO b
unsafeWith Vector a
v forall a b. (a -> b) -> a -> b
$! \Ptr a
p -> forall a. Storable a => Ptr a -> Int -> IO a
peekElemOff Ptr a
p Int
k
                                       b
y <- Int -> a -> m b
f Int
k a
x
                                       forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$! forall a. IO a -> a
inlinePerformIO forall a b. (a -> b) -> a -> b
$! forall a b. Storable a => Vector a -> (Ptr a -> IO b) -> IO b
unsafeWith Vector b
w' forall a b. (a -> b) -> a -> b
$! \Ptr b
q -> forall a. Storable a => Ptr a -> Int -> a -> IO ()
pokeElemOff Ptr b
q Int
k b
y
              | Bool
otherwise            = do
                                       a
x <- forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$! forall a. IO a -> a
inlinePerformIO forall a b. (a -> b) -> a -> b
$! forall a b. Storable a => Vector a -> (Ptr a -> IO b) -> IO b
unsafeWith Vector a
v forall a b. (a -> b) -> a -> b
$! \Ptr a
p -> forall a. Storable a => Ptr a -> Int -> IO a
peekElemOff Ptr a
p Int
k
                                       b
y <- Int -> a -> m b
f Int
k a
x
                                       ()
_ <- forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$! forall a. IO a -> a
inlinePerformIO forall a b. (a -> b) -> a -> b
$! forall a b. Storable a => Vector a -> (Ptr a -> IO b) -> IO b
unsafeWith Vector b
w' forall a b. (a -> b) -> a -> b
$! \Ptr b
q -> forall a. Storable a => Ptr a -> Int -> a -> IO ()
pokeElemOff Ptr b
q Int
k b
y
                                       Vector b -> Int -> Int -> m ()
mapVectorM' Vector b
w' (Int
kforall a. Num a => a -> a -> a
+Int
1) Int
t
{-# INLINE mapVectorWithIndexM #-}

-- | monadic map over Vectors with the zero-indexed index passed to the mapping function
mapVectorWithIndexM_ :: (Storable a, Monad m) => (Int -> a -> m ()) -> Vector a -> m ()
mapVectorWithIndexM_ :: forall a (m :: * -> *).
(Storable a, Monad m) =>
(Int -> a -> m ()) -> Vector a -> m ()
mapVectorWithIndexM_ Int -> a -> m ()
f Vector a
v = do
    Int -> Int -> m ()
mapVectorM' Int
0 (forall t. Storable t => Vector t -> Int
dim Vector a
v forall a. Num a => a -> a -> a
-Int
1)
    where mapVectorM' :: Int -> Int -> m ()
mapVectorM' !Int
k !Int
t
              | Int
k forall a. Eq a => a -> a -> Bool
== Int
t            = do
                                    a
x <- forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$! forall a. IO a -> a
inlinePerformIO forall a b. (a -> b) -> a -> b
$! forall a b. Storable a => Vector a -> (Ptr a -> IO b) -> IO b
unsafeWith Vector a
v forall a b. (a -> b) -> a -> b
$! \Ptr a
p -> forall a. Storable a => Ptr a -> Int -> IO a
peekElemOff Ptr a
p Int
k
                                    Int -> a -> m ()
f Int
k a
x
              | Bool
otherwise         = do
                                    a
x <- forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$! forall a. IO a -> a
inlinePerformIO forall a b. (a -> b) -> a -> b
$! forall a b. Storable a => Vector a -> (Ptr a -> IO b) -> IO b
unsafeWith Vector a
v forall a b. (a -> b) -> a -> b
$! \Ptr a
p -> forall a. Storable a => Ptr a -> Int -> IO a
peekElemOff Ptr a
p Int
k
                                    ()
_ <- Int -> a -> m ()
f Int
k a
x
                                    Int -> Int -> m ()
mapVectorM' (Int
kforall a. Num a => a -> a -> a
+Int
1) Int
t
{-# INLINE mapVectorWithIndexM_ #-}


mapVectorWithIndex :: (Storable a, Storable b) => (Int -> a -> b) -> Vector a -> Vector b
--mapVectorWithIndex g = head . mapVectorWithIndexM (\a b -> [g a b])
mapVectorWithIndex :: forall a b.
(Storable a, Storable b) =>
(Int -> a -> b) -> Vector a -> Vector b
mapVectorWithIndex Int -> a -> b
f Vector a
v = forall a. IO a -> a
unsafePerformIO forall a b. (a -> b) -> a -> b
$ do
    Vector b
w <- forall a. Storable a => Int -> IO (Vector a)
createVector (forall t. Storable t => Vector t -> Int
dim Vector a
v)
    forall a b. Storable a => Vector a -> (Ptr a -> IO b) -> IO b
unsafeWith Vector a
v forall a b. (a -> b) -> a -> b
$ \Ptr a
p ->
        forall a b. Storable a => Vector a -> (Ptr a -> IO b) -> IO b
unsafeWith Vector b
w forall a b. (a -> b) -> a -> b
$ \Ptr b
q -> do
            let go :: Int -> IO ()
go (-1) = forall (m :: * -> *) a. Monad m => a -> m a
return ()
                go !Int
k = do a
x <- forall a. Storable a => Ptr a -> Int -> IO a
peekElemOff Ptr a
p Int
k
                           forall a. Storable a => Ptr a -> Int -> a -> IO ()
pokeElemOff      Ptr b
q Int
k (Int -> a -> b
f Int
k a
x)
                           Int -> IO ()
go (Int
kforall a. Num a => a -> a -> a
-Int
1)
            Int -> IO ()
go (forall t. Storable t => Vector t -> Int
dim Vector a
v forall a. Num a => a -> a -> a
-Int
1)
    forall (m :: * -> *) a. Monad m => a -> m a
return Vector b
w
{-# INLINE mapVectorWithIndex #-}

--------------------------------------------------------------------------------



-- a 64K cache, with a Double taking 13 bytes in Bytestring,
-- implies a chunk size of 5041
chunk :: Int
chunk :: Int
chunk = Int
5000

chunks :: Int -> [Int]
chunks :: Int -> [Int]
chunks Int
d = let c :: Int
c = Int
d forall a. Integral a => a -> a -> a
`div` Int
chunk
               m :: Int
m = Int
d forall a. Integral a => a -> a -> a
`mod` Int
chunk
           in if Int
m forall a. Eq a => a -> a -> Bool
/= Int
0 then forall a. [a] -> [a]
reverse (Int
mforall a. a -> [a] -> [a]
:(forall a. Int -> a -> [a]
replicate Int
c Int
chunk)) else (forall a. Int -> a -> [a]
replicate Int
c Int
chunk)

putVector :: (Storable t, Binary t) => Vector t -> Data.Binary.Put.PutM ()
putVector :: forall t. (Storable t, Binary t) => Vector t -> PutM ()
putVector Vector t
v = forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ forall t. Binary t => t -> PutM ()
put forall a b. (a -> b) -> a -> b
$! forall a. Storable a => Vector a -> [a]
toList Vector t
v

getVector :: (Storable a, Binary a) => Int -> Get (Vector a)
getVector :: forall a. (Storable a, Binary a) => Int -> Get (Vector a)
getVector Int
d = do
              [a]
xs <- forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
d forall t. Binary t => Get t
get
              forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$! forall a. Storable a => [a] -> Vector a
fromList [a]
xs

--------------------------------------------------------------------------------

toByteString :: Storable t => Vector t -> BS.ByteString
toByteString :: forall t. Storable t => Vector t -> ByteString
toByteString Vector t
v = ForeignPtr Word8 -> Int -> Int -> ByteString
BS.PS (forall a b. ForeignPtr a -> ForeignPtr b
castForeignPtr ForeignPtr t
fp) (Int
szforall a. Num a => a -> a -> a
*Int
o) (Int
sz forall a. Num a => a -> a -> a
* forall t. Storable t => Vector t -> Int
dim Vector t
v)
  where
    (ForeignPtr t
fp,Int
o,Int
_n) = forall a. Storable a => Vector a -> (ForeignPtr a, Int, Int)
unsafeToForeignPtr Vector t
v
    sz :: Int
sz = forall a. Storable a => a -> Int
sizeOf (Vector t
vforall t. Storable t => Vector t -> Int -> t
@>Int
0)


fromByteString :: Storable t => BS.ByteString -> Vector t
fromByteString :: forall t. Storable t => ByteString -> Vector t
fromByteString (BS.PS ForeignPtr Word8
fp Int
o Int
n) = Vector t
r
  where
    r :: Vector t
r = forall a. Storable a => ForeignPtr a -> Int -> Int -> Vector a
unsafeFromForeignPtr (forall a b. ForeignPtr a -> ForeignPtr b
castForeignPtr (forall a. (Ptr a -> Ptr a) -> ForeignPtr a -> ForeignPtr a
updPtr (forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
o) ForeignPtr Word8
fp)) Int
0 Int
n'
    n' :: Int
n' = Int
n forall a. Integral a => a -> a -> a
`div` Int
sz
    sz :: Int
sz = forall a. Storable a => a -> Int
sizeOf (Vector t
rforall t. Storable t => Vector t -> Int -> t
@>Int
0)

--------------------------------------------------------------------------------

instance (Binary a, Storable a) => Binary (Vector a) where

    put :: Vector a -> PutM ()
put Vector a
v = do
            let d :: Int
d = forall t. Storable t => Vector t -> Int
dim Vector a
v
            forall t. Binary t => t -> PutM ()
put Int
d
            forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ forall t. (Storable t, Binary t) => Vector t -> PutM ()
putVector forall a b. (a -> b) -> a -> b
$! forall t. Storable t => [Int] -> Vector t -> [Vector t]
takesV (Int -> [Int]
chunks Int
d) Vector a
v

    -- put = put . v2bs

    get :: Get (Vector a)
get = do
          Int
d <- forall t. Binary t => Get t
get
          [Vector a]
vs <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall a. (Storable a, Binary a) => Int -> Get (Vector a)
getVector forall a b. (a -> b) -> a -> b
$ Int -> [Int]
chunks Int
d
          forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$! forall t. Storable t => [Vector t] -> Vector t
vjoin [Vector a]
vs

    -- get = fmap bs2v get



-------------------------------------------------------------------

{- | creates a Vector of the specified length using the supplied function to
     to map the index to the value at that index.

@> buildVector 4 fromIntegral
4 |> [0.0,1.0,2.0,3.0]@

-}
buildVector :: Storable a => Int -> (Int -> a) -> Vector a
buildVector :: forall a. Storable a => Int -> (Int -> a) -> Vector a
buildVector Int
len Int -> a
f =
    forall a. Storable a => [a] -> Vector a
fromList forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map Int -> a
f [Int
0 .. (Int
len forall a. Num a => a -> a -> a
- Int
1)]


-- | zip for Vectors
zipVector :: (Storable a, Storable b, Storable (a,b)) => Vector a -> Vector b -> Vector (a,b)
zipVector :: forall a b.
(Storable a, Storable b, Storable (a, b)) =>
Vector a -> Vector b -> Vector (a, b)
zipVector = forall a b c.
(Storable a, Storable b, Storable c) =>
(a -> b -> c) -> Vector a -> Vector b -> Vector c
zipVectorWith (,)

-- | unzip for Vectors
unzipVector :: (Storable a, Storable b, Storable (a,b)) => Vector (a,b) -> (Vector a,Vector b)
unzipVector :: forall a b.
(Storable a, Storable b, Storable (a, b)) =>
Vector (a, b) -> (Vector a, Vector b)
unzipVector = forall a b c d.
(Storable (a, b), Storable c, Storable d) =>
((a, b) -> (c, d)) -> Vector (a, b) -> (Vector c, Vector d)
unzipVectorWith forall a. a -> a
id

-------------------------------------------------------------------