{-|
Module      : Z.Data.Vector.FlatSet
Description : Fast set based on sorted vector
Copyright   : (c) Dong Han, 2017-2019
              (c) Tao He, 2018-2019
License     : BSD
Maintainer  : winterland1989@gmail.com
Stability   : experimental
Portability : non-portable

This module provides a simple value set based on sorted vector and binary search. It's particularly
suitable for small sized value collections such as deserializing intermediate representation.
But can also used in various place where insertion and deletion is rare but require fast elem.

-}

module Z.Data.Vector.FlatSet
  ( -- * FlatSet backed by sorted vector
    FlatSet, sortedValues, size, null, empty, map'
  , pack, packN, packR, packRN
  , unpack, unpackR, packVector, packVectorR
  , elem
  , delete
  , insert
  , merge
    -- * search on vectors
  , binarySearch
  ) where

import           Control.DeepSeq
import           Control.Monad
import           Control.Monad.ST
import qualified Data.Primitive.SmallArray  as A
import qualified Data.Semigroup             as Semigroup
import qualified Data.Monoid                as Monoid
import qualified Z.Data.Vector.Base         as V
import qualified Z.Data.Vector.Extra        as V
import qualified Z.Data.Vector.Sort         as V
import qualified Z.Data.Text.Print          as T
import           Data.Bits                   (unsafeShiftR)
import           Data.Data
import           Prelude hiding (elem, null)
import           Test.QuickCheck.Arbitrary (Arbitrary(..), CoArbitrary(..))

--------------------------------------------------------------------------------

