content: {-# LANGUAGE BangPatterns #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# LANGUAGE MonoLocalBinds #-}
  
  -- ---------------------------------------------------------------------------
  -- |
  -- Module      : Data.Vector.Algorithms.Search
  -- Copyright   : (c) 2009-2010 Dan Doel
  -- 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
         , Comparison
         ) where
  
  import Prelude hiding (read, length)
  
  import Control.Monad.Primitive
  
  import Data.Bits
  
  import Data.Vector.Generic.Mutable
  
  import Data.Vector.Algorithms.Common (shiftRI, Comparison)
  
  ------------------------------------------------------------------------------------
  -- LIQUID API Specifications -------------------------------------------------------
  ------------------------------------------------------------------------------------
  
  {-@ binarySearch, binarySearchL, binarySearchR 
        :: (PrimMonad m, MVector v e, Ord e) 
        => vec:(v (PrimState m) e) -> e -> m (AOkIdx vec)
    @-}
  
  {-@ binarySearchBy, binarySearchLBy, binarySearchRBy 
        :: (PrimMonad m, MVector v e) 
        => Comparison e -> vec:(v (PrimState m) e) -> e -> m (AOkIdx vec) 
    @-}
  
  {-@ binarySearchByBounds, binarySearchLByBounds, binarySearchRByBounds 
        :: (PrimMonad m, MVector v e)
        => Comparison e -> vec:(v (PrimState m) e) -> e 
        -> l:{v:Nat | v <= (vsize vec)}
        -> u:{v:Nat | (l <= v && v <= (vsize vec))}
        -> m {v:Int | (l <= v && v <= u)}
    @-}
  
  
  -------------------------------------------------------------------------------------
  
  -- | 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 = binarySearchBy 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 cmp veck e = binarySearchByBounds cmp veck e 0 (length veck)
  {-# 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 cmp vec e lo hi = loop (hi - lo) lo hi  
   where
    {- LIQUID WITNESS -}
   loop (twit :: Int) !l !u
     | u <= l    = return l
     | otherwise = do e' <- unsafeRead vec k
                      case cmp e' e of
                        LT -> loop (u - (k+1)) (k+1) u
                        EQ -> return  k
                        GT -> loop (k -l)      l     k
    where k = (u + l) `shiftRI` 1
  {-# 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 = binarySearchLBy 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 cmp vec e = binarySearchLByBounds cmp vec e 0 (length 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 cmp vec e lo hi = loop (hi - lo) lo hi
   where
    {- LIQUID WITNESS -}
   loop (twit :: Int) !l !u
     | u <= l    = return l
     | otherwise = do e' <- unsafeRead vec k
                      case cmp e' e of
                        LT -> loop (u - (k+1)) (k+1) u
                        _  -> loop (k - l)     l     k
    where k = (u + l) `shiftRI` 1
  {-# 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 = binarySearchRBy 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 cmp vec e = binarySearchRByBounds cmp vec e 0 (length 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 cmp vec e lo hi = loop (hi - lo) lo hi
   where
    {- LIQUID WITNESS -}
   loop (twit :: Int) !l !u
     | u <= l    = return l
     | otherwise = do e' <- unsafeRead vec k
                      case cmp e' e of
                        GT -> loop (k - l)      l     k
                        _  -> loop (u - (k+1))  (k+1) u
    where k = (u + l) `shiftRI` 1
  {-# INLINE binarySearchRByBounds #-}
  