{-# LANGUAGE MultiParamTypeClasses      #-}
{-# LANGUAGE MagicHash                  #-}
module Basement.Alg.PrimArray
    ( Indexable, index
    , findIndexElem
    , revFindIndexElem
    , findIndexPredicate
    , revFindIndexPredicate
    , foldl
    , foldr
    , foldl1
    , all
    , any
    , filter
    ) where

import           GHC.Types
import           GHC.Prim
import           Basement.Alg.Class
import           Basement.Compat.Base
import           Basement.Numerical.Additive
import           Basement.Numerical.Multiplicative
import           Basement.Types.OffsetSize
import           Basement.PrimType
import           Basement.Monad

findIndexElem :: (Indexable container ty, Eq ty) => ty -> container -> Offset ty -> Offset ty -> Offset ty
findIndexElem :: forall container ty.
(Indexable container ty, Eq ty) =>
ty -> container -> Offset ty -> Offset ty -> Offset ty
findIndexElem ty
ty container
ba Offset ty
startIndex Offset ty
endIndex = Offset ty -> Offset ty
loop Offset ty
startIndex
  where
    loop :: Offset ty -> Offset ty
loop !Offset ty
i
        | Offset ty
i Offset ty -> Offset ty -> Bool
forall a. Ord a => a -> a -> Bool
>= Offset ty
endIndex    = Offset ty
forall {ty}. Offset ty
sentinel
        | container -> Offset ty -> ty
forall container ty.
Indexable container ty =>
container -> Offset ty -> ty
index container
ba Offset ty
i ty -> ty -> Bool
forall a. Eq a => a -> a -> Bool
== ty
ty = Offset ty
i
        | Bool
otherwise        = Offset ty -> Offset ty
loop (Offset ty
iOffset ty -> Offset ty -> Offset ty
forall a. Additive a => a -> a -> a
+Offset ty
1)
{-# INLINE findIndexElem #-}

revFindIndexElem :: (Indexable container ty, Eq ty) => ty -> container -> Offset ty -> Offset ty -> Offset ty
revFindIndexElem :: forall container ty.
(Indexable container ty, Eq ty) =>
ty -> container -> Offset ty -> Offset ty -> Offset ty
revFindIndexElem ty
ty container
ba Offset ty
startIndex Offset ty
endIndex = Offset ty -> Offset ty
loop Offset ty
endIndex
  where
    loop :: Offset ty -> Offset ty
loop !Offset ty
iplus1
        | Offset ty
iplus1 Offset ty -> Offset ty -> Bool
forall a. Ord a => a -> a -> Bool
<= Offset ty
startIndex = Offset ty
forall {ty}. Offset ty
sentinel
        | container -> Offset ty -> ty
forall container ty.
Indexable container ty =>
container -> Offset ty -> ty
index container
ba Offset ty
i ty -> ty -> Bool
forall a. Eq a => a -> a -> Bool
== ty
ty     = Offset ty
i
        | Bool
otherwise            = Offset ty -> Offset ty
loop Offset ty
i
      where !i :: Offset ty
i = Offset ty
iplus1 Offset ty -> CountOf ty -> Offset ty
forall ty. Offset ty -> CountOf ty -> Offset ty
`offsetMinusE` CountOf ty
1
{-# INLINE revFindIndexElem #-}

findIndexPredicate :: Indexable container ty => (ty -> Bool) -> container -> Offset ty -> Offset ty -> Offset ty
findIndexPredicate :: forall container ty.
Indexable container ty =>
(ty -> Bool) -> container -> Offset ty -> Offset ty -> Offset ty
findIndexPredicate ty -> Bool
predicate container
ba Offset ty
startIndex Offset ty
endIndex = Offset ty -> Offset ty
loop Offset ty
startIndex
  where
    loop :: Offset ty -> Offset ty
loop !Offset ty
i
        | Offset ty
i Offset ty -> Offset ty -> Bool
forall a. Ord a => a -> a -> Bool
>= Offset ty
endIndex          = Offset ty
forall {ty}. Offset ty
sentinel
        | ty -> Bool
predicate (container -> Offset ty -> ty
forall container ty.
Indexable container ty =>
container -> Offset ty -> ty
index container
ba Offset ty
i) = Offset ty
i
        | Bool
otherwise              = Offset ty -> Offset ty
loop (Offset ty
iOffset ty -> Offset ty -> Offset ty
forall a. Additive a => a -> a -> a
+Offset ty
1)
{-# INLINE findIndexPredicate #-}

revFindIndexPredicate :: Indexable container ty => (ty -> Bool) -> container -> Offset ty -> Offset ty -> Offset ty
revFindIndexPredicate :: forall container ty.
Indexable container ty =>
(ty -> Bool) -> container -> Offset ty -> Offset ty -> Offset ty
revFindIndexPredicate ty -> Bool
predicate container
ba Offset ty
startIndex Offset ty
endIndex = Offset ty -> Offset ty
loop Offset ty
endIndex
  where
    loop :: Offset ty -> Offset ty
loop !Offset ty
iplus1
        | Offset ty
iplus1 Offset ty -> Offset ty -> Bool
forall a. Ord a => a -> a -> Bool
<= Offset ty
startIndex   = Offset ty
forall {ty}. Offset ty
sentinel
        | ty -> Bool
predicate (container -> Offset ty -> ty
forall container ty.
Indexable container ty =>
container -> Offset ty -> ty
index container
ba Offset ty
i) = Offset ty
i
        | Bool
otherwise              = Offset ty -> Offset ty
loop Offset ty
i
      where !i :: Offset ty
i = Offset ty
iplus1 Offset ty -> CountOf ty -> Offset ty
forall ty. Offset ty -> CountOf ty -> Offset ty
`offsetMinusE` CountOf ty
1
{-# INLINE revFindIndexPredicate #-}

foldl :: Indexable container ty => (a -> ty -> a) -> a -> container -> Offset ty -> Offset ty -> a
foldl :: forall container ty a.
Indexable container ty =>
(a -> ty -> a) -> a -> container -> Offset ty -> Offset ty -> a
foldl a -> ty -> a
f !a
initialAcc container
ba !Offset ty
startIndex !Offset ty
endIndex = Offset ty -> a -> a
loop Offset ty
startIndex a
initialAcc
  where
    loop :: Offset ty -> a -> a
loop !Offset ty
i !a
acc
        | Offset ty
i Offset ty -> Offset ty -> Bool
forall a. Eq a => a -> a -> Bool
== Offset ty
endIndex = a
acc
        | Bool
otherwise     = Offset ty -> a -> a
loop (Offset ty
iOffset ty -> Offset ty -> Offset ty
forall a. Additive a => a -> a -> a
+Offset ty
1) (a -> ty -> a
f a
acc (container -> Offset ty -> ty
forall container ty.
Indexable container ty =>
container -> Offset ty -> ty
index container
ba Offset ty
i))
{-# INLINE foldl #-}

foldr :: Indexable container ty => (ty -> a -> a) -> a -> container -> Offset ty -> Offset ty -> a
foldr :: forall container ty a.
Indexable container ty =>
(ty -> a -> a) -> a -> container -> Offset ty -> Offset ty -> a
foldr ty -> a -> a
f !a
initialAcc container
ba Offset ty
startIndex Offset ty
endIndex = Offset ty -> a
loop Offset ty
startIndex
  where
    loop :: Offset ty -> a
loop !Offset ty
i
        | Offset ty
i Offset ty -> Offset ty -> Bool
forall a. Eq a => a -> a -> Bool
== Offset ty
endIndex = a
initialAcc
        | Bool
otherwise     = container -> Offset ty -> ty
forall container ty.
Indexable container ty =>
container -> Offset ty -> ty
index container
ba Offset ty
i ty -> a -> a
`f` Offset ty -> a
loop (Offset ty
iOffset ty -> Offset ty -> Offset ty
forall a. Additive a => a -> a -> a
+Offset ty
1)
{-# INLINE foldr #-}

foldl1 :: Indexable container ty => (ty -> ty -> ty) -> container -> Offset ty -> Offset ty -> ty
foldl1 :: forall container ty.
Indexable container ty =>
(ty -> ty -> ty) -> container -> Offset ty -> Offset ty -> ty
foldl1 ty -> ty -> ty
f container
ba Offset ty
startIndex Offset ty
endIndex = Offset ty -> ty -> ty
loop (Offset ty
startIndexOffset ty -> Offset ty -> Offset ty
forall a. Additive a => a -> a -> a
+Offset ty
1) (container -> Offset ty -> ty
forall container ty.
Indexable container ty =>
container -> Offset ty -> ty
index container
ba Offset ty
startIndex)
  where
    loop :: Offset ty -> ty -> ty
loop !Offset ty
i !ty
acc
        | Offset ty
i Offset ty -> Offset ty -> Bool
forall a. Eq a => a -> a -> Bool
== Offset ty
endIndex = ty
acc
        | Bool
otherwise     = Offset ty -> ty -> ty
loop (Offset ty
iOffset ty -> Offset ty -> Offset ty
forall a. Additive a => a -> a -> a
+Offset ty
1) (ty -> ty -> ty
f ty
acc (container -> Offset ty -> ty
forall container ty.
Indexable container ty =>
container -> Offset ty -> ty
index container
ba Offset ty
i))
{-# INLINE foldl1 #-}

filter :: (PrimMonad prim, PrimType ty, Indexable container ty)
       => (ty -> Bool) -> MutableByteArray# (PrimState prim) 
       -> container -> Offset ty -> Offset ty -> prim (CountOf ty)
filter :: forall (prim :: * -> *) ty container.
(PrimMonad prim, PrimType ty, Indexable container ty) =>
(ty -> Bool)
-> MutableByteArray# (PrimState prim)
-> container
-> Offset ty
-> Offset ty
-> prim (CountOf ty)
filter ty -> Bool
predicate MutableByteArray# (PrimState prim)
dst container
src Offset ty
start Offset ty
end = Offset ty -> Offset ty -> prim (CountOf ty)
loop Offset ty
forall a. Additive a => a
azero Offset ty
start
  where
    loop :: Offset ty -> Offset ty -> prim (CountOf ty)
loop !Offset ty
d !Offset ty
s
        | Offset ty
s Offset ty -> Offset ty -> Bool
forall a. Eq a => a -> a -> Bool
== Offset ty
end    = CountOf ty -> prim (CountOf ty)
forall a. a -> prim a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Offset ty -> CountOf ty
forall a. Offset a -> CountOf a
offsetAsSize Offset ty
d)
        | ty -> Bool
predicate ty
v = MutableByteArray# (PrimState prim) -> Offset ty -> ty -> prim ()
forall ty (prim :: * -> *).
(PrimType ty, PrimMonad prim) =>
MutableByteArray# (PrimState prim) -> Offset ty -> ty -> prim ()
primMbaWrite MutableByteArray# (PrimState prim)
dst Offset ty
d ty
v prim () -> prim (CountOf ty) -> prim (CountOf ty)
forall a b. prim a -> prim b -> prim b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Offset ty -> Offset ty -> prim (CountOf ty)
loop (Offset ty
dOffset ty -> Offset ty -> Offset ty
forall a. Additive a => a -> a -> a
+Int -> Offset ty
forall ty. Int -> Offset ty
Offset Int
1) (Offset ty
sOffset ty -> Offset ty -> Offset ty
forall a. Additive a => a -> a -> a
+Int -> Offset ty
forall ty. Int -> Offset ty
Offset Int
1)
        | Bool
otherwise   = Offset ty -> Offset ty -> prim (CountOf ty)
loop Offset ty
d (Offset ty
sOffset ty -> Offset ty -> Offset ty
forall a. Additive a => a -> a -> a
+Int -> Offset ty
forall ty. Int -> Offset ty
Offset Int
1)
      where
        v :: ty
v = container -> Offset ty -> ty
forall container ty.
Indexable container ty =>
container -> Offset ty -> ty
index container
src Offset ty
s
{-# INLINE filter #-}

all :: Indexable container ty => (ty -> Bool) -> container -> Offset ty -> Offset ty -> Bool
all :: forall container ty.
Indexable container ty =>
(ty -> Bool) -> container -> Offset ty -> Offset ty -> Bool
all ty -> Bool
predicate container
ba Offset ty
start Offset ty
end = Offset ty -> Bool
loop Offset ty
start
  where
    loop :: Offset ty -> Bool
loop !Offset ty
i
        | Offset ty
i Offset ty -> Offset ty -> Bool
forall a. Eq a => a -> a -> Bool
== Offset ty
end               = Bool
True
        | ty -> Bool
predicate (container -> Offset ty -> ty
forall container ty.
Indexable container ty =>
container -> Offset ty -> ty
index container
ba Offset ty
i) = Offset ty -> Bool
loop (Offset ty
iOffset ty -> Offset ty -> Offset ty
forall a. Additive a => a -> a -> a
+Offset ty
1)
        | Bool
otherwise              = Bool
False
{-# INLINE all #-}

any :: Indexable container ty => (ty -> Bool) -> container -> Offset ty -> Offset ty -> Bool
any :: forall container ty.
Indexable container ty =>
(ty -> Bool) -> container -> Offset ty -> Offset ty -> Bool
any ty -> Bool
predicate container
ba Offset ty
start Offset ty
end = Offset ty -> Bool
loop Offset ty
start
  where
    loop :: Offset ty -> Bool
loop !Offset ty
i
        | Offset ty
i Offset ty -> Offset ty -> Bool
forall a. Eq a => a -> a -> Bool
== Offset ty
end               = Bool
False
        | ty -> Bool
predicate (container -> Offset ty -> ty
forall container ty.
Indexable container ty =>
container -> Offset ty -> ty
index container
ba Offset ty
i) = Bool
True
        | Bool
otherwise              = Offset ty -> Bool
loop (Offset ty
iOffset ty -> Offset ty -> Offset ty
forall a. Additive a => a -> a -> a
+Offset ty
1)
{-# INLINE any #-}