Fix #814: add argSet and fromArgSet (#817)

Co-authored-by: David Feuer <[email protected]>

Author: josephcsible

containers-tests/tests/map-properties.hs

@@ -19,6 +19,7 @@ import Data.Monoid
 import Data.Maybe hiding (mapMaybe)
 import qualified Data.Maybe as Maybe (mapMaybe)
 import Data.Ord
+import Data.Semigroup (Arg(..))
 import Data.Function
 import qualified Data.Foldable as Foldable
 #if MIN_VERSION_base(4,10,0)
@@ -99,7 +100,9 @@ main = defaultMain $ testGroup "map-properties"
          , testCase "keys" test_keys
          , testCase "assocs" test_assocs
          , testCase "keysSet" test_keysSet
+         , testCase "argSet" test_argSet
          , testCase "fromSet" test_fromSet
+         , testCase "fromArgSet" test_fromArgSet
          , testCase "toList" test_toList
          , testCase "fromList" test_fromList
          , testCase "fromListWith" test_fromListWith
@@ -238,7 +241,9 @@ main = defaultMain $ testGroup "map-properties"
          , testProperty "bifoldl'"             prop_bifoldl'
 #endif
          , testProperty "keysSet"              prop_keysSet
+         , testProperty "argSet"               prop_argSet
          , testProperty "fromSet"              prop_fromSet
+         , testProperty "fromArgSet"           prop_fromArgSet
          , testProperty "takeWhileAntitone"    prop_takeWhileAntitone
          , testProperty "dropWhileAntitone"    prop_dropWhileAntitone
          , testProperty "spanAntitone"         prop_spanAntitone
@@ -711,11 +716,21 @@ test_keysSet = do
     keysSet (fromList [(5,"a"), (3,"b")]) @?= Set.fromList [3,5]
     keysSet (empty :: UMap) @?= Set.empty
 
+test_argSet :: Assertion
+test_argSet = do
+    argSet (fromList [(5,"a"), (3,"b")]) @?= Set.fromList [Arg 3 "b",Arg 5 "a"]
+    argSet (empty :: UMap) @?= Set.empty
+
 test_fromSet :: Assertion
 test_fromSet = do
    fromSet (\k -> replicate k 'a') (Set.fromList [3, 5]) @?= fromList [(5,"aaaaa"), (3,"aaa")]
    fromSet undefined Set.empty @?= (empty :: IMap)
 
+test_fromArgSet :: Assertion
+test_fromArgSet = do
+   fromArgSet (Set.fromList [Arg 3 "aaa", Arg 5 "aaaaa"]) @?= fromList [(5,"aaaaa"), (3,"aaa")]
+   fromArgSet Set.empty @?= (empty :: IMap)
+
 ----------------------------------------------------------------
 -- Lists
 
@@ -1556,7 +1571,16 @@ prop_keysSet :: [(Int, Int)] -> Bool
 prop_keysSet xs =
   keysSet (fromList xs) == Set.fromList (List.map fst xs)
 
+prop_argSet :: [(Int, Int)] -> Bool
+prop_argSet xs =
+  argSet (fromList xs) == Set.fromList (List.map (uncurry Arg) xs)
+
 prop_fromSet :: [(Int, Int)] -> Bool
 prop_fromSet ys =
   let xs = List.nubBy ((==) `on` fst) ys
   in fromSet (\k -> fromJust $ List.lookup k xs) (Set.fromList $ List.map fst xs) == fromList xs
+
+prop_fromArgSet :: [(Int, Int)] -> Bool
+prop_fromArgSet ys =
+  let xs = List.nubBy ((==) `on` fst) ys
+  in fromArgSet (Set.fromList $ List.map (uncurry Arg) xs) == fromList xs

containers/src/Data/Map/Internal.hs

