{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE TypeFamilies #-}

-- ---------------------------------------------------------------------------
-- |
-- Module      : Data.Vector.Algorithms.Search
-- Copyright   : (c) 2009-2015 Dan Doel, 2015 Tim Baumann
-- Maintainer  : Dan Doel <dan.doel@gmail.com>
-- Stability   : Experimental
-- Portability : Non-portable (bang patterns)
--
-- This module implements several methods of searching for indicies to insert
-- elements into a sorted vector.

module Data.Vector.Algorithms.Search
       ( binarySearch
       , binarySearchBy
       , binarySearchByBounds
       , binarySearchL
       , binarySearchLBy
       , binarySearchLByBounds
       , binarySearchR
       , binarySearchRBy
       , binarySearchRByBounds
       , binarySearchP
       , binarySearchPBounds
       , gallopingSearchLeftP
       , gallopingSearchLeftPBounds
       , gallopingSearchRightP
       , gallopingSearchRightPBounds
       , Comparison
       ) where

import Prelude hiding (read, length)

import Control.Monad.Primitive

import Data.Bits

import Data.Vector.Generic.Mutable

import Data.Vector.Algorithms.Common (Comparison, midPoint)

-- | Finds an index in a given sorted vector at which the given element could
-- be inserted while maintaining the sortedness of the vector.
binarySearch :: (PrimMonad m, MVector v e, Ord e)
             => v (PrimState m) e -> e -> m Int
binarySearch :: forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e, Ord e) =>
v (PrimState m) e -> e -> m Int
binarySearch = Comparison e -> v (PrimState m) e -> e -> m Int
forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e -> v (PrimState m) e -> e -> m Int
binarySearchBy Comparison e
forall a. Ord a => a -> a -> Ordering
compare
{-# INLINE binarySearch #-}

-- | Finds an index in a given vector, which must be sorted with respect to the
-- given comparison function, at which the given element could be inserted while
-- preserving the vector's sortedness.
binarySearchBy :: (PrimMonad m, MVector v e)
               => Comparison e -> v (PrimState m) e -> e -> m Int
binarySearchBy :: forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e -> v (PrimState m) e -> e -> m Int
binarySearchBy Comparison e
cmp v (PrimState m) e
vec e
e = Comparison e -> v (PrimState m) e -> e -> Int -> Int -> m Int
forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e -> v (PrimState m) e -> e -> Int -> Int -> m Int
binarySearchByBounds Comparison e
cmp v (PrimState m) e
vec e
e Int
0 (v (PrimState m) e -> Int
forall (v :: * -> * -> *) a s. MVector v a => v s a -> Int
length v (PrimState m) e
vec)
{-# INLINE binarySearchBy #-}

-- | Given a vector sorted with respect to a given comparison function in indices
-- in [l,u), finds an index in [l,u] at which the given element could be inserted
-- while preserving sortedness.
binarySearchByBounds :: (PrimMonad m, MVector v e)
                     => Comparison e -> v (PrimState m) e -> e -> Int -> Int -> m Int
binarySearchByBounds :: forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e -> v (PrimState m) e -> e -> Int -> Int -> m Int
binarySearchByBounds Comparison e
cmp v (PrimState m) e
vec e
e = Int -> Int -> m Int
loop
 where
 loop :: Int -> Int -> m Int
loop !Int
l !Int
u
   | Int
u Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
l    = Int -> m Int
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Int
l
   | Bool
otherwise = do e
e' <- v (PrimState m) e -> Int -> m e
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
unsafeRead v (PrimState m) e
vec Int
k
                    case Comparison e
cmp e
e' e
e of
                      Ordering
LT -> Int -> Int -> m Int
loop (Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) Int
u
                      Ordering
EQ -> Int -> m Int
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Int
k
                      Ordering
GT -> Int -> Int -> m Int
loop Int
l     Int
k
  where k :: Int
k = Int -> Int -> Int
midPoint Int
u Int
l
{-# INLINE binarySearchByBounds #-}

-- | Finds the lowest index in a given sorted vector at which the given element
-- could be inserted while maintaining the sortedness.
binarySearchL :: (PrimMonad m, MVector v e, Ord e) => v (PrimState m) e -> e -> m Int
binarySearchL :: forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e, Ord e) =>
v (PrimState m) e -> e -> m Int
binarySearchL = Comparison e -> v (PrimState m) e -> e -> m Int
forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e -> v (PrimState m) e -> e -> m Int
binarySearchLBy Comparison e
forall a. Ord a => a -> a -> Ordering
compare
{-# INLINE binarySearchL #-}

-- | Finds the lowest index in a given vector, which must be sorted with respect to
-- the given comparison function, at which the given element could be inserted
-- while preserving the sortedness.
binarySearchLBy :: (PrimMonad m, MVector v e)
                => Comparison e -> v (PrimState m) e -> e -> m Int
binarySearchLBy :: forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e -> v (PrimState m) e -> e -> m Int
binarySearchLBy Comparison e
cmp v (PrimState m) e
vec e
e = Comparison e -> v (PrimState m) e -> e -> Int -> Int -> m Int
forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e -> v (PrimState m) e -> e -> Int -> Int -> m Int
binarySearchLByBounds Comparison e
cmp v (PrimState m) e
vec e
e Int
0 (v (PrimState m) e -> Int
forall (v :: * -> * -> *) a s. MVector v a => v s a -> Int
length v (PrimState m) e
vec)
{-# INLINE binarySearchLBy #-}

-- | Given a vector sorted with respect to a given comparison function on indices
-- in [l,u), finds the lowest index in [l,u] at which the given element could be
-- inserted while preserving sortedness.
binarySearchLByBounds :: (PrimMonad m, MVector v e)
                      => Comparison e -> v (PrimState m) e -> e -> Int -> Int -> m Int
binarySearchLByBounds :: forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e -> v (PrimState m) e -> e -> Int -> Int -> m Int
binarySearchLByBounds Comparison e
cmp v (PrimState m) e
vec e
e = (e -> Bool) -> v (PrimState m) e -> Int -> Int -> m Int
forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
(e -> Bool) -> v (PrimState m) e -> Int -> Int -> m Int
binarySearchPBounds e -> Bool
p v (PrimState m) e
vec
 where p :: e -> Bool
p e
e' = case Comparison e
cmp e
e' e
e of Ordering
LT -> Bool
False ; Ordering
_ -> Bool
True
{-# INLINE binarySearchLByBounds #-}

-- | Finds the greatest index in a given sorted vector at which the given element
-- could be inserted while maintaining sortedness.
binarySearchR :: (PrimMonad m, MVector v e, Ord e) => v (PrimState m) e -> e -> m Int
binarySearchR :: forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e, Ord e) =>
v (PrimState m) e -> e -> m Int
binarySearchR = Comparison e -> v (PrimState m) e -> e -> m Int
forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e -> v (PrimState m) e -> e -> m Int
binarySearchRBy Comparison e
forall a. Ord a => a -> a -> Ordering
compare
{-# INLINE binarySearchR #-}

-- | Finds the greatest index in a given vector, which must be sorted with respect to
-- the given comparison function, at which the given element could be inserted
-- while preserving the sortedness.
binarySearchRBy :: (PrimMonad m, MVector v e)
                => Comparison e -> v (PrimState m) e -> e -> m Int
binarySearchRBy :: forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e -> v (PrimState m) e -> e -> m Int
binarySearchRBy Comparison e
cmp v (PrimState m) e
vec e
e = Comparison e -> v (PrimState m) e -> e -> Int -> Int -> m Int
forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e -> v (PrimState m) e -> e -> Int -> Int -> m Int
binarySearchRByBounds Comparison e
cmp v (PrimState m) e
vec e
e Int
0 (v (PrimState m) e -> Int
forall (v :: * -> * -> *) a s. MVector v a => v s a -> Int
length v (PrimState m) e
vec)
{-# INLINE binarySearchRBy #-}

-- | Given a vector sorted with respect to the given comparison function on indices
-- in [l,u), finds the greatest index in [l,u] at which the given element could be
-- inserted while preserving sortedness.
binarySearchRByBounds :: (PrimMonad m, MVector v e)
                      => Comparison e -> v (PrimState m) e -> e -> Int -> Int -> m Int
binarySearchRByBounds :: forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e -> v (PrimState m) e -> e -> Int -> Int -> m Int
binarySearchRByBounds Comparison e
cmp v (PrimState m) e
vec e
e = (e -> Bool) -> v (PrimState m) e -> Int -> Int -> m Int
forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
(e -> Bool) -> v (PrimState m) e -> Int -> Int -> m Int
binarySearchPBounds e -> Bool
p v (PrimState m) e
vec
 where p :: e -> Bool
p e
e' = case Comparison e
cmp e
e' e
e of Ordering
GT -> Bool
True ; Ordering
_ -> Bool
False
{-# INLINE binarySearchRByBounds #-}

-- | Given a predicate that is guaranteed to be monotone on the given vector,
-- finds the first index at which the predicate returns True, or the length of
-- the array if the predicate is false for the entire array.
binarySearchP :: (PrimMonad m, MVector v e) => (e -> Bool) -> v (PrimState m) e -> m Int
binarySearchP :: forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
(e -> Bool) -> v (PrimState m) e -> m Int
binarySearchP e -> Bool
p v (PrimState m) e
vec = (e -> Bool) -> v (PrimState m) e -> Int -> Int -> m Int
forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
(e -> Bool) -> v (PrimState m) e -> Int -> Int -> m Int
binarySearchPBounds e -> Bool
p v (PrimState m) e
vec Int
0 (v (PrimState m) e -> Int
forall (v :: * -> * -> *) a s. MVector v a => v s a -> Int
length v (PrimState m) e
vec)
{-# INLINE binarySearchP #-}

-- | Given a predicate that is guaranteed to be monotone on the indices [l,u) in
-- a given vector, finds the index in [l,u] at which the predicate turns from
-- False to True (yielding u if the entire interval is False).
binarySearchPBounds :: (PrimMonad m, MVector v e)
                    => (e -> Bool) -> v (PrimState m) e -> Int -> Int -> m Int
binarySearchPBounds :: forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
(e -> Bool) -> v (PrimState m) e -> Int -> Int -> m Int
binarySearchPBounds e -> Bool
p v (PrimState m) e
vec = Int -> Int -> m Int
loop
 where
 loop :: Int -> Int -> m Int
loop !Int
l !Int
u
   | Int
u Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
l    = Int -> m Int
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Int
l
   | Bool
otherwise = v (PrimState m) e -> Int -> m e
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
unsafeRead v (PrimState m) e
vec Int
k m e -> (e -> m Int) -> m Int
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \e
e -> if e -> Bool
p e
e then Int -> Int -> m Int
loop Int
l Int
k else Int -> Int -> m Int
loop (Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) Int
u
  where k :: Int
k = Int -> Int -> Int
midPoint Int
u Int
l
{-# INLINE binarySearchPBounds #-}

-- | Given a predicate that is guaranteed to be monotone on the vector elements
-- in order, finds the index at which the predicate turns from False to True.
-- The length of the vector is returned if the predicate is False for the entire
-- vector.
--
-- Begins searching at the start of the vector, in increasing steps of size 2^n.
gallopingSearchLeftP
  :: (PrimMonad m, MVector v e) => (e -> Bool) -> v (PrimState m) e -> m Int
gallopingSearchLeftP :: forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
(e -> Bool) -> v (PrimState m) e -> m Int
gallopingSearchLeftP e -> Bool
p v (PrimState m) e
vec = (e -> Bool) -> v (PrimState m) e -> Int -> Int -> m Int
forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
(e -> Bool) -> v (PrimState m) e -> Int -> Int -> m Int
gallopingSearchLeftPBounds e -> Bool
p v (PrimState m) e
vec Int
0 (v (PrimState m) e -> Int
forall (v :: * -> * -> *) a s. MVector v a => v s a -> Int
length v (PrimState m) e
vec)
{-# INLINE gallopingSearchLeftP #-}

-- | Given a predicate that is guaranteed to be monotone on the vector elements
-- in order, finds the index at which the predicate turns from False to True.
-- The length of the vector is returned if the predicate is False for the entire
-- vector.
--
-- Begins searching at the end of the vector, in increasing steps of size 2^n.
gallopingSearchRightP
  :: (PrimMonad m, MVector v e) => (e -> Bool) -> v (PrimState m) e -> m Int
gallopingSearchRightP :: forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
(e -> Bool) -> v (PrimState m) e -> m Int
gallopingSearchRightP e -> Bool
p v (PrimState m) e
vec = (e -> Bool) -> v (PrimState m) e -> Int -> Int -> m Int
forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
(e -> Bool) -> v (PrimState m) e -> Int -> Int -> m Int
gallopingSearchRightPBounds e -> Bool
p v (PrimState m) e
vec Int
0 (v (PrimState m) e -> Int
forall (v :: * -> * -> *) a s. MVector v a => v s a -> Int
length v (PrimState m) e
vec)
{-# INLINE gallopingSearchRightP #-}

-- | Given a predicate that is guaranteed to be monotone on the indices [l,u) in
-- a given vector, finds the index in [l,u] at which the predicate turns from
-- False to True (yielding u if the entire interval is False).
-- Begins searching at l, going right in increasing (2^n)-steps.
gallopingSearchLeftPBounds :: (PrimMonad m, MVector v e)
                           => (e -> Bool)
                           -> v (PrimState m) e
                           -> Int -- ^ l
                           -> Int -- ^ u
                           -> m Int
gallopingSearchLeftPBounds :: forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
(e -> Bool) -> v (PrimState m) e -> Int -> Int -> m Int
gallopingSearchLeftPBounds e -> Bool
p v (PrimState m) e
vec Int
l Int
u
  | Int
u Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
l    = Int -> m Int
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Int
l
  | Bool
otherwise = do e
x <- v (PrimState m) e -> Int -> m e
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
unsafeRead v (PrimState m) e
vec Int
l
                   if e -> Bool
p e
x then Int -> m Int
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Int
l else Int -> Int -> Int -> m Int
iter (Int
lInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) Int
l Int
2
 where
 binSearch :: Int -> Int -> m Int
binSearch = (e -> Bool) -> v (PrimState m) e -> Int -> Int -> m Int
forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
(e -> Bool) -> v (PrimState m) e -> Int -> Int -> m Int
binarySearchPBounds e -> Bool
p v (PrimState m) e
vec
 iter :: Int -> Int -> Int -> m Int
iter !Int
i !Int
j !Int
_stepSize | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
u Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1 = do
   e
x <- v (PrimState m) e -> Int -> m e
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
unsafeRead v (PrimState m) e
vec (Int
uInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)
   if e -> Bool
p e
x then Int -> Int -> m Int
binSearch (Int
jInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) (Int
uInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) else Int -> m Int
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Int
u
 iter !Int
i !Int
j !Int
stepSize = do
   e
x <- v (PrimState m) e -> Int -> m e
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
unsafeRead v (PrimState m) e
vec Int
i
   if e -> Bool
p e
x then Int -> Int -> m Int
binSearch (Int
jInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) Int
i else Int -> Int -> Int -> m Int
iter (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
stepSize) Int
i (Int
2Int -> Int -> Int
forall a. Num a => a -> a -> a
*Int
stepSize)
{-# INLINE gallopingSearchLeftPBounds #-}

-- | Given a predicate that is guaranteed to be monotone on the indices [l,u) in
-- a given vector, finds the index in [l,u] at which the predicate turns from
-- False to True (yielding u if the entire interval is False).
-- Begins searching at u, going left in increasing (2^n)-steps.
gallopingSearchRightPBounds :: (PrimMonad m, MVector v e)
                            => (e -> Bool)
                            -> v (PrimState m) e
                            -> Int -- ^ l
                            -> Int -- ^ u
                            -> m Int
gallopingSearchRightPBounds :: forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
(e -> Bool) -> v (PrimState m) e -> Int -> Int -> m Int
gallopingSearchRightPBounds e -> Bool
p v (PrimState m) e
vec Int
l Int
u
  | Int
u Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
l    = Int -> m Int
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Int
l
  | Bool
otherwise = Int -> Int -> Int -> m Int
iter (Int
uInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) (Int
uInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) (-Int
1)
 where
 binSearch :: Int -> Int -> m Int
binSearch = (e -> Bool) -> v (PrimState m) e -> Int -> Int -> m Int
forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
(e -> Bool) -> v (PrimState m) e -> Int -> Int -> m Int
binarySearchPBounds e -> Bool
p v (PrimState m) e
vec
 iter :: Int -> Int -> Int -> m Int
iter !Int
i !Int
j !Int
_stepSize | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
l = do
   e
x <- v (PrimState m) e -> Int -> m e
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
unsafeRead v (PrimState m) e
vec Int
l
   if e -> Bool
p e
x then Int -> m Int
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Int
l else Int -> Int -> m Int
binSearch (Int
lInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) Int
j
 iter !Int
i !Int
j !Int
stepSize = do
   e
x <- v (PrimState m) e -> Int -> m e
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
unsafeRead v (PrimState m) e
vec Int
i
   if e -> Bool
p e
x then Int -> Int -> Int -> m Int
iter (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
stepSize) Int
i (Int
2Int -> Int -> Int
forall a. Num a => a -> a -> a
*Int
stepSize) else Int -> Int -> m Int
binSearch (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) Int
j
{-# INLINE gallopingSearchRightPBounds #-}