Add RULES for nub functions (#517)
authorDavid Feuer <David.Feuer@gmail.com>
Sun, 28 Jan 2018 01:10:50 +0000 (20:10 -0500)
committerGitHub <noreply@github.com>
Sun, 28 Jan 2018 01:10:50 +0000 (20:10 -0500)
* Add rewrite rules to allow the nub functions in `ListUtils`
  to participate in fold/build fusion.

* For the sake of simplicity, define `nubOrd` and `nubInt` in terms
  of `nubOrdOn` and `nubIntOn`.

Data/Containers/ListUtils.hs
tests/listutils-properties.hs

index f8885f3..81b7da3 100644 (file)
@@ -1,6 +1,7 @@
 {-# LANGUAGE CPP #-}
-#if !defined(TESTING) && __GLASGOW_HASKELL__ >= 703
-{-# LANGUAGE Safe #-}
+{-# LANGUAGE BangPatterns #-}
+#if __GLASGOW_HASKELL__ >= 703
+{-# LANGUAGE Trustworthy #-}
 #endif
 
 -----------------------------------------------------------------------------
@@ -21,45 +22,160 @@ module Data.Containers.ListUtils (
        nubIntOn
        ) where
 
+import Data.Set (Set)
 import qualified Data.Set as Set
 import qualified Data.IntSet as IntSet
+import Data.IntSet (IntSet)
+#ifdef __GLASGOW_HASKELL__
+import GHC.Exts ( build )
+#endif
 
--- | /O(n log n)/. The 'nubOrd' function removes duplicate elements from a list.
--- In particular, it keeps only the first occurrence of each element. By using a 'Set' internally
--- it has better asymptotics than the standard 'nub' function.
-nubOrd :: (Ord a) => [a] -> [a]
-nubOrd = go Set.empty
-  where
-    go _ [] = []
-    go s (x:xs) = if x `Set.member` s then go s xs
-                                      else x : go (Set.insert x s) xs
+-- *** Ord-based nubbing ***
 
--- | The `nubOrdOn` function behaves just like `nubOrd` except it performs comparisons not on the
--- original datatype, but a user-specified projection from that datatype.
-nubOrdOn :: (Ord b) => (a -> b) -> [a] -> [a]
-nubOrdOn f = go Set.empty
+
+-- | \( O(n \log n \). The @nubOrd@ function removes duplicate elements from a list.
+-- In particular, it keeps only the first occurrence of each element. By using a
+-- 'Set' internally it has better asymptotics than the standard 'Data.List.nub'
+-- function.
+--
+-- ==== Strictness
+--
+-- @nubOrd@ is strict in the elements of the list.
+--
+-- ==== Efficiency note
+--
+-- When applicable, it is almost always better to use 'nubInt' or 'nubIntOn' instead
+-- of this function. For example, the best way to nub a list of characters is
+--
+-- @ nubIntOn fromEnum xs @
+nubOrd :: Ord a => [a] -> [a]
+nubOrd = nubOrdOn id
+{-# INLINE nubOrd #-}
+
+-- | The @nubOrdOn@ function behaves just like 'nubOrd' except it performs
+-- comparisons not on the original datatype, but a user-specified projection
+-- from that datatype.
+--
+-- ==== Strictness
+--
+-- @nubOrdOn@ is strict in the values of the function applied to the
+-- elements of the list.
+nubOrdOn :: Ord b => (a -> b) -> [a] -> [a]
+-- For some reason we need to write an explicit lambda here to allow this
+-- to inline when only applied to a function.
+nubOrdOn f = \xs -> nubOrdOnExcluding f Set.empty xs
+{-# INLINE nubOrdOn #-}
+
+-- Splitting nubOrdOn like this means that we don't have to worry about
+-- matching specifically on Set.empty in the rewrite-back rule.
+nubOrdOnExcluding :: Ord b => (a -> b) -> Set b -> [a] -> [a]
+nubOrdOnExcluding f = go
   where
     go _ [] = []
-    go s (x:xs) = let fx = f x
-                  in if fx `Set.member` s then go s xs
-                                          else x : go (Set.insert fx s) xs
+    go s (x:xs)
+      | fx `Set.member` s = go s xs
+      | otherwise = x : go (Set.insert fx s) xs
+      where !fx = f x
+
+#ifdef __GLASGOW_HASKELL__
+-- We want this inlinable to specialize to the necessary Ord instance.
+{-# INLINABLE [1] nubOrdOnExcluding #-}
+
+{-# RULES
+-- Rewrite to a fusible form.
+"nubOrdOn" [~1] forall f as s. nubOrdOnExcluding  f s as =
+  build (\c n -> foldr (nubOrdOnFB f c) (constNubOn n) as s)
+
+-- Rewrite back to a plain form
+"nubOrdOnList" [1] forall f as s.
+    foldr (nubOrdOnFB f (:)) (constNubOn []) as s =
+       nubOrdOnExcluding f s as
+ #-}
+
+nubOrdOnFB :: Ord b
+           => (a -> b)
+           -> (a -> r -> r)
+           -> a
+           -> (Set b -> r)
+           -> Set b
+           -> r
+nubOrdOnFB f c x r s
+  | fx `Set.member` s = r s
+  | otherwise = x `c` r (Set.insert fx s)
+  where !fx = f x
+{-# INLINABLE [0] nubOrdOnFB #-}
 
--- | /O(n min(n,W))/. The 'nubInt' function removes duplicate elements from a list.
--- In particular, it keeps only the first occurrence of each element. By using an 'IntSet' internally
--- it has better asymptotics than the standard 'nub' function.
+constNubOn :: a -> b -> a
+constNubOn x _ = x
+{-# INLINE [0] constNubOn #-}
+#endif
+
+
+-- *** Int-based nubbing ***
+
+
+-- | \( O(n \min(n,W)) \). The @nubInt@ function removes duplicate 'Int'
+-- values from a list. In particular, it keeps only the first occurrence
+-- of each element. By using an 'IntSet' internally, it attains better
+-- asymptotics than the standard 'Data.List.nub' function.
+--
+-- See also 'nubIntOn', a more widely applicable generalization.
+--
+-- ==== Strictness
+--
+-- @nubInt@ is strict in the elements of the list.
 nubInt :: [Int] -> [Int]
-nubInt = go IntSet.empty
-  where
-    go _ [] = []
-    go s (x:xs) = if x `IntSet.member` s then go s xs
-                                         else x : go (IntSet.insert x s) xs
+nubInt = nubIntOn id
+{-# INLINE nubInt #-}
 
--- | The `nubIntOn` function behaves just like 'nubInt' except it performs comparisons not on the
--- original datatype, but a user-specified projection from that datatype to 'Int'.
+-- | The @nubIntOn@ function behaves just like 'nubInt' except it performs
+-- comparisons not on the original datatype, but a user-specified projection
+-- from that datatype.
+--
+-- ==== Strictness
+--
+-- @nubIntOn@ is strict in the values of the function applied to the
+-- elements of the list.
 nubIntOn :: (a -> Int) -> [a] -> [a]
-nubIntOn f = go IntSet.empty
+-- For some reason we need to write an explicit lambda here to allow this
+-- to inline when only applied to a function.
+nubIntOn f = \xs -> nubIntOnExcluding f IntSet.empty xs
+{-# INLINE nubIntOn #-}
+
+-- Splitting nubIntOn like this means that we don't have to worry about
+-- matching specifically on IntSet.empty in the rewrite-back rule.
+nubIntOnExcluding :: (a -> Int) -> IntSet -> [a] -> [a]
+nubIntOnExcluding f = go
   where
     go _ [] = []
-    go s (x:xs) = let fx = f x
-                  in if fx `IntSet.member` s then go s xs
-                                             else x : go (IntSet.insert fx s) xs
+    go s (x:xs)
+      | fx `IntSet.member` s = go s xs
+      | otherwise = x : go (IntSet.insert fx s) xs
+      where !fx = f x
+
+#ifdef __GLASGOW_HASKELL__
+-- We don't mark this INLINABLE because it doesn't seem obviously useful
+-- to inline it anywhere; the elements the function operates on are actually
+-- pulled from a list and installed in a list; the situation is very different
+-- when fusion occurs. In this case, we let GHC make the call.
+{-# NOINLINE [1] nubIntOnExcluding #-}
+
+{-# RULES
+"nubIntOn" [~1] forall f as s. nubIntOnExcluding  f s as =
+  build (\c n -> foldr (nubIntOnFB f c) (constNubOn n) as s)
+"nubIntOnList" [1] forall f as s. foldr (nubIntOnFB f (:)) (constNubOn []) as s =
+  nubIntOnExcluding f s as
+ #-}
+
+nubIntOnFB :: (a -> Int)
+           -> (a -> r -> r)
+           -> a
+           -> (IntSet -> r)
+           -> IntSet
+           -> r
+nubIntOnFB f c x r s
+  | fx `IntSet.member` s = r s
+  | otherwise = x `c` r (IntSet.insert fx s)
+  where !fx = f x
+{-# INLINABLE [0] nubIntOnFB #-}
+#endif
index 658c2d5..055c626 100644 (file)
@@ -4,24 +4,57 @@ import Data.List (nub, nubBy)
 import Data.Containers.ListUtils
 import Test.Framework
 import Test.Framework.Providers.QuickCheck2
+import Test.QuickCheck (Property, (===))
+import Test.QuickCheck.Function (Fun, apply)
+import Test.QuickCheck.Poly (A, OrdA, B, OrdB, C)
 
 main :: IO ()
 main = defaultMain
-         [   testProperty "nubOrd"     prop_nubOrd
-           , testProperty "nubOrdOn"   prop_nubOrdOn
-           , testProperty "nubInt"     prop_nubInt
-           , testProperty "nubIntOn"     prop_nubIntOn
+         [ testProperty "nubOrd" prop_nubOrd
+         , testProperty "nubOrdOn" prop_nubOrdOn
+         , testProperty "nubOrdOn fusion" prop_nubOrdOnFusion
+         , testProperty "nubInt" prop_nubInt
+         , testProperty "nubIntOn" prop_nubIntOn
+         , testProperty "nubIntOn fusion" prop_nubIntOnFusion
          ]
 
 
-prop_nubOrd :: [Int] -> Bool
-prop_nubOrd xs = nubOrd xs == nub xs
+prop_nubOrd :: [OrdA] -> Property
+prop_nubOrd xs = nubOrd xs === nub xs
 
-prop_nubInt :: [Int] -> Bool
-prop_nubInt xs = nubInt xs == nub xs
+prop_nubInt :: [Int] -> Property
+prop_nubInt xs = nubInt xs === nub xs
 
-prop_nubOrdOn :: [(Int,Int)] -> Bool
-prop_nubOrdOn xs = nubOrdOn snd xs == nubBy (\x y -> snd x == snd y) xs
+prop_nubOrdOn :: Fun A OrdB -> [A] -> Property
+prop_nubOrdOn f' xs =
+  nubOrdOn f xs === nubBy (\x y -> f x == f y) xs
+  where f = apply f'
 
-prop_nubIntOn :: [(Int,Int)] -> Bool
-prop_nubIntOn xs = nubIntOn snd xs == nubBy (\x y -> snd x == snd y) xs
+prop_nubIntOn :: Fun A Int -> [A] -> Property
+prop_nubIntOn f' xs =
+  nubIntOn f xs === nubBy (\x y -> f x == f y) xs
+  where f = apply f'
+
+prop_nubOrdOnFusion :: Fun B C
+                    -> Fun B OrdB
+                    -> Fun A B
+                    -> [A] -> Property
+prop_nubOrdOnFusion f' g' h' xs =
+  (map f . nubOrdOn g . map h $ xs)
+    === (map f . nubBy (\x y -> g x == g y) . map h $ xs)
+  where
+    f = apply f'
+    g = apply g'
+    h = apply h'
+
+prop_nubIntOnFusion :: Fun B C
+                    -> Fun B Int
+                    -> Fun A B
+                    -> [A] -> Property
+prop_nubIntOnFusion f' g' h' xs =
+  (map f . nubIntOn g . map h $ xs)
+    === (map f . nubBy (\x y -> g x == g y) . map h $ xs)
+  where
+    f = apply f'
+    g = apply g'
+    h = apply h'