1{-# LANGUAGE MultiParamTypeClasses      #-}
2{-# LANGUAGE MagicHash                  #-}
3module Basement.Alg.PrimArray
4    ( Indexable, index
5    , findIndexElem
6    , revFindIndexElem
7    , findIndexPredicate
8    , revFindIndexPredicate
9    , foldl
10    , foldr
11    , foldl1
12    , all
13    , any
14    , filter
15    ) where
16
17import           GHC.Types
18import           GHC.Prim
19import           Basement.Alg.Class
20import           Basement.Compat.Base
21import           Basement.Numerical.Additive
22import           Basement.Numerical.Multiplicative
23import           Basement.Types.OffsetSize
24import           Basement.PrimType
25import           Basement.Monad
26
27findIndexElem :: (Indexable container ty, Eq ty) => ty -> container -> Offset ty -> Offset ty -> Offset ty
28findIndexElem ty ba startIndex endIndex = loop startIndex
29  where
30    loop !i
31        | i >= endIndex    = sentinel
32        | index ba i == ty = i
33        | otherwise        = loop (i+1)
34{-# INLINE findIndexElem #-}
35
36revFindIndexElem :: (Indexable container ty, Eq ty) => ty -> container -> Offset ty -> Offset ty -> Offset ty
37revFindIndexElem ty ba startIndex endIndex = loop endIndex
38  where
39    loop !iplus1
40        | iplus1 <= startIndex = sentinel
41        | index ba i == ty     = i
42        | otherwise            = loop i
43      where !i = iplus1 `offsetMinusE` 1
44{-# INLINE revFindIndexElem #-}
45
46findIndexPredicate :: Indexable container ty => (ty -> Bool) -> container -> Offset ty -> Offset ty -> Offset ty
47findIndexPredicate predicate ba startIndex endIndex = loop startIndex
48  where
49    loop !i
50        | i >= endIndex          = sentinel
51        | predicate (index ba i) = i
52        | otherwise              = loop (i+1)
53{-# INLINE findIndexPredicate #-}
54
55revFindIndexPredicate :: Indexable container ty => (ty -> Bool) -> container -> Offset ty -> Offset ty -> Offset ty
56revFindIndexPredicate predicate ba startIndex endIndex = loop endIndex
57  where
58    loop !iplus1
59        | iplus1 <= startIndex   = sentinel
60        | predicate (index ba i) = i
61        | otherwise              = loop i
62      where !i = iplus1 `offsetMinusE` 1
63{-# INLINE revFindIndexPredicate #-}
64
65foldl :: Indexable container ty => (a -> ty -> a) -> a -> container -> Offset ty -> Offset ty -> a
66foldl f !initialAcc ba !startIndex !endIndex = loop startIndex initialAcc
67  where
68    loop !i !acc
69        | i == endIndex = acc
70        | otherwise     = loop (i+1) (f acc (index ba i))
71{-# INLINE foldl #-}
72
73foldr :: Indexable container ty => (ty -> a -> a) -> a -> container -> Offset ty -> Offset ty -> a
74foldr f !initialAcc ba startIndex endIndex = loop startIndex
75  where
76    loop !i
77        | i == endIndex = initialAcc
78        | otherwise     = index ba i `f` loop (i+1)
79{-# INLINE foldr #-}
80
81foldl1 :: Indexable container ty => (ty -> ty -> ty) -> container -> Offset ty -> Offset ty -> ty
82foldl1 f ba startIndex endIndex = loop (startIndex+1) (index ba startIndex)
83  where
84    loop !i !acc
85        | i == endIndex = acc
86        | otherwise     = loop (i+1) (f acc (index ba i))
87{-# INLINE foldl1 #-}
88
89filter :: (PrimMonad prim, PrimType ty, Indexable container ty)
90       => (ty -> Bool) -> MutableByteArray# (PrimState prim)
91       -> container -> Offset ty -> Offset ty -> prim (CountOf ty)
92filter predicate dst src start end = loop azero start
93  where
94    loop !d !s
95        | s == end    = pure (offsetAsSize d)
96        | predicate v = primMbaWrite dst d v >> loop (d+Offset 1) (s+Offset 1)
97        | otherwise   = loop d (s+Offset 1)
98      where
99        v = index src s
100{-# INLINE filter #-}
101
102all :: Indexable container ty => (ty -> Bool) -> container -> Offset ty -> Offset ty -> Bool
103all predicate ba start end = loop start
104  where
105    loop !i
106        | i == end               = True
107        | predicate (index ba i) = loop (i+1)
108        | otherwise              = False
109{-# INLINE all #-}
110
111any :: Indexable container ty => (ty -> Bool) -> container -> Offset ty -> Offset ty -> Bool
112any predicate ba start end = loop start
113  where
114    loop !i
115        | i == end               = False
116        | predicate (index ba i) = True
117        | otherwise              = loop (i+1)
118{-# INLINE any #-}
119