1 {-# LANGUAGE BangPatterns, CPP #-}
2 module LookupGE_Map where
4 import Data.Map.Internal
6 lookupGE1 :: Ord k => k -> Map k a -> Maybe (k,a)
7 lookupGE1 k m =
8 case splitLookup k m of
9 (_,Just v,_) -> Just (k,v)
10 (_,Nothing,r) -> findMinMaybe r
11 {-# INLINABLE lookupGE1 #-}
13 lookupGE2 :: Ord k => k -> Map k a -> Maybe (k,a)
14 lookupGE2 = go
15 where
16 go !_ Tip = Nothing
17 go !k (Bin _ kx x l r) =
18 case compare k kx of
19 LT -> case go k l of
20 Nothing -> Just (kx,x)
21 ret -> ret
22 GT -> go k r
23 EQ -> Just (kx,x)
24 {-# INLINABLE lookupGE2 #-}
26 lookupGE3 :: Ord k => k -> Map k a -> Maybe (k,a)
27 lookupGE3 = go Nothing
28 where
29 go def !_ Tip = def
30 go def !k (Bin _ kx x l r) =
31 case compare k kx of
32 LT -> go (Just (kx,x)) k l
33 GT -> go def k r
34 EQ -> Just (kx,x)
35 {-# INLINABLE lookupGE3 #-}
37 lookupGE4 :: Ord k => k -> Map k a -> Maybe (k,a)
38 lookupGE4 k = k `seq` goNothing
39 where
40 goNothing Tip = Nothing
41 goNothing (Bin _ kx x l r) = case compare k kx of
42 LT -> goJust kx x l
43 EQ -> Just (kx, x)
44 GT -> goNothing r
46 goJust ky y Tip = Just (ky, y)
47 goJust ky y (Bin _ kx x l r) = case compare k kx of
48 LT -> goJust kx x l
49 EQ -> Just (kx, x)
50 GT -> goJust ky y r
51 {-# INLINABLE lookupGE4 #-}
53 -------------------------------------------------------------------------------
54 -- Utilities
55 -------------------------------------------------------------------------------
57 findMinMaybe :: Map k a -> Maybe (k,a)
58 findMinMaybe (Bin _ kx x Tip _) = Just (kx,x)
59 findMinMaybe (Bin _ _ _ l _) = findMinMaybe l
60 findMinMaybe Tip = Nothing
62 #ifdef TESTING
63 -------------------------------------------------------------------------------
64 -- Properties:
65 -------------------------------------------------------------------------------
67 prop_lookupGE12 :: Int -> [Int] -> Bool
68 prop_lookupGE12 x xs = case fromList \$ zip xs xs of m -> lookupGE1 x m == lookupGE2 x m
70 prop_lookupGE13 :: Int -> [Int] -> Bool
71 prop_lookupGE13 x xs = case fromList \$ zip xs xs of m -> lookupGE1 x m == lookupGE3 x m
73 prop_lookupGE14 :: Int -> [Int] -> Bool
74 prop_lookupGE14 x xs = case fromList \$ zip xs xs of m -> lookupGE1 x m == lookupGE4 x m
75 #endif