newtype FlatSet v = FlatSet { forall v. FlatSet v -> Vector v
sortedValues :: V.Vector v }
    deriving (Int -> FlatSet v -> ShowS
forall v. Show v => Int -> FlatSet v -> ShowS
forall v. Show v => [FlatSet v] -> ShowS
forall v. Show v => FlatSet v -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [FlatSet v] -> ShowS
$cshowList :: forall v. Show v => [FlatSet v] -> ShowS
show :: FlatSet v -> String
$cshow :: forall v. Show v => FlatSet v -> String
showsPrec :: Int -> FlatSet v -> ShowS
$cshowsPrec :: forall v. Show v => Int -> FlatSet v -> ShowS
Show, FlatSet v -> FlatSet v -> Bool
forall v. Eq v => FlatSet v -> FlatSet v -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: FlatSet v -> FlatSet v -> Bool
$c/= :: forall v. Eq v => FlatSet v -> FlatSet v -> Bool
== :: FlatSet v -> FlatSet v -> Bool
$c== :: forall v. Eq v => FlatSet v -> FlatSet v -> Bool
Eq, FlatSet v -> FlatSet v -> Bool
FlatSet v -> FlatSet v -> Ordering
FlatSet v -> FlatSet v -> FlatSet v
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
forall {v}. Ord v => Eq (FlatSet v)
forall v. Ord v => FlatSet v -> FlatSet v -> Bool
forall v. Ord v => FlatSet v -> FlatSet v -> Ordering
forall v. Ord v => FlatSet v -> FlatSet v -> FlatSet v
min :: FlatSet v -> FlatSet v -> FlatSet v
$cmin :: forall v. Ord v => FlatSet v -> FlatSet v -> FlatSet v
max :: FlatSet v -> FlatSet v -> FlatSet v
$cmax :: forall v. Ord v => FlatSet v -> FlatSet v -> FlatSet v
>= :: FlatSet v -> FlatSet v -> Bool
$c>= :: forall v. Ord v => FlatSet v -> FlatSet v -> Bool
> :: FlatSet v -> FlatSet v -> Bool
$c> :: forall v. Ord v => FlatSet v -> FlatSet v -> Bool
<= :: FlatSet v -> FlatSet v -> Bool
$c<= :: forall v. Ord v => FlatSet v -> FlatSet v -> Bool
< :: FlatSet v -> FlatSet v -> Bool
$c< :: forall v. Ord v => FlatSet v -> FlatSet v -> Bool
compare :: FlatSet v -> FlatSet v -> Ordering
$ccompare :: forall v. Ord v => FlatSet v -> FlatSet v -> Ordering
Ord, Typeable, forall a. Eq a => a -> FlatSet a -> Bool
forall a. Num a => FlatSet a -> a
forall a. Ord a => FlatSet a -> a
forall m. Monoid m => FlatSet m -> m
forall a. FlatSet a -> Bool
forall a. FlatSet a -> Int
forall a. FlatSet a -> [a]
forall a. (a -> a -> a) -> FlatSet a -> a
forall m a. Monoid m => (a -> m) -> FlatSet a -> m
forall b a. (b -> a -> b) -> b -> FlatSet a -> b
forall a b. (a -> b -> b) -> b -> FlatSet a -> b
forall (t :: * -> *).
(forall m. Monoid m => t m -> m)
-> (forall m a. Monoid m => (a -> m) -> t a -> m)
-> (forall m a. Monoid m => (a -> m) -> t a -> m)
-> (forall a b. (a -> b -> b) -> b -> t a -> b)
-> (forall a b. (a -> b -> b) -> b -> t a -> b)
-> (forall b a. (b -> a -> b) -> b -> t a -> b)
-> (forall b a. (b -> a -> b) -> b -> t a -> b)
-> (forall a. (a -> a -> a) -> t a -> a)
-> (forall a. (a -> a -> a) -> t a -> a)
-> (forall a. t a -> [a])
-> (forall a. t a -> Bool)
-> (forall a. t a -> Int)
-> (forall a. Eq a => a -> t a -> Bool)
-> (forall a. Ord a => t a -> a)
-> (forall a. Ord a => t a -> a)
-> (forall a. Num a => t a -> a)
-> (forall a. Num a => t a -> a)
-> Foldable t
product :: forall a. Num a => FlatSet a -> a
$cproduct :: forall a. Num a => FlatSet a -> a
sum :: forall a. Num a => FlatSet a -> a
$csum :: forall a. Num a => FlatSet a -> a
minimum :: forall a. Ord a => FlatSet a -> a
$cminimum :: forall a. Ord a => FlatSet a -> a
maximum :: forall a. Ord a => FlatSet a -> a
$cmaximum :: forall a. Ord a => FlatSet a -> a
elem :: forall a. Eq a => a -> FlatSet a -> Bool
$celem :: forall a. Eq a => a -> FlatSet a -> Bool
length :: forall a. FlatSet a -> Int
$clength :: forall a. FlatSet a -> Int
null :: forall a. FlatSet a -> Bool
$cnull :: forall a. FlatSet a -> Bool
toList :: forall a. FlatSet a -> [a]
$ctoList :: forall a. FlatSet a -> [a]
foldl1 :: forall a. (a -> a -> a) -> FlatSet a -> a
$cfoldl1 :: forall a. (a -> a -> a) -> FlatSet a -> a
foldr1 :: forall a. (a -> a -> a) -> FlatSet a -> a
$cfoldr1 :: forall a. (a -> a -> a) -> FlatSet a -> a
foldl' :: forall b a. (b -> a -> b) -> b -> FlatSet a -> b
$cfoldl' :: forall b a. (b -> a -> b) -> b -> FlatSet a -> b
foldl :: forall b a. (b -> a -> b) -> b -> FlatSet a -> b
$cfoldl :: forall b a. (b -> a -> b) -> b -> FlatSet a -> b
foldr' :: forall a b. (a -> b -> b) -> b -> FlatSet a -> b
$cfoldr' :: forall a b. (a -> b -> b) -> b -> FlatSet a -> b
foldr :: forall a b. (a -> b -> b) -> b -> FlatSet a -> b
$cfoldr :: forall a b. (a -> b -> b) -> b -> FlatSet a -> b
foldMap' :: forall m a. Monoid m => (a -> m) -> FlatSet a -> m
$cfoldMap' :: forall m a. Monoid m => (a -> m) -> FlatSet a -> m
foldMap :: forall m a. Monoid m => (a -> m) -> FlatSet a -> m
$cfoldMap :: forall m a. Monoid m => (a -> m) -> FlatSet a -> m
fold :: forall m. Monoid m => FlatSet m -> m
$cfold :: forall m. Monoid m => FlatSet m -> m
Foldable)

