-- |
-- Module      : Foundation.Array.Indexed
-- License     : BSD-style
-- Maintainer  : Vincent Hanquez <vincent@snarc.org>
-- Stability   : experimental
-- Portability : portable
--
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE CPP                   #-}

#if MIN_VERSION_base(4,9,0)
{-# LANGUAGE TypeApplications      #-}
{-# LANGUAGE UndecidableInstances  #-}
#endif

module Foundation.Collection.Indexed
    ( IndexedCollection(..)
    ) where

import           Basement.Compat.Base
import           Basement.Numerical.Additive ((+))
import           Basement.Types.OffsetSize
import           Foundation.Collection.Element
import qualified Data.List
import qualified Basement.Block as BLK
import qualified Basement.UArray as UV
import qualified Basement.BoxedArray as BA
import qualified Basement.Exception as A
import qualified Basement.String as S

#if MIN_VERSION_base(4,9,0)
import qualified Basement.Sized.Block as BLKN
import qualified Basement.Sized.List  as LN
import           Basement.Nat
#endif

-- | Collection of elements that can indexed by int
class IndexedCollection c where
    (!) :: c -> Offset (Element c) -> Maybe (Element c)
    findIndex :: (Element c -> Bool) -> c -> Maybe (Offset (Element c))

instance IndexedCollection [a] where
    (!) [a]
l (Offset Int
n)
        | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0     = Maybe (Element [a])
forall a. Maybe a
Nothing
        | Bool
otherwise = case Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
Data.List.drop Int
n [a]
l of
                        []  -> Maybe (Element [a])
forall a. Maybe a
Nothing
                        a
x:[a]
_ -> a -> Maybe a
forall a. a -> Maybe a
Just a
x
    findIndex :: (Element [a] -> Bool) -> [a] -> Maybe (Offset (Element [a]))
findIndex Element [a] -> Bool
predicate = (Int -> Offset a) -> Maybe Int -> Maybe (Offset a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Int -> Offset a
forall ty. Int -> Offset ty
Offset (Maybe Int -> Maybe (Offset a))
-> ([a] -> Maybe Int) -> [a] -> Maybe (Offset a)
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. (a -> Bool) -> [a] -> Maybe Int
forall a. (a -> Bool) -> [a] -> Maybe Int
Data.List.findIndex a -> Bool
Element [a] -> Bool
predicate

instance UV.PrimType ty => IndexedCollection (BLK.Block ty) where
    (!) Block ty
l Offset (Element (Block ty))
n
        | Offset ty -> CountOf ty -> Bool
forall ty. Offset ty -> CountOf ty -> Bool
A.isOutOfBound Offset ty
Offset (Element (Block ty))
n (Block ty -> CountOf ty
forall ty. PrimType ty => Block ty -> CountOf ty
BLK.length Block ty
l) = Maybe (Element (Block ty))
forall a. Maybe a
Nothing
        | Bool
otherwise                       = ty -> Maybe ty
forall a. a -> Maybe a
Just (ty -> Maybe ty) -> ty -> Maybe ty
forall a b. (a -> b) -> a -> b
$ Block ty -> Offset ty -> ty
forall ty. PrimType ty => Block ty -> Offset ty -> ty
BLK.index Block ty
l Offset ty
Offset (Element (Block ty))
n
    findIndex :: (Element (Block ty) -> Bool)
-> Block ty -> Maybe (Offset (Element (Block ty)))
findIndex Element (Block ty) -> Bool
predicate Block ty
c = Offset ty -> Maybe (Offset ty)
loop Offset ty
0
      where
        !len :: CountOf ty
len = Block ty -> CountOf ty
forall ty. PrimType ty => Block ty -> CountOf ty
BLK.length Block ty
c
        loop :: Offset ty -> Maybe (Offset ty)
loop Offset ty
i
            | Offset ty
i Offset ty -> CountOf ty -> Bool
forall ty. Offset ty -> CountOf ty -> Bool
.==# CountOf ty
len                      = Maybe (Offset ty)
forall a. Maybe a
Nothing
            | Element (Block ty) -> Bool
predicate (Block ty -> Offset ty -> ty
forall ty. PrimType ty => Block ty -> Offset ty -> ty
BLK.unsafeIndex Block ty
c Offset ty
i) = Offset ty -> Maybe (Offset ty)
forall a. a -> Maybe a
Just Offset ty
i
            | Bool
otherwise                       = Offset ty -> Maybe (Offset ty)
loop (Offset ty
i Offset ty -> Offset ty -> Offset ty
forall a. Additive a => a -> a -> a
+ Offset ty
1)

instance UV.PrimType ty => IndexedCollection (UV.UArray ty) where
    (!) UArray ty
l Offset (Element (UArray ty))
n
        | Offset ty -> CountOf ty -> Bool
forall ty. Offset ty -> CountOf ty -> Bool
A.isOutOfBound Offset ty
Offset (Element (UArray ty))
n (UArray ty -> CountOf ty
forall ty. UArray ty -> CountOf ty
UV.length UArray ty
l) = Maybe (Element (UArray ty))
forall a. Maybe a
Nothing
        | Bool
otherwise                          = ty -> Maybe ty
forall a. a -> Maybe a
Just (ty -> Maybe ty) -> ty -> Maybe ty
forall a b. (a -> b) -> a -> b
$ UArray ty -> Offset ty -> ty
forall ty. PrimType ty => UArray ty -> Offset ty -> ty
UV.index UArray ty
l Offset ty
Offset (Element (UArray ty))
n
    findIndex :: (Element (UArray ty) -> Bool)
-> UArray ty -> Maybe (Offset (Element (UArray ty)))
findIndex Element (UArray ty) -> Bool
predicate UArray ty
c = Offset ty -> Maybe (Offset ty)
loop Offset ty
0
      where
        !len :: CountOf ty
len = UArray ty -> CountOf ty
forall ty. UArray ty -> CountOf ty
UV.length UArray ty
c
        loop :: Offset ty -> Maybe (Offset ty)
loop Offset ty
i
            | Offset ty
i Offset ty -> CountOf ty -> Bool
forall ty. Offset ty -> CountOf ty -> Bool
.==# CountOf ty
len                     = Maybe (Offset ty)
forall a. Maybe a
Nothing
            | Element (UArray ty) -> Bool
predicate (UArray ty -> Offset ty -> ty
forall ty. PrimType ty => UArray ty -> Offset ty -> ty
UV.unsafeIndex UArray ty
c Offset ty
i) = Offset ty -> Maybe (Offset ty)
forall a. a -> Maybe a
Just Offset ty
i
            | Bool
otherwise                      = Offset ty -> Maybe (Offset ty)
loop (Offset ty
i Offset ty -> Offset ty -> Offset ty
forall a. Additive a => a -> a -> a
+ Offset ty
1)

instance IndexedCollection (BA.Array ty) where
    (!) Array ty
l Offset (Element (Array ty))
n
        | Offset ty -> CountOf ty -> Bool
forall ty. Offset ty -> CountOf ty -> Bool
A.isOutOfBound Offset ty
Offset (Element (Array ty))
n (Array ty -> CountOf ty
forall a. Array a -> CountOf a
BA.length Array ty
l) = Maybe (Element (Array ty))
forall a. Maybe a
Nothing
        | Bool
otherwise                          = ty -> Maybe ty
forall a. a -> Maybe a
Just (ty -> Maybe ty) -> ty -> Maybe ty
forall a b. (a -> b) -> a -> b
$ Array ty -> Offset ty -> ty
forall ty. Array ty -> Offset ty -> ty
BA.index Array ty
l Offset ty
Offset (Element (Array ty))
n
    findIndex :: (Element (Array ty) -> Bool)
-> Array ty -> Maybe (Offset (Element (Array ty)))
findIndex Element (Array ty) -> Bool
predicate Array ty
c = Offset ty -> Maybe (Offset ty)
loop Offset ty
0
      where
        !len :: CountOf ty
len = Array ty -> CountOf ty
forall a. Array a -> CountOf a
BA.length Array ty
c
        loop :: Offset ty -> Maybe (Offset ty)
loop Offset ty
i
            | Offset ty
i Offset ty -> CountOf ty -> Bool
forall ty. Offset ty -> CountOf ty -> Bool
.==# CountOf ty
len = Maybe (Offset ty)
forall a. Maybe a
Nothing
            | Bool
otherwise  =
                if Element (Array ty) -> Bool
predicate (Array ty -> Offset ty -> ty
forall ty. Array ty -> Offset ty -> ty
BA.unsafeIndex Array ty
c Offset ty
i) then Offset ty -> Maybe (Offset ty)
forall a. a -> Maybe a
Just Offset ty
i else Offset ty -> Maybe (Offset ty)
loop (Offset ty
i Offset ty -> Offset ty -> Offset ty
forall a. Additive a => a -> a -> a
+ Offset ty
1)

instance IndexedCollection S.String where
    (!) = String -> Offset Char -> Maybe Char
String -> Offset (Element String) -> Maybe (Element String)
S.index
    findIndex :: (Element String -> Bool)
-> String -> Maybe (Offset (Element String))
findIndex = (Char -> Bool) -> String -> Maybe (Offset Char)
(Element String -> Bool)
-> String -> Maybe (Offset (Element String))
S.findIndex

#if MIN_VERSION_base(4,9,0)
instance (NatWithinBound Int n, KnownNat n) => IndexedCollection (LN.ListN n a) where
    (!) ListN n a
c Offset (Element (ListN n a))
off
        | Offset a -> CountOf a -> Bool
forall ty. Offset ty -> CountOf ty -> Bool
A.isOutOfBound Offset a
Offset (Element (ListN n a))
off (ListN n a -> CountOf a
forall a (n :: Nat).
(KnownNat n, NatWithinBound Int n) =>
ListN n a -> CountOf a
LN.length ListN n a
c) = Maybe (Element (ListN n a))
forall a. Maybe a
Nothing
        | Bool
otherwise                        = a -> Maybe a
forall a. a -> Maybe a
Just (a -> Maybe a) -> a -> Maybe a
forall a b. (a -> b) -> a -> b
$ ListN n a -> Offset a -> a
forall (n :: Nat) ty. ListN n ty -> Offset ty -> ty
LN.index ListN n a
c Offset a
Offset (Element (ListN n a))
off
    findIndex :: (Element (ListN n a) -> Bool)
-> ListN n a -> Maybe (Offset (Element (ListN n a)))
findIndex Element (ListN n a) -> Bool
predicate ListN n a
c = Offset a -> Maybe (Offset a)
loop Offset a
0
      where
        !len :: CountOf a
len = ListN n a -> CountOf a
forall a (n :: Nat).
(KnownNat n, NatWithinBound Int n) =>
ListN n a -> CountOf a
LN.length ListN n a
c
        loop :: Offset a -> Maybe (Offset a)
loop Offset a
i
            | Offset a
i Offset a -> CountOf a -> Bool
forall ty. Offset ty -> CountOf ty -> Bool
.==# CountOf a
len               = Maybe (Offset a)
forall a. Maybe a
Nothing
            | Element (ListN n a) -> Bool
predicate (ListN n a -> Offset a -> a
forall (n :: Nat) ty. ListN n ty -> Offset ty -> ty
LN.index ListN n a
c Offset a
i) = Offset a -> Maybe (Offset a)
forall a. a -> Maybe a
Just Offset a
i
            | Bool
otherwise                = Offset a -> Maybe (Offset a)
loop (Offset a
i Offset a -> Offset a -> Offset a
forall a. Additive a => a -> a -> a
+ Offset a
1)

instance (NatWithinBound (CountOf ty) n, KnownNat n, UV.PrimType ty) => IndexedCollection (BLKN.BlockN n ty) where
    (!) BlockN n ty
c Offset (Element (BlockN n ty))
off
        | Offset ty -> CountOf ty -> Bool
forall ty. Offset ty -> CountOf ty -> Bool
A.isOutOfBound Offset ty
Offset (Element (BlockN n ty))
off (BlockN n ty -> CountOf ty
forall (n :: Nat) ty.
(KnownNat n, Countable ty n) =>
BlockN n ty -> CountOf ty
BLKN.length BlockN n ty
c) = Maybe (Element (BlockN n ty))
forall a. Maybe a
Nothing
        | Bool
otherwise                          = ty -> Maybe ty
forall a. a -> Maybe a
Just (ty -> Maybe ty) -> ty -> Maybe ty
forall a b. (a -> b) -> a -> b
$ BlockN n ty -> Offset ty -> ty
forall i (n :: Nat) ty.
PrimType ty =>
BlockN n ty -> Offset ty -> ty
BLKN.index BlockN n ty
c Offset ty
Offset (Element (BlockN n ty))
off
    findIndex :: (Element (BlockN n ty) -> Bool)
-> BlockN n ty -> Maybe (Offset (Element (BlockN n ty)))
findIndex Element (BlockN n ty) -> Bool
predicate BlockN n ty
c = Offset ty -> Maybe (Offset ty)
loop Offset ty
0
      where
        !len :: CountOf ty
len = BlockN n ty -> CountOf ty
forall (n :: Nat) ty.
(KnownNat n, Countable ty n) =>
BlockN n ty -> CountOf ty
BLKN.length BlockN n ty
c
        loop :: Offset ty -> Maybe (Offset ty)
loop Offset ty
i
            | Offset ty
i Offset ty -> CountOf ty -> Bool
forall ty. Offset ty -> CountOf ty -> Bool
.==# CountOf ty
len                 = Maybe (Offset ty)
forall a. Maybe a
Nothing
            | Element (BlockN n ty) -> Bool
predicate (BlockN n ty -> Offset ty -> ty
forall i (n :: Nat) ty.
PrimType ty =>
BlockN n ty -> Offset ty -> ty
BLKN.index BlockN n ty
c Offset ty
i) = Offset ty -> Maybe (Offset ty)
forall a. a -> Maybe a
Just Offset ty
i
            | Bool
otherwise                  = Offset ty -> Maybe (Offset ty)
loop (Offset ty
i Offset ty -> Offset ty -> Offset ty
forall a. Additive a => a -> a -> a
+ Offset ty
1)
#endif