Improve {Set,Map}.fromAscList and friends

containers-tests/benchmarks/Map.hs

@@ -21,7 +21,7 @@ main = do
         m_even = M.fromAscList elems_even :: M.Map Int Int
         m_odd = M.fromAscList elems_odd :: M.Map Int Int
     evaluate $ rnf [m, m_even, m_odd]
-    evaluate $ rnf elems_rev
+    evaluate $ rnf [elems_rev, elems_asc, elems_desc]
     defaultMain
         [ bench "lookup absent" $ whnf (lookup evens) m_odd
         , bench "lookup present" $ whnf (lookup evens) m_even
@@ -89,7 +89,10 @@ main = do
         , bench "split" $ whnf (M.split (bound `div` 2)) m
         , bench "fromList" $ whnf M.fromList elems
         , bench "fromList-desc" $ whnf M.fromList (reverse elems)
-        , bench "fromAscList" $ whnf M.fromAscList elems
+        , bench "fromAscList" $ whnf M.fromAscList elems_asc
+        , bench "fromAscList:fusion" $ whnf (\n -> M.fromAscList [(i `div` 2, i) | i <- [1..n]]) bound
+        , bench "fromDescList" $ whnf M.fromDescList elems_desc
+        , bench "fromDescList:fusion" $ whnf (\n -> M.fromDescList [(i `div` 2, i) | i <- [n,n-1..1]]) bound
         , bench "fromDistinctAscList" $ whnf M.fromDistinctAscList elems
         , bench "fromDistinctAscList:fusion" $ whnf (\n -> M.fromDistinctAscList [(i,i) | i <- [1..n]]) bound
         , bench "fromDistinctDescList" $ whnf M.fromDistinctDescList elems_rev
@@ -102,6 +105,10 @@ main = do
     elems_even = zip evens evens
     elems_odd = zip odds odds
     elems_rev = reverse elems
+    keys_asc = map (`div` 2) [1..bound]
+    elems_asc = zip keys_asc values
+    keys_desc = map (`div` 2) [bound,bound-1..1]
+    elems_desc = zip keys_desc values
     keys = [1..bound]
     evens = [2,4..bound]
     odds = [1,3..bound]

containers-tests/benchmarks/Set.hs

@@ -14,7 +14,7 @@ main = do
         s_odd = S.fromAscList elems_odd :: S.Set Int
         strings_s = S.fromList strings
     evaluate $ rnf [s, s_even, s_odd]
-    evaluate $ rnf elems_rev
+    evaluate $ rnf [elems_rev, elems_asc, elems_desc]
     defaultMain
         [ bench "member" $ whnf (member elems) s
         , bench "insert" $ whnf (ins elems) S.empty
@@ -33,9 +33,12 @@ main = do
         , bench "intersection" $ whnf (S.intersection s) s_even
         , bench "fromList" $ whnf S.fromList elems
         , bench "fromList-desc" $ whnf S.fromList (reverse elems)
-        , bench "fromAscList" $ whnf S.fromAscList elems
+        , bench "fromAscList" $ whnf S.fromAscList elems_asc
         , bench "fromDistinctAscList" $ whnf S.fromDistinctAscList elems
         , bench "fromDistinctAscList:fusion" $ whnf (\n -> S.fromDistinctAscList [1..n]) bound
+        , bench "fromAscList:fusion" $ whnf (\n -> S.fromAscList [i `div` 2 | i <- [1..n]]) bound
+        , bench "fromDescList" $ whnf S.fromDescList elems_desc
+        , bench "fromDescList:fusion" $ whnf (\n -> S.fromDescList [i `div` 2 | i <- [n,n-1..1]]) bound
         , bench "fromDistinctDescList" $ whnf S.fromDistinctDescList elems_rev
         , bench "fromDistinctDescList:fusion" $ whnf (\n -> S.fromDistinctDescList [n,n-1..1]) bound
         , bench "disjoint:false" $ whnf (S.disjoint s) s_even
@@ -62,6 +65,8 @@ main = do
     elems_even = [2,4..bound]
     elems_odd = [1,3..bound]
     elems_rev = reverse elems
+    elems_asc = map (`div` 2) [1..bound]
+    elems_desc = map (`div` 2) [bound,bound-1..1]
     strings = map show elems
 
 member :: [Int] -> S.Set Int -> Int

containers-tests/tests/map-properties.hs

@@ -28,6 +28,8 @@ import qualified Prelude
 
 import Data.List (nub,sort)
 import qualified Data.List as List
+import Data.List.NonEmpty (NonEmpty(..))
+import qualified Data.List.NonEmpty as NE
 import qualified Data.Set as Set
 import Test.Tasty
 import Test.Tasty.HUnit
@@ -184,9 +186,13 @@ main = defaultMain $ testGroup "map-properties"
          , testProperty "unionWithKeyMerge"   prop_unionWithKeyMerge
          , testProperty "mergeWithKey model"   prop_mergeWithKeyModel
          , testProperty "mergeA effects"       prop_mergeA_effects
-         , testProperty "fromAscList"          prop_ordered
+         , testProperty "fromAscList"          prop_fromAscList
+         , testProperty "fromAscListWith"      prop_fromAscListWith
+         , testProperty "fromAscListWithKey"   prop_fromAscListWithKey
          , testProperty "fromDistinctAscList"  prop_fromDistinctAscList
-         , testProperty "fromDescList"         prop_rev_ordered
+         , testProperty "fromDescList"         prop_fromDescList
+         , testProperty "fromDescListWith"     prop_fromDescListWith
+         , testProperty "fromDescListWithKey"  prop_fromDescListWithKey
          , testProperty "fromDistinctDescList" prop_fromDistinctDescList
          , testProperty "fromList then toList" prop_list
          , testProperty "toDescList"           prop_descList
@@ -1222,31 +1228,68 @@ prop_mergeA_effects xs ys
 
 ----------------------------------------------------------------
 
-prop_ordered :: Property
-prop_ordered
-  = forAll (choose (5,100)) $ \n ->
-    let xs = [(x,()) | x <- [0..n::Int]]
-    in fromAscList xs == fromList xs
-
-prop_rev_ordered :: Property
-prop_rev_ordered
-  = forAll (choose (5,100)) $ \n ->
-    let xs = [(x,()) | x <- [0..n::Int]]
-    in fromDescList (reverse xs) == fromList xs
+-- fromAscListWith, fromAscListWithKey, fromDescListWith, fromDescListWithKey
+-- all effectively perform a left fold over adjacent elements in the input list
+-- using some function as long as the keys are equal.
+--
+-- The property tests for these functions compare the result against the
+-- sequence we would get if we used NE.groupBy instead. We use Magma to check
+-- the fold direction (left, not right) and the order of arguments to the fold
+-- function (new then old).
+
+data Magma a
+  = Inj a
+  | Magma a :* Magma a
+  deriving (Eq, Show)
+
+groupByK :: Eq k => [(k, a)] -> [(k, NonEmpty a)]
+groupByK =
+    List.map (\ys -> (fst (NE.head ys), NE.map snd ys)) .
+    NE.groupBy ((==) `on` fst)
 
 prop_list :: [Int] -> Bool
 prop_list xs = (sort (nub xs) == [x | (x,()) <- toList (fromList [(x,()) | x <- xs])])
 
 prop_descList :: [Int] -> Bool
 prop_descList xs = (reverse (sort (nub xs)) == [x | (x,()) <- toDescList (fromList [(x,()) | x <- xs])])
 
+prop_fromDescList :: [(Int, A)] -> Property
+prop_fromDescList xs =
+    valid t .&&.
+    toList t === reverse nub_last_down_sort_xs
+  where
+    down_sort_xs = List.sortBy (comparing (Down . fst)) xs
+    t = fromDescList down_sort_xs
+    nub_last_down_sort_xs = List.map NE.last $ NE.groupBy ((==) `on` fst) down_sort_xs
+
+prop_fromDescListWith :: [(Int, A)] -> Property
+prop_fromDescListWith ys =
+    valid t .&&.
+    toList t === reverse combined_down_sort_xs
+  where
+    xs = [(kx, Inj x) | (kx,x) <- ys]
+    down_sort_xs = List.sortBy (comparing (Down . fst)) xs
+    t = fromDescListWith (:*) down_sort_xs
+    combined_down_sort_xs = [(kx, Foldable.foldl1 (flip (:*)) x) | (kx,x) <- groupByK down_sort_xs]
+
+prop_fromDescListWithKey :: [(Int, A)] -> Property
+prop_fromDescListWithKey ys =
+    valid t .&&.
+    toList t === reverse combined_down_sort_xs
+  where
+    xs = [(kx, Inj (Left x)) | (kx,x) <- ys]
+    down_sort_xs = List.sortBy (comparing (Down . fst)) xs
+    t = fromDescListWithKey (\kx (Inj (Left x)) acc -> Inj (Right (kx,x)) :* acc) down_sort_xs
+    combined_down_sort_xs = [ (kx, Foldable.foldl1 (\acc (Inj (Left x)) -> Inj (Right (kx,x)) :* acc) xs)
+                            | (kx,xs) <- groupByK down_sort_xs ]
+
 prop_fromDistinctDescList :: [(Int, A)] -> Property
 prop_fromDistinctDescList xs =
     valid t .&&.
-    toList t === nub_sort_xs
+    toList t === reverse nub_down_sort_xs
   where
-    t = fromDistinctDescList (reverse nub_sort_xs)
-    nub_sort_xs = List.map List.head $ List.groupBy ((==) `on` fst) $ List.sortBy (comparing fst) xs
+    t = fromDistinctDescList nub_down_sort_xs
+    nub_down_sort_xs = List.map NE.last $ NE.groupBy ((==) `on` fst) $ List.sortBy (comparing (Down . fst)) xs
 
 prop_ascDescList :: [Int] -> Bool
 prop_ascDescList xs = toAscList m == reverse (toDescList m)
@@ -1259,6 +1302,36 @@ prop_fromList xs
            t == List.foldr (uncurry insert) empty (zip xs xs)
   where sort_xs = sort xs
 
+prop_fromAscList :: [(Int, A)] -> Property
+prop_fromAscList xs =
+    valid t .&&.
+    toList t === nub_last_sort_xs
+  where
+    sort_xs = List.sortBy (comparing fst) xs
+    t = fromAscList sort_xs
+    nub_last_sort_xs = List.map NE.last $ NE.groupBy ((==) `on` fst) sort_xs
+
+prop_fromAscListWith :: [(Int, A)] -> Property
+prop_fromAscListWith ys =
+    valid t .&&.
+    toList t === combined_sort_xs
+  where
+    xs = [(kx, Inj x) | (kx,x) <- ys]
+    sort_xs = List.sortBy (comparing fst) xs
+    t = fromAscListWith (:*) sort_xs
+    combined_sort_xs = [(kx, Foldable.foldl1 (flip (:*)) x) | (kx,x) <- groupByK sort_xs]
+
+prop_fromAscListWithKey :: [(Int, A)] -> Property
+prop_fromAscListWithKey ys =
+    valid t .&&.
+    toList t === combined_sort_xs
+  where
+    xs = [(kx, Inj (Left x)) | (kx,x) <- ys]
+    sort_xs = List.sortBy (comparing fst) xs
+    t = fromAscListWithKey (\kx (Inj (Left x)) acc -> Inj (Right (kx,x)) :* acc) sort_xs
+    combined_sort_xs = [ (kx, Foldable.foldl1 (\acc (Inj (Left x)) -> Inj (Right (kx,x)) :* acc) xs)
+                       | (kx,xs) <- groupByK sort_xs ]
+
 prop_fromDistinctAscList :: [(Int, A)] -> Property
 prop_fromDistinctAscList xs =
     valid t .&&.

containers-tests/tests/set-properties.hs

@@ -1,6 +1,6 @@
 {-# LANGUAGE CPP #-}
 import qualified Data.IntSet as IntSet
-import Data.List (nub,sort)
+import Data.List (nub, sort, sortBy)
 import qualified Data.List as List
 import Data.Monoid (mempty)
 import Data.Maybe
@@ -15,6 +15,7 @@ import Control.Monad.Trans.Class
 import Control.Monad (liftM, liftM3)
 import Data.Functor.Identity
 import Data.Foldable (all)
+import Data.Ord (Down(..), comparing)
 import Control.Applicative (liftA2)
 
 #if __GLASGOW_HASKELL__ >= 806
@@ -67,8 +68,9 @@ main = defaultMain $ testGroup "set-properties"
                    , testProperty "prop_DescList" prop_DescList
                    , testProperty "prop_AscDescList" prop_AscDescList
                    , testProperty "prop_fromList" prop_fromList
+                   , testProperty "prop_fromAscList" prop_fromAscList
                    , testProperty "prop_fromDistinctAscList" prop_fromDistinctAscList
-                   , testProperty "prop_fromListDesc" prop_fromListDesc
+                   , testProperty "prop_fromDescList" prop_fromDescList
                    , testProperty "prop_fromDistinctDescList" prop_fromDistinctDescList
                    , testProperty "prop_isProperSubsetOf" prop_isProperSubsetOf
                    , testProperty "prop_isProperSubsetOf2" prop_isProperSubsetOf2
@@ -515,10 +517,18 @@ prop_AscDescList xs = toAscList s == reverse (toDescList s)
 
 prop_fromList :: [Int] -> Property
 prop_fromList xs =
-           t === fromAscList sort_xs .&&.
+           valid t .&&.
            t === List.foldr insert empty xs
   where t = fromList xs
-        sort_xs = sort xs
+
+prop_fromAscList :: [Int] -> Property
+prop_fromAscList xs =
+    valid t .&&.
+    toList t === nub_sort_xs
+  where
+    sort_xs = sort xs
+    t = fromAscList sort_xs
+    nub_sort_xs = List.map List.head $ List.group sort_xs
 
 prop_fromDistinctAscList :: [Int] -> Property
 prop_fromDistinctAscList xs =
@@ -528,22 +538,22 @@ prop_fromDistinctAscList xs =
     t = fromDistinctAscList nub_sort_xs
     nub_sort_xs = List.map List.head $ List.group $ sort xs
 
-prop_fromListDesc :: [Int] -> Property
-prop_fromListDesc xs =
-           t === fromDescList sort_xs .&&.
-           t === fromDistinctDescList nub_sort_xs .&&.
-           t === List.foldr insert empty xs
-  where t = fromList xs
-        sort_xs = reverse (sort xs)
-        nub_sort_xs = List.map List.head $ List.group sort_xs
+prop_fromDescList :: [Int] -> Property
+prop_fromDescList xs =
+    valid t .&&.
+    toList t === reverse nub_down_sort_xs
+  where
+    down_sort_xs = sortBy (comparing Down) xs
+    t = fromDescList down_sort_xs
+    nub_down_sort_xs = List.map List.head $ List.group down_sort_xs
 
 prop_fromDistinctDescList :: [Int] -> Property
 prop_fromDistinctDescList xs =
     valid t .&&.
-    toList t === nub_sort_xs
+    toList t === reverse nub_down_sort_xs
   where
-    t = fromDistinctDescList (reverse nub_sort_xs)
-    nub_sort_xs = List.map List.head $ List.group $ sort xs
+    t = fromDistinctDescList nub_down_sort_xs
+    nub_down_sort_xs = List.map List.head $ List.group $ sortBy (comparing Down) xs
 
 {--------------------------------------------------------------------
   Set operations are like IntSet operations