Fix #818: add unsafe functions for converting between Set and Map

Author: josephcsible

containers-tests/tests/map-properties.hs

@@ -104,6 +104,10 @@ main = defaultMain $ testGroup "map-properties"
          , testCase "argSet" test_argSet
          , testCase "fromSet" test_fromSet
          , testCase "fromArgSet" test_fromArgSet
+         , testCase "unsafeSet" test_unsafeSet
+         , testCase "unsafeSetA" test_unsafeSetA
+         , testCase "unsafeFromSet" test_unsafeFromSet
+         , testCase "unsafeFromSetA" test_unsafeFromSetA
          , testCase "toList" test_toList
          , testCase "fromList" test_fromList
          , testCase "fromListWith" test_fromListWith
@@ -252,6 +256,10 @@ main = defaultMain $ testGroup "map-properties"
          , testProperty "argSet"               prop_argSet
          , testProperty "fromSet"              prop_fromSet
          , testProperty "fromArgSet"           prop_fromArgSet
+         , testProperty "unsafeSet"            prop_unsafeSet
+         , testProperty "unsafeSetA"           prop_unsafeSetA
+         , testProperty "unsafeFromSet"        prop_unsafeFromSet
+         , testProperty "unsafeFromSetA"       prop_unsafeFromSetA
          , testProperty "takeWhileAntitone"    prop_takeWhileAntitone
          , testProperty "dropWhileAntitone"    prop_dropWhileAntitone
          , testProperty "spanAntitone"         prop_spanAntitone
@@ -707,6 +715,26 @@ test_fromArgSet = do
    fromArgSet (Set.fromList [Arg 3 "aaa", Arg 5 "aaaaa"]) @?= fromList [(5,"aaaaa"), (3,"aaa")]
    fromArgSet Set.empty @?= (empty :: IMap)
 
+test_unsafeSet :: Assertion
+test_unsafeSet = do
+   unsafeSet (,) (fromList [(5,"a"), (3,"b")]) @?= Set.fromList [(5,"a"), (3,"b")]
+   unsafeSet undefined (empty :: UMap) @?= (Set.empty :: Set.Set Int)
+
+test_unsafeSetA :: Assertion
+test_unsafeSetA = do
+   unsafeSetA (\x y -> Just (x,y)) (fromList [(5,"a"), (3,"b")]) @?= Just (Set.fromList [(5,"a"), (3,"b")])
+   unsafeSetA undefined (empty :: UMap) @?= Identity (Set.empty :: Set.Set Int)
+
+test_unsafeFromSet :: Assertion
+test_unsafeFromSet = do
+   unsafeFromSet (\k -> (k+1, replicate k 'a')) (Set.fromList [2, 4]) @?= fromList [(5,"aaaa"), (3,"aa")]
+   unsafeFromSet undefined Set.empty @?= (empty :: IMap)
+
+test_unsafeFromSetA :: Assertion
+test_unsafeFromSetA = do
+   unsafeFromSetA (\k -> Just (k+1, replicate k 'a')) (Set.fromList [2, 4]) @?= Just (fromList [(5,"aaaa"), (3,"aa")])
+   unsafeFromSetA undefined Set.empty @?= Identity (empty :: IMap)
+
 ----------------------------------------------------------------
 -- Lists
 
@@ -1672,6 +1700,26 @@ prop_fromArgSet ys =
   let xs = List.nubBy ((==) `on` fst) ys
   in fromArgSet (Set.fromList $ List.map (uncurry Arg) xs) == fromList xs
 
+prop_unsafeSet :: [(Int, Int)] -> Bool
+prop_unsafeSet ys =
+  let xs = List.nubBy ((==) `on` fst) ys
+  in unsafeSet (,) (fromList xs) == Set.fromList xs
+
+prop_unsafeSetA :: [(Int, Int)] -> Bool
+prop_unsafeSetA ys =
+  let xs = List.nubBy ((==) `on` fst) ys
+  in unsafeSetA (\x y -> Identity (x,y)) (fromList xs) == Identity (Set.fromList xs)
+
+prop_unsafeFromSet :: [(Int, Int)] -> Bool
+prop_unsafeFromSet ys =
+  let xs = List.nubBy ((==) `on` fst) ys
+  in unsafeFromSet id (Set.fromList xs) == fromList xs
+
+prop_unsafeFromSetA :: [(Int, Int)] -> Bool
+prop_unsafeFromSetA ys =
+  let xs = List.nubBy ((==) `on` fst) ys
+  in unsafeFromSetA Identity (Set.fromList xs) == Identity (fromList xs)
+
 prop_eq :: Map Int A -> Map Int A -> Property
 prop_eq m1 m2 = (m1 == m2) === (toList m1 == toList m2)
 

containers/src/Data/Map/Internal.hs

@@ -270,6 +270,10 @@ module Data.Map.Internal (
     , argSet
     , fromSet
     , fromArgSet
+    , unsafeSet
+    , unsafeSetA
+    , unsafeFromSet
+    , unsafeFromSetA
 
     -- ** Lists
     , toList
@@ -3504,6 +3508,58 @@ 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)
 
