Add more indexed and unsafe functions for Data.Map
authorDavid Feuer <David.Feuer@gmail.com>
Thu, 25 Aug 2016 21:44:51 +0000 (17:44 -0400)
committerDavid Feuer <David.Feuer@gmail.com>
Tue, 30 Aug 2016 20:30:20 +0000 (16:30 -0400)
* Offer `take`, `drop`, and `splitAt` by index.

* Offer 'takeWhileAntitone`, `dropWhileAntitone`, and `spanAntitone`.

Data/Map/Base.hs
Data/Map/Lazy.hs
Data/Map/Strict.hs
Data/Map/Strict/Internal.hs
tests/map-properties.hs

index e1c176b..9ebae06 100644 (file)
@@ -3,7 +3,7 @@
 #if __GLASGOW_HASKELL__
 {-# LANGUAGE DeriveDataTypeable, StandaloneDeriving #-}
 #endif
-#if !defined(TESTING) && __GLASGOW_HASKELL__ >= 703
+#if __GLASGOW_HASKELL__ >= 703
 {-# LANGUAGE Trustworthy #-}
 #endif
 #if __GLASGOW_HASKELL__ >= 708
@@ -280,6 +280,11 @@ module Data.Map.Base (
     -- * Filter
     , filter
     , filterWithKey
+
+    , takeWhileAntitone
+    , dropWhileAntitone
+    , spanAntitone
+
     , restrictKeys
     , withoutKeys
     , partition
@@ -304,6 +309,9 @@ module Data.Map.Base (
     , elemAt
     , updateAt
     , deleteAt
+    , take
+    , drop
+    , splitAt
 
     -- * Min\/Max
     , findMin
@@ -370,7 +378,7 @@ import Control.DeepSeq (NFData(rnf))
 import Data.Bits (shiftL, shiftR)
 import qualified Data.Foldable as Foldable
 import Data.Typeable
-import Prelude hiding (lookup, map, filter, foldr, foldl, null)
+import Prelude hiding (lookup, map, filter, foldr, foldl, null, splitAt, take, drop)
 
 import qualified Data.Set.Base as Set
 import Data.Set.Base (Set)
@@ -1444,6 +1452,66 @@ elemAt i (Bin _ kx x l r)
   where
     sizeL = size l
 
+-- | Take a given number of entries in key order, beginning
+-- with the smallest keys.
+--
+-- @
+-- take n = 'fromDistinctAscList' . 'Prelude.take' n . 'toAscList'
+-- @
+
+take :: Int -> Map k a -> Map k 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 _ kx x l r) =
+      case compare i sizeL of
+        LT -> go i l
+        GT -> link kx x l (go (i - sizeL - 1) r)
+        EQ -> l
+      where sizeL = size l
+
+-- | Drop a given number of entries in key order, beginning
+-- with the smallest keys.
+--
+-- @
+-- drop n = 'fromDistinctAscList' . 'Prelude.drop' n . 'toAscList'
+-- @
+drop :: Int -> Map k a -> Map k 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 _ kx x l r) =
+      case compare i sizeL of
+        LT -> link kx x (go i l) r
+        GT -> go (i - sizeL - 1) r
+        EQ -> insertMin kx x r
+      where sizeL = size l
+
+-- | /O(log n)/. Split a map at a particular index.
+--
+-- @
+-- splitAt !n !xs = ('take' n xs, 'drop' n xs)
+-- @
+splitAt :: Int -> Map k a -> (Map k a, Map k 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 _ kx x l r)
+      = case compare i sizeL of
+          LT -> case go i l of
+                  ll :*: lr -> ll :*: link kx x lr r
+          GT -> case go (i - sizeL - 1) r of
+                  rl :*: rr -> link kx x l rl :*: rr
+          EQ -> l :*: insertMin kx x r
+      where sizeL = size l
+
 -- | /O(log n)/. Update the element at /index/, i.e. by its zero-based index in
 -- the sequence sorted by keys. If the /index/ is out of range (less than zero,
 -- greater or equal to 'size' of the map), 'error' is called.
@@ -2597,6 +2665,58 @@ filterWithKeyA p t@(Bin _ kx x l r) =
       | otherwise = link kx x pl pr
     combine False pl pr = link2 pl pr
 
+-- | /O(log n)/. Take while a predicate on the keys holds.
+-- The user is responsible for ensuring that for all keys @j@ and @k@ in the map,
+-- @j \< k ==\> p j \>= p k@. See note at 'spanAntitone'.
+--
+-- @
+-- takeWhileAntitone p = 'fromDistinctAscList' . 'Data.List.takeWhile' (p . fst) . 'toList'
+-- takeWhileAntitone p = 'filterWithKey' (\k _ -> p k)
+-- @
+
+takeWhileAntitone :: (k -> Bool) -> Map k a -> Map k a
+takeWhileAntitone _ Tip = Tip
+takeWhileAntitone p (Bin _ kx x l r)
+  | p kx = link kx x l (takeWhileAntitone p r)
+  | otherwise = takeWhileAntitone p l
+
+-- | /O(log n)/. Drop while a predicate on the keys holds.
+-- The user is responsible for ensuring that for all keys @j@ and @k@ in the map,
+-- @j \< k ==\> p j \>= p k@. See note at 'spanAntitone'.
+--
+-- @
+-- dropWhileAntitone p = 'fromDistinctAscList' . 'Data.List.dropWhile' (p . fst) . 'toList'
+-- dropWhileAntitone p = 'filterWithKey' (\k -> not (p k))
+-- @
+
+dropWhileAntitone :: (k -> Bool) -> Map k a -> Map k a
+dropWhileAntitone _ Tip = Tip
+dropWhileAntitone p (Bin _ kx x l r)
+  | p kx = dropWhileAntitone p r
+  | otherwise = link kx x (dropWhileAntitone p l) r
+
+-- | /O(log n)/. Divide a map at the point where a predicate on the keys stops holding.
+-- The user is responsible for ensuring that for all keys @j@ and @k@ in the map,
+-- @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 map
+-- at some /unspecified/ point where the predicate switches from holding to not
+-- holding (where the predicate is seen to hold before the first key and to fail
+-- after the last key).
+
+spanAntitone :: (k -> Bool) -> Map k a -> (Map k a, Map k a)
+spanAntitone p0 m = toPair (go p0 m)
+  where
+    go _ Tip = Tip :*: Tip
+    go p (Bin _ kx x l r)
+      | p kx = let u :*: v = go p r in link kx x l u :*: v
+      | otherwise = let u :*: v = go p l in u :*: link kx x v r
+
 -- | /O(n)/. Partition the map according to a predicate. The first
 -- map contains all elements that satisfy the predicate, the second all
 -- elements that fail the predicate. See also 'split'.
index 617ae17..4f3e25f 100644 (file)
@@ -1,5 +1,5 @@
 {-# LANGUAGE CPP #-}
-#if !defined(TESTING) && __GLASGOW_HASKELL__ >= 703
+#if __GLASGOW_HASKELL__ >= 703
 {-# LANGUAGE Safe #-}
 #endif
 
@@ -63,11 +63,7 @@ module Data.Map.Lazy (
     -- $strictness
 
     -- * Map type
-#if !defined(TESTING)
     Map              -- instance Eq,Show,Read
-#else
-    Map(..)          -- instance Eq,Show,Read
-#endif
 
     -- * Operators
     , (!), (\\)
@@ -126,7 +122,7 @@ module Data.Map.Lazy (
     -- ** General combining functions
     -- | See "Data.Map.Lazy.Merge"
 
-    -- ** Deprecated general combining function
+    -- ** Unsafe general combining function
 
     , mergeWithKey
 
@@ -188,6 +184,9 @@ module Data.Map.Lazy (
     , withoutKeys
     , partition
     , partitionWithKey
+    , takeWhileAntitone
+    , dropWhileAntitone
+    , spanAntitone
 
     , mapMaybe
     , mapMaybeWithKey
@@ -208,6 +207,9 @@ module Data.Map.Lazy (
     , elemAt
     , updateAt
     , deleteAt
+    , take
+    , drop
+    , splitAt
 
     -- * Min\/Max
     , findMin
@@ -229,18 +231,10 @@ module Data.Map.Lazy (
     , showTree
     , showTreeWith
     , valid
-
-#if defined(TESTING)
-    -- * Internals
-    , bin
-    , balanced
-    , link
-    , link2
-#endif
-
     ) where
 
 import Data.Map.Base as M
+import Prelude ()
 
 -- $strictness
 --
index 78b24df..b3774c5 100644 (file)
@@ -193,6 +193,10 @@ module Data.Map.Strict
     , partition
     , partitionWithKey
 
+    , takeWhileAntitone
+    , dropWhileAntitone
+    , spanAntitone
+
     , mapMaybe
     , mapMaybeWithKey
     , mapEither
@@ -212,6 +216,9 @@ module Data.Map.Strict
     , elemAt
     , updateAt
     , deleteAt
+    , take
+    , drop
+    , splitAt
 
     -- * Min\/Max
     , findMin
index 89d93ef..0b0a34e 100644 (file)
@@ -247,6 +247,9 @@ module Data.Map.Strict.Internal
     , withoutKeys
     , partition
     , partitionWithKey
+    , takeWhileAntitone
+    , dropWhileAntitone
+    , spanAntitone
 
     , mapMaybe
     , mapMaybeWithKey
@@ -267,6 +270,9 @@ module Data.Map.Strict.Internal
     , elemAt
     , updateAt
     , deleteAt
+    , take
+    , drop
+    , splitAt
 
     -- * Min\/Max
     , findMin
@@ -295,7 +301,7 @@ module Data.Map.Strict.Internal
     , link2
     ) where
 
-import Prelude hiding (lookup,map,filter,foldr,foldl,null)
+import Prelude hiding (lookup,map,filter,foldr,foldl,null,take,drop,splitAt)
 
 import Data.Map.Base
   ( Map (..)
@@ -332,6 +338,8 @@ import Data.Map.Base
   , deleteMin
   , deleteMax
   , difference
+  , drop
+  , dropWhileAntitone
   , filter
   , filterWithKey
   , findIndex
@@ -378,9 +386,13 @@ import Data.Map.Base
   , showTree
   , showTreeWith
   , size
+  , spanAntitone
   , split
+  , splitAt
   , splitLookup
   , splitRoot
+  , take
+  , takeWhileAntitone
   , toList
   , toAscList
   , toDescList
index 065b85d..dc115c0 100644 (file)
@@ -1,11 +1,13 @@
 {-# LANGUAGE CPP #-}
 
 #ifdef STRICT
-import Data.Map.Strict.Internal as Data.Map
+import Data.Map.Strict as Data.Map
+import Data.Map.Strict.Merge
 #else
 import Data.Map.Lazy as Data.Map
 import Data.Map.Lazy.Merge
 #endif
+import Data.Map.Base (Map (..), balanced, link2, link, bin)
 
 import Control.Applicative (Const(Const, getConst), pure, (<$>), (<*>))
 import Data.Functor.Identity (Identity(runIdentity))
@@ -14,8 +16,8 @@ import Data.Maybe hiding (mapMaybe)
 import qualified Data.Maybe as Maybe (mapMaybe)
 import Data.Ord
 import Data.Function
-import Prelude hiding (lookup, null, map, filter, foldr, foldl)
-import qualified Prelude (map)
+import Prelude hiding (lookup, null, map, filter, foldr, foldl, take, drop, splitAt)
+import qualified Prelude
 
 import Data.List (nub,sort)
 import qualified Data.List as List
@@ -213,6 +215,12 @@ main = defaultMain
          , testProperty "foldl'"               prop_foldl'
          , testProperty "keysSet"              prop_keysSet
          , testProperty "fromSet"              prop_fromSet
+         , testProperty "takeWhileAntitone"    prop_takeWhileAntitone
+         , testProperty "dropWhileAntitone"    prop_dropWhileAntitone
+         , testProperty "spanAntitone"         prop_spanAntitone
+         , testProperty "take"                 prop_take
+         , testProperty "drop"                 prop_drop
+         , testProperty "splitAt"              prop_splitAt
          ]
 
 {--------------------------------------------------------------------
@@ -1257,6 +1265,50 @@ prop_filter p ys = length ys > 0 ==>
       m  = fromList xs
   in  filter (apply p) m == fromList (List.filter (apply p . snd) xs)
 
+prop_take :: Int -> Map Int Int -> Property
+prop_take n xs = valid taken .&&.
+                 taken === fromDistinctAscList (List.take n (toList xs))
+  where
+    taken = take n xs
+
+prop_drop :: Int -> Map Int Int -> Property
+prop_drop n xs = valid dropped .&&.
+                 dropped === fromDistinctAscList (List.drop n (toList xs))
+  where
+    dropped = drop n xs
+
+prop_splitAt :: Int -> Map Int 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, Int)] -> Property
+prop_takeWhileAntitone xs' = valid tw .&&. (tw === filterWithKey (\k _ -> isLeft k) xs)
+  where
+    xs = fromList xs'
+    tw = takeWhileAntitone isLeft xs
+
+prop_dropWhileAntitone :: [(Either Int Int, Int)] -> Property
+prop_dropWhileAntitone xs' = valid tw .&&. (tw === filterWithKey (\k _ -> not (isLeft k)) xs)
+  where
+    xs = fromList xs'
+    tw = dropWhileAntitone isLeft xs
+
+prop_spanAntitone :: [(Either Int 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
+
 prop_partition :: Fun Int Bool -> [(Int, Int)] -> Property
 prop_partition p ys = length ys > 0 ==>
   let xs = List.nubBy ((==) `on` fst) ys