module Data.Array.Repa.Internals.Forcing
( toVector
, toList
, force, forceWith
, force2, forceWith2)
where
import Data.Array.Repa.Internals.EvalChunked
import Data.Array.Repa.Internals.EvalCursored
import Data.Array.Repa.Internals.Elt
import Data.Array.Repa.Internals.Base
import Data.Array.Repa.Index
import Data.Array.Repa.Shape as S
import qualified Data.Vector.Unboxed as V
import qualified Data.Vector.Unboxed.Mutable as VM
import Data.Vector.Unboxed (Vector)
import System.IO.Unsafe
stage = "Data.Array.Repa.Internals.Forcing"
toVector
:: (Shape sh, Elt a)
=> Array sh a
-> Vector a
toVector arr
= case force arr of
Array _ [Region _ (GenManifest vec)] -> vec
_ -> error $ stage ++ ".toVector: force failed"
toList :: (Shape sh, Elt a)
=> Array sh a
-> [a]
toList arr
= V.toList $ toVector arr
force :: (Shape sh, Elt a)
=> Array sh a -> Array sh a
force arr
= unsafePerformIO
$ do (sh, vec) <- forceIO arr
return $ sh `seq` vec `seq`
Array sh [Region RangeAll (GenManifest vec)]
where forceIO arr'
= case arr' of
Array sh [Region RangeAll (GenManifest vec)]
-> return (sh, vec)
Array sh _
-> do mvec <- VM.unsafeNew (S.size sh)
forceWith (VM.unsafeWrite mvec) arr'
vec <- V.unsafeFreeze mvec
return (sh, vec)
forceWith
:: (Shape sh, Elt a)
=> (Int -> a -> IO ())
-> Array sh a
-> IO ()
forceWith !update arr@(Array sh _)
= fillChunkedP
(S.size sh)
update
(\ix -> arr `unsafeIndex` fromIndex sh ix)
force2 :: Elt a => Array DIM2 a -> Array DIM2 a
force2 arr
= unsafePerformIO
$ do (sh, vec) <- forceIO2 arr
return $ sh `seq` vec `seq`
Array sh [Region RangeAll (GenManifest vec)]
where forceIO2 arr'
= arr' `deepSeqArray`
case arr' of
Array sh [Region RangeAll (GenManifest vec)]
-> return (sh, vec)
Array sh _
-> do mvec <- VM.new (S.size sh)
forceWith2 (VM.unsafeWrite mvec) arr'
vec <- V.unsafeFreeze mvec
return (sh, vec)
forceWith2
:: Elt a
=> (Int -> a -> IO ())
-> Array DIM2 a
-> IO ()
forceWith2 !write arr
= arr `deepSeqArray`
case arr of
Array _ [Region RangeAll (GenManifest _)]
-> forceWith write arr
Array sh [r1]
-> do fillRegion2P write sh r1
Array sh [r1, r2]
-> do fillRegion2P write sh r1
fillRegion2P write sh r2
Array sh regions
-> do mapM_ (fillRegion2P write sh) regions
fillRegion2P
:: Elt a
=> (Int -> a -> IO ())
-> DIM2
-> Region DIM2 a
-> IO ()
fillRegion2P write sh@(_ :. height :. width) (Region range gen)
= write `seq` height `seq` width `seq`
case range of
RangeAll
-> fillRect2 write sh gen
(Rect (Z :. 0 :. 0)
(Z :. height 1 :. width 1))
RangeRects _ [r1]
-> do fillRect2 write sh gen r1
RangeRects _ [r1, r2]
-> do fillRect2 write sh gen r1
fillRect2 write sh gen r2
RangeRects _ [r1, r2, r3]
-> do fillRect2 write sh gen r1
fillRect2 write sh gen r2
fillRect2 write sh gen r3
RangeRects _ [r1, r2, r3, r4]
-> do fillRect2 write sh gen r1
fillRect2 write sh gen r2
fillRect2 write sh gen r3
fillRect2 write sh gen r4
RangeRects _ rects
-> mapM_ (fillRect2 write sh gen) rects
fillRect2
:: Elt a
=> (Int -> a -> IO ())
-> DIM2
-> Generator DIM2 a
-> Rect DIM2
-> IO ()
fillRect2 write sh@(_ :. _ :. width) gen (Rect (Z :. y0 :. x0) (Z :. y1 :. x1))
= write `seq` width `seq` y0 `seq` x0 `seq` y1 `seq` x1 `seq`
case gen of
GenManifest vec
-> fillCursoredBlock2P write
id addDim (\ix -> vec `V.unsafeIndex` toIndex sh ix)
width x0 y0 x1 y1
GenCursor makeCursor shiftCursor loadElem
-> fillCursoredBlock2P write
makeCursor shiftCursor loadElem
width x0 y0 x1 y1