+-- | \(O(n)\). Build a set from the elements in a map and a function which for each
+-- element computes its value. The function must preserve the relative ordering
+-- of the keys. /The precondition is not checked./
+--
+-- > unsafeSet id (fromList [(5,"a"), (3,"b")]) == Data.Set.fromList [(3,"b"),(5,"a")]
+-- > unsafeSet undefined empty == Data.Set.empty
+
+unsafeSet :: (k -> a -> b) -> Map k a -> Set.Set b
+unsafeSet f = go
+  where
+    go Tip = Set.Tip
+    go (Bin sz kx x l r) = Set.Bin sz (f kx x) (go l) (go r)
+
+-- | \(O(n)\). Build a set from the elements in a map and a function which for each
+-- element computes its value inside an 'Applicative'. The function must preserve
+-- the relative ordering of the keys. /The precondition is not checked./
+--
+-- > unsafeSetA Identity (fromList [(5,"a"), (3,"b")]) == Data.Set.fromList [(3,"b"),(5,"a")]
+-- > unsafeSetA undefined empty == Identity (Data.Set.empty)
+
+unsafeSetA :: Applicative t => (k -> a -> t b) -> Map k a -> t (Set.Set b)
+unsafeSetA f = go
+  where
+    go Tip = pure Set.Tip
+    go (Bin sz kx x l r) = liftA3 (Set.Bin sz) (f kx x) (go l) (go r)
+
+-- | \(O(n)\). Build a map from a set of elements and a function which for each
+-- element computes its key and value. The function must preserve the relative
+-- ordering of the elements. /The precondition is not checked./
+--
+-- > unsafeFromSet id (Data.Set.fromList [(3,"aaa"), (5,"aaaaa")]) == fromList [(5,"aaaaa"), (3,"aaa")]
+-- > unsafeFromSet undefined Data.Set.empty == empty
+
+unsafeFromSet :: (b -> (k, a)) -> Set.Set b -> Map k a
+unsafeFromSet f = go
+  where
+    go Set.Tip = Tip
+    go (Set.Bin sz x l r) = uncurry (Bin sz) (f x) (go l) (go r)
+
+-- | \(O(n)\). Build a map from a set of elements and a function which for each
+-- element computes its key and value inside an 'Applicative'. The function must
+-- preserve the relative ordering of the elements. /The precondition is not checked./
+--
+-- > unsafeFromSetA Identity (Data.Set.fromList [(3,"aaa"), (5,"aaaaa")]) == Identity (fromList [(5,"aaaaa"), (3,"aaa")])
+-- > unsafeFromSetA undefined Data.Set.empty == Identity empty
+
+unsafeFromSetA :: Applicative t => (b -> t (k, a)) -> Set.Set b -> t (Map k a)
+unsafeFromSetA f = go
+  where
+    go Set.Tip = pure Tip
+    go (Set.Bin sz x l r) = liftA3 (uncurry (Bin sz)) (f x) (go l) (go r)
+
 {--------------------------------------------------------------------
   Lists
 --------------------------------------------------------------------}

containers/src/Data/Map/Lazy.hs

@@ -99,6 +99,8 @@ module Data.Map.Lazy (
     , singleton
     , fromSet
     , fromArgSet
+    , unsafeFromSet
+    , unsafeFromSetA
 
     -- ** From Unordered Lists
     , fromList
@@ -218,6 +220,8 @@ module Data.Map.Lazy (
     , assocs
     , keysSet
     , argSet
+    , unsafeSet
+    , unsafeSetA
 
     -- ** Lists
     , toList

containers/src/Data/Map/Strict.hs

@@ -114,6 +114,8 @@ module Data.Map.Strict
     , singleton
     , fromSet
     , fromArgSet
+    , unsafeFromSet
+    , unsafeFromSetA
 
     -- ** From Unordered Lists
     , fromList
@@ -233,6 +235,8 @@ module Data.Map.Strict
     , assocs
     , keysSet
     , argSet
+    , unsafeSet
+    , unsafeSetA
 
     -- ** Lists
     , toList

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

@@ -232,6 +232,10 @@ module Data.Map.Strict.Internal
     , argSet
     , fromSet
     , fromArgSet
+    , unsafeSet
+    , unsafeSetA
+    , unsafeFromSet
+    , unsafeFromSetA
 
     -- ** Lists
     , toList
@@ -420,6 +424,8 @@ import Data.Map.Internal
   , toDescList
   , union
   , unions
+  , unsafeSet
+  , unsafeSetA
   , withoutKeys )
 
 import Data.Map.Internal.Debug (valid)
@@ -1476,6 +1482,33 @@ 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)
 
+-- | \(O(n)\). Build a map from a set of elements and a function which for each
+-- element computes its key and value. The function must preserve the relative
+-- ordering of the elements. /The precondition is not checked./
+--
+-- > unsafeFromSet id (Data.Set.fromList [(3,"aaa"), (5,"aaaaa")]) == fromList [(5,"aaaaa"), (3,"aaa")]
+-- > unsafeFromSet undefined Data.Set.empty == empty
+
+unsafeFromSet :: (b -> (k, a)) -> Set.Set b -> Map k a
+unsafeFromSet f = go
+  where
+    go Set.Tip = Tip
+    go (Set.Bin sz x l r) = case f x of
+      (k,!v) -> Bin sz k v (go l) (go r)
+
+-- | \(O(n)\). Build a map from a set of elements and a function which for each
+-- element computes its key and value inside an 'Applicative'. The function must
+-- preserve the relative ordering of the elements. /The precondition is not checked./
+--
+-- > unsafeFromSetA Identity (Data.Set.fromList [(3,"aaa"), (5,"aaaaa")]) == Identity (fromList [(5,"aaaaa"), (3,"aaa")])
+-- > unsafeFromSetA undefined Data.Set.empty == Identity empty
+
+unsafeFromSetA :: Applicative t => (b -> t (k, a)) -> Set.Set b -> t (Map k a)
+unsafeFromSetA f = go
+  where
+    go Set.Tip = pure Tip
+    go (Set.Bin sz x l r) = liftA3 (\(k,!v) -> Bin sz k v) (f x) (go l) (go r)
+
 {--------------------------------------------------------------------
   Lists
 --------------------------------------------------------------------}