{-# LANGUAGE CPP #-}
module Map_FindGE where

import Data.Map.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 :: Ord k => k -> Map k a -> Maybe (k,a)
findGreaterEqual1 k m =
    case splitLookup k m of
        (_,Just v,_)  -> Just (k,v)
        (_,Nothing,r) -> findMinMaybe r 

findGreaterEqual2 :: Ord k => k -> Map k a -> Maybe (k,a)
findGreaterEqual2 = go
  where
    STRICT_1_OF_2(go)
    go _ Tip = Nothing
    go k (Bin _ kx x l r) =
        case compare k kx of
            LT -> case go k l of
                    Nothing -> Just (kx,x)
                    ret -> ret
            GT -> go k r
            EQ -> Just (kx,x)

findGreaterEqual3 :: Ord k => k -> Map k a -> Maybe (k,a)
findGreaterEqual3 = go Nothing
  where
    STRICT_2_OF_3(go)
    go def _ Tip = def
    go def k (Bin _ kx x l r) =
        case compare k kx of
            LT -> go (Just (kx,x)) k l
            GT -> go def k r
            EQ -> Just (kx,x)

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

findGreater1 :: Ord k => k -> Map k a -> Maybe (k,a)
findGreater1 k m = findMinMaybe (snd (split k m)) 

findGreater3 :: Ord k => k -> Map k a -> Maybe (k,a)
findGreater3 = go Nothing
  where
    STRICT_2_OF_3(go)
    go def _ Tip = def
    go def k (Bin _ kx x l r) =
        case compare k kx of
            LT -> go (Just (kx,x)) k l
            GT -> go def k r
            EQ -> def

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

findMinMaybe :: Map k a -> Maybe (k,a)
findMinMaybe (Bin _ kx x Tip _)  = Just (kx,x)
findMinMaybe (Bin _ _  _ l _)    = findMinMaybe l
findMinMaybe Tip                 = Nothing

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

type IMap = Map Int 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_findGreater13 :: Int -> IMap -> Bool
prop_findGreater13 k t = findGreater1 k t == findGreater3 k t

-- copy/paste from map-properties.hs
{--------------------------------------------------------------------
  Arbitrary, reasonably balanced trees
--------------------------------------------------------------------}
instance (Enum k,Arbitrary a) => Arbitrary (Map k a) where
  arbitrary = sized (arbtree 0 maxkey)
    where maxkey = 10^5

          arbtree :: (Enum k, Arbitrary a) => Int -> Int -> Int -> Gen (Map k a)
          arbtree lo hi n = do t <- gentree lo hi n
                               if balanced t then return t else arbtree lo hi n
            where gentree lo hi n
                    | n <= 0        = return Tip
                    | lo >= hi      = return Tip
                    | otherwise     = do{ x  <- arbitrary
                                        ; i  <- choose (lo,hi)
                                        ; m  <- choose (1,70)
                                        ; let (ml,mr)  | m==(1::Int)= (1,2)
                                                       | m==2       = (2,1)
                                                       | m==3       = (1,1)
                                                       | otherwise  = (2,2)
                                        ; l  <- gentree lo (i-1) (n `div` ml)
                                        ; r  <- gentree (i+1) hi (n `div` mr)
                                        ; return (bin (toEnum i) x l r)
                                        }