Add new index-based and unsafe Set functions
authorDavid Feuer <David.Feuer@gmail.com>
Wed, 31 Aug 2016 02:48:12 +0000 (22:48 -0400)
committerDavid Feuer <David.Feuer@gmail.com>
Wed, 31 Aug 2016 02:49:33 +0000 (22:49 -0400)
These match the ones just added for maps.

Data/Set.hs
Data/Set/Base.hs
changelog.md
tests/set-properties.hs

index a17f6b5..07b566d 100644 (file)
@@ -90,6 +90,9 @@ module Data.Set (
 
             -- * Filter
             , S.filter
+            , takeWhileAntitone
+            , dropWhileAntitone
+            , spanAntitone
             , partition
             , split
             , splitMember
@@ -100,6 +103,9 @@ module Data.Set (
             , findIndex
             , elemAt
             , deleteAt
+            , S.take
+            , S.drop
+            , S.splitAt
 
             -- * Map
             , S.map
index c8a3688..881e8ce 100644 (file)
@@ -139,6 +139,9 @@ module Data.Set.Base (
 
             -- * Filter
             , filter
+            , takeWhileAntitone
+            , dropWhileAntitone
+            , spanAntitone
             , partition
             , split
             , splitMember
@@ -149,6 +152,9 @@ module Data.Set.Base (
             , findIndex
             , elemAt
             , deleteAt
+            , take
+            , drop
+            , splitAt
 
             -- * Map
             , map
@@ -200,7 +206,7 @@ module Data.Set.Base (
             , merge
             ) where
 
-import Prelude hiding (filter,foldl,foldr,null,map)
+import Prelude hiding (filter,foldl,foldr,null,map,take,drop,splitAt)
 import qualified Data.List as List
 import Data.Bits (shiftL, shiftR)
 #if !MIN_VERSION_base(4,8,0)
@@ -1166,6 +1172,117 @@ deleteAt !i t =
       where
         sizeL = size l
 
+-- | Take a given number of elements in order, beginning
+-- with the smallest ones.
+--
+-- @
+-- take n = 'fromDistinctAscList' . 'Prelude.take' n . 'toAscList'
+-- @
+take :: Int -> Set a -> Set a
+take i m | i >= size m = m
+take i0 m0 = go i0 m0
+  where
+    go i !_ | i <= 0 = Tip
+    go !_ Tip = Tip
+    go i (Bin _ x l r) =
+      case compare i sizeL of
+        LT -> go i l
+        GT -> link x l (go (i - sizeL - 1) r)
+        EQ -> l
+      where sizeL = size l
+
+-- | Drop a given number of elements in order, beginning
+-- with the smallest ones.
+--
+-- @
+-- drop n = 'fromDistinctAscList' . 'Prelude.drop' n . 'toAscList'
+-- @
+drop :: Int -> Set a -> Set a
+drop i m | i >= size m = Tip
+drop i0 m0 = go i0 m0
+  where
+    go i m | i <= 0 = m
+    go !_ Tip = Tip
+    go i (Bin _ x l r) =
+      case compare i sizeL of
+        LT -> link x (go i l) r
+        GT -> go (i - sizeL - 1) r
+        EQ -> insertMin x r
+      where sizeL = size l
+
+-- | /O(log n)/. Split a set at a particular index.
+--
+-- @
+-- splitAt !n !xs = ('take' n xs, 'drop' n xs)
+-- @
+splitAt :: Int -> Set a -> (Set a, Set a)
+splitAt i0 m0
+  | i0 >= size m0 = (m0, Tip)
+  | otherwise = toPair $ go i0 m0
+  where
+    go i m | i <= 0 = Tip :*: m
+    go !_ Tip = Tip :*: Tip
+    go i (Bin _ x l r)
+      = case compare i sizeL of
+          LT -> case go i l of
+                  ll :*: lr -> ll :*: link x lr r
+          GT -> case go (i - sizeL - 1) r of
+                  rl :*: rr -> link x l rl :*: rr
+          EQ -> l :*: insertMin x r
+      where sizeL = size l
+
+-- | /O(log n)/. Take while a predicate on the elements holds.
+-- The user is responsible for ensuring that for all elements @j@ and @k@ in the set,
+-- @j \< k ==\> p j \>= p k@. See note at 'spanAntitone'.
+--
+-- @
+-- takeWhileAntitone p = 'fromDistinctAscList' . 'Data.List.takeWhile' p . 'toList'
+-- takeWhileAntitone p = 'filter' p
+-- @
+
+takeWhileAntitone :: (a -> Bool) -> Set a -> Set a
+takeWhileAntitone _ Tip = Tip
+takeWhileAntitone p (Bin _ x l r)
+  | p x = link x l (takeWhileAntitone p r)
+  | otherwise = takeWhileAntitone p l
+
+-- | /O(log n)/. Drop while a predicate on the elements holds.
+-- The user is responsible for ensuring that for all elements @j@ and @k@ in the set,
+-- @j \< k ==\> p j \>= p k@. See note at 'spanAntitone'.
+--
+-- @
+-- dropWhileAntitone p = 'fromDistinctAscList' . 'Data.List.dropWhile' p . 'toList'
+-- dropWhileAntitone p = 'filter' (not . p)
+-- @
+
+dropWhileAntitone :: (a -> Bool) -> Set a -> Set a
+dropWhileAntitone _ Tip = Tip
+dropWhileAntitone p (Bin _ x l r)
+  | p x = dropWhileAntitone p r
+  | otherwise = link x (dropWhileAntitone p l) r
+
+-- | /O(log n)/. Divide a set at the point where a predicate on the elements stops holding.
+-- The user is responsible for ensuring that for all elements @j@ and @k@ in the set,
+-- @j \< k ==\> p j \>= p k@.
+--
+-- @
+-- spanAntitone p xs = ('takeWhileAntitone' p xs, 'dropWhileAntitone' p xs)
+-- spanAntitone p xs = partition p xs
+-- @
+--
+-- Note: if @p@ is not actually antitone, then @spanAntitone@ will split the set
+-- at some /unspecified/ point where the predicate switches from holding to not
+-- holding (where the predicate is seen to hold before the first element and to fail
+-- after the last element).
+
+spanAntitone :: (a -> Bool) -> Set a -> (Set a, Set a)
+spanAntitone p0 m = toPair (go p0 m)
+  where
+    go _ Tip = Tip :*: Tip
+    go p (Bin _ x l r)
+      | p x = let u :*: v = go p r in link x l u :*: v
+      | otherwise = let u :*: v = go p l in u :*: link x v r
+
 
 {--------------------------------------------------------------------
   Utility functions that maintain the balance properties of the tree.
index b21e18f..8abdb54 100644 (file)
@@ -30,8 +30,8 @@
     and `Data.IntMap`.
 
   * Add `take`, `drop`, `splitAt`, `takeWhileAntitone`, `dropWhileAntitone`,
-    and `spanAntitone` for `Data.Map`. Thanks to Cale Gibbard for suggesting
-    these.
+    and `spanAntitone` for `Data.Map` and `Data.Set`. Thanks to Cale Gibbard
+    for suggesting these.
 
   * Add `merge`, `mergeA`, and associated merge tactics for `Data.Map`.
     Many thanks to Cale Gibbard, Ryan Trinkle, and Dan Doel for
index afe7298..5a48397 100644 (file)
@@ -5,7 +5,7 @@ import qualified Data.List as List
 import Data.Monoid (mempty)
 import Data.Maybe
 import Data.Set
-import Prelude hiding (lookup, null, map, filter, foldr, foldl, all)
+import Prelude hiding (lookup, null, map, filter, foldr, foldl, all, take, drop, splitAt)
 import Test.Framework
 import Test.Framework.Providers.HUnit
 import Test.Framework.Providers.QuickCheck2
@@ -85,6 +85,12 @@ main = defaultMain [ testCase "lookupLT" test_lookupLT
                    , testProperty "prop_splitRoot" prop_splitRoot
                    , testProperty "prop_partition" prop_partition
                    , testProperty "prop_filter" prop_filter
+                   , testProperty "takeWhileAntitone"    prop_takeWhileAntitone
+                   , testProperty "dropWhileAntitone"    prop_dropWhileAntitone
+                   , testProperty "spanAntitone"         prop_spanAntitone
+                   , testProperty "take"                 prop_take
+                   , testProperty "drop"                 prop_drop
+                   , testProperty "splitAt"              prop_splitAt
                    ]
 
 -- A type with a peculiar Eq instance designed to make sure keys
@@ -548,3 +554,47 @@ prop_partition s i = case partition odd s of
 
 prop_filter :: Set Int -> Int -> Bool
 prop_filter s i = partition odd s == (filter odd s, filter even s)
+
+prop_take :: Int -> Set Int -> Property
+prop_take n xs = valid taken .&&.
+                 taken === fromDistinctAscList (List.take n (toList xs))
+  where
+    taken = take n xs
+
+prop_drop :: Int -> Set Int -> Property
+prop_drop n xs = valid dropped .&&.
+                 dropped === fromDistinctAscList (List.drop n (toList xs))
+  where
+    dropped = drop n xs
+
+prop_splitAt :: Int -> Set Int -> Property
+prop_splitAt n xs = valid taken .&&.
+                    valid dropped .&&.
+                    taken === take n xs .&&.
+                    dropped === drop n xs
+  where
+    (taken, dropped) = splitAt n xs
+
+prop_takeWhileAntitone :: [Either Int Int] -> Property
+prop_takeWhileAntitone xs' = valid tw .&&. tw === filter isLeft xs
+  where
+    xs = fromList xs'
+    tw = takeWhileAntitone isLeft xs
+
+prop_dropWhileAntitone :: [Either Int Int] -> Property
+prop_dropWhileAntitone xs' = valid tw .&&. tw === filter (not . isLeft) xs
+  where
+    xs = fromList xs'
+    tw = dropWhileAntitone isLeft xs
+
+prop_spanAntitone :: [Either Int Int] -> Property
+prop_spanAntitone xs' = valid tw .&&. valid dw
+                        .&&. tw === takeWhileAntitone isLeft xs
+                        .&&. dw === dropWhileAntitone isLeft xs
+  where
+    xs = fromList xs'
+    (tw, dw) = spanAntitone isLeft xs
+
+isLeft :: Either a b -> Bool
+isLeft (Left _) = True
+isLeft _ = False