diff --git a/Data/Set.hs b/Data/Set.hs index fd8c8b96e..f2cf1ea49 100644 --- a/Data/Set.hs +++ b/Data/Set.hs @@ -63,6 +63,7 @@ module Data.Set ( , size , member , notMember + , lookup , lookupLT , lookupGT , lookupLE @@ -146,6 +147,7 @@ module Data.Set ( ) where import Data.Set.Base as S +import Prelude () -- $strictness -- diff --git a/Data/Set/Base.hs b/Data/Set/Base.hs index 92bfc1dcc..fc1bc5ff3 100644 --- a/Data/Set/Base.hs +++ b/Data/Set/Base.hs @@ -114,6 +114,7 @@ module Data.Set.Base ( , size , member , notMember + , lookup , lookupLT , lookupGT , lookupLE @@ -194,7 +195,7 @@ module Data.Set.Base ( , merge ) where -import Prelude hiding (filter,foldl,foldr,null,map) +import Prelude hiding (filter,foldl,foldr,null,map,lookup) import qualified Data.List as List import Data.Bits (shiftL, shiftR) #if !MIN_VERSION_base(4,8,0) @@ -370,6 +371,21 @@ notMember a t = not $ member a t {-# INLINE notMember #-} #endif +-- | /O(log n)/. Find the given element and return the copy contained in the set. +lookup :: Ord a => a -> Set a -> Maybe a +lookup = go + where + go !_ Tip = Nothing + go x (Bin _ y l r) = case compare x y of + LT -> go x l + GT -> go x r + EQ -> Just y +#if __GLASGOW_HASKELL__ +{-# INLINABLE lookup #-} +#else +{-# INLINE lookup #-} +#endif + -- | /O(log n)/. Find largest element smaller than the given one. -- -- > lookupLT 3 (fromList [3, 5]) == Nothing diff --git a/tests/set-properties.hs b/tests/set-properties.hs index 694437c0f..5a083d496 100644 --- a/tests/set-properties.hs +++ b/tests/set-properties.hs @@ -12,7 +12,8 @@ import Test.HUnit hiding (Test, Testable) import Test.QuickCheck main :: IO () -main = defaultMain [ testCase "lookupLT" test_lookupLT +main = defaultMain [ testCase "lookup" test_lookup + , testCase "lookupLT" test_lookupLT , testCase "lookupGT" test_lookupGT , testCase "lookupLE" test_lookupLE , testCase "lookupGE" test_lookupGE @@ -24,6 +25,7 @@ main = defaultMain [ testCase "lookupLT" test_lookupLT , testProperty "prop_Single" prop_Single , testProperty "prop_Member" prop_Member , testProperty "prop_NotMember" prop_NotMember + , testProperty "prop_Lookup" prop_Lookup , testProperty "prop_LookupLT" prop_LookupLT , testProperty "prop_LookupGT" prop_LookupGT , testProperty "prop_LookupLE" prop_LookupLE @@ -73,6 +75,11 @@ main = defaultMain [ testCase "lookupLT" test_lookupLT -- Unit tests ---------------------------------------------------------------- +test_lookup :: Assertion +test_lookup = do + lookup 3 (fromList [3, 5]) @?= Just 3 + lookup 4 (fromList [3, 5]) @?= Nothing + test_lookupLT :: Assertion test_lookupLT = do lookupLT 3 (fromList [3, 5]) @?= Nothing @@ -192,6 +199,9 @@ test_LookupSomething lookup' cmp xs = filter_odd [_] = [] filter_odd (_ : o : xs) = o : filter_odd xs +prop_Lookup :: [Int] -> Bool +prop_Lookup = test_LookupSomething lookup (==) + prop_LookupLT :: [Int] -> Bool prop_LookupLT = test_LookupSomething lookupLT (<)