diff --git a/dense-int-set.cabal b/dense-int-set.cabal --- a/dense-int-set.cabal +++ b/dense-int-set.cabal @@ -1,5 +1,5 @@ name: dense-int-set -version: 0.2 +version: 0.3 category: Data synopsis: Dense int-set homepage: https://github.com/metrix-ai/dense-int-set @@ -24,7 +24,23 @@ base >=4.11 && <5, cereal >=0.5.5 && <0.6, cereal-vector >=0.2.0.1 && <0.3, - deferred-folds >=0.9.7 && <0.10, + deferred-folds >=0.9.9 && <0.10, hashable >=1 && <2, vector >=0.12 && <0.13, vector-algorithms >=0.7.0.4 && <0.8 + +test-suite test + type: exitcode-stdio-1.0 + hs-source-dirs: test + default-extensions: Arrows, BangPatterns, ConstraintKinds, DataKinds, DefaultSignatures, DeriveDataTypeable, DeriveFoldable, DeriveFunctor, DeriveGeneric, DeriveTraversable, EmptyDataDecls, FlexibleContexts, FlexibleInstances, FunctionalDependencies, GADTs, GeneralizedNewtypeDeriving, LambdaCase, LiberalTypeSynonyms, MagicHash, MultiParamTypeClasses, MultiWayIf, NoImplicitPrelude, NoMonomorphismRestriction, OverloadedLists, OverloadedStrings, PatternGuards, PatternSynonyms, ParallelListComp, QuasiQuotes, RankNTypes, RecordWildCards, ScopedTypeVariables, StandaloneDeriving, TemplateHaskell, TupleSections, TypeApplications, TypeFamilies, TypeOperators, UnboxedTuples + default-language: Haskell2010 + main-is: + Main.hs + build-depends: + dense-int-set, + QuickCheck >=2.8.1 && <3, + quickcheck-instances >=0.3.11 && <0.4, + rerebase <2, + tasty >=0.12 && <2, + tasty-hunit >=0.9 && <0.11, + tasty-quickcheck >=0.9 && <0.11 diff --git a/library/DenseIntSet.hs b/library/DenseIntSet.hs --- a/library/DenseIntSet.hs +++ b/library/DenseIntSet.hs @@ -6,11 +6,13 @@ foldable, topValueIndices, filteredIndices, + invert, -- *** Composition - intersection, - union, + intersections, + unions, -- ** Accessors - size, + capacity, + population, lookup, -- *** Vectors presentElementsVector, @@ -20,10 +22,6 @@ presentElementsUnfoldr, absentElementsUnfoldr, vectorElementsUnfoldr, - -- * Composition - DenseIntSetComposition, - compose, - composeList, ) where @@ -50,7 +48,7 @@ the instances are provided for "DenseIntSetComposition", which is open for interpretation of how to compose. -} -newtype DenseIntSet = DenseIntSet (UnboxedVector Word64) +data DenseIntSet = DenseIntSet {-# UNPACK #-} !Int {-# UNPACK #-} !(UnboxedVector Word64) deriving instance Eq DenseIntSet @@ -59,14 +57,16 @@ instance Show DenseIntSet where show = show . toList -deriving instance Serialize DenseIntSet +deriving instance Generic DenseIntSet +instance Serialize DenseIntSet + instance Hashable DenseIntSet where - hashWithSalt salt (DenseIntSet vec) = hashWithSalt salt (GenericVector.toList vec) + hashWithSalt salt (DenseIntSet capacity vec) = hashWithSalt (hashWithSalt salt capacity) (GenericVector.toList vec) instance IsList DenseIntSet where type Item DenseIntSet = Int - fromList list = foldable (foldl' max 0 list) list + fromList list = foldable (succ (foldl' max 0 list)) list toList = toList . presentElementsUnfoldr @@ -74,41 +74,46 @@ ------------------------- {-| -Given a maximum int, construct from a foldable of ints, which are smaller or equal to it. +Given a maximum int, construct from a foldable of ints, which are smaller or equal to it and larger than 0. It is your responsibility to ensure that the values match this contract. -} foldable :: Foldable foldable => Int -> foldable Int -> DenseIntSet -foldable size foldable = let - !wordsAmount = divCeiling size 64 - in DenseIntSet $ runST $ do +foldable capacity foldable = let + !wordsAmount = divCeiling capacity 64 + in DenseIntSet capacity $ runST $ do indexSetMVec <- MutableGenericVector.new wordsAmount forM_ foldable $ \ index -> let (wordIndex, bitIndex) = divMod index 64 in MutableGenericVector.modify indexSetMVec (flip setBit bitIndex) wordIndex GenericVector.unsafeFreeze indexSetMVec -{-| -Interpret a composition as an intersection of sets. --} -intersection :: DenseIntSetComposition -> DenseIntSet -intersection = zipWords (.&.) maxBound +{-| Intersect multiple sets -} +intersections :: [DenseIntSet] -> DenseIntSet +intersections list = if null list + then DenseIntSet 0 mempty + else let + cap = foldl1' min (fmap capacity list) + vecs = fmap (\ (DenseIntSet _ vec) -> vec) list + in compositions (.&.) maxBound cap vecs -{-| -Interpret a composition as a union of sets. --} -union :: DenseIntSetComposition -> DenseIntSet -union = zipWords (.|.) 0 +{-| Unite multiple sets -} +unions :: [DenseIntSet] -> DenseIntSet +unions list = let + cap = foldl' max 0 (fmap capacity list) + vecs = fmap (\ (DenseIntSet _ vec) -> vec) list + in compositions (.|.) 0 cap vecs -zipWords :: (Word64 -> Word64 -> Word64) -> Word64 -> DenseIntSetComposition -> DenseIntSet -zipWords append empty (DenseIntSetComposition minLength vecs) = DenseIntSet $ runST $ let - wordIndexUnfoldr = Unfoldr.intsInRange 0 (pred minLength) +compositions :: (Word64 -> Word64 -> Word64) -> Word64 -> Int -> [UnboxedVector Word64] -> DenseIntSet +compositions append empty capacity vecs = DenseIntSet capacity $ runST $ let + wordsAmount = divCeiling capacity 64 + wordIndexUnfoldr = Unfoldr.intsInRange 0 (pred wordsAmount) vecVec = Vector.fromList vecs vecUnfoldr = Unfoldr.foldable vecVec - wordUnfoldrAt wordIndex = fmap (flip GenericVector.unsafeIndex wordIndex) vecUnfoldr + wordUnfoldrAt wordIndex = vecUnfoldr >>= Unfoldr.foldable . flip (GenericVector.!?) wordIndex finalWordAt = foldr append empty . wordUnfoldrAt in do - indexSetMVec <- MutableGenericVector.new minLength + indexSetMVec <- MutableGenericVector.new wordsAmount forM_ wordIndexUnfoldr $ \ index -> MutableGenericVector.unsafeWrite indexSetMVec index (finalWordAt index) GenericVector.unsafeFreeze indexSetMVec @@ -128,39 +133,60 @@ (_, index) <- MutableGenericVector.unsafeRead pairMVec pairIndex let (wordIndex, bitIndex) = divMod index 64 MutableGenericVector.modify indexSetMVec (flip setBit bitIndex) wordIndex - DenseIntSet <$> GenericVector.unsafeFreeze indexSetMVec + DenseIntSet valuesAmount <$> GenericVector.unsafeFreeze indexSetMVec {-| Select the indices of vector elements, which match the predicate. -} filteredIndices :: GenericVector.Vector vector a => (a -> Bool) -> vector a -> DenseIntSet -filteredIndices predicate valueVec = DenseIntSet $ let +filteredIndices predicate valueVec = let valuesAmount = GenericVector.length valueVec wordsAmount = divCeiling valuesAmount 64 indexUnfoldr = do (index, a) <- Unfoldr.vectorWithIndices valueVec guard (predicate a) return (divMod index 64) - in runST $ do + in DenseIntSet valuesAmount $ runST $ do indexSetMVec <- MutableGenericVector.new wordsAmount forM_ indexUnfoldr $ \ (wordIndex, bitIndex) -> MutableGenericVector.modify indexSetMVec (flip setBit bitIndex) wordIndex GenericVector.unsafeFreeze indexSetMVec +{-| +Invert the set. +-} +invert :: DenseIntSet -> DenseIntSet +invert (DenseIntSet capacity vec) = DenseIntSet capacity $ let + invertedVec = GenericVector.map complement vec + (lastWordIndex, claimedBitsOfLastWord) = divMod capacity 64 + unclaimedBitsOfLastWord = 64 - claimedBitsOfLastWord + in if claimedBitsOfLastWord == 0 + then invertedVec + else runST $ do + mv <- GenericVector.unsafeThaw invertedVec + flip (MutableGenericVector.modify mv) lastWordIndex $ flip shiftR unclaimedBitsOfLastWord . flip shiftL unclaimedBitsOfLastWord + GenericVector.unsafeFreeze mv + -- * Accessors ------------------------- {-| +/O(1)/. Size of the value space. +-} +capacity :: DenseIntSet -> Int +capacity (DenseIntSet x _) = x + +{-| /O(log n)/. Count the amount of present elements in the set. -} -size :: DenseIntSet -> Int -size (DenseIntSet vec) = getSum (foldMap (Sum . popCount) (Unfoldr.vector vec)) +population :: DenseIntSet -> Int +population (DenseIntSet _ vec) = getSum (foldMap (Sum . popCount) (Unfoldr.vector vec)) {-| /O(1)/. Check whether an int is a member of the set. -} lookup :: Int -> DenseIntSet -> Bool -lookup index (DenseIntSet vec) = let +lookup index (DenseIntSet _ vec) = let (wordIndex, bitIndex) = divMod index 64 in vec GenericVector.!? wordIndex & maybe False (flip testBit bitIndex) @@ -173,10 +199,10 @@ -} presentElementsVector :: GenericVector.Vector vector element => DenseIntSet -> (Int -> element) -> vector element presentElementsVector intSet intToIndex = let - sizeVal = size intSet + vectorPop = population intSet unfoldr = Unfoldr.zipWithIndex (presentElementsUnfoldr intSet) in runST $ do - mv <- MutableGenericVector.unsafeNew sizeVal + mv <- MutableGenericVector.unsafeNew vectorPop forM_ unfoldr $ \ (index, element) -> MutableGenericVector.unsafeWrite mv index (intToIndex element) GenericVector.unsafeFreeze mv @@ -184,13 +210,11 @@ Construct a vector, which maps from the original ints into their indices amongst the ones present in the set. -} indexVector :: GenericVector.Vector vector (Maybe index) => DenseIntSet -> (Int -> index) -> vector (Maybe index) -indexVector set@(DenseIntSet setVec) intToIndex = let - indexVecSize = GenericVector.length setVec * 64 - in runST $ do - v <- MutableGenericVector.replicate indexVecSize Nothing - forM_ (Unfoldr.zipWithIndex (presentElementsUnfoldr set)) $ \ (index, element) -> do - MutableGenericVector.unsafeWrite v element (Just (intToIndex index)) - GenericVector.unsafeFreeze v +indexVector set@(DenseIntSet capacity setVec) intToIndex = runST $ do + v <- MutableGenericVector.replicate capacity Nothing + forM_ (Unfoldr.zipWithIndex (presentElementsUnfoldr set)) $ \ (index, element) -> do + MutableGenericVector.unsafeWrite v element (Just (intToIndex index)) + GenericVector.unsafeFreeze v {-| Filter a vector, leaving only the entries, under the indices, which are in the set. @@ -199,9 +223,9 @@ -} filterVector :: GenericVector.Vector vector a => DenseIntSet -> vector a -> vector a filterVector set vector = let - !newVectorSize = size set + !newVectorPop = population set in runST $ do - newVector <- MutableGenericVector.unsafeNew newVectorSize + newVector <- MutableGenericVector.unsafeNew newVectorPop forM_ (Unfoldr.zipWithIndex (vectorElementsUnfoldr vector set)) $ \ (newIndex, a) -> do MutableGenericVector.unsafeWrite newVector newIndex a GenericVector.unsafeFreeze newVector @@ -214,7 +238,7 @@ Unfold the present elements. -} presentElementsUnfoldr :: DenseIntSet -> Unfoldr Int -presentElementsUnfoldr (DenseIntSet vec) = do +presentElementsUnfoldr (DenseIntSet capacity vec) = do (wordIndex, word) <- Unfoldr.vectorWithIndices vec bitIndex <- Unfoldr.setBitIndices word return (wordIndex * 64 + bitIndex) @@ -223,10 +247,20 @@ Unfold the absent elements. -} absentElementsUnfoldr :: DenseIntSet -> Unfoldr Int -absentElementsUnfoldr (DenseIntSet vec) = do - (wordIndex, word) <- Unfoldr.vectorWithIndices vec - bitIndex <- Unfoldr.unsetBitIndices word - return (wordIndex * 64 + bitIndex) +absentElementsUnfoldr (DenseIntSet capacity vec) = let + isLastWordIndex = let + !maxWordIndex = pred (divCeiling capacity 64) + in \ wordIndex -> wordIndex == maxWordIndex + in do + (wordIndex, word) <- Unfoldr.vectorWithIndices vec + let + wordElemIndex = wordIndex * 64 + in if isLastWordIndex wordIndex + then let + !bitsLeft = capacity - wordElemIndex + predicate bitIndex = bitIndex < bitsLeft + in fmap (wordElemIndex +) (Unfoldr.takeWhile predicate (Unfoldr.unsetBitIndices word)) + else fmap (wordElemIndex +) (Unfoldr.unsetBitIndices word) {-| Unfold the elements of a vector by indices in the set. @@ -236,40 +270,3 @@ vectorElementsUnfoldr :: (GenericVector.Vector vector a) => vector a -> DenseIntSet -> Unfoldr a vectorElementsUnfoldr vec = fmap (vec GenericVector.!) . presentElementsUnfoldr - --- * Composition -------------------------- - -{-| -Abstraction over the composition of sets, -which is cheap to append and can be used for interpreted merging of sets. --} -data DenseIntSetComposition = DenseIntSetComposition !Int [UnboxedVector Word64] - -instance Semigroup DenseIntSetComposition where - (<>) (DenseIntSetComposition leftMinLength leftVecs) = - if null leftVecs - then id - else \ (DenseIntSetComposition rightMinLength rightVecs) -> if null rightVecs - then DenseIntSetComposition leftMinLength leftVecs - else DenseIntSetComposition (min leftMinLength rightMinLength) (leftVecs <> rightVecs) - -instance Monoid DenseIntSetComposition where - mempty = DenseIntSetComposition 0 [] - mappend = (<>) - -{-| -Lift a set into composition. --} -compose :: DenseIntSet -> DenseIntSetComposition -compose (DenseIntSet vec) = DenseIntSetComposition (UnboxedVector.length vec) (pure vec) - -{-| -Lift a list of sets into composition. --} -composeList :: [DenseIntSet] -> DenseIntSetComposition -composeList list = if null list - then mempty - else let - unboxedVec = fmap (\ (DenseIntSet x) -> x) list - in DenseIntSetComposition (foldr1 min (fmap UnboxedVector.length unboxedVec)) unboxedVec diff --git a/test/Main.hs b/test/Main.hs new file mode 100644 --- /dev/null +++ b/test/Main.hs @@ -0,0 +1,37 @@ +module Main where + +import Prelude hiding (toList, choose) +import GHC.Exts +import Test.QuickCheck.Instances +import Test.Tasty +import Test.Tasty.Runners +import Test.Tasty.HUnit +import Test.Tasty.QuickCheck +import qualified Test.QuickCheck as QuickCheck +import qualified Test.QuickCheck.Property as QuickCheck +import qualified DenseIntSet +import qualified Data.HashSet as HashSet + + +main = + defaultMain $ + testGroup "All" $ let + setGen = listOf (choose (0, 99)) <&> HashSet.fromList + listSetGen = setGen <&> HashSet.toList + nonEmptyListOf gen = do + length <- choose (1, 99) + replicateM length gen + in [ + testProperty "List roundtrip" $ forAll listSetGen $ \ list -> + sort list === sort (toList (fromList @DenseIntSet.DenseIntSet list)) + , + testProperty "Inversion" $ forAll listSetGen $ \ list -> let + invertedList = enumFromTo 0 (foldl' max 0 list) \\ list + in sort invertedList === sort (toList (DenseIntSet.invert (fromList list))) + , + testProperty "Union" $ forAll (listOf listSetGen) $ \ lists -> + sort (foldr union [] lists) === sort (toList (DenseIntSet.unions (fmap fromList lists))) + , + testProperty "Intersection" $ forAll (nonEmptyListOf listSetGen) $ \ lists -> + sort (foldl1' intersect lists) === sort (toList (DenseIntSet.intersections (fmap fromList lists))) + ]