Add lookup{LT,LE,GT,GE}, find{Min,Max}

src/Data/Map.purs

@@ -10,6 +10,12 @@ module Data.Map
   , checkValid
   , insert
   , lookup
+  , lookupLE
+  , lookupLT
+  , lookupGE
+  , lookupGT
+  , findMin
+  , findMax
   , fromFoldable
   , fromFoldableWith
   , toList
@@ -32,7 +38,7 @@ import Prelude
 
 import Data.Foldable (foldl, foldMap, foldr, class Foldable)
 import Data.List (List(..), length, nub)
-import Data.Maybe (Maybe(..), maybe, isJust)
+import Data.Maybe (Maybe(..), maybe, isJust, fromMaybe)
 import Data.Monoid (class Monoid)
 import Data.Traversable (traverse, class Traversable)
 import Data.Tuple (Tuple(..), uncurry, snd)
@@ -139,6 +145,67 @@ lookup = unsafePartial \k tree ->
                 _ , GT -> lookup k right
                 _ , _  -> lookup k mid
 
+
+-- | Lookup a value for the specified key, or the greatest one less than it
+lookupLE :: forall k v. (Ord k) => k -> Map k v -> Maybe { key :: k, value :: v }
+lookupLE _ Leaf = Nothing
+lookupLE k (Two left k1 v1 right) = case compare k k1 of
+  EQ -> Just { key: k1, value: v1 }
+  GT -> Just $ fromMaybe { key: k1, value: v1 } $ lookupLE k right
+  LT -> lookupLE k left
+lookupLE k (Three left k1 v1 mid k2 v2 right) = case compare k k2 of
+  EQ -> Just { key: k2, value: v2 }
+  GT -> Just $ fromMaybe { key: k2, value: v2 } $ lookupLE k right
+  LT -> lookupLE k $ Two left k1 v1 mid
+
+-- | Lookup a value for the greatest key less than the specified key
+lookupLT :: forall k v. (Ord k) => k -> Map k v -> Maybe { key :: k, value :: v }
+lookupLT _ Leaf = Nothing
+lookupLT k (Two left k1 v1 right) = case compare k k1 of
+  EQ -> findMax left
+  GT -> Just $ fromMaybe { key: k1, value: v1 } $ lookupLT k right
+  LT -> lookupLT k left
+lookupLT k (Three left k1 v1 mid k2 v2 right) = case compare k k2 of
+  EQ -> findMax $ Two left k1 v1 mid
+  GT -> Just $ fromMaybe { key: k2, value: v2 } $ lookupLT k right
+  LT -> lookupLT k $ Two left k1 v1 mid
+
+-- | Lookup a value for the specified key, or the least one greater than it
+lookupGE :: forall k v. (Ord k) => k -> Map k v -> Maybe { key :: k, value :: v }
+lookupGE _ Leaf = Nothing
+lookupGE k (Two left k1 v1 right) = case compare k k1 of
+  EQ -> Just { key: k1, value: v1 }
+  LT -> Just $ fromMaybe { key: k1, value: v1 } $ lookupGE k left
+  GT -> lookupGE k right
+lookupGE k (Three left k1 v1 mid k2 v2 right) = case compare k k1 of
+  EQ -> Just { key: k1, value: v1 }
+  LT -> Just $ fromMaybe { key: k1, value: v1 } $ lookupGE k left
+  GT -> lookupGE k $ Two mid k2 v2 right
+
+-- | Lookup a value for the least key greater than the specified key
+lookupGT :: forall k v. (Ord k) => k -> Map k v -> Maybe { key :: k, value :: v }
+lookupGT _ Leaf = Nothing
+lookupGT k (Two left k1 v1 right) = case compare k k1 of
+  EQ -> findMin right
+  LT -> Just $ fromMaybe { key: k1, value: v1 } $ lookupGT k left
+  GT -> lookupGT k right
+lookupGT k (Three left k1 v1 mid k2 v2 right) = case compare k k1 of
+  EQ -> findMin $ Two mid k2 v2 right
+  LT -> Just $ fromMaybe { key: k1, value: v1 } $ lookupGT k left
+  GT -> lookupGT k $ Two mid k2 v2 right
+
+-- | Returns the pair with the greatest key
+findMax :: forall k v. Map k v -> Maybe { key :: k, value :: v }
+findMax Leaf = Nothing
+findMax (Two _ k1 v1 right) = Just $ fromMaybe { key: k1, value: v1 } $ findMax right
+findMax (Three _ _ _ _ k2 v2 right) = Just $ fromMaybe { key: k2, value: v2 } $ findMax right
+
+-- | Returns the pair with the least key
+findMin :: forall k v. Map k v -> Maybe { key :: k, value :: v }
+findMin Leaf = Nothing
+findMin (Two left k1 v1 _) = Just $ fromMaybe { key: k1, value: v1 } $ findMin left
+findMin (Three left k1 v1 _ _ _ _) = Just $ fromMaybe { key: k1, value: v1 } $ findMin left
+
 -- | Test if a key is a member of a map
 member :: forall k v. Ord k => k -> Map k v -> Boolean
 member k m = isJust (k `lookup` m)

