ansatz for improved powerset

Author: Johannes Waldmann

containers/src/Data/Set/Internal.hs

@@ -247,6 +247,9 @@ import Data.Functor.Identity (Identity)
 import qualified Data.Foldable as Foldable
 import Control.DeepSeq (NFData(rnf))
 
+import qualified Data.Array as A
+import Data.Bits
+
 import Utils.Containers.Internal.StrictPair
 import Utils.Containers.Internal.PtrEquality
 
@@ -1824,6 +1827,58 @@ splitRoot orig =
 --
 -- @since 0.5.11
 
+
+powerSet :: Set a -> Set (Set a)
+powerSet xs =
+  let !w = length xs
+      !u = A.listArray (0, w-1) $ toList xs
+      !v = generateA (0, 2^w -1) $ \ m ->
+        if m == 0
+        then Tip
+        else let ST up med lo = splitBits m
+             in  bin (u A.! (w - 1 - med))
+                     (v A.! up) (v A.! lo)
+      make !begin !s =
+        if s == 0 then Tip
+        else let !sl = div (s-1) 2; !sr = s - 1 - sl
+             in  bin (v A.! forw w (begin + sl))
+                     (make begin sl)
+                     (make (begin + sl+1) sr)
+  in  make 0 (2^w)
+
+generateA :: A.Ix i => (i,i) -> (i -> a) -> A.Array i a
+generateA bnd f = A.listArray bnd $ fmap f $ A.range bnd
+
+forw :: Int -> Int -> Int
+forw !width !n = fr_go_branch (bit $ width-1) n 0
+
+fr_go_branch :: Int -> Int -> Int -> Int
+fr_go_branch !topmask !n !set =
+  if n == 0 then set
+  else if 0 == ((n-1) .&. topmask)
+  then fr_go_branch (shiftR topmask 1)
+         (n-1) (set .|. topmask)
+  else fr_go_branch (shiftR topmask 1)
+         (n .&. complement topmask) set
+
+data StrictTriple = ST !Int !Int !Int
+
+-- | return bitmask for upper half,
+-- index of middle bit, bitmask for lower half
+splitBits :: Int -> StrictTriple
+splitBits m | m > 0 =
+  let go 0 !x = x; go k !x = go (k-1) (clearLowest x)
+      !up_med = go (div (popCount m) 2) m
+      !lo = xor m up_med
+      !up = clearLowest up_med
+      !med = xor up_med up
+  in  ST up (countTrailingZeros med) lo
+
+clearLowest :: Int -> Int
+clearLowest m | m > 0 = m .&. (m-1)
+
+
+
 -- Proof of complexity: step executes n times. At the ith step,
 -- "insertMin x `mapMonotonic` pxs" takes O(2^i log i) time since pxs has size
 -- 2^i - 1 and we insertMin into its elements which are sets of size <= i.
@@ -1834,8 +1889,8 @@ splitRoot orig =
 -- = O(log n * \sum_{i=1}^{n-1} 2^i)
 -- = O(2^n log n)
 
-powerSet :: Set a -> Set (Set a)
-powerSet xs0 = insertMin empty (foldr' step Tip xs0) where
+powerSet_orig :: Set a -> Set (Set a)
+powerSet_orig xs0 = insertMin empty (foldr' step Tip xs0) where
   step x pxs = insertMin (singleton x) (insertMin x `mapMonotonic` pxs) `glue` pxs
 
 -- | \(O(nm)\). Calculate the Cartesian product of two sets.