{-# LANGUAGE CPP #-}
module IntMap_FindGE where

import Prelude hiding (null)
import Data.IntMap.Base
import Test.QuickCheck

#define STRICT_1_OF_2(fn) fn arg _ | arg `seq` False = undefined
#define STRICT_1_OF_3(fn) fn arg _ _ | arg `seq` False = undefined
#define STRICT_2_OF_3(fn) fn _ arg _ | arg `seq` False = undefined
#define STRICT_2_OF_4(fn) fn _ arg _ _ | arg `seq` False = undefined

-------------------------------------------------------------------------------
-- findGreaterEqual variants
-------------------------------------------------------------------------------

findGreaterEqual1 :: Key -> IntMap a -> Maybe (Key,a)
findGreaterEqual1 k m =
    case splitLookup k m of
        (_,Just v,_)  -> Just (k,v)
        (_,Nothing,r) -> findMinMaybe r 

findGreaterEqual2 :: Key -> IntMap a -> Maybe (Key,a)
findGreaterEqual2 k t = case t of
    Bin _ m l r | m < 0 -> if k >= 0
      then go l
      else case go r of
        Nothing -> Just $ findMin l
        justx -> justx
    _ -> go t
  where
    go (Bin p m l r)
      | nomatch k p m = if k < p
        then Just $ findMin l
        else Nothing
      | zero k m = case go l of
        Nothing -> Just $ findMin r
        justx -> justx
      | otherwise = go r
    go (Tip ky y)
      | k > ky = Nothing
      | otherwise = Just (ky, y)
    go Nil = Nothing

findGreaterEqual3 :: Key -> IntMap a -> Maybe (Key,a)
findGreaterEqual3 k t = k `seq` case t of
    Bin _ m l r | m < 0 -> if k >= 0
      then go Nothing l
      else go (Just (findMin l)) r
    _ -> go Nothing t
  where
    go def (Bin p m l r)
      | nomatch k p m = if k < p then Just $ findMin l else def
      | zero k m  = go (Just $ findMin r) l
      | otherwise = go def r
    go def (Tip ky y)
      | k > ky    = def
      | otherwise = Just (ky, y)
    go def Nil  = def

-------------------------------------------------------------------------------
-- findGreater variants
-------------------------------------------------------------------------------

findGreater1 :: Key -> IntMap a -> Maybe (Key,a)
findGreater1 k m = findMinMaybe (snd (split k m)) 

findGreater3 :: Key -> IntMap a -> Maybe (Key,a)
findGreater3 k t = k `seq` case t of
    Bin _ m l r | m < 0 -> if k >= 0
      then go Nothing l
      else go (Just (findMin l)) r
    _ -> go Nothing t
  where
    go def (Bin p m l r)
      | nomatch k p m = if k < p then Just $ findMin l else def
      | zero k m  = go (Just $ findMin r) l
      | otherwise = go def r
    go def (Tip ky y)
      | ky > k    = Just (ky, y)
      | otherwise = def
    go def Nil  = def


-------------------------------------------------------------------------------
-- Utilities
-------------------------------------------------------------------------------

-- | /O(log n)/. The minimal key of the map.
findMinMaybe :: IntMap a -> Maybe (Key, a)
findMinMaybe m
  | null m = Nothing
  | otherwise = Just (findMin m)

-------------------------------------------------------------------------------
-- Properties:
-------------------------------------------------------------------------------

type IMap = IntMap Int

prop_findGreaterEqual12 :: Int -> IMap -> Bool
prop_findGreaterEqual12 k t = findGreaterEqual1 k t == findGreaterEqual2 k t

prop_findGreaterEqual13 :: Int -> IMap -> Bool
prop_findGreaterEqual13 k t = findGreaterEqual1 k t == findGreaterEqual3 k t

--prop_findGreaterEqual14 :: Int -> IMap -> Bool
--prop_findGreaterEqual14 k t = findGreaterEqual1 k t == findGreaterEqual4 k t

prop_findGreater13 :: Int -> IMap -> Bool
prop_findGreater13 k t = findGreater1 k t == findGreater3 k t

--prop_findGreater14 :: Int -> IMap -> Bool
--prop_findGreater14 k t = findGreater1 k t == findGreater4 k t


-- copy/paste from intmap-properties.hs
instance Arbitrary a => Arbitrary (IntMap a) where
  arbitrary = do{ ks <- arbitrary
                ; xs <- arbitrary
                ; return (fromList (zip xs ks))
                }