module Data.Repa.Array.Material.Nested
( N (..)
, Name (..)
, Array (..)
, U.Unbox
, fromLists
, fromListss
, mapElems
, slices
, concats
, segment
, segmentOn
, dice
, diceSep
, trims
, trimEnds
, trimStarts
, ragspose3)
where
import Data.Repa.Array.Meta.Delayed as A
import Data.Repa.Array.Meta.Window as A
import Data.Repa.Array.Generic.Index as A
import Data.Repa.Array.Material.Unboxed as A
import Data.Repa.Array.Material.Foreign.Base as A
import Data.Repa.Array.Internals.Bulk as A
import Data.Repa.Array.Internals.Target as A
import Data.Repa.Eval.Stream as A
import Data.Repa.Stream as S
import qualified Data.Repa.Vector.Generic as G
import qualified Data.Repa.Vector.Unboxed as U
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Fusion.Stream as S
import qualified Data.Vector.Mutable as VM
import qualified Data.Vector as VV
import Control.Monad.ST
import Control.Monad
import GHC.Exts hiding (fromList)
import Prelude as P
import Prelude hiding (concat)
#include "repa-array.h"
data N = Nested
{ nestedLength :: !Int }
deriving instance Eq N
deriving instance Show N
instance Layout N where
data Name N = N
type Index N = Int
name = N
create N len = Nested len
extent (Nested len) = len
toIndex _ ix = ix
fromIndex _ ix = ix
deriving instance Eq (Name N)
deriving instance Show (Name N)
instance (BulkI l a, Windowable l a)
=> Bulk N (Array l a) where
data Array N (Array l a)
= NArray
{ nArrayStarts :: !(U.Vector Int)
, nArrayLengths :: !(U.Vector Int)
, nArrayElems :: !(Array l a)
}
layout (NArray starts _lengths _elems)
= Nested (U.length starts)
index (NArray starts lengths elems) ix
= window (starts `U.unsafeIndex` ix)
(lengths `U.unsafeIndex` ix)
elems
deriving instance Show (Array l a)
=> Show (Array N (Array l a))
instance (Bulk l a, Target l a, Index l ~ Int)
=> Target N (Array l a) where
data Buffer N (Array l a)
= NBuffer !(VM.IOVector (Array l a))
unsafeNewBuffer (Nested n)
= NBuffer `liftM` VM.unsafeNew n
unsafeReadBuffer (NBuffer mv) i
= VM.unsafeRead mv i
unsafeWriteBuffer (NBuffer mv) i !x
= VM.unsafeWrite mv i x
unsafeGrowBuffer (NBuffer mv) x
= NBuffer `liftM` VM.unsafeGrow mv x
unsafeSliceBuffer i n (NBuffer mv)
= return $ NBuffer (VM.unsafeSlice i n mv)
touchBuffer _
= return ()
bufferLayout (NBuffer mv)
= Nested $ VM.length mv
unsafeFreezeBuffer (NBuffer mvec)
= do
!(vec :: VV.Vector (Array l a)) <- VV.unsafeFreeze mvec
let !(lengths :: U.Vector Int) = U.convert $ VV.map A.length vec
let !(starts :: U.Vector Int) = U.unsafeInit $ U.scanl (+) 0 lengths
let !(I# lenElems) = U.sum lengths
let !(I# lenArrs) = VV.length vec
!bufElems <- unsafeNewBuffer (create name (I# lenElems))
let loop_freeze !iDst !iSrcArr
| I# iSrcArr >= I# lenArrs
= return ()
| otherwise
= do let !arrSrc = VV.unsafeIndex vec (I# iSrcArr)
let !(I# lenSrc) = A.length arrSrc
let loop_freeze_copy iDst' iSrc'
| I# iSrc' >= I# lenSrc
= loop_freeze iDst' (iSrcArr +# 1#)
| otherwise
= do let !x = A.index arrSrc (I# iSrc')
unsafeWriteBuffer bufElems (I# iDst') x
loop_freeze_copy (iDst' +# 1#) (iSrc' +# 1#)
loop_freeze_copy iDst 0#
loop_freeze 0# 0#
!arrElems <- unsafeFreezeBuffer bufElems
return $ NArray starts lengths arrElems
instance (BulkI l a, Windowable l a)
=> Windowable N (Array l a) where
window start len (NArray starts lengths elems)
= NArray (U.unsafeSlice start len starts)
(U.unsafeSlice start len lengths)
elems
fromLists
:: TargetI l a
=> Name l -> [[a]] -> Array N (Array l a)
fromLists nDst xss
= let xs = concat xss
elems = fromList nDst xs
lengths = U.fromList $ P.map P.length xss
starts = U.unsafeInit $ U.scanl (+) 0 lengths
in NArray starts lengths elems
fromListss
:: TargetI l a
=> Name l -> [[[a]]] -> Array N (Array N (Array l a))
fromListss nDst xs
= let xs1 = concat xs
xs2 = concat xs1
elems = fromList nDst xs2
lengths1 = U.fromList $ P.map P.length xs
starts1 = U.unsafeInit $ U.scanl (+) 0 lengths1
lengths2 = U.fromList $ P.map P.length xs1
starts2 = U.unsafeInit $ U.scanl (+) 0 lengths2
in NArray starts1 lengths1
$ NArray starts2 lengths2
$ elems
mapElems :: (Array l1 a -> Array l2 b)
-> Array N (Array l1 a)
-> Array N (Array l2 b)
mapElems f (NArray starts lengths elems)
= NArray starts lengths (f elems)
slices :: Array F Int
-> Array F Int
-> Array l a
-> Array N (Array l a)
slices (FArray starts) (FArray lens) !elems
= NArray (VV.convert starts)
(VV.convert lens)
elems
concats :: Array N (Array N (Array l a))
-> Array N (Array l a)
concats (NArray starts1 lengths1 (NArray starts2 lengths2 elems))
= let
!starts2' = U.extract (U.unsafeIndex starts2)
$ U.zip starts1 lengths1
!lengths2' = U.extract (U.unsafeIndex lengths2)
$ U.zip starts1 lengths1
in NArray starts2' lengths2' elems
segment :: (BulkI l a, U.Unbox a)
=> (a -> Bool)
-> (a -> Bool)
-> Array l a
-> Array N (Array l a)
segment pStart pEnd !elems
= let len = size (extent $ layout elems)
(starts, lens)
= U.findSegments pStart pEnd
$ U.generate len (\ix -> index elems ix)
in NArray starts lens elems
segmentOn
:: (BulkI l a, U.Unbox a)
=> (a -> Bool)
-> Array l a
-> Array N (Array l a)
segmentOn !pEnd !arr
= segment (const True) pEnd arr
dice :: (BulkI l a, Windowable l a, U.Unbox a)
=> (a -> Bool)
-> (a -> Bool)
-> (a -> Bool)
-> (a -> Bool)
-> Array l a
-> Array N (Array N (Array l a))
dice pStart1 pEnd1 pStart2 pEnd2 !arr
= let lenArr = size (extent $ layout arr)
(starts1, lens1) = U.findSegments pStart1 pEnd1
$ U.generate lenArr (index arr)
pStart2' arr'
= pStart2 $ index arr' 0
pEnd2' arr'
= pEnd2 $ index arr' (size (extent $ layout arr') 1)
!lenArrInner = U.length starts1
!arrInner = NArray starts1 lens1 arr
(starts2, lens2) = U.findSegmentsFrom pStart2' pEnd2'
lenArrInner (index arrInner)
in NArray starts2 lens2 arrInner
diceSep :: (BulkI l a, Eq a)
=> a
-> a
-> Array l a
-> Array N (Array N (Array l a))
diceSep !xEndCol !xEndRow !arr
= let (startsLensCol, startsLensRow)
= runST
$ G.unstreamToVector2
$ S.diceSepS (== xEndCol) (== xEndRow)
$ S.liftStream
$ streamOfArray arr
(startsCol, endsCol) = U.unzip startsLensCol
(startsRow, endsRow) = U.unzip startsLensRow
in NArray startsRow endsRow $ NArray startsCol endsCol arr
trims :: BulkI l a
=> (a -> Bool)
-> Array N (Array l a)
-> Array N (Array l a)
trims pTrim (NArray starts lengths elems)
= let
loop_trimEnds !start !len
| len == 0 = (start, len)
| pTrim (elems `index` (start + len 1))
= loop_trimEnds start (len 1)
| otherwise = loop_trimStarts start len
loop_trimStarts !start !len
| len == 0 = (start, len)
| pTrim (elems `index` (start + len 1))
= loop_trimStarts (start + 1) (len 1)
| otherwise = (start, len)
(starts', lengths')
= U.unzip $ U.zipWith loop_trimEnds starts lengths
in NArray starts' lengths' elems
trimEnds :: BulkI l a
=> (a -> Bool)
-> Array N (Array l a)
-> Array N (Array l a)
trimEnds pTrim (NArray starts lengths elems)
= let
loop_trimEnds !start !len
| len == 0 = 0
| pTrim (elems `index` (start + len 1))
= loop_trimEnds start (len 1)
| otherwise = len
lengths' = U.zipWith loop_trimEnds starts lengths
in NArray starts lengths' elems
trimStarts :: BulkI l a
=> (a -> Bool)
-> Array N (Array l a)
-> Array N (Array l a)
trimStarts pTrim (NArray starts lengths elems)
= let
loop_trimStarts !start !len
| len == 0 = (start, len)
| pTrim (elems `index` (start + len 1))
= loop_trimStarts (start + 1) (len 1)
| otherwise = (start, len)
(starts', lengths')
= U.unzip $ U.zipWith loop_trimStarts starts lengths
in NArray starts' lengths' elems
ragspose3 :: Array N (Array N (Array l a))
-> Array N (Array N (Array l a))
ragspose3 (NArray starts1 lengths1 (NArray starts2 lengths2 elems))
= let
startStops1 = U.zipWith (\s l -> (s, s + l)) starts1 lengths1
(ixs', lengths1') = U.ratchet startStops1
starts2' = U.map (U.unsafeIndex starts2) ixs'
lengths2' = U.map (U.unsafeIndex lengths2) ixs'
starts1' = U.unsafeInit $ U.scanl (+) 0 lengths1'
in NArray starts1' lengths1' (NArray starts2' lengths2' elems)