Word64 based implementation for resource multiset

src/RMultiSet.hs

@@ -17,24 +17,25 @@ module RMultiSet
   , toSet
   ) where
 
-import qualified Data.Vector.Unboxed as V
-import qualified Data.Vector.Unboxed.Mutable as VM
 import qualified Data.Set as S
 import Data.Aeson
+import Data.Bits
+import Data.Word
+import Data.List (foldl')
 
 import Startups.Base
 
 newtype ResourceSet
-    = ResourceSet { getMultiSet :: V.Vector Int }
+    = ResourceSet Word64
     deriving (Eq, Ord)
 
 instance Show ResourceSet where
     show = show . toList
 
 instance Monoid ResourceSet where
-    mempty = ResourceSet (V.replicate 7 0)
+    mempty = ResourceSet 0
     {-# INLINE mempty #-}
-    mappend (ResourceSet a) (ResourceSet b) = ResourceSet (V.zipWith (+) a b)
+    mappend (ResourceSet a) (ResourceSet b) = ResourceSet (a + b)
     {-# INLINE mappend #-}
 
 singleton :: Resource -> ResourceSet
@@ -52,8 +53,12 @@ toList :: ResourceSet -> [Resource]
 toList = concatMap (uncurry (flip replicate)) . toOccurList
 {-# INLINE toList #-}
 
+occur :: Resource -> ResourceSet -> Int
+occur r (ResourceSet s) = fromIntegral ((s `shiftR` (fromEnum r * 8)) .&. 0xff)
+{-# INLINE occur #-}
+
 toOccurList :: ResourceSet -> [(Resource, Int)]
-toOccurList = zip resources . V.toList . getMultiSet
+toOccurList s = [ (r, occur r s) | r <- resources ]
 {-# INLINE toOccurList #-}
 
 fromOccurList :: [(Resource, Int)] -> ResourceSet
@@ -65,35 +70,50 @@ foldMap f = Prelude.foldMap f . toList
 {-# INLINE foldMap #-}
 
 insertMany :: Resource -> Int -> ResourceSet -> ResourceSet
-insertMany r n (ResourceSet s) = ResourceSet (V.modify im s)
-  where
-    im t = VM.unsafeModify t (+n) (fromEnum r)
+insertMany r n (ResourceSet s) = ResourceSet (s + (fromIntegral n `shiftL` (fromEnum r * 8)))
 {-# INLINE insertMany #-}
 
 insert :: Resource -> ResourceSet -> ResourceSet
-insert r (ResourceSet s) = ResourceSet (V.modify im s)
-  where
-    im t = VM.unsafeModify t (+1) (fromEnum r)
+insert r (ResourceSet s) = ResourceSet (s + (1 `shiftL` (fromEnum r * 8)))
 {-# INLINE insert #-}
 
 isSubsetOf :: ResourceSet -> ResourceSet -> Bool
 isSubsetOf (ResourceSet s1) (ResourceSet s2) =
-    not (V.any (< 0) (V.zipWith (-) s2 s1))
+    all check [0..6]
+  where
+    check n =
+      let !mask = 0xff `shiftL` (n * 8)
+      in  s1 .&. mask <= s2 .&. mask
 {-# INLINE isSubsetOf #-}
 
 member :: Resource -> ResourceSet -> Bool
-member r (ResourceSet s) = s V.! fromEnum r > 0
+member r (ResourceSet s) = (s `shiftR` (fromEnum r * 8)) .&. 0xff > 0
 {-# INLINE member #-}
 
 delete :: Resource -> ResourceSet -> ResourceSet
-delete r (ResourceSet s) = ResourceSet (V.modify im s)
+delete r (ResourceSet s) = ResourceSet modified
   where
-    im t = VM.unsafeModify t (\x -> max 0 (x-1)) (fromEnum r)
+    !idx = fromEnum r * 8
+    !mask = 0xff `shiftL` idx
+    !imask = complement mask
+    !cur = (s `shiftR` idx) .&. 0xff
+    !new = if cur == 0
+             then 0
+             else cur - 1
+    !modified = (s .&. imask) .|. (new `shiftL` idx)
 {-# INLINE delete #-}
 
 difference :: ResourceSet -> ResourceSet -> ResourceSet
 difference (ResourceSet s1) (ResourceSet s2) =
-    ResourceSet (V.zipWith (\a b -> max 0 (a - b)) s1 s2)
+    ResourceSet (foldl' foo 0 [0..6])
+  where
+    foo cur i =
+      let !mask = 0xff `shiftL` (i * 8)
+          !m1 = s1 .&. mask
+          !m2 = s2 .&. mask
+      in  if m2 >= m1
+            then cur
+            else cur .|. (m1 - m2)
 {-# INLINE difference #-}
 
 toSet :: ResourceSet -> S.Set Resource