test/Test/Data/Map.purs

@@ -8,7 +8,7 @@ import Control.Monad.Eff.Console (log, CONSOLE)
 import Control.Monad.Eff.Exception (EXCEPTION)
 import Control.Monad.Eff.Random (RANDOM)
 
-import Data.Foldable (foldl, for_)
+import Data.Foldable (foldl, for_, all)
 import Data.Function (on)
 import Data.List (List(..), groupBy, length, nubBy, sortBy, singleton)
 import Data.Map as M
@@ -112,6 +112,9 @@ smallKey k = k
 number :: Int -> Int
 number n = n
 
+smallKeyToNumberMap :: M.Map SmallKey Int -> M.Map SmallKey Int
+smallKeyToNumberMap m = m
+
 mapTests :: forall eff. Eff (console :: CONSOLE, random :: RANDOM, err :: EXCEPTION | eff) Unit
 mapTests = do
 
@@ -247,3 +250,53 @@ mapTests = do
   quickCheck $ \xs ->
     let xs' = nubBy ((==) `on` fst) xs
     in  M.size (M.fromList xs') == length (xs' :: List (Tuple SmallKey Int))
+
+  log "lookupLE result is correct"
+  quickCheck $ \k (TestMap m) -> case M.lookupLE k (smallKeyToNumberMap m) of
+    Nothing -> all (_ > k) $ M.keys m
+    Just { key: k1, value: v } -> let
+      isCloserKey k2 = k1 < k2 && k2 < k
+      isLTwhenEQexists = k1 < k && M.member k m
+      in   k1 <= k
+        && all (not <<< isCloserKey) (M.keys m)
+        && not isLTwhenEQexists
+        && M.lookup k1 m == Just v
+
+  log "lookupGE result is correct"
+  quickCheck $ \k (TestMap m) -> case M.lookupGE k (smallKeyToNumberMap m) of
+    Nothing -> all (_ < k) $ M.keys m
+    Just { key: k1, value: v } -> let
+      isCloserKey k2 = k < k2 && k2 < k1
+      isGTwhenEQexists = k < k1 && M.member k m
+      in   k1 >= k
+        && all (not <<< isCloserKey) (M.keys m)
+        && not isGTwhenEQexists
+        && M.lookup k1 m == Just v
+
+  log "lookupLT result is correct"
+  quickCheck $ \k (TestMap m) -> case M.lookupLT k (smallKeyToNumberMap m) of
+    Nothing -> all (_ >= k) $ M.keys m
+    Just { key: k1, value: v } -> let
+      isCloserKey k2 = k1 < k2 && k2 < k
+      in   k1 < k
+        && all (not <<< isCloserKey) (M.keys m)
+        && M.lookup k1 m == Just v
+
+  log "lookupGT result is correct"
+  quickCheck $ \k (TestMap m) -> case M.lookupGT k (smallKeyToNumberMap m) of
+    Nothing -> all (_ <= k) $ M.keys m
+    Just { key: k1, value: v } -> let
+      isCloserKey k2 = k < k2 && k2 < k1
+      in   k1 > k
+        && all (not <<< isCloserKey) (M.keys m)
+        && M.lookup k1 m == Just v
+
+  log "findMin result is correct"
+  quickCheck $ \(TestMap m) -> case M.findMin (smallKeyToNumberMap m) of
+    Nothing -> M.isEmpty m
+    Just { key: k, value: v } -> M.lookup k m == Just v && all (_ >= k) (M.keys m)
+
+  log "findMax result is correct"
+  quickCheck $ \(TestMap m) -> case M.findMax (smallKeyToNumberMap m) of
+    Nothing -> M.isEmpty m
+    Just { key: k, value: v } -> M.lookup k m == Just v && all (_ <= k) (M.keys m)