1-- |
2-- Module      : Foundation.Array.Chunked.Unboxed
3-- License     : BSD-style -- Maintainer  : Alfredo Di Napoli <alfredo.dinapoli@gmail.com>
4-- Stability   : experimental
5-- Portability : portable
6--
7-- Simple array-of-arrays abstraction
8--
9{-# LANGUAGE MagicHash #-}
10{-# LANGUAGE BangPatterns #-}
11{-# LANGUAGE ExistentialQuantification #-}
12{-# LANGUAGE RankNTypes #-}
13{-# LANGUAGE ScopedTypeVariables #-}
14{-# LANGUAGE ViewPatterns #-}
15module Foundation.Array.Chunked.Unboxed
16    ( ChunkedUArray
17    ) where
18
19import           Data.Typeable
20import           Control.Arrow ((***))
21import           Basement.BoxedArray (Array)
22import qualified Basement.BoxedArray as A
23import           Basement.Exception
24import           Basement.UArray (UArray)
25import qualified Basement.UArray as U
26import           Basement.Compat.Bifunctor
27import           Basement.Compat.Semigroup
28import           Basement.Compat.Base
29import           Basement.Types.OffsetSize
30import           Basement.PrimType
31import           GHC.ST
32
33import           Foundation.Numerical
34import           Foundation.Primitive
35import qualified Foundation.Collection as C
36
37
38newtype ChunkedUArray ty = ChunkedUArray (Array (UArray ty))
39                      deriving (Show, Ord, Typeable)
40
41instance PrimType ty => Eq (ChunkedUArray ty) where
42  (==) = equal
43instance NormalForm (ChunkedUArray ty) where
44    toNormalForm (ChunkedUArray spine) = toNormalForm spine
45
46instance Semigroup (ChunkedUArray a) where
47    (<>) = append
48instance Monoid (ChunkedUArray a) where
49    mempty  = empty
50    mappend = append
51    mconcat = concat
52
53type instance C.Element (ChunkedUArray ty) = ty
54
55instance PrimType ty => IsList (ChunkedUArray ty) where
56    type Item (ChunkedUArray ty) = ty
57    fromList = vFromList
58    toList = vToList
59
60instance PrimType ty => C.Foldable (ChunkedUArray ty) where
61    foldl' = foldl'
62    foldr = foldr
63    -- Use the default foldr' instance
64
65instance PrimType ty => C.Collection (ChunkedUArray ty) where
66    null = null
67    length = length
68    elem   = elem
69    minimum = minimum
70    maximum = maximum
71    all p (ChunkedUArray cua) = A.all (U.all p) cua
72    any p (ChunkedUArray cua) = A.any (U.any p) cua
73
74instance PrimType ty => C.Sequential (ChunkedUArray ty) where
75    take = take
76    drop = drop
77    splitAt = splitAt
78    revTake = revTake
79    revDrop = revDrop
80    splitOn = splitOn
81    break = break
82    breakEnd = breakEnd
83    intersperse = intersperse
84    filter = filter
85    reverse = reverse
86    unsnoc = unsnoc
87    uncons = uncons
88    snoc = snoc
89    cons = cons
90    find = find
91    sortBy = sortBy
92    singleton = fromList . (:[])
93    replicate n = fromList . C.replicate n
94
95instance PrimType ty => C.IndexedCollection (ChunkedUArray ty) where
96    (!) l n
97        | isOutOfBound n (length l) = Nothing
98        | otherwise                     = Just $ index l n
99    findIndex predicate c = loop 0
100      where
101        !len = length c
102        loop i
103            | i .==# len = Nothing
104            | otherwise  =
105                if predicate (unsafeIndex c i) then Just i else Nothing
106
107empty :: ChunkedUArray ty
108empty = ChunkedUArray A.empty
109
110append :: ChunkedUArray ty -> ChunkedUArray ty -> ChunkedUArray ty
111append (ChunkedUArray a1) (ChunkedUArray a2) = ChunkedUArray (mappend a1 a2)
112
113concat :: [ChunkedUArray ty] -> ChunkedUArray ty
114concat x = ChunkedUArray (mconcat $ fmap (\(ChunkedUArray spine) -> spine) x)
115
116vFromList :: PrimType ty => [ty] -> ChunkedUArray ty
117vFromList l = ChunkedUArray $ A.singleton $ fromList l
118
119vToList :: PrimType ty => ChunkedUArray ty -> [ty]
120vToList (ChunkedUArray a) = mconcat $ toList $ toList <$> a
121
122null :: PrimType ty => ChunkedUArray ty -> Bool
123null (ChunkedUArray array) =
124    C.null array || allNulls 0
125  where
126    !len = A.length array
127    allNulls !idx
128      | idx .==# len = True
129      | otherwise    = C.null (array `A.unsafeIndex` idx) && allNulls (idx + 1)
130
131-- | Returns the length of this `ChunkedUArray`, by summing each inner length.
132-- Complexity: O(n) where `n` is the number of chunks, as U.length u is O(1).
133length :: PrimType ty => ChunkedUArray ty -> CountOf ty
134length (ChunkedUArray array) = C.foldl' (\acc l -> acc + U.length l) 0 array
135
136-- | Returns `True` if the given element is contained in the `ChunkedUArray`.
137-- Complexity: O(n) where `n` is the number of chunks, as U.length u is O(1).
138elem :: PrimType ty => ty -> ChunkedUArray ty -> Bool
139elem el (ChunkedUArray array) = loop 0
140  where
141    !len = A.length array
142    loop i
143        | i .==# len = False
144        | otherwise  =
145            case C.elem el (A.unsafeIndex array i) of
146                True  -> True
147                False -> loop (i+1)
148
149-- | Fold a `ChunkedUArray' leftwards strictly. Implemented internally using a double
150-- fold on the nested Array structure. Other folds implemented analogously.
151foldl' :: PrimType ty => (a -> ty -> a) -> a -> ChunkedUArray ty -> a
152foldl' f initialAcc (ChunkedUArray cua) = A.foldl' (U.foldl' f) initialAcc cua
153
154foldr :: PrimType ty => (ty -> a -> a) -> a -> ChunkedUArray ty -> a
155foldr f initialAcc (ChunkedUArray cua) = A.foldr (flip $ U.foldr f) initialAcc cua
156
157minimum :: (Ord ty, PrimType ty) => C.NonEmpty (ChunkedUArray ty) -> ty
158minimum cua = foldl' min (unsafeIndex cua' 0) (drop 1 cua')
159  where
160    cua' = C.getNonEmpty cua
161
162maximum :: (Ord ty, PrimType ty) => C.NonEmpty (ChunkedUArray ty) -> ty
163maximum cua = foldl' max (unsafeIndex cua' 0) (drop 1 cua')
164  where
165    cua' = C.getNonEmpty cua
166
167-- | Equality between `ChunkedUArray`.
168-- This function is fiddly to write as is not enough to compare for
169-- equality the inner `UArray`(s), we need an element-by-element
170-- comparison.
171equal :: PrimType ty => ChunkedUArray ty -> ChunkedUArray ty -> Bool
172equal ca1 ca2 =
173    len1 == len2 && go 0
174  where
175    len1 = length ca1
176    len2 = length ca2
177
178    go !x
179      | x .==# len1 = True
180      | otherwise   = (ca1 `unsafeIndex` x == ca2 `unsafeIndex` x) && go (x + 1)
181
182-- given an offset express in element of ty, return the offset in array in the spine,
183-- plus the relative offset in element on this array
184findPos :: PrimType ty => Offset ty -> ChunkedUArray ty -> Maybe (Offset (UArray ty), Offset ty)
185findPos absOfs (ChunkedUArray array)
186    | A.null array = Nothing
187    | otherwise    = loop absOfs 0
188  where
189    !len = A.length array
190    loop relOfs outerI
191        | outerI .==# len = Nothing -- haven't found what to do
192        | relOfs == 0     = Just (outerI, 0)
193        | otherwise       =
194            let !innera   = A.unsafeIndex array outerI
195                !innerLen = U.length innera
196             in case removeArraySize relOfs innerLen of
197                        Nothing      -> Just (outerI, relOfs)
198                        Just relOfs' -> loop relOfs' (outerI + 1)
199
200splitChunk :: Offset (UArray ty) -> ChunkedUArray ty -> (ChunkedUArray ty, ChunkedUArray ty)
201splitChunk ofs (ChunkedUArray c) = (ChunkedUArray *** ChunkedUArray) $ A.splitAt (offsetAsSize ofs) c
202
203take :: PrimType ty => CountOf ty -> ChunkedUArray ty -> ChunkedUArray ty
204take n c@(ChunkedUArray spine)
205    | n <= 0    = empty
206    | otherwise =
207        case findPos (sizeAsOffset n) c of
208            Nothing       -> c
209            Just (ofs, 0) -> ChunkedUArray (A.take (offsetAsSize ofs) spine)
210            Just (ofs, r) ->
211                let uarr = A.unsafeIndex spine ofs
212                 in ChunkedUArray (A.take (offsetAsSize ofs) spine `A.snoc` U.take (offsetAsSize r) uarr)
213
214drop :: PrimType ty => CountOf ty -> ChunkedUArray ty -> ChunkedUArray ty
215drop n c@(ChunkedUArray spine)
216    | n <= 0    = c
217    | otherwise =
218        case findPos (sizeAsOffset n) c of
219            Nothing       -> empty
220            Just (ofs, 0) -> ChunkedUArray (A.drop (offsetAsSize ofs) spine)
221            Just (ofs, r) ->
222                let uarr = A.unsafeIndex spine ofs
223                 in ChunkedUArray (U.drop (offsetAsSize r) uarr `A.cons` A.drop (offsetAsSize ofs+1) spine)
224
225splitAt :: PrimType ty => CountOf ty -> ChunkedUArray ty -> (ChunkedUArray ty, ChunkedUArray ty)
226splitAt n c@(ChunkedUArray spine)
227    | n <= 0    = (empty, c)
228    | otherwise =
229        case findPos (sizeAsOffset n) c of
230            Nothing       -> (c, empty)
231            Just (ofs, 0) -> splitChunk ofs c
232            Just (ofs, offsetAsSize -> r) ->
233                let uarr = A.unsafeIndex spine ofs
234                 in ( ChunkedUArray (A.take (offsetAsSize ofs) spine `A.snoc` U.take r uarr)
235                    , ChunkedUArray (U.drop r uarr `A.cons` A.drop (offsetAsSize ofs+1) spine)
236                    )
237
238revTake :: PrimType ty => CountOf ty -> ChunkedUArray ty -> ChunkedUArray ty
239revTake n c = case length c - n of
240    Nothing -> c
241    Just elems -> drop elems c
242
243revDrop :: PrimType ty => CountOf ty -> ChunkedUArray ty -> ChunkedUArray ty
244revDrop n c = case length c - n of
245    Nothing -> empty
246    Just keepElems -> take keepElems c
247
248-- TODO: Improve implementation.
249splitOn :: PrimType ty => (ty -> Bool) -> ChunkedUArray ty -> [ChunkedUArray ty]
250splitOn p = fmap fromList . C.splitOn p . toList
251
252-- TODO: Improve implementation.
253break :: PrimType ty => (ty -> Bool) -> ChunkedUArray ty -> (ChunkedUArray ty, ChunkedUArray ty)
254break p = bimap fromList fromList . C.break p . toList
255
256-- TODO: Improve implementation.
257breakEnd :: PrimType ty => (ty -> Bool) -> ChunkedUArray ty -> (ChunkedUArray ty, ChunkedUArray ty)
258breakEnd p = bimap fromList fromList . C.breakEnd p . toList
259
260-- TODO: Improve implementation.
261intersperse :: PrimType ty => ty -> ChunkedUArray ty -> ChunkedUArray ty
262intersperse el = fromList . C.intersperse el . toList
263
264-- TODO: Improve implementation.
265reverse :: PrimType ty => ChunkedUArray ty -> ChunkedUArray ty
266reverse = fromList . C.reverse . toList
267
268-- TODO: Improve implementation.
269filter :: PrimType ty => (ty -> Bool) -> ChunkedUArray ty -> ChunkedUArray ty
270filter p = fromList . C.filter p . toList
271
272-- TODO: Improve implementation.
273unsnoc :: PrimType ty => ChunkedUArray ty -> Maybe (ChunkedUArray ty, ty)
274unsnoc v = first fromList <$> (C.unsnoc $ toList v)
275
276-- TODO: Improve implementation.
277uncons :: PrimType ty => ChunkedUArray ty -> Maybe (ty, ChunkedUArray ty)
278uncons v = second fromList <$> (C.uncons $ toList v)
279
280cons :: PrimType ty => ty -> ChunkedUArray ty -> ChunkedUArray ty
281cons el (ChunkedUArray inner) = ChunkedUArray $ runST $ do
282  let newLen = C.length inner + 1
283  newArray   <- A.new newLen
284  let single = fromList [el]
285  A.unsafeWrite newArray 0 single
286  A.unsafeCopyAtRO newArray (Offset 1) inner (Offset 0) (C.length inner)
287  A.unsafeFreeze newArray
288
289snoc :: PrimType ty => ChunkedUArray ty -> ty -> ChunkedUArray ty
290snoc (ChunkedUArray spine) el = ChunkedUArray $ runST $ do
291  newArray  <- A.new (A.length spine + 1)
292  let single = U.singleton el
293  A.unsafeCopyAtRO newArray (Offset 0) spine (Offset 0) (C.length spine)
294  A.unsafeWrite newArray (sizeAsOffset $ A.length spine) single
295  A.unsafeFreeze newArray
296
297-- TODO optimise
298find :: PrimType ty => (ty -> Bool) -> ChunkedUArray ty -> Maybe ty
299find fn v = loop 0
300  where
301    len = length v
302    loop !idx
303      | idx .==# len = Nothing
304      | otherwise    =
305        let currentElem = v `unsafeIndex` idx
306        in case fn currentElem of
307          True  -> Just currentElem
308          False -> loop (idx + 1)
309
310-- TODO: Improve implementation.
311sortBy :: PrimType ty => (ty -> ty -> Ordering) -> ChunkedUArray ty -> ChunkedUArray ty
312sortBy p = fromList . C.sortBy p . toList
313
314index :: PrimType ty => ChunkedUArray ty -> Offset ty -> ty
315index array n
316    | isOutOfBound n len = outOfBound OOB_Index n len
317    | otherwise          = unsafeIndex array n
318  where len = length array
319{-# INLINE index #-}
320
321unsafeIndex :: PrimType ty => ChunkedUArray ty -> Offset ty -> ty
322unsafeIndex (ChunkedUArray array) idx = go (A.unsafeIndex array 0) 0 idx
323  where
324    go u globalIndex 0 = case C.null u of
325      -- Skip empty chunks.
326      True  -> go (A.unsafeIndex array (globalIndex + 1)) (globalIndex + 1) 0
327      False -> U.unsafeIndex u 0
328    go u !globalIndex !i
329      -- Skip empty chunks.
330      | C.null u  = go (A.unsafeIndex array (globalIndex + 1)) (globalIndex + 1) i
331      | otherwise =
332          case removeArraySize i (U.length u) of
333              Just i' -> go (A.unsafeIndex array (globalIndex + 1)) (globalIndex + 1) i'
334              Nothing -> U.unsafeIndex u i
335
336{-# INLINE unsafeIndex #-}
337
338removeArraySize :: Offset ty -> CountOf ty -> Maybe (Offset ty)
339removeArraySize (Offset ty) (CountOf s)
340    | ty >= s   = Just (Offset (ty - s))
341    | otherwise = Nothing
342