    {-# 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' <- 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 cmp 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 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 p x then return l else iter (l+1) l 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
       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 p x then binSearch (j+1) (u-1) else return u
     iter !Int
    i !Int
    j !Int
    stepSize = do
       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 p x then binSearch (j+1) i else iter (i+stepSize) i (2*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
       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 p x then return l else binSearch (l+1) j
     iter !Int
    i !Int
    j !Int
    stepSize = do
       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 p x then iter (i+stepSize) i (2*stepSize) else binSearch (i+1) j
    {-# INLINE gallopingSearchRightPBounds #-}