instance T.Print v => T.Print (FlatSet v) where
    {-# INLINE toUTF8BuilderP #-}
    toUTF8BuilderP :: Int -> FlatSet v -> Builder ()
toUTF8BuilderP Int
p (FlatSet Vector v
vec) = Bool -> Builder () -> Builder ()
T.parenWhen (Int
p forall a. Ord a => a -> a -> Bool
> Int
10) forall a b. (a -> b) -> a -> b
$ do
        Builder ()
"FlatSet{"
        forall (v :: * -> *) a.
Vec v a =>
Builder () -> (a -> Builder ()) -> v a -> Builder ()
T.intercalateVec Builder ()
T.comma forall a. Print a => a -> Builder ()
T.toUTF8Builder Vector v
vec
        Char -> Builder ()
T.char7 Char
'}'

instance Ord v => Semigroup.Semigroup (FlatSet v) where
    {-# INLINE (<>) #-}
    <> :: FlatSet v -> FlatSet v -> FlatSet v
(<>) = forall v. Ord v => FlatSet v -> FlatSet v -> FlatSet v
merge

instance Ord v => Monoid.Monoid (FlatSet v) where
    {-# INLINE mappend #-}
    mappend :: FlatSet v -> FlatSet v -> FlatSet v
mappend = forall a. Semigroup a => a -> a -> a
(<>)
    {-# INLINE mempty #-}
    mempty :: FlatSet v
mempty = forall v. FlatSet v
empty

instance NFData v => NFData (FlatSet v) where
    {-# INLINE rnf #-}
    rnf :: FlatSet v -> ()
rnf (FlatSet Vector v
vs) = forall a. NFData a => a -> ()
rnf Vector v
vs

instance (Ord v, Arbitrary v) => Arbitrary (FlatSet v) where
    arbitrary :: Gen (FlatSet v)
arbitrary = forall v. Ord v => [v] -> FlatSet v
pack forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. Arbitrary a => Gen a
arbitrary
    shrink :: FlatSet v -> [FlatSet v]
shrink FlatSet v
v = forall v. Ord v => [v] -> FlatSet v
pack forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. Arbitrary a => a -> [a]
shrink (forall a. FlatSet a -> [a]
unpack FlatSet v
v)

instance (CoArbitrary v) => CoArbitrary (FlatSet v) where
    coarbitrary :: forall b. FlatSet v -> Gen b -> Gen b
coarbitrary = forall a b. CoArbitrary a => a -> Gen b -> Gen b
coarbitrary forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. FlatSet a -> [a]
unpack

size :: FlatSet v -> Int
{-# INLINE size #-}
size :: forall a. FlatSet a -> Int
size = forall (v :: * -> *) a. Vec v a => v a -> Int
V.length forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall v. FlatSet v -> Vector v
sortedValues

null :: FlatSet v -> Bool
{-# INLINE null #-}
null :: forall a. FlatSet a -> Bool
null = forall (v :: * -> *) a. Vec v a => v a -> Bool
V.null forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall v. FlatSet v -> Vector v
sortedValues

-- | Mapping values of within a set, the result size may change if there're duplicated values
-- after mapping.
map' :: forall v. Ord v => (v -> v) -> FlatSet v -> FlatSet v
{-# INLINE map' #-}
map' :: forall v. Ord v => (v -> v) -> FlatSet v -> FlatSet v
map' v -> v
f (FlatSet Vector v
vs) = forall v. Ord v => Vector v -> FlatSet v
packVector (forall (u :: * -> *) (v :: * -> *) a b.
(Vec u a, Vec v b) =>
(a -> b) -> u a -> v b
V.map' v -> v
f Vector v
vs :: V.Vector v)

-- | /O(1)/ empty flat set.
empty :: FlatSet v
{-# NOINLINE empty #-}
empty :: forall v. FlatSet v
empty = forall v. Vector v -> FlatSet v
FlatSet forall (v :: * -> *) a. Vec v a => v a
V.empty

-- | /O(N*logN)/ Pack list of values, on duplication prefer left one.
pack :: Ord v => [v] -> FlatSet v
{-# INLINABLE pack #-}
pack :: forall v. Ord v => [v] -> FlatSet v
pack [v]
vs = forall v. Vector v -> FlatSet v
FlatSet (forall (v :: * -> *) a. Vec v a => (a -> a -> Bool) -> v a -> v a
V.mergeDupAdjacentLeft forall a. Eq a => a -> a -> Bool
(==) (forall (v :: * -> *) a. (Vec v a, Ord a) => v a -> v a
V.mergeSort (forall (v :: * -> *) a. Vec v a => [a] -> v a
V.pack [v]
vs)))

-- | /O(N*logN)/ Pack list of values with suggested size, on duplication prefer left one.
packN :: Ord v => Int -> [v] -> FlatSet v
{-# INLINABLE packN #-}
packN :: forall v. Ord v => Int -> [v] -> FlatSet v
packN Int
n [v]
vs = forall v. Vector v -> FlatSet v
FlatSet (forall (v :: * -> *) a. Vec v a => (a -> a -> Bool) -> v a -> v a
V.mergeDupAdjacentLeft forall a. Eq a => a -> a -> Bool
(==) (forall (v :: * -> *) a. (Vec v a, Ord a) => v a -> v a
V.mergeSort (forall (v :: * -> *) a. Vec v a => Int -> [a] -> v a
V.packN Int
n [v]
vs)))

-- | /O(N*logN)/ Pack list of values, on duplication prefer right one.
packR :: Ord v => [v] -> FlatSet v
{-# INLINABLE packR #-}
packR :: forall v. Ord v => [v] -> FlatSet v
packR [v]
vs = forall v. Vector v -> FlatSet v
FlatSet (forall (v :: * -> *) a. Vec v a => (a -> a -> Bool) -> v a -> v a
V.mergeDupAdjacentRight forall a. Eq a => a -> a -> Bool
(==) (forall (v :: * -> *) a. (Vec v a, Ord a) => v a -> v a
V.mergeSort (forall (v :: * -> *) a. Vec v a => [a] -> v a
V.pack [v]
vs)))

-- | /O(N*logN)/ Pack list of values with suggested size, on duplication prefer right one.
packRN :: Ord v => Int -> [v] -> FlatSet v
{-# INLINABLE packRN #-}
packRN :: forall v. Ord v => Int -> [v] -> FlatSet v
packRN Int
n [v]
vs = forall v. Vector v -> FlatSet v
FlatSet (forall (v :: * -> *) a. Vec v a => (a -> a -> Bool) -> v a -> v a
V.mergeDupAdjacentRight forall a. Eq a => a -> a -> Bool
(==) (forall (v :: * -> *) a. (Vec v a, Ord a) => v a -> v a
V.mergeSort (forall (v :: * -> *) a. Vec v a => Int -> [a] -> v a
V.packN Int
n [v]
vs)))

-- | /O(N)/ Unpack a set of values to a list s in ascending order.
--
-- This function works with @foldr/build@ fusion in base.
unpack :: FlatSet v -> [v]
{-# INLINE unpack #-}
unpack :: forall a. FlatSet a -> [a]
unpack = forall (v :: * -> *) a. Vec v a => v a -> [a]
V.unpack forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall v. FlatSet v -> Vector v
sortedValues

-- | /O(N)/ Unpack a set of values to a list s in descending order.
--
-- This function works with @foldr/build@ fusion in base.
unpackR :: FlatSet v -> [v]
{-# INLINE unpackR #-}
unpackR :: forall a. FlatSet a -> [a]
unpackR = forall (v :: * -> *) a. Vec v a => v a -> [a]
V.unpackR forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall v. FlatSet v -> Vector v
sortedValues

-- | /O(N*logN)/ Pack vector of values, on duplication prefer left one.
packVector :: Ord v => V.Vector v -> FlatSet v
{-# INLINABLE packVector #-}
packVector :: forall v. Ord v => Vector v -> FlatSet v
packVector Vector v
vs = forall v. Vector v -> FlatSet v
FlatSet (forall (v :: * -> *) a. Vec v a => (a -> a -> Bool) -> v a -> v a
V.mergeDupAdjacentLeft forall a. Eq a => a -> a -> Bool
(==) (forall (v :: * -> *) a. (Vec v a, Ord a) => v a -> v a
V.mergeSort Vector v
vs))

-- | /O(N*logN)/ Pack vector of values, on duplication prefer right one.
packVectorR :: Ord v => V.Vector v -> FlatSet v
{-# INLINABLE packVectorR #-}
packVectorR :: forall v. Ord v => Vector v -> FlatSet v
packVectorR Vector v
vs = forall v. Vector v -> FlatSet v
FlatSet (forall (v :: * -> *) a. Vec v a => (a -> a -> Bool) -> v a -> v a
V.mergeDupAdjacentRight forall a. Eq a => a -> a -> Bool
(==) (forall (v :: * -> *) a. (Vec v a, Ord a) => v a -> v a
V.mergeSort Vector v
vs))

-- | /O(logN)/ Binary search on flat set.
elem :: Ord v => v -> FlatSet v -> Bool
{-# INLINABLE elem #-}
elem :: forall v. Ord v => v -> FlatSet v -> Bool
elem v
v (FlatSet Vector v
vec) = case forall v. Ord v => Vector v -> v -> Either Int Int
binarySearch Vector v
vec v
v of Left Int
_ -> Bool
False
                                                  Either Int Int
_      -> Bool
True
-- | /O(N)/ Insert new value into set.
insert :: Ord v => v -> FlatSet v -> FlatSet v
{-# INLINABLE insert #-}
insert :: forall v. Ord v => v -> FlatSet v -> FlatSet v
insert v
v m :: FlatSet v
m@(FlatSet Vector v
vec) =
    case forall v. Ord v => Vector v -> v -> Either Int Int
binarySearch Vector v
vec v
v of
        Left Int
i -> forall v. Vector v -> FlatSet v
FlatSet (forall (v :: * -> *) a.
(Vec v a, HasCallStack) =>
v a -> Int -> a -> v a
V.unsafeInsertIndex Vector v
vec Int
i v
v)
        Right Int
_ -> FlatSet v
m

-- | /O(N)/ Delete a value from set.
delete :: Ord v => v -> FlatSet v -> FlatSet v
{-# INLINABLE delete #-}
delete :: forall v. Ord v => v -> FlatSet v -> FlatSet v
delete v
v m :: FlatSet v
m@(FlatSet Vector v
vec) =
    case forall v. Ord v => Vector v -> v -> Either Int Int
binarySearch Vector v
vec v
v of
        Left Int
_ -> FlatSet v
m
        Right Int
i -> forall v. Vector v -> FlatSet v
FlatSet (forall (v :: * -> *) a.
(Vec v a, HasCallStack) =>
v a -> Int -> v a
V.unsafeDeleteIndex Vector v
vec Int
i)

-- | /O(n+m)/ Merge two 'FlatSet', prefer right value on value duplication.
merge :: forall v . Ord v => FlatSet v -> FlatSet v -> FlatSet v
{-# INLINABLE merge #-}
merge :: forall v. Ord v => FlatSet v -> FlatSet v -> FlatSet v
merge fmL :: FlatSet v
fmL@(FlatSet (V.Vector SmallArray v
arrL Int
sL Int
lL)) fmR :: FlatSet v
fmR@(FlatSet (V.Vector SmallArray v
arrR Int
sR Int
lR))
    | forall a. FlatSet a -> Bool
null FlatSet v
fmL = FlatSet v
fmR
    | forall a. FlatSet a -> Bool
null FlatSet v
fmR = FlatSet v
fmL
    | Bool
otherwise = forall v. Vector v -> FlatSet v
FlatSet (forall (v :: * -> *) a.
(Vec v a, HasCallStack) =>
Int -> (forall s. MArr (IArray v) s a -> ST s Int) -> v a
V.createN (Int
lLforall a. Num a => a -> a -> a
+Int
lR) (forall s. Int -> Int -> Int -> SmallMutableArray s v -> ST s Int
go Int
sL Int
sR Int
0))
  where
    endL :: Int
endL = Int
sL forall a. Num a => a -> a -> a
+ Int
lL
    endR :: Int
endR = Int
sR forall a. Num a => a -> a -> a
+ Int
lR
    go :: Int -> Int -> Int -> A.SmallMutableArray s v -> ST s Int
    go :: forall s. Int -> Int -> Int -> SmallMutableArray s v -> ST s Int
go !Int
i !Int
j !Int
k SmallMutableArray s v
marr
        | Int
i forall a. Ord a => a -> a -> Bool
>= Int
endL = do
            forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a
-> Int -> SmallArray a -> Int -> Int -> m ()
A.copySmallArray SmallMutableArray s v
marr Int
k SmallArray v
arrR Int
j (Int
lRforall a. Num a => a -> a -> a
-Int
j)
            forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$! Int
kforall a. Num a => a -> a -> a
+Int
lRforall a. Num a => a -> a -> a
-Int
j
        | Int
j forall a. Ord a => a -> a -> Bool
>= Int
endR = do
            forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a
-> Int -> SmallArray a -> Int -> Int -> m ()
A.copySmallArray SmallMutableArray s v
marr Int
k SmallArray v
arrL Int
i (Int
lLforall a. Num a => a -> a -> a
-Int
i)
            forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$! Int
kforall a. Num a => a -> a -> a
+Int
lLforall a. Num a => a -> a -> a
-Int
i
        | Bool
otherwise = do
            v
vL <- SmallArray v
arrL forall (m :: * -> *) a. Monad m => SmallArray a -> Int -> m a
`A.indexSmallArrayM` Int
i
            v
vR <- SmallArray v
arrR forall (m :: * -> *) a. Monad m => SmallArray a -> Int -> m a
`A.indexSmallArrayM` Int
j
            case v
vL forall a. Ord a => a -> a -> Ordering
`compare` v
vR of Ordering
LT -> do forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
A.writeSmallArray SmallMutableArray s v
marr Int
k v
vL
                                             forall s. Int -> Int -> Int -> SmallMutableArray s v -> ST s Int
go (Int
iforall a. Num a => a -> a -> a
+Int
1) Int
j (Int
kforall a. Num a => a -> a -> a
+Int
1) SmallMutableArray s v
marr
                                    Ordering
EQ -> do forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
A.writeSmallArray SmallMutableArray s v
marr Int
k v
vR
                                             forall s. Int -> Int -> Int -> SmallMutableArray s v -> ST s Int
go (Int
iforall a. Num a => a -> a -> a
+Int
1) (Int
jforall a. Num a => a -> a -> a
+Int
1) (Int
kforall a. Num a => a -> a -> a
+Int
1) SmallMutableArray s v
marr
                                    Ordering
_  -> do forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
A.writeSmallArray SmallMutableArray s v
marr Int
k v
vR
                                             forall s. Int -> Int -> Int -> SmallMutableArray s v -> ST s Int
go Int
i (Int
jforall a. Num a => a -> a -> a
+Int
1) (Int
kforall a. Num a => a -> a -> a
+Int
1) SmallMutableArray s v
marr

--------------------------------------------------------------------------------

-- | Find the value's index in the vector, if value exists return 'Right',
-- otherwise 'Left', i.e. the insert index
--
-- This function only works on ascending sorted vectors.
binarySearch :: Ord v => V.Vector v -> v -> Either Int Int
{-# INLINABLE binarySearch #-}
binarySearch :: forall v. Ord v => Vector v -> v -> Either Int Int
binarySearch (V.Vector SmallArray v
_ Int
_ Int
0) v
_   = forall a b. a -> Either a b
Left Int
0
binarySearch (V.Vector SmallArray v
arr Int
s0 Int
l) !v
v' = Int -> Int -> Either Int Int
go Int
s0 (Int
s0forall a. Num a => a -> a -> a
+Int
lforall a. Num a => a -> a -> a
-Int
1)
  where
    go :: Int -> Int -> Either Int Int
go !Int
s !Int
e
        | Int
s forall a. Eq a => a -> a -> Bool
== Int
e =
            let v :: v
v = SmallArray v
arr forall a. SmallArray a -> Int -> a
`A.indexSmallArray` Int
s
            in case v
v' forall a. Ord a => a -> a -> Ordering
`compare` v
v of Ordering
LT -> forall a b. a -> Either a b
Left (Int
sforall a. Num a => a -> a -> a
-Int
s0)
                                      Ordering
GT -> let !s' :: Int
s' = Int
sforall a. Num a => a -> a -> a
+Int
1 in forall a b. a -> Either a b
Left (Int
s'forall a. Num a => a -> a -> a
-Int
s0)
                                      Ordering
_  -> forall a b. b -> Either a b
Right (Int
sforall a. Num a => a -> a -> a
-Int
s0)
        | Int
s forall a. Ord a => a -> a -> Bool
>  Int
e = forall a b. a -> Either a b
Left Int
s
        | Bool
otherwise =
            let !mid :: Int
mid = (Int
sforall a. Num a => a -> a -> a
+Int
e) forall a. Bits a => a -> Int -> a
`unsafeShiftR` Int
1
                v :: v
v = SmallArray v
arr forall a. SmallArray a -> Int -> a
`A.indexSmallArray` Int
mid
            in case v
v' forall a. Ord a => a -> a -> Ordering
`compare` v
v of Ordering
LT -> Int -> Int -> Either Int Int
go Int
s (Int
midforall a. Num a => a -> a -> a
-Int
1)
                                      Ordering
GT -> Int -> Int -> Either Int Int
go (Int
midforall a. Num a => a -> a -> a
+Int
1) Int
e
                                      Ordering
_  -> forall a b. b -> Either a b
Right (Int
midforall a. Num a => a -> a -> a
-Int
s0)