@@ -264,7 +264,9 @@ module Data.Map.Internal (
     , keys
     , assocs
     , keysSet
+    , argSet
     , fromSet
+    , fromArgSet
 
     -- ** Lists
     , toList
@@ -370,7 +372,7 @@ import Data.Functor.Identity (Identity (..))
 import Control.Applicative (liftA3)
 import Data.Functor.Classes
 import Data.Semigroup (stimesIdempotentMonoid)
-import Data.Semigroup (Semigroup(stimes))
+import Data.Semigroup (Arg(..), Semigroup(stimes))
 #if !(MIN_VERSION_base(4,11,0))
 import Data.Semigroup (Semigroup((<>)))
 #endif
@@ -3362,6 +3364,15 @@ keysSet :: Map k a -> Set.Set k
 keysSet Tip = Set.Tip
 keysSet (Bin sz kx _ l r) = Set.Bin sz kx (keysSet l) (keysSet r)
 
+-- | \(O(n)\). The set of all elements of the map contained in 'Arg's.
+--
+-- > argSet (fromList [(5,"a"), (3,"b")]) == Data.Set.fromList [Arg 3 "b",Arg 5 "a"]
+-- > argSet empty == Data.Set.empty
+
+argSet :: Map k a -> Set.Set (Arg k a)
+argSet Tip = Set.Tip
+argSet (Bin sz kx x l r) = Set.Bin sz (Arg kx x) (argSet l) (argSet r)
+
 -- | \(O(n)\). Build a map from a set of keys and a function which for each key
 -- computes its value.
 --
@@ -3372,6 +3383,15 @@ fromSet :: (k -> a) -> Set.Set k -> Map k a
 fromSet _ Set.Tip = Tip
 fromSet f (Set.Bin sz x l r) = Bin sz x (f x) (fromSet f l) (fromSet f r)
 
+-- | /O(n)/. Build a map from a set of elements contained inside 'Arg's.
+--
+-- > fromArgSet (Data.Set.fromList [Arg 3 "aaa", Arg 5 "aaaaa"]) == fromList [(5,"aaaaa"), (3,"aaa")]
+-- > fromArgSet Data.Set.empty == empty
+
+fromArgSet :: Set.Set (Arg k a) -> Map k a
+fromArgSet Set.Tip = Tip
+fromArgSet (Set.Bin sz (Arg x v) l r) = Bin sz x v (fromArgSet l) (fromArgSet r)
+
 {--------------------------------------------------------------------
   Lists
 --------------------------------------------------------------------}

containers/src/Data/Map/Lazy.hs

@@ -92,6 +92,7 @@ module Data.Map.Lazy (
     , empty
     , singleton
     , fromSet
+    , fromArgSet
 
     -- ** From Unordered Lists
     , fromList
@@ -207,6 +208,7 @@ module Data.Map.Lazy (
     , keys
     , assocs
     , keysSet
+    , argSet
 
     -- ** Lists
     , toList

containers/src/Data/Map/Strict.hs

@@ -108,6 +108,7 @@ module Data.Map.Strict
     , empty
     , singleton
     , fromSet
+    , fromArgSet
 
     -- ** From Unordered Lists
     , fromList
@@ -223,6 +224,7 @@ module Data.Map.Strict
     , keys
     , assocs
     , keysSet
+    , argSet
 
     -- ** Lists
     , toList

containers/src/Data/Map/Strict/Internal.hs

@@ -226,7 +226,9 @@ module Data.Map.Strict.Internal
     , keys
     , assocs
     , keysSet
+    , argSet
     , fromSet
+    , fromArgSet
 
     -- ** Lists
     , toList
@@ -327,6 +329,7 @@ import Data.Map.Internal
   , (!)
   , (!?)
   , (\\)
+  , argSet
   , assocs
   , atKeyImpl
   , atKeyPlain
@@ -413,6 +416,7 @@ import Data.Map.Internal.DeprecatedShowTree (showTree, showTreeWith)
 import Data.Map.Internal.Debug (valid)
 
 import Control.Applicative (Const (..), liftA3)
+import Data.Semigroup (Arg (..))
 import qualified Data.Set.Internal as Set
 import qualified Data.Map.Internal as L
 import Utils.Containers.Internal.StrictPair
@@ -1467,6 +1471,15 @@ fromSet :: (k -> a) -> Set.Set k -> Map k a
 fromSet _ Set.Tip = Tip
 fromSet f (Set.Bin sz x l r) = case f x of v -> v `seq` Bin sz x v (fromSet f l) (fromSet f r)
 
+-- | /O(n)/. Build a map from a set of elements contained inside 'Arg's.
+--
+-- > fromArgSet (Data.Set.fromList [Arg 3 "aaa", Arg 5 "aaaaa"]) == fromList [(5,"aaaaa"), (3,"aaa")]
+-- > fromArgSet Data.Set.empty == empty
+
+fromArgSet :: Set.Set (Arg k a) -> Map k a
+fromArgSet Set.Tip = Tip
+fromArgSet (Set.Bin sz (Arg x v) l r) = v `seq` Bin sz x v (fromArgSet l) (fromArgSet r)
+
 {--------------------------------------------------------------------
   Lists
 --------------------------------------------------------------------}