module Foundation.Array.Chunked.Unboxed
( ChunkedUArray
) where
import qualified Data.List
import Data.Typeable
import Foundation.Array.Boxed (Array)
import qualified Foundation.Array.Boxed as A
import Foundation.Array.Common
import Foundation.Array.Unboxed (UArray)
import qualified Foundation.Array.Unboxed as U
import Foundation.Class.Bifunctor
import qualified Foundation.Collection as C
import Foundation.Internal.Base
import Foundation.Internal.Types
import Foundation.Numerical
import Foundation.Primitive.Monad
import Foundation.Primitive.Types
import GHC.ST
import qualified Prelude as P
data ChunkedUArray ty = ChunkedUArray (Array (UArray ty))
deriving (Show, Ord, Typeable)
instance PrimType ty => Eq (ChunkedUArray ty) where
(==) = equal
instance Monoid (ChunkedUArray a) where
mempty = empty
mappend = append
mconcat = concat
type instance C.Element (ChunkedUArray ty) = ty
instance PrimType ty => IsList (ChunkedUArray ty) where
type Item (ChunkedUArray ty) = ty
fromList = vFromList
toList = vToList
instance PrimType ty => C.Collection (ChunkedUArray ty) where
null = null
length = length
elem = elem
minimum = minimum
maximum = maximum
all p = Data.List.all p . toList
any p = Data.List.any p . toList
instance PrimType ty => C.Sequential (ChunkedUArray ty) where
take = take
drop = drop
revTake = revTake
revDrop = revDrop
splitOn = splitOn
break = break
intersperse = intersperse
filter = filter
reverse = reverse
unsnoc = unsnoc
uncons = uncons
snoc = snoc
cons = cons
find = find
sortBy = sortBy
singleton = fromList . (:[])
instance PrimType ty => C.IndexedCollection (ChunkedUArray ty) where
(!) l n
| n < 0 || n >= length l = Nothing
| otherwise = Just $ index l n
findIndex predicate c = loop 0
where
!len = length c
loop i
| i == len = Nothing
| otherwise =
if predicate (unsafeIndex c i) then Just i else Nothing
empty :: ChunkedUArray ty
empty = ChunkedUArray (A.empty)
append :: ChunkedUArray ty -> ChunkedUArray ty -> ChunkedUArray ty
append (ChunkedUArray a1) (ChunkedUArray a2) = ChunkedUArray $ runST $ do
let a1Size@(Size a1len) = Size $ C.length a1
let a2Size = Size $ C.length a2
a <- A.new (a1Size + a2Size)
A.thaw a1 >>= \a1' -> A.copyAt a (Offset 0) a1' (Offset 0) a1Size
A.thaw a2 >>= \a2' -> A.copyAt a (Offset a1len) a2' (Offset 0) a2Size
A.unsafeFreeze a
concat :: [ChunkedUArray ty] -> ChunkedUArray ty
concat x = C.foldl' append mempty x
vFromList :: PrimType ty => [ty] -> ChunkedUArray ty
vFromList l = ChunkedUArray array
where
array = runST $ do
a <- A.new (Size 1)
A.unsafeWrite a 0 (fromList l)
A.unsafeFreeze a
vToList :: PrimType ty => ChunkedUArray ty -> [ty]
vToList (ChunkedUArray a) = mconcat $ toList $ toList <$> a
null :: PrimType ty => ChunkedUArray ty -> Bool
null (ChunkedUArray array) =
let len = C.length array
in C.null array || allNulls 0 len
where
allNulls !idx len
| idx == len = True
| otherwise = C.null (array `A.unsafeIndex` idx) && allNulls (idx + 1) len
length :: PrimType ty => ChunkedUArray ty -> Int
length (ChunkedUArray array) = C.foldl' (\acc l -> acc + C.length l) 0 array
elem :: PrimType ty => ty -> ChunkedUArray ty -> Bool
elem el array = go 0
where
len = C.length array
go !currentIndex = case currentIndex < len of
True -> case el == array `unsafeIndex` currentIndex of
True -> True
False -> go (currentIndex + 1)
False -> False
minimum :: (Ord ty, PrimType ty) => C.NonEmpty (ChunkedUArray ty) -> ty
minimum = Data.List.minimum . toList . C.getNonEmpty
maximum :: (Ord ty, PrimType ty) => C.NonEmpty (ChunkedUArray ty) -> ty
maximum = Data.List.maximum . toList . C.getNonEmpty
equal :: PrimType ty => ChunkedUArray ty -> ChunkedUArray ty -> Bool
equal ca1 ca2 = len1 == len2 && deepEqual
where
len1 = C.length ca1
len2 = C.length ca2
deepEqual :: Bool
deepEqual = go 0 0
go !x !y
| x == len1 && y == len2 = True
| otherwise =
(ca1 `unsafeIndex` x == ca2 `unsafeIndex` y) && go (x + 1) (y + 1)
take :: PrimType ty => Int -> ChunkedUArray ty -> ChunkedUArray ty
take nbElems v@(ChunkedUArray inner)
| nbElems <= 0 = empty
| C.null v = empty
| nbElems >= C.length v = v
| otherwise =
let newSize = Size requiredChunks
in ChunkedUArray $ runST (A.new newSize >>= iter inner nbElems)
where
requiredChunks = loop 0 nbElems
where
loop !idx !remaining
| remaining <= 0 = idx
| otherwise =
let vec = inner `A.unsafeIndex` idx
l = U.length vec
in loop (idx + 1) (remaining l)
iter :: (PrimType ty, PrimMonad prim)
=> Array (UArray ty)
-> Int
-> A.MArray (UArray ty) (PrimState prim)
-> prim (Array (UArray ty))
iter inner0 elems finalVector = loop 0 elems
where
loop !currentIndex !remainingElems
| remainingElems <= 0 || currentIndex >= C.length inner0 = A.unsafeFreeze finalVector
| otherwise =
let chunk = inner0 `A.unsafeIndex` currentIndex
chunkLen = C.length chunk
in case C.null chunk of
True -> loop (currentIndex + 1) remainingElems
False -> case chunkLen <= remainingElems of
True -> do
A.unsafeWrite finalVector currentIndex chunk
loop (currentIndex + 1) (remainingElems chunkLen)
False -> do
nc <- do
newChunk <- U.new (Size remainingElems)
U.unsafeCopyAtRO newChunk (Offset 0) chunk (Offset 0) (Size remainingElems)
U.unsafeFreeze newChunk
A.unsafeWrite finalVector currentIndex nc
A.freeze finalVector
drop :: PrimType ty => Int -> ChunkedUArray ty -> ChunkedUArray ty
drop nbElems v@(ChunkedUArray inner)
| nbElems >= C.length v = empty
| nbElems <= 0 = v
| C.null v = empty
| otherwise =
let newSize = Size (C.length inner chunksToSkip)
in ChunkedUArray $ runST (A.new newSize >>= iter inner nbElems)
where
chunksToSkip = loop 0 nbElems
where
loop !idx !remaining =
let vec = inner `A.unsafeIndex` idx
l = U.length vec
slack = remaining l
in case slack of
x | x == 0 -> idx + 1
x | x < 0 -> idx
_ -> loop (idx + 1) slack
iter :: (PrimType ty, PrimMonad prim)
=> Array (UArray ty)
-> Int
-> A.MArray (UArray ty) (PrimState prim)
-> prim (Array (UArray ty))
iter inner0 elems finalVector = loop 0 elems
where
loop !currentIndex !remainingElems
| remainingElems <= 0 = do
A.unsafeCopyAtRO finalVector (Offset 0) inner0 (Offset currentIndex) (Size $ C.length inner0 currentIndex)
A.freeze finalVector
| otherwise =
let chunk = inner0 `A.unsafeIndex` currentIndex
chunkLen = C.length chunk
slack = chunkLen P.- remainingElems
in case chunkLen <= remainingElems of
True -> do
loop (currentIndex + 1) (remainingElems chunkLen)
False -> do
nc <- do
newChunk <- U.new (Size slack)
U.unsafeCopyAtRO newChunk (Offset 0) chunk (Offset remainingElems) (Size slack)
U.unsafeFreeze newChunk
A.unsafeWrite finalVector 0 nc
let !nextIdx = currentIndex + 1
A.unsafeCopyAtRO finalVector (Offset 1) inner0 (Offset nextIdx) (Size $ C.length inner0 nextIdx)
A.freeze finalVector
revTake :: PrimType ty => Int -> ChunkedUArray ty -> ChunkedUArray ty
revTake x = fromList . C.revTake x . toList
revDrop :: PrimType ty => Int -> ChunkedUArray ty -> ChunkedUArray ty
revDrop x = fromList . C.revDrop x . toList
splitOn :: PrimType ty => (ty -> Bool) -> ChunkedUArray ty -> [ChunkedUArray ty]
splitOn p = fmap fromList . C.splitOn p . toList
break :: PrimType ty => (ty -> Bool) -> ChunkedUArray ty -> (ChunkedUArray ty, ChunkedUArray ty)
break p = bimap fromList fromList . C.break p . toList
intersperse :: PrimType ty => ty -> ChunkedUArray ty -> ChunkedUArray ty
intersperse el = fromList . C.intersperse el . toList
reverse :: PrimType ty => ChunkedUArray ty -> ChunkedUArray ty
reverse = fromList . C.reverse . toList
filter :: PrimType ty => (ty -> Bool) -> ChunkedUArray ty -> ChunkedUArray ty
filter p = fromList . C.filter p . toList
unsnoc :: PrimType ty => ChunkedUArray ty -> Maybe (ChunkedUArray ty, ty)
unsnoc v = first fromList <$> (C.unsnoc $ toList v)
uncons :: PrimType ty => ChunkedUArray ty -> Maybe (ty, ChunkedUArray ty)
uncons v = second fromList <$> (C.uncons $ toList v)
cons :: PrimType ty => ty -> ChunkedUArray ty -> ChunkedUArray ty
cons el (ChunkedUArray inner) = ChunkedUArray $ runST $ do
let newLen = (Size $ C.length inner + 1)
newArray <- A.new newLen
let single = fromList [el]
A.unsafeWrite newArray 0 single
A.unsafeCopyAtRO newArray (Offset 1) inner (Offset 0) (Size $ C.length inner)
A.unsafeFreeze newArray
snoc :: PrimType ty => ChunkedUArray ty -> ty -> ChunkedUArray ty
snoc (ChunkedUArray inner) el = ChunkedUArray $ runST $ do
newArray <- A.new (Size $ C.length inner + 1)
let single = fromList [el]
A.unsafeCopyAtRO newArray (Offset 0) inner (Offset 0) (Size $ C.length inner)
A.unsafeWrite newArray (C.length inner) single
A.unsafeFreeze newArray
find :: PrimType ty => (ty -> Bool) -> ChunkedUArray ty -> Maybe ty
find fn v = loop 0 (C.length v)
where
loop !idx len
| idx >= len = Nothing
| otherwise =
let currentElem = v `unsafeIndex` idx
in case fn currentElem of
True -> Just currentElem
False -> loop (idx + 1) len
sortBy :: PrimType ty => (ty -> ty -> Ordering) -> ChunkedUArray ty -> ChunkedUArray ty
sortBy p = fromList . C.sortBy p . toList
index :: PrimType ty => ChunkedUArray ty -> Int -> ty
index array n
| n < 0 || n >= len = throw (OutOfBound OOB_Index n len)
| otherwise = unsafeIndex array n
where len = C.length array
unsafeIndex :: PrimType ty => ChunkedUArray ty -> Int -> ty
unsafeIndex (ChunkedUArray array) idx = go (A.unsafeIndex array 0) 0 idx
where
go u globalIndex 0 = case C.null u of
True -> go (A.unsafeIndex array (globalIndex + 1)) (globalIndex + 1) 0
False -> U.unsafeIndex u 0
go u !globalIndex !i
| C.null u = go (A.unsafeIndex array (globalIndex + 1)) (globalIndex + 1) i
| otherwise = case i (C.length u) of
i' | i' >= 0 -> go (A.unsafeIndex array (globalIndex + 1)) (globalIndex + 1) i'
_ -> U.unsafeIndex u i