diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,35 @@ +Copyright (c) 2008, David Himmelstrup +All rights reserved. + +Redistribution and use in source and binary forms, +with or without modification, are permitted provided +that the following conditions are met: + + * Redistributions of source code must retain + the above copyright notice, this list of + conditions and the following disclaimer. + * Redistributions in binary form must reproduce + the above copyright notice, this list of + conditions and the following disclaimer in the + documentation and/or other materials provided + with the distribution. + * Neither the name of David Himmelstrup nor the + names of other contributors may be used to + endorse or promote products derived from this + software without specific prior written permission. + + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND +CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, +INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES +OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER +OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, +BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING +NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE +USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY +OF SUCH DAMAGE. diff --git a/Setup.hs b/Setup.hs new file mode 100644 --- /dev/null +++ b/Setup.hs @@ -0,0 +1,2 @@ +import Distribution.Simple +main = defaultMain diff --git a/compact-map.cabal b/compact-map.cabal new file mode 100644 --- /dev/null +++ b/compact-map.cabal @@ -0,0 +1,24 @@ +Name: compact-map +Version: 2008.11.8 +Author: David Himmelstrup +Maintainer: David Himmelstrup +Copyright: 2008 David Himmelstrup +Build-Type: Simple +Build-Depends: base, bytestring, binary, array, containers +ghc-prof-options: -auto-all +Exposed-Modules: Data.CompactMap +Other-Modules: Data.CompactMap.Types + Data.CompactMap.Index + Data.CompactMap.Buffer + Data.CompactMap.MemoryMap + Data.CompactMap.Fetch +Hs-Source-Dirs: src +Extensions: CPP +License: BSD3 +License-file: LICENSE +Tested-with: GHC ==6.8.3 +Category: Data +Synopsis: Compact Data.Map implementation using Data.Binary +Description: + This library attempts to provide a memory efficient alternative to + Data.Map. diff --git a/src/Data/CompactMap.hs b/src/Data/CompactMap.hs new file mode 100644 --- /dev/null +++ b/src/Data/CompactMap.hs @@ -0,0 +1,1031 @@ +{-# LANGUAGE NoBangPatterns, CPP, DeriveDataTypeable, ScopedTypeVariables #-} +----------------------------------------------------------------------------- +-- | +-- Module : Data.CompactMap +-- Copyright : (c) David Himmelstrup 2008 +-- License : BSD-style +-- Maintainer : libraries@haskell.org +-- Stability : experimental +-- Portability : portable +-- +-- An efficient implementation of maps from keys to values (dictionaries). +-- +-- Since many function names (but not the type name) clash with +-- "Prelude" names, this module is usually imported @qualified@, e.g. +-- +-- > import Data.CompactMap (Map) +-- > import qualified Data.CompactMap as Map +-- +-- The implementation of 'Map' is based on /size balanced/ binary trees (or +-- trees of /bounded balance/) as described by: +-- +-- * Stephen Adams, \"/Efficient sets: a balancing act/\", +-- Journal of Functional Programming 3(4):553-562, October 1993, +-- . +-- +-- * J. Nievergelt and E.M. Reingold, +-- \"/Binary search trees of bounded balance/\", +-- SIAM journal of computing 2(1), March 1973. +-- +-- Note that the implementation is /left-biased/ -- the elements of a +-- first argument are always preferred to the second, for example in +-- 'union' or 'insert'. +-- +-- Operation comments contain the operation time complexity in +-- the Big-O notation . +----------------------------------------------------------------------------- +module Data.CompactMap + ( -- * Map type + Map -- instance Eq,Show,Read + + -- * Operators + , (!) --, (\\) + + + -- * Query + , null + , size + , member + , notMember + , lookup + , findWithDefault + + -- * Construction + , empty + , singleton + + -- ** Insertion + , insert + , insertWith, insertWithKey, insertLookupWithKey + + -- ** Delete\/Update + , delete + , adjust + , adjustWithKey + , update + , updateWithKey + , updateLookupWithKey + , alter + + -- * Combine + , union + , unionWith + , unionWithKey + , unions + , unionsWith{- + + -- ** Difference + , difference + , differenceWith + , differenceWithKey + + -- ** Intersection + , intersection + , intersectionWith + , intersectionWithKey-} + + -- * Traversal + -- ** Map + , map + , mapWithKey{- + , mapAccum + , mapAccumWithKey-} + , mapKeys + , mapKeysWith + , mapKeysMonotonic + + -- ** Fold + , fold + , foldWithKey + + -- * Conversion + , elems + , keys + , keysSet + , assocs + + -- ** Lists + , toList + , fromList + , fromListWith + , fromListWithKey + + -- ** Ordered lists + , toAscList + , fromAscList + , fromAscListWith + , fromAscListWithKey + , fromDistinctAscList + + -- * Filter + , filter + , filterWithKey + , partition + , partitionWithKey + + , mapMaybe + , mapMaybeWithKey + , mapEither + , mapEitherWithKey{- + + , split + , splitLookup + + -- * Submap + , isSubmapOf, isSubmapOfBy + , isProperSubmapOf, isProperSubmapOfBy + + -- * Indexed + , lookupIndex + , findIndex + , elemAt + , updateAt + , deleteAt + -- * Min\/Max + , findMin + , findMax + , deleteMin + , deleteMax + , deleteFindMin + , deleteFindMax + , updateMin + , updateMax + , updateMinWithKey + , updateMaxWithKey + , minView + , maxView + , minViewWithKey + , maxViewWithKey + + -- * Debugging + , showTree + , showTreeWith + , valid-} + ) where + + +import Data.Monoid (Monoid(..)) +import Control.Concurrent +import Data.IORef +import Data.Binary +import Data.Typeable +import Data.List (foldl') +import System.IO.Unsafe + +import qualified Data.Maybe as Maybe +import Data.Maybe (isJust) +import Foreign (nullPtr) +import Text.Read hiding (get) +import Control.Monad +import qualified Data.CompactMap.Index as Index +import qualified Data.CompactMap.Types as Types +import qualified Data.Array.IArray as IArray +import qualified Data.Set as Set + +import qualified Data.ByteString as Strict +import qualified Data.ByteString.Lazy as Lazy + +import Prelude hiding (null,lookup,map,filter) +import qualified Prelude + +data Range = Range Int Int deriving Show + +-- | A Map from keys @k@ to values @a@. +data Map k a = Empty + | Existing + { index :: !(MVar Types.Index) + , uniq :: {-# UNPACK #-} !(IORef Int) + , range :: ![Range] + , mapSize :: {-# UNPACK #-} !Int + } deriving (Typeable) + +{-------------------------------------------------------------------- + Instances +--------------------------------------------------------------------} + +instance (Eq k, Eq a, Binary k, Binary a) => Eq (Map k a) where + m1 == m2 = toList m1 == toList m2 + +instance (Ord k, Ord a, Binary k, Binary a) => Ord (Map k a) where + m1 `compare` m2 = toList m1 `compare` toList m2 + +instance (Binary k, Binary a, Show k, Show a) => Show (Map k a) where + showsPrec d m = showParen (d > 10) $ + showString "fromList " . shows (toList m) + +instance (Ord k, Binary k, Binary a, Read k, Read a) => Read (Map k a) where +#ifdef __GLASGOW_HASKELL__ + readPrec = parens $ prec 10 $ do + Ident "fromList" <- lexP + xs <- readPrec + return (fromList xs) + + readListPrec = readListPrecDefault +#else + readsPrec p = readParen (p > 10) $ \ r -> do + ("fromList",s) <- lex r + (xs,t) <- reads s + return (fromList xs,t) +#endif + +instance Binary (Map k a) where + put Empty = put (0::Int) + put Existing{index=index,range=range,mapSize=mapSize} = + let a = unsafePerformIO $ withMVar index $ Index.listKeyPointers + in do put mapSize + forM_ (IArray.elems a) $ \ptr -> + do let ls = unsafePerformIO $ Index.getDataFromPointer ptr + case findValue range ls of + Nothing -> return () + Just val -> do let key = unsafePerformIO $ Index.getKeyFromPointer ptr + put (key,val) + let x = unsafePerformIO $ do withMVar index Index.touchIndex + return () + x `seq` return () + get = do n <- get + ls <- replicateM n get + unsafePerformIO $ + do idx <- Index.newIndex + forM_ ls $ \(k,v) -> do keyCursor <- Index.newKeyCursor (Types.indexBuffer idx) (Lazy.fromChunks [k]) + Index.insertLargestKeyCursor idx keyCursor + dataCursor <- Index.newDataCursor (Types.indexBuffer idx) 0 (Just (Lazy.fromChunks [v])) + Index.pushNewDataCursor keyCursor dataCursor + --Index.insertBS idx (decodeStrict k :: k) 0 (Just (Lazy.fromChunks [v])) + uniq <- newIORef 1 + index <- newMVar idx + return $ return $ Existing{index=index,uniq=uniq,range=addToRange 0 [],mapSize=n} + +instance (Ord k, Binary k, Binary a) => Monoid (Map k a) where + mempty = empty + mappend = union + mconcat = unions + + +{-------------------------------------------------------------------- + Methods +--------------------------------------------------------------------} + +infixl 9 ! -- ,\\ + +-- | /O(log n)/. Find the value at a key. +-- Calls 'error' when the element can not be found. +-- +-- > fromList [(5,'a'), (3,'b')] ! 1 Error: element not in the map +-- > fromList [(5,'a'), (3,'b')] ! 5 == 'a' + +(!) :: (Ord k, Binary k, Binary a) => Map k a -> k -> a +m ! k = case lookup k m of + Nothing -> error "element not in the map" + Just x -> x + + +-- | /O(1)/. Is the map empty? +-- +-- > Data.Map.null (empty) == True +-- > Data.Map.null (singleton 1 'a') == False +null :: Map k a -> Bool +null m = size m == 0 + +-- | /O(1)/. The number of elements in the map. +-- +-- > size empty == 0 +-- > size (singleton 1 'a') == 1 +-- > size (fromList([(1,'a'), (2,'c'), (3,'b')])) == 3 +size :: Map k a -> Int +size Empty = 0 +size Existing{mapSize=size} = size + +-- | /O(log n)/. Is the key a member of the map? See also 'notMember'. +-- +-- > member 5 (fromList [(5,'a'), (3,'b')]) == True +-- > member 1 (fromList [(5,'a'), (3,'b')]) == False +member :: (Ord k, Binary k) => k -> Map k a -> Bool +member k Empty = False +member k Existing{index=index,range=range} + = unsafePerformIO $ withMVar index $ \idx -> + do ls <- Index.lookupList idx k + return $ isJust $ findValue range ls + +-- | /O(log n)/. Is the key not a member of the map? See also 'member'. +-- +-- > notMember 5 (fromList [(5,'a'), (3,'b')]) == False +-- > notMember 1 (fromList [(5,'a'), (3,'b')]) == True +notMember :: (Ord k, Binary k) => k -> Map k a -> Bool +notMember k m = not (member k m) + +-- | /O(log n)/. Lookup the value at a key in the map. +-- +-- The function will return the corresponding value as @('Just' value)@, +-- or 'Nothing' if the key isn't in the map. +-- +-- An example of using @lookup@: +-- +-- > import Prelude hiding (lookup) +-- > import Data.Map +-- > +-- > employeeDept = fromList([("John","Sales"), ("Bob","IT")]) +-- > deptCountry = fromList([("IT","USA"), ("Sales","France")]) +-- > countryCurrency = fromList([("USA", "Dollar"), ("France", "Euro")]) +-- > +-- > employeeCurrency :: String -> Maybe String +-- > employeeCurrency name = do +-- > dept <- lookup name employeeDept +-- > country <- lookup dept deptCountry +-- > lookup country countryCurrency +-- > +-- > main = do +-- > putStrLn $ "John's currency: " ++ (show (employeeCurrency "John")) +-- > putStrLn $ "Pete's currency: " ++ (show (employeeCurrency "Pete")) +-- +-- The output of this program: +-- +-- > John's currency: Just "Euro" +-- > Pete's currency: Nothing +lookup :: (Ord k, Binary k, Binary a) => k -> Map k a -> Maybe a +lookup k Empty = Nothing +lookup k Existing{index=index,range=range} + = unsafePerformIO $ withMVar index $ \idx -> + do ls <- Index.lookupList idx k + return $ case findValue range ls of + Nothing -> Nothing + Just bs -> Just (decodeStrict bs) + +-- | /O(log n)/. The expression @('findWithDefault' def k map)@ returns +-- the value at key @k@ or returns default value @def@ +-- when the key is not in the map. +-- +-- > findWithDefault 'x' 1 (fromList [(5,'a'), (3,'b')]) == 'x' +-- > findWithDefault 'x' 5 (fromList [(5,'a'), (3,'b')]) == 'a' +findWithDefault :: (Ord k, Binary k, Binary a) => a -> k -> Map k a -> a +findWithDefault def k m = case lookup k m of + Nothing -> def + Just x -> x + +-- | /O(1)/. The empty map. +-- +-- > empty == fromList [] +-- > size empty == 0 +empty :: Map k a +empty = Empty + +-- | /O(1)/. A map with a single element. +-- +-- > singleton 1 'a' == fromList [(1, 'a')] +-- > size (singleton 1 'a') == 1 +singleton :: (Ord k, Binary k, Binary a) => k -> a -> Map k a +singleton k a = insert k a empty + +-- | /O(log n)/. Insert a new key and value in the map. +-- If the key is already present in the map, the associated value is +-- replaced with the supplied value. 'insert' is equivalent to +-- @'insertWith' 'const'@. +-- +-- > insert 5 'x' (fromList [(5,'a'), (3,'b')]) == fromList [(3, 'b'), (5, 'x')] +-- > insert 7 'x' (fromList [(5,'a'), (3,'b')]) == fromList [(3, 'b'), (5, 'a'), (7, 'x')] +-- > insert 5 'x' empty == singleton 5 'x' +insert :: (Ord k, Binary k, Binary a) => k -> a -> Map k a -> Map k a +insert k a m + = unsafePerformIO $ + withExisting m $ \Existing{index=index,uniq=uniq,range=range,mapSize=mapSize} -> + withMVar index $ \idx -> + do u <- readIORef uniq + modifyIORef uniq succ + ls <- Index.insert idx k u (Just a) + let newSize | haveOldValue range ls = mapSize + | otherwise = mapSize+1 + return Existing{index=index,uniq=uniq,range=addToRange u range,mapSize=newSize} + +-- | /O(log n)/. Insert with a function, combining new value and old value. +-- @'insertWith' f key value mp@ +-- will insert the pair (key, value) into @mp@ if key does +-- not exist in the map. If the key does exist, the function will +-- insert the pair @(key, f new_value old_value)@. +-- +-- > insertWith (++) 5 "xxx" (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "xxxa")] +-- > insertWith (++) 7 "xxx" (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a"), (7, "xxx")] +-- > insertWith (++) 5 "xxx" empty == singleton 5 "xxx" +insertWith :: (Ord k, Binary k, Binary a) => (a -> a -> a) -> k -> a -> Map k a -> Map k a +insertWith f k x m + = insertWithKey (\_ x' y' -> f x' y') k x m + +-- | /O(log n)/. Insert with a function, combining key, new value and old value. +-- @'insertWithKey' f key value mp@ +-- will insert the pair (key, value) into @mp@ if key does +-- not exist in the map. If the key does exist, the function will +-- insert the pair @(key,f key new_value old_value)@. +-- Note that the key passed to f is the same key passed to 'insertWithKey'. +-- +-- > let f key new_value old_value = (show key) ++ ":" ++ new_value ++ "|" ++ old_value +-- > insertWithKey f 5 "xxx" (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "5:xxx|a")] +-- > insertWithKey f 7 "xxx" (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a"), (7, "xxx")] +-- > insertWithKey f 5 "xxx" empty == singleton 5 "xxx" +insertWithKey :: (Ord k, Binary k, Binary a) => (k -> a -> a -> a) -> k -> a -> Map k a -> Map k a +insertWithKey f kx x Empty = singleton kx x +insertWithKey f kx x Existing{index=index,uniq=uniq,range=range,mapSize=mapSize} + = unsafePerformIO $ withMVar index $ \idx -> + do u <- readIORef uniq + modifyIORef uniq succ + keyCursor <- Index.insertKey idx kx + ls <- Index.getDataFromPointer keyCursor + let oldVal = findValue range ls + newVal = case oldVal of + Nothing -> x + Just old -> f kx x (decodeStrict old) + newSize = if isJust oldVal then mapSize else mapSize + 1 + dataCursor <- Index.newDataCursor (Types.indexBuffer idx) u (Just $ encode newVal) + Index.pushNewDataCursor keyCursor dataCursor + return $ Existing{index=index,uniq=uniq,range=addToRange u range,mapSize=newSize} + +-- | /O(log n)/. Combines insert operation with old value retrieval. +-- The expression (@'insertLookupWithKey' f k x map@) +-- is a pair where the first element is equal to (@'lookup' k map@) +-- and the second element equal to (@'insertWithKey' f k x map@). +-- +-- > let f key new_value old_value = (show key) ++ ":" ++ new_value ++ "|" ++ old_value +-- > insertLookupWithKey f 5 "xxx" (fromList [(5,"a"), (3,"b")]) == (Just "a", fromList [(3, "b"), (5, "5:xxx|a")]) +-- > insertLookupWithKey f 7 "xxx" (fromList [(5,"a"), (3,"b")]) == (Nothing, fromList [(3, "b"), (5, "a"), (7, "xxx")]) +-- > insertLookupWithKey f 5 "xxx" empty == (Nothing, singleton 5 "xxx") +-- +-- This is how to define @insertLookup@ using @insertLookupWithKey@: +-- +-- > let insertLookup kx x t = insertLookupWithKey (\_ a _ -> a) kx x t +-- > insertLookup 5 "x" (fromList [(5,"a"), (3,"b")]) == (Just "a", fromList [(3, "b"), (5, "x")]) +-- > insertLookup 7 "x" (fromList [(5,"a"), (3,"b")]) == (Nothing, fromList [(3, "b"), (5, "a"), (7, "x")]) +insertLookupWithKey :: (Ord k, Binary k, Binary a) => (k -> a -> a -> a) -> k -> a -> Map k a -> (Maybe a,Map k a) +insertLookupWithKey f k a Empty = (Nothing, singleton k a) +insertLookupWithKey f k a Existing{index=index,uniq=uniq,range=range,mapSize=mapSize} + = unsafePerformIO $ withMVar index $ \idx -> + do u <- readIORef uniq + modifyIORef uniq succ + keyCursor <- Index.insertKey idx k + ls <- Index.getDataFromPointer keyCursor + let oldVal = fmap decodeStrict $ findValue range ls + newVal = case oldVal of + Nothing -> a + Just old -> f k a old + newSize = if isJust oldVal then mapSize else mapSize + 1 + dataCursor <- Index.newDataCursor (Types.indexBuffer idx) u (Just $ encode newVal) + Index.pushNewDataCursor keyCursor dataCursor + return $ (oldVal, Existing{index=index,uniq=uniq,range=addToRange u range,mapSize=newSize}) + + + +-- | /O(log n)/. Delete a key and its value from the map. When the key is not +-- a member of the map, the original map is returned. +-- +-- > delete 5 (fromList [(5,"a"), (3,"b")]) == singleton 3 "b" +-- > delete 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")] +-- > delete 5 empty == empty +delete :: (Ord k, Binary k) => k -> Map k a -> Map k a +delete k Empty = Empty +delete k Existing{index=index,uniq=uniq,range=range,mapSize=mapSize} + = unsafePerformIO $ withMVar index $ \idx -> + do u <- readIORef uniq + modifyIORef uniq succ + ls <- Index.insert idx k u (Nothing :: Maybe ()) + let newSize | haveOldValue range ls = mapSize-1 + | otherwise = mapSize + return Existing{index=index,uniq=uniq,range=addToRange u range,mapSize=newSize} + +-- | /O(log n)/. Update a value at a specific key with the result of the provided function. +-- When the key is not +-- a member of the map, the original map is returned. +-- +-- > adjust ("new " ++) 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "new a")] +-- > adjust ("new " ++) 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")] +-- > adjust ("new " ++) 7 empty == empty +adjust :: (Ord k, Binary k, Binary a) => (a -> a) -> k -> Map k a -> Map k a +adjust f k m + = adjustWithKey (\_ x -> f x) k m + +-- | /O(log n)/. Adjust a value at a specific key. When the key is not +-- a member of the map, the original map is returned. +-- +-- > let f key x = (show key) ++ ":new " ++ x +-- > adjustWithKey f 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "5:new a")] +-- > adjustWithKey f 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")] +-- > adjustWithKey f 7 empty == empty +adjustWithKey :: (Ord k, Binary k, Binary a) => (k -> a -> a) -> k -> Map k a -> Map k a +adjustWithKey f k m + = updateWithKey (\k' x' -> Just (f k' x')) k m + +-- | /O(log n)/. The expression (@'update' f k map@) updates the value @x@ +-- at @k@ (if it is in the map). If (@f x@) is 'Nothing', the element is +-- deleted. If it is (@'Just' y@), the key @k@ is bound to the new value @y@. +-- +-- > let f x = if x == "a" then Just "new a" else Nothing +-- > update f 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "new a")] +-- > update f 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")] +-- > update f 3 (fromList [(5,"a"), (3,"b")]) == singleton 5 "a" +update :: (Ord k, Binary k, Binary a) => (a -> Maybe a) -> k -> Map k a -> Map k a +update f k m + = updateWithKey (\_ x -> f x) k m + +-- | /O(log n)/. The expression (@'updateWithKey' f k map@) updates the +-- value @x@ at @k@ (if it is in the map). If (@f k x@) is 'Nothing', +-- the element is deleted. If it is (@'Just' y@), the key @k@ is bound +-- to the new value @y@. +-- +-- > let f k x = if x == "a" then Just ((show k) ++ ":new a") else Nothing +-- > updateWithKey f 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "5:new a")] +-- > updateWithKey f 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")] +-- > updateWithKey f 3 (fromList [(5,"a"), (3,"b")]) == singleton 5 "a" +updateWithKey :: (Ord k, Binary k, Binary a) => (k -> a -> Maybe a) -> k -> Map k a -> Map k a +updateWithKey f k m = snd (updateLookupWithKey f k m) + +-- | /O(log n)/. Lookup and update. See also 'updateWithKey'. +-- The function returns changed value, if it is updated. +-- Returns the original key value if the map entry is deleted. +-- +-- > let f k x = if x == "a" then Just ((show k) ++ ":new a") else Nothing +-- > updateLookupWithKey f 5 (fromList [(5,"a"), (3,"b")]) == (Just "5:new a", fromList [(3, "b"), (5, "5:new a")]) +-- > updateLookupWithKey f 7 (fromList [(5,"a"), (3,"b")]) == (Nothing, fromList [(3, "b"), (5, "a")]) +-- > updateLookupWithKey f 3 (fromList [(5,"a"), (3,"b")]) == (Just "b", singleton 5 "a") +updateLookupWithKey :: (Ord k, Binary k, Binary a) => (k -> a -> Maybe a) -> k -> Map k a -> (Maybe a,Map k a) +updateLookupWithKey f k Empty = (Nothing,Empty) +updateLookupWithKey f k m@Existing{index=index,uniq=uniq,range=range,mapSize=mapSize} + = unsafePerformIO $ withMVar index $ \idx -> + do ls <- Index.lookupList idx k + case fmap decodeStrict $ findValue range ls of + Nothing -> return (Nothing, m) + Just val -> do let newVal = f k val + u <- readIORef uniq + modifyIORef uniq succ + Index.insert idx k u newVal + let newSize = case isJust newVal of + False -> mapSize-1 + True -> mapSize + return (newVal `mplus` Just val, Existing{index=index,uniq=uniq,range=addToRange u range,mapSize=newSize}) + + +-- | /O(log n)/. The expression (@'alter' f k map@) alters the value @x@ at @k@, or absence thereof. +-- 'alter' can be used to insert, delete, or update a value in a 'Map'. +-- In short : @'lookup' k ('alter' f k m) = f ('lookup' k m)@. +-- +-- > let f _ = Nothing +-- > alter f 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")] +-- > alter f 5 (fromList [(5,"a"), (3,"b")]) == singleton 3 "b" +-- > +-- > let f _ = Just "c" +-- > alter f 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a"), (7, "c")] +-- > alter f 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "c")] +alter :: (Ord k, Binary k, Binary a) => (Maybe a -> Maybe a) -> k -> Map k a -> Map k a +alter f k m + = case f (lookup k m) of + Nothing -> delete k m + Just val -> insert k val m + +-- | /O(log n*m)/. +-- The expression (@'union' t1 t2@) takes the left-biased union of @t1@ and @t2@. +-- It prefers @t1@ when duplicate keys are encountered, +-- i.e. (@'union' == 'unionWith' 'const'@). +-- +-- > union (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == fromList [(3, "b"), (5, "a"), (7, "C")] +union :: (Ord k, Binary k, Binary a) => Map k a -> Map k a -> Map k a +union = unionWith const + +-- | /O(log n*m)/. Union with a combining function. +-- +-- > unionWith (++) (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == fromList [(3, "b"), (5, "aA"), (7, "C")] +unionWith :: (Ord k, Binary k, Binary a) => (a -> a -> a) -> Map k a -> Map k a -> Map k a +unionWith f m1 m2 + = unionWithKey (\_ x y -> f x y) m1 m2 + + +-- | /O(log n*m)/. +-- Union with a combining function. +-- +-- > let f key new_value old_value = (show key) ++ ":" ++ new_value ++ "|" ++ old_value +-- > unionWithKey f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == fromList [(3, "b"), (5, "5:a|A"), (7, "C")] +unionWithKey :: (Ord k, Binary k, Binary a) => (k -> a -> a -> a) -> Map k a -> Map k a -> Map k a +unionWithKey f t1 t2 = foldl' (\m (k,v) -> insertWithKey f k v m) t2 (toList t1) + +-- | The union of a list of maps: +-- (@'unions' == 'Prelude.foldl' 'union' 'empty'@). +-- +-- > unions [(fromList [(5, "a"), (3, "b")]), (fromList [(5, "A"), (7, "C")]), (fromList [(5, "A3"), (3, "B3")])] +-- > == fromList [(3, "b"), (5, "a"), (7, "C")] +-- > unions [(fromList [(5, "A3"), (3, "B3")]), (fromList [(5, "A"), (7, "C")]), (fromList [(5, "a"), (3, "b")])] +-- > == fromList [(3, "B3"), (5, "A3"), (7, "C")] +unions :: (Ord k, Binary k, Binary a) => [Map k a] -> Map k a +unions ts + = foldl' union empty ts + +-- | The union of a list of maps, with a combining operation: +-- (@'unionsWith' f == 'Prelude.foldl' ('unionWith' f) 'empty'@). +-- +-- > unionsWith (++) [(fromList [(5, "a"), (3, "b")]), (fromList [(5, "A"), (7, "C")]), (fromList [(5, "A3"), (3, "B3")])] +-- > == fromList [(3, "bB3"), (5, "aAA3"), (7, "C")] +unionsWith :: (Ord k, Binary k, Binary a) => (a -> a -> a) -> [Map k a] -> Map k a +unionsWith f ts + = foldl' (unionWith f) empty ts + +-- | /O(n)/. Map a function over all values in the map. +-- +-- > map (++ "x") (fromList [(5,"a"), (3,"b")]) == fromList [(3, "bx"), (5, "ax")] +map :: (Ord k, Binary k, Binary a, Binary b) => (a -> b) -> Map k a -> Map k b +map f m + = mapWithKey (\_ x -> f x) m + +-- | /O(n)/. Map a function over all values in the map. +-- +-- > let f key x = (show key) ++ ":" ++ x +-- > mapWithKey f (fromList [(5,"a"), (3,"b")]) == fromList [(3, "3:b"), (5, "5:a")] +mapWithKey :: (Ord k, Binary k, Binary a, Binary b) => (k -> a -> b) -> Map k a -> Map k b +mapWithKey f m = fromDistinctAscList [ (k, f k x) | (k,x) <- toList m ] + +-- | /O(n*log n)/. +-- @'mapKeys' f s@ is the map obtained by applying @f@ to each key of @s@. +-- +-- The size of the result may be smaller if @f@ maps two or more distinct +-- keys to the same new key. In this case the value at the smallest of +-- these keys is retained. +-- +-- > mapKeys (+ 1) (fromList [(5,"a"), (3,"b")]) == fromList [(4, "b"), (6, "a")] +-- > mapKeys (\ _ -> 1) (fromList [(1,"b"), (2,"a"), (3,"d"), (4,"c")]) == singleton 1 "c" +-- > mapKeys (\ _ -> 3) (fromList [(1,"b"), (2,"a"), (3,"d"), (4,"c")]) == singleton 3 "c" +mapKeys :: (Ord k2,Binary k1,Binary k2,Binary a) => (k1->k2) -> Map k1 a -> Map k2 a +mapKeys = mapKeysWith (\x _ -> x) + +-- | /O(n*log n)/. +-- @'mapKeysWith' c f s@ is the map obtained by applying @f@ to each key of @s@. +-- +-- The size of the result may be smaller if @f@ maps two or more distinct +-- keys to the same new key. In this case the associated values will be +-- combined using @c@. +-- +-- > mapKeysWith (++) (\ _ -> 1) (fromList [(1,"b"), (2,"a"), (3,"d"), (4,"c")]) == singleton 1 "cdab" +-- > mapKeysWith (++) (\ _ -> 3) (fromList [(1,"b"), (2,"a"), (3,"d"), (4,"c")]) == singleton 3 "cdab" +mapKeysWith :: (Ord k2, Binary k1, Binary k2, Binary a) => (a -> a -> a) -> (k1->k2) -> Map k1 a -> Map k2 a +mapKeysWith c f m = fromListWith c [ (f x,y) | (x,y) <- toList m ] + +-- | /O(n)/. +-- @'mapKeysMonotonic' f s == 'mapKeys' f s@, but works only when @f@ +-- is strictly monotonic. +-- That is, for any values @x@ and @y@, if @x@ < @y@ then @f x@ < @f y@. +-- /The precondition is not checked./ +-- Semi-formally, we have: +-- +-- > and [x < y ==> f x < f y | x <- ls, y <- ls] +-- > ==> mapKeysMonotonic f s == mapKeys f s +-- > where ls = keys s +-- +-- This means that @f@ maps distinct original keys to distinct resulting keys. +-- This function has better performance than 'mapKeys'. +-- +-- > mapKeysMonotonic (\ k -> k * 2) (fromList [(5,"a"), (3,"b")]) == fromList [(6, "b"), (10, "a")] +-- > valid (mapKeysMonotonic (\ k -> k * 2) (fromList [(5,"a"), (3,"b")])) == True +-- > valid (mapKeysMonotonic (\ _ -> 1) (fromList [(5,"a"), (3,"b")])) == False +mapKeysMonotonic :: (Binary k1, Binary k2, Binary a) => (k1->k2) -> Map k1 a -> Map k2 a +mapKeysMonotonic f m = fromDistinctAscList [ (f x, y) | (x,y) <- toList m ] + +-- | /O(n)/. Fold the values in the map, such that +-- @'fold' f z == 'Prelude.foldr' f z . 'elems'@. +-- For example, +-- +-- > elems map = fold (:) [] map +-- +-- > let f a len = len + (length a) +-- > fold f 0 (fromList [(5,"a"), (3,"bbb")]) == 4 +fold :: (Binary k, Binary a) => (a -> b -> b) -> b -> Map k a -> b +fold f z m + = foldWithKey (\_ x' z' -> f x' z') z m + +-- | /O(n)/. Fold the keys and values in the map, such that +-- @'foldWithKey' f z == 'Prelude.foldr' ('uncurry' f) z . 'toAscList'@. +-- For example, +-- +-- > keys map = foldWithKey (\k x ks -> k:ks) [] map +-- +-- > let f k a result = result ++ "(" ++ (show k) ++ ":" ++ a ++ ")" +-- > foldWithKey f "Map: " (fromList [(5,"a"), (3,"b")]) == "Map: (5:a)(3:b)" +foldWithKey :: (Binary k, Binary a) => (k -> a -> b -> b) -> b -> Map k a -> b +foldWithKey f z = Prelude.foldr (uncurry f) z . toList + + +-- | /O(n)/. +-- Return all elements of the map in the ascending order of their keys. +-- +-- > elems (fromList [(5,"a"), (3,"b")]) == ["b","a"] +-- > elems empty == [] +elems :: (Binary k, Binary a) => Map k a -> [a] +elems = Prelude.map snd . toList + +-- | /O(n)/. Return all keys of the map in ascending order. +-- +-- > keys (fromList [(5,"a"), (3,"b")]) == [3,5] +-- > keys empty == [] +keys :: (Binary k, Binary a) => Map k a -> [k] +keys = Prelude.map fst . toList + +-- | /O(n)/. The set of all keys of the map. +-- +-- > keysSet (fromList [(5,"a"), (3,"b")]) == Data.Set.fromList [3,5] +-- > keysSet empty == Data.Set.empty +keysSet :: (Ord k, Binary k, Binary a) => Map k a -> Set.Set k +keysSet m = Set.fromDistinctAscList (keys m) + +-- | /O(n)/. Return all key\/value pairs in the map in ascending key order. +-- +-- > assocs (fromList [(5,"a"), (3,"b")]) == [(3,"b"), (5,"a")] +-- > assocs empty == [] +assocs :: (Binary k, Binary a) => Map k a -> [(k,a)] +assocs m + = toList m + + + +{-# SPECIALISE fromList :: (Binary a) => [(Strict.ByteString,a)] -> Map Strict.ByteString a #-} +{-# SPECIALISE fromList :: (Binary a) => [(Int,a)] -> Map Int a #-} +-- | /O(n*log n)/. Build a map from a list of key\/value pairs. See also 'fromAscList'. +-- If the list contains more than one value for the same key, the last value +-- for the key is retained. +-- +-- > fromList [] == empty +-- > fromList [(5,"a"), (3,"b"), (5, "c")] == fromList [(5,"c"), (3,"b")] +-- > fromList [(5,"c"), (3,"b"), (5, "a")] == fromList [(5,"a"), (3,"b")] +fromList :: (Ord k, Binary k, Binary a) => [(k,a)] -> Map k a +fromList [] = Empty +fromList ls + = unsafePerformIO $ + do idx <- Index.newIndex + let loop n _ | n `seq` False = undefined + loop n [] = return n + loop n ((k,v):rs) + = do keyCursor <- Index.insertKey idx k + oldData <- Index.peekKeyCursorData keyCursor + newData <- Index.newDataCursor (Types.indexBuffer idx) 0 (Just (encode v)) + Index.pushNewDataCursor keyCursor newData + loop (if oldData==nullPtr then n+1 else n) rs + size <- loop 0 ls + uniq <- newIORef 1 + index <- newMVar idx + return $ Existing{index=index,uniq=uniq,range=addToRange 0 [],mapSize=size} + +-- | /O(n*log n)/. Build a map from a list of key\/value pairs with a combining function. See also 'fromAscListWith'. +-- +-- > fromListWith (++) [(5,"a"), (5,"b"), (3,"b"), (3,"a"), (5,"a")] == fromList [(3, "ab"), (5, "aba")] +-- > fromListWith (++) [] == empty +fromListWith :: (Ord k, Binary k, Binary a) => (a -> a -> a) -> [(k,a)] -> Map k a +fromListWith f xs + = fromListWithKey (\_ x y -> f x y) xs + +-- | /O(n*log n)/. Build a map from a list of key\/value pairs with a combining function. See also 'fromAscListWithKey'. +-- +-- > let f k a1 a2 = (show k) ++ a1 ++ a2 +-- > fromListWithKey f [(5,"a"), (5,"b"), (3,"b"), (3,"a"), (5,"a")] == fromList [(3, "3ab"), (5, "5a5ba")] +-- > fromListWithKey f [] == empty +fromListWithKey :: (Ord k, Binary k, Binary a) => (k -> a -> a -> a) -> [(k,a)] -> Map k a +fromListWithKey f [] = empty +fromListWithKey f ls + = unsafePerformIO $ + do idx <- Index.newIndex + let loop n _ | n `seq` False = undefined + loop n [] = return n + loop n ((k,v):rs) + = do keyCursor <- Index.insertKey idx k + oldData <- Index.getDataFromPointer keyCursor + let newVal = case oldData of + ((_,Just old):_) -> f k v (decodeStrict old) + _ -> v + newData <- Index.newDataCursor (Types.indexBuffer idx) 0 (Just (encode newVal)) + Index.pushNewDataCursor keyCursor newData + loop (if Prelude.null oldData then n+1 else n) rs + size <- loop 0 ls + uniq <- newIORef 1 + index <- newMVar idx + return $ Existing{index=index,uniq=uniq,range=addToRange 0 [],mapSize=size} + + +-- | /O(n)/. Convert to a list of key\/value pairs. +-- +-- > toList (fromList [(5,"a"), (3,"b")]) == [(3,"b"), (5,"a")] +-- > toList empty == [] +toList :: (Binary k, Binary a) => Map k a -> [(k,a)] +toList Empty = [] +toList Existing{index=index,range=range} + = unsafePerformIO $ + do keys <- withMVar index $ Index.listKeyPointers + let loop [] = withMVar index Index.touchIndex >> return [] + loop (keyCursor:xs) + = unsafeInterleaveIO $ + do ls <- Index.getDataFromPointer keyCursor + case findValue range ls of + Nothing -> loop xs + Just bs -> do key <- Index.getKeyFromPointer keyCursor + let ckey = Strict.copy key + cbs = Strict.copy bs + let pair = (decodeStrict ckey, decodeStrict cbs) + ckey `seq` cbs `seq` liftM (pair:) (loop xs) + loop (IArray.elems keys) + + +-- | /O(n)/. Convert to an ascending list. +-- +-- > toAscList (fromList [(5,"a"), (3,"b")]) == [(3,"b"), (5,"a")] +toAscList :: (Binary k, Binary a) =>Map k a -> [(k,a)] +toAscList = toList + +-- | /O(n)/. Build a map from an ascending list in linear time. +-- /The precondition (input list is ascending) is not checked./ +-- +-- > fromAscList [(3,"b"), (5,"a")] == fromList [(3, "b"), (5, "a")] +-- > fromAscList [(3,"b"), (5,"a"), (5,"b")] == fromList [(3, "b"), (5, "b")] +-- > valid (fromAscList [(3,"b"), (5,"a"), (5,"b")]) == True +-- > valid (fromAscList [(5,"a"), (3,"b"), (5,"b")]) == False +fromAscList :: (Eq k, Binary k, Binary a) => [(k,a)] -> Map k a +fromAscList xs + = fromAscListWithKey (\_ x _ -> x) xs + +-- | /O(n)/. Build a map from an ascending list in linear time with a combining function for equal keys. +-- /The precondition (input list is ascending) is not checked./ +-- +-- > fromAscListWith (++) [(3,"b"), (5,"a"), (5,"b")] == fromList [(3, "b"), (5, "ba")] +-- > valid (fromAscListWith (++) [(3,"b"), (5,"a"), (5,"b")]) == True +-- > valid (fromAscListWith (++) [(5,"a"), (3,"b"), (5,"b")]) == False +fromAscListWith :: (Eq k, Binary k, Binary a) => (a -> a -> a) -> [(k,a)] -> Map k a +fromAscListWith f xs + = fromAscListWithKey (\_ x y -> f x y) xs + +-- | /O(n)/. Build a map from an ascending list in linear time with a +-- combining function for equal keys. +-- /The precondition (input list is ascending) is not checked./ +-- +-- > let f k a1 a2 = (show k) ++ ":" ++ a1 ++ a2 +-- > fromAscListWithKey f [(3,"b"), (5,"a"), (5,"b"), (5,"b")] == fromList [(3, "b"), (5, "5:b5:ba")] +-- > valid (fromAscListWithKey f [(3,"b"), (5,"a"), (5,"b"), (5,"b")]) == True +-- > valid (fromAscListWithKey f [(5,"a"), (3,"b"), (5,"b"), (5,"b")]) == False +fromAscListWithKey :: (Eq k, Binary k, Binary a) => (k -> a -> a -> a) -> [(k,a)] -> Map k a +fromAscListWithKey f xs + = fromDistinctAscList (combineEq f xs) + where + -- [combineEq f xs] combines equal elements with function [f] in an ordered list [xs] + combineEq _ xs' + = case xs' of + [] -> [] + [x] -> [x] + (x:xx) -> combineEq' x xx + + combineEq' z [] = [z] + combineEq' z@(kz,zz) (x@(kx,xx):xs') + | kx==kz = let yy = f kx xx zz in combineEq' (kx,yy) xs' + | otherwise = z:combineEq' x xs' + + +-- | /O(n)/. Build a map from an ascending list of distinct elements in linear time. +-- /The precondition is not checked./ +-- +-- > fromDistinctAscList [(3,"b"), (5,"a")] == fromList [(3, "b"), (5, "a")] +-- > valid (fromDistinctAscList [(3,"b"), (5,"a")]) == True +-- > valid (fromDistinctAscList [(3,"b"), (5,"a"), (5,"b")]) == False +fromDistinctAscList :: (Binary k, Binary a) => [(k,a)] -> Map k a +fromDistinctAscList [] = Empty +fromDistinctAscList ls + = unsafePerformIO $ + do idx <- Index.newIndex + n <- foldM (\s (k,v) -> do keyCursor <- Index.insertLargestKey idx k + dataCursor <- Index.newDataCursor (Types.indexBuffer idx) 0 (Just $ encode v) + Index.pushNewDataCursor keyCursor dataCursor + return $! s+1) 0 ls + index <- newMVar idx + uniq <- newIORef 1 + return Existing{index=index,uniq=uniq,range=addToRange 0 [],mapSize=n} + + +-- | /O(n)/. Filter all values that satisfy the predicate. +-- +-- > filter (> "a") (fromList [(5,"a"), (3,"b")]) == singleton 3 "b" +-- > filter (> "x") (fromList [(5,"a"), (3,"b")]) == empty +-- > filter (< "a") (fromList [(5,"a"), (3,"b")]) == empty +filter :: (Ord k, Binary k, Binary a) => (a -> Bool) -> Map k a -> Map k a +filter p m + = filterWithKey (\_ x -> p x) m + +-- FIXME: optimize this. +-- | /O(n)/. Filter all keys\/values that satisfy the predicate. +-- +-- > filterWithKey (\k _ -> k > 4) (fromList [(5,"a"), (3,"b")]) == singleton 5 "a" +filterWithKey :: (Ord k, Binary k, Binary a) => (k -> a -> Bool) -> Map k a -> Map k a +filterWithKey p m = fromDistinctAscList [ (k, v) | (k,v) <- toList m, p k v ] + +-- | /O(n)/. Partition the map according to a predicate. The first +-- map contains all elements that satisfy the predicate, the second all +-- elements that fail the predicate. See also 'split'. +-- +-- > partition (> "a") (fromList [(5,"a"), (3,"b")]) == (singleton 3 "b", singleton 5 "a") +-- > partition (< "x") (fromList [(5,"a"), (3,"b")]) == (fromList [(3, "b"), (5, "a")], empty) +-- > partition (> "x") (fromList [(5,"a"), (3,"b")]) == (empty, fromList [(3, "b"), (5, "a")]) +partition :: (Ord k, Binary k, Binary a) => (a -> Bool) -> Map k a -> (Map k a,Map k a) +partition p = partitionWithKey (\_ -> p) + +-- | /O(n)/. Partition the map according to a predicate. The first +-- map contains all elements that satisfy the predicate, the second all +-- elements that fail the predicate. See also 'split'. +-- +-- > partitionWithKey (\ k _ -> k > 3) (fromList [(5,"a"), (3,"b")]) == (singleton 5 "a", singleton 3 "b") +-- > partitionWithKey (\ k _ -> k < 7) (fromList [(5,"a"), (3,"b")]) == (fromList [(3, "b"), (5, "a")], empty) +-- > partitionWithKey (\ k _ -> k > 7) (fromList [(5,"a"), (3,"b")]) == (empty, fromList [(3, "b"), (5, "a")]) +partitionWithKey :: (Ord k, Binary k, Binary a) => (k -> a -> Bool) -> Map k a -> (Map k a,Map k a) +partitionWithKey p m = mapEitherWithKey (\k x -> if p k x then Left x else Right x) m + +-- | /O(n)/. Map values and collect the 'Just' results. +-- +-- > let f x = if x == "a" then Just "new a" else Nothing +-- > mapMaybe f (fromList [(5,"a"), (3,"b")]) == singleton 5 "new a" +mapMaybe :: (Ord k, Binary k, Binary a, Binary b) => (a -> Maybe b) -> Map k a -> Map k b +mapMaybe f m + = mapMaybeWithKey (\_ x -> f x) m + +-- | /O(n)/. Map keys\/values and collect the 'Just' results. +-- +-- > let f k _ = if k < 5 then Just ("key : " ++ (show k)) else Nothing +-- > mapMaybeWithKey f (fromList [(5,"a"), (3,"b")]) == singleton 3 "key : 3" +mapMaybeWithKey :: (Ord k, Binary k, Binary a, Binary b) => (k -> a -> Maybe b) -> Map k a -> Map k b +mapMaybeWithKey f m = fromDistinctAscList [ (k, v) | (k,x) <- toList m, Just v <- [f k x] ] + + +-- | /O(n)/. Map values and separate the 'Left' and 'Right' results. +-- +-- > let f a = if a < "c" then Left a else Right a +-- > mapEither f (fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")]) +-- > == (fromList [(3,"b"), (5,"a")], fromList [(1,"x"), (7,"z")]) +-- > +-- > mapEither (\ a -> Right a) (fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")]) +-- > == (empty, fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")]) + +mapEither :: (Ord k, Binary k, Binary a, Binary b, Binary c) => (a -> Either b c) -> Map k a -> (Map k b, Map k c) +mapEither f m + = mapEitherWithKey (\_ x -> f x) m + +-- The key doesn't change. Don't re-encode it. Copy bytestring instead. +-- | /O(n)/. Map keys\/values and separate the 'Left' and 'Right' results. +-- +-- > let f k a = if k < 5 then Left (k * 2) else Right (a ++ a) +-- > mapEitherWithKey f (fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")]) +-- > == (fromList [(1,2), (3,6)], fromList [(5,"aa"), (7,"zz")]) +-- > +-- > mapEitherWithKey (\_ a -> Right a) (fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")]) +-- > == (empty, fromList [(1,"x"), (3,"b"), (5,"a"), (7,"z")]) +mapEitherWithKey :: (Ord k, Binary k, Binary a, Binary c, Binary b) => + (k -> a -> Either b c) -> Map k a -> (Map k b, Map k c) +mapEitherWithKey f m + = unsafePerformIO $ + do idxL <- Index.newIndex + idxR <- Index.newIndex + (s1,s2) <- foldM (\(s1,s2) (k,v) -> s1 `seq` s2 `seq` + do let cond = f k v + (idx,v',s1',s2') = case cond of + Left v' -> (idxL,encode v',s1+1,s2) + Right v' -> (idxR,encode v',s1,s2+1) + keyCursor <- Index.insertLargestKey idx k + dataCursor <- Index.newDataCursor (Types.indexBuffer idx) 0 (Just v') + Index.pushNewDataCursor keyCursor dataCursor + return $! (s1',s2')) (0,0) (toList m) + indexL <- newMVar idxL + indexR <- newMVar idxR + uniqL <- newIORef 1 + uniqR <- newIORef 1 + return $ (Existing{index=indexL,uniq=uniqR,range=addToRange 0 [],mapSize=s1} + ,Existing{index=indexR,uniq=uniqL,range=addToRange 0 [],mapSize=s2}) + +{-------------------------------------------------------------------- + Utilities +--------------------------------------------------------------------} + +decodeStrict bs = decode (Lazy.fromChunks [bs]) + +haveOldValue range ls + = isJust (findValue range ls) + +withExisting Empty fn + = do idx <- newMVar =<< Index.newIndex + uniq <- newIORef 0 + fn (Existing idx uniq [] 0) +withExisting m fn + = fn m + +findValue range [] = Nothing +findValue range ((uniqId, value):rs) + | uniqId `isInRange` range = value + | otherwise = findValue range rs + + +isInRange :: Int -> [Range] -> Bool +isInRange i [] = False +isInRange i (Range x y:rs) + | i > x = False + | i < y = isInRange i rs + | otherwise = True + +addToRange :: Int -> [Range] -> [Range] +addToRange i [] = [Range i i] +addToRange i (Range x y:rs) + = merge (Range i i:Range x y:rs) + +merge [] = [] +merge [x] = [x] +merge (Range x y:Range a b:rs) + | y == a+1 = merge (Range x b:rs) + | otherwise = Range x y:merge (Range a b:rs) + + diff --git a/src/Data/CompactMap/Buffer.hs b/src/Data/CompactMap/Buffer.hs new file mode 100644 --- /dev/null +++ b/src/Data/CompactMap/Buffer.hs @@ -0,0 +1,52 @@ +{-# OPTIONS -fglasgow-exts -fbang-patterns #-} +module Data.CompactMap.Buffer where + + +import Foreign (Ptr,Storable(..),plusPtr, castPtr) + +import Data.CompactMap.MemoryMap +import Data.CompactMap.Types + +import Foreign.ForeignPtr (ForeignPtr, withForeignPtr, touchForeignPtr, castForeignPtr) +import Foreign.Concurrent +import Data.IORef + + +newBuffer :: Int -> IO Buffer +newBuffer initSize + = do aligned <- alignSize initSize + dataPtr <- mmap aligned [Read,Write] [Anonymous,Private,NoReserve] + fptr <- newIORef =<< newForeignPtr dataPtr (munmap dataPtr initSize) + old <- newIORef [] + posRef <- newFastMutInt 0 + size <- newFastMutInt aligned + return $ Buffer{ bufferData = fptr + , bufferOld = old + , bufferPos = posRef + , bufferSize = size } + +withBytes :: Buffer -> Int -> (Ptr a -> IO b) -> IO b +withBytes !Buffer{bufferPos=bufferPos,bufferData=bufferData,bufferSize=bufferSize,bufferOld=bufferOld} !bytesNeeded fn + = do !currentPos <- readFastMutInt bufferPos + !currentSize <- readFastMutInt bufferSize + !oldPtr <- readIORef bufferData + if currentSize >= currentPos + bytesNeeded + then do writeFastMutInt bufferPos (currentPos+bytesNeeded) + withForeignPtr (castForeignPtr oldPtr) $ \ptr -> fn $! (ptr `plusPtr` currentPos) + else do let minSize = max bytesNeeded currentSize + newSize = minSize + minSize `div` 4 -- Add 25% to the buffer. + aligned <- alignSize newSize + --putStrLn $ "Expanding from " ++ show currentSize ++ " to " ++ show aligned + !newPtr <- mmap aligned [Read,Write] [Anonymous,Private,NoReserve] + fptr <- newForeignPtr newPtr (munmap newPtr aligned) + writeIORef bufferData fptr + modifyIORef bufferOld (oldPtr:) + writeFastMutInt bufferPos bytesNeeded + writeFastMutInt bufferSize aligned + fn $! (castPtr newPtr) + +touchBuffer :: Buffer -> IO () +touchBuffer buffer + = do touchForeignPtr =<< readIORef (bufferData buffer) + ls <- readIORef (bufferOld buffer) + mapM_ touchForeignPtr ls diff --git a/src/Data/CompactMap/Fetch.hs b/src/Data/CompactMap/Fetch.hs new file mode 100644 --- /dev/null +++ b/src/Data/CompactMap/Fetch.hs @@ -0,0 +1,54 @@ +{-# OPTIONS -fglasgow-exts -fbang-patterns #-} +module Data.CompactMap.Fetch where + + +import Data.CompactMap.Types +import Data.CompactMap.Buffer + +import Foreign +import GHC.Ptr + +import qualified Data.ByteString.Internal as B +import qualified Data.ByteString.Unsafe as B +import qualified Data.ByteString as B +import qualified Data.ByteString.Lazy as LBS + +import Data.Binary +--import Data.BinaryLinear + +sizeOfInt :: Int +sizeOfInt = sizeOf (0::Int) + +{-# INLINE [2] getElement #-} +getElement :: (Binary a) => Ptr () -> IO a +getElement ptr + = do size <- peek (castPtr ptr) :: IO Int + extractElement (ptr `plusPtr` (sizeOfInt * 1)) size + +extractElement :: Binary a => Ptr () -> Int -> IO a +extractElement !ptr !size +-- = return $! decode (castPtr ptr) + = do bs <- B.unsafePackCStringLen (castPtr ptr, size) + return $! decode (LBS.fromChunks [bs]) + +{-# RULES "extractElement/ByteString" extractElement = extractElementBS #-} +extractElementBS :: Ptr () -> Int -> IO B.ByteString +extractElementBS ptr !size + = let n = sizeOf (0::Int) + Ptr addr# = ptr `plusPtr` n + in B.unsafePackAddressLen (size-n) addr# +{- +extractRawString :: DiskSet RawString -> Int -> IO RawString +extractRawString !(DiskSet {tPosition=pos, tData=dat}) !n + = do posPtr <- bufferPtr pos + datPtr <- bufferPtr dat + (from, size) <- getElemDimensions posPtr n + let n = sizeOf (0::Int) + return $! RawString (size-n) (castPtr datPtr `plusPtr` (n+from)) +-} +{- +{-# RULES "extractElement/RawString" extractElement = extractElementRaw #-} +extractElementRaw :: Ptr () -> Int -> IO RawString +extractElementRaw ptr size = let n = sizeOf (0::Int) + in return $! RawString (size-n) (castPtr ptr `plusPtr` n) +-} diff --git a/src/Data/CompactMap/Index.hs b/src/Data/CompactMap/Index.hs new file mode 100644 --- /dev/null +++ b/src/Data/CompactMap/Index.hs @@ -0,0 +1,547 @@ +{-# OPTIONS -fglasgow-exts #-} +{-# LANGUAGE NoMonomorphismRestriction, BangPatterns, EmptyDataDecls #-} +module Data.CompactMap.Index where + +import Foreign hiding (rotateL,rotateR) +import Foreign.C (CStringLen) +import Foreign.Storable +import Control.Monad +import System.Exit +import GHC.Exts +import Data.Maybe +import Text.Printf +import Data.Int +import Numeric +import System.IO.Unsafe + +import Data.Array.IArray +import Data.Array.IO +import Data.Array.Unboxed + +import Data.Binary + +import qualified Data.ByteString as Strict +import qualified Data.ByteString.Unsafe as Strict +import qualified Data.ByteString.Internal as Strict +import qualified Data.ByteString.Lazy as Lazy + +import Data.CompactMap.Buffer +import Data.CompactMap.Types +import Data.CompactMap.Fetch + +import GHC.Exts (addr2Int#) + +import Prelude hiding (Either(..)) + +type Tag = Int + + +-- ints are native GHC ints. +{- KeyCursor +void *dataPointer; +int keyLen; +void *key; +-} +{- DataCursor +void *next; +int tag; +word8 isJust; +int dataLen; +void *data; +-} + +peekKeyCursorData :: Ptr KeyCursor -> IO (Ptr DataCursor) +peekKeyCursorData ptr + = peek (castPtr ptr) + +peekKeyCursorKey :: Ptr KeyCursor -> IO Strict.ByteString +peekKeyCursorKey ptr = do len <- peek (ptr `plusPtr` ptrSize) + Strict.unsafePackCStringLen (ptr `plusPtr` (ptrSize+intSize), len) + +pokeKeyCursorData :: Ptr KeyCursor -> Ptr DataCursor -> IO () +pokeKeyCursorData ptr dataPtr + = poke (castPtr ptr) dataPtr + +newKeyCursor :: Buffer -> Lazy.ByteString -> IO (Ptr KeyCursor) +newKeyCursor buffer keyE + = withBytes buffer (intSize*2 + keyLen) $ \keyPtr -> + do poke (castPtr keyPtr) nullPtr + putByteString (keyPtr `plusPtr` intSize) keyE keyLen + return keyPtr + where keyLen = fromIntegral $ Lazy.length keyE + +newBinaryKeyCursor :: (Binary a) => Buffer -> a -> IO (Ptr KeyCursor) +newBinaryKeyCursor !buffer !key + = newKeyCursor buffer (encode key) + +pushNewDataCursor :: Ptr KeyCursor -> Ptr DataCursor -> IO () +pushNewDataCursor keyCursor dataCursor + = do oldData <- peekKeyCursorData keyCursor + pokeDataCursorNext dataCursor oldData + pokeKeyCursorData keyCursor dataCursor + +peekDataCursorNext :: Ptr DataCursor -> IO (Ptr DataCursor) +peekDataCursorNext ptr = peek (castPtr ptr) + +peekDataCursorTag :: Ptr DataCursor -> IO Int +peekDataCursorTag ptr = peek (ptr `plusPtr` ptrSize) + +peekDataCursorData :: Ptr DataCursor -> IO (Maybe Strict.ByteString) +peekDataCursorData ptr + = do isJ <- peek (ptr `plusPtr` (ptrSize+intSize)) + case isJ == (1::Word8) of + False -> do return Nothing + True -> do len <- peek (ptr `plusPtr` (ptrSize+intSize+1)) + bs <- Strict.unsafePackCStringLen (ptr `plusPtr` (intSize+intSize+1+intSize),len) + return (Just bs) + +pokeDataCursorNext :: Ptr DataCursor -> Ptr DataCursor -> IO () +pokeDataCursorNext ptr next + = poke (castPtr ptr) next + +newDataCursor :: Buffer -> Tag -> Maybe Lazy.ByteString -> IO (Ptr DataCursor) +newDataCursor !buffer !tag !mbString + = do let !bsLen = fromIntegral $ maybe 0 Lazy.length mbString + !ext = if isJust mbString then intSize else 0 + withBytes buffer (ptrSize+intSize+1+bsLen+ext) $ \ !ptr -> + do poke (castPtr ptr) (nullPtr :: Ptr DataCursor) + poke (ptr `plusPtr` ptrSize) tag + case mbString of + Nothing -> do poke (ptr `plusPtr` (ptrSize+intSize)) (0::Word8) + Just bs -> do poke (ptr `plusPtr` (ptrSize+intSize)) (1::Word8) + putByteString (ptr `plusPtr` (ptrSize+intSize+1)) bs bsLen + return ptr + + +intToPtr i = nullPtr `plusPtr` i +ptrToInt (Ptr addr#) = I# (addr2Int# addr#) + +extractTop = extractField 0 +extractSize = fmap ptrToInt . extractField 1 +extractElemIdx = fmap castPtr . extractField 2 +extractLeft = extractField 3 +extractRight = extractField 4 + +putTop ptr val = if ptr == nullPtr then return () else putField 0 ptr val +putSize p s = putField 1 p (intToPtr s) +--putElemIdx p e = putField 2 p (castPtr e) +putLeft = putField 3 +putRight :: Ptr IndexItem -> Ptr IndexItem -> IO () +putRight = putField 4 + +--getElement set e = return $ set IntMap.! e +{- +ppHex n = "0x" ++ (drop (length hex) zeroes) ++ hex + where hex = showHex n "" + zeroes = replicate ((sizeOf nullPtr) * 2) '0' +-} +data Direction = Left | Right | Stop + +{-# INLINE walkTree #-} +walkTree start move + = let loop n = do keyCursor <- extractElemIdx n + dir <- move keyCursor + case dir of + Left -> extractLeft n >>= \left -> + if left == nullPtr + then return (Left, n) else loop left + Right -> extractRight n >>= \right -> + if right == nullPtr + then return (Right, n) else loop right + Stop -> return (Stop, n) + in loop start + + +{-# INLINE lookupNearest #-} +lookupNearest :: (Ord a, Binary a) => Ptr IndexItem + -> a -> IO (Direction, Ptr IndexItem) +lookupNearest start e + = walkTree start $ \keyCursor -> + do idxElem <- getElement (keyCursor `plusPtr` intSize) + case compare e idxElem of + LT -> return Left + GT -> return Right + EQ -> return Stop + +{-# INLINE lookupLargest #-} +lookupLargest :: Ptr IndexItem + -> IO (Direction, Ptr IndexItem) +lookupLargest start + = walkTree start $ \_ -> return Right + + +putByteString :: Ptr () -> Lazy.ByteString -> Int -> IO () +putByteString dst lbs len + = do poke (castPtr dst) len + let loop !ptr [] = return () + loop !ptr (chunk:cs) = do Strict.unsafeUseAsCString chunk $ \cstr -> + copyArray ptr cstr (Strict.length chunk) + loop (ptr `plusPtr` Strict.length chunk) cs + loop (dst `plusPtr` intSize) (Lazy.toChunks lbs) + + +intSize :: Int +intSize = sizeOf (undefined::Int) + +ptrSize :: Int +ptrSize = sizeOf (undefined::Ptr ()) + +insert :: (Ord k, Binary k, Binary a) => Index -> k -> Tag -> Maybe a -> IO [(Tag,Maybe Strict.ByteString)] +insert idx key tag mbVal + = insertBS idx key tag (fmap encode mbVal) + +{- +insertWith :: (Ord k, Binary k, Binary a) => Index -> k -> Tag -> (Ptr DataCursor -> IO (Maybe a)) -> IO [(Tag,Maybe Strict.ByteString)] +insertWith idx key tag genVal + = insertWithBS idx key tag (\dataCursor -> do val <- genVal dataCursor; return $ fmap encode val) +-} +{- SPECIALISE insertBS :: Index -> Strict.ByteString -> Tag -> Maybe Lazy.ByteString -> IO [(Tag,Maybe Strict.ByteString)] -} +insertBS :: (Ord k, Binary k) => Index -> k -> Tag -> Maybe Lazy.ByteString -> IO [(Tag,Maybe Strict.ByteString)] +insertBS idx key tag mbVal + = insertWithBS idx key tag (\_ -> return mbVal) + +{- SPECIALISE insertWithBS :: Index -> Strict.ByteString -> Tag -> (Ptr DataCursor -> IO (Maybe Lazy.ByteString)) -> IO [(Tag,Maybe Strict.ByteString)] -} +insertWithBS :: (Ord k, Binary k) => Index -> k -> Tag -> (Ptr DataCursor -> IO (Maybe Lazy.ByteString)) -> IO [(Tag,Maybe Strict.ByteString)] +insertWithBS (Index orig buffer) key tag genVal + = do keyCursor <- insertKey (Index orig buffer) key + oldData <- peekKeyCursorData keyCursor -- Get the old data item + dataPtr <- newDataCursor buffer tag =<< genVal oldData + pushNewDataCursor keyCursor dataPtr + if oldData == nullPtr + then return [] + else fetchAllElts oldData + +{-# INLINE insertKey #-} +insertKey :: (Ord k, Binary k) => Index -> k -> IO (Ptr KeyCursor) +insertKey (Index orig buffer) key + = insertPrim (lookupNearest orig key) orig buffer (newBinaryKeyCursor buffer key) + +{-# INLINE insertLargestKey #-} +insertLargestKey :: (Binary k) => Index -> k -> IO (Ptr KeyCursor) +insertLargestKey (Index orig buffer) key + = insertPrim (lookupLargest orig) orig buffer (newBinaryKeyCursor buffer key) + +insertLargestKeyCursor :: Index -> Ptr KeyCursor -> IO () +insertLargestKeyCursor (Index orig buffer) keyCursor + = do insertPrim (lookupLargest orig) orig buffer (return keyCursor) + return () + +lookupKey :: (Ord k, Binary k) => Index -> k -> IO (Maybe (Ptr KeyCursor)) +lookupKey (Index orig buffer) key + = do (dir,pos) <- lookupNearest orig key + case dir of + Stop -> fmap Just $ extractElemIdx pos + _ -> return Nothing + +lookupList :: (Ord k, Binary k) => Index -> k -> IO [(Tag,Maybe Strict.ByteString)] +lookupList idx key + = do mbKey <- lookupKey idx key + case mbKey of + Nothing -> return [] + Just key -> fetchAllElts =<< peekKeyCursorData key + +fetchAllElts :: Ptr DataCursor -> IO [(Tag,Maybe Strict.ByteString)] +fetchAllElts ptr | ptr == nullPtr = return [] +fetchAllElts ptr + = unsafeInterleaveIO $ + do next <- peekDataCursorNext ptr + tag <- peekDataCursorTag ptr + mbData <- peekDataCursorData ptr + liftM ((tag,mbData):) (fetchAllElts next) + +indexItemSize :: Int +indexItemSize = sizeOf (undefined :: IndexItem) + +{- + Insert a key in the map. Return pointer to the old key if it exists. +-} +{- SPECIALISE insertPrim :: IO (Direction,Ptr IndexItem) -> Ptr IndexItem -> Buffer -> IO (Ptr KeyCursor) -> IO (Ptr KeyCursor) -} +{-# INLINE insertPrim #-} +insertPrim :: (IO (Direction,Ptr IndexItem)) -> Ptr IndexItem -> Buffer -> IO (Ptr KeyCursor) -> IO (Ptr KeyCursor) +insertPrim getPos !orig !buffer genIdx + = do size <- getSize orig + if size==0 -- We need a special case for size=0 /-: + then do eIdx <- genIdx + poke orig (IndexItem nullPtr (intToPtr 1) eIdx nullPtr nullPtr) + return eIdx + else do (dir,pos) <- getPos + case dir of + Right -> withBytes buffer indexItemSize $ \ptr -> + do eIdx <- genIdx + poke ptr (IndexItem pos (intToPtr 1) eIdx nullPtr nullPtr) + putRight pos ptr + balanceTree pos + return eIdx + Left -> withBytes buffer indexItemSize $ \ptr -> + do eIdx <- genIdx + poke ptr (IndexItem pos (intToPtr 1) eIdx nullPtr nullPtr) + putLeft pos ptr + balanceTree pos + return eIdx + Stop -> extractElemIdx pos + + + +listKeyPointers :: Index -> IO (UArray Int (Ptr KeyCursor)) +listKeyPointers (Index orig buffer) + = do size <- getSize orig + a <- newArray_ (0,size-1) :: IO (IOUArray Int (Ptr KeyCursor)) + let loop n ptr | ptr == nullPtr = return () + loop n ptr = do left <- extractLeft ptr + right <- extractRight ptr + leftSize <- getSize left + key <- extractElemIdx ptr + writeArray a (leftSize+n) key + loop (n) left + loop (leftSize+1+n) right + unless (size==0) $ loop (0::Int) orig + unsafeFreeze a + +getKeyFromPointer :: Ptr KeyCursor -> IO Strict.ByteString +getKeyFromPointer ptr + = peekKeyCursorKey ptr + +getDataFromPointer :: Ptr KeyCursor -> IO [(Tag, Maybe Strict.ByteString)] +getDataFromPointer ptr + = do dataPtr <- peekKeyCursorData ptr + fetchAllElts dataPtr + + +{- +getAllElements fn (Index buffer) + = do elems <- readBufferPos buffer + indices <- if elems == 0 then return id else sumIndex 0 buffer + vals <- forM (indices []) $ \idx -> sumIndex idx buffer + return $ concatMap ($ []) vals + where sumIndex = foldIndex sumNode sumLeaf + sumNode _ _ idx left right = return $ fn idx left right + sumLeaf = return id + +getSortedElements = getAllElements (\idx left right -> left . (idx:) . right) +getReverseElements = getAllElements (\idx left right -> right . (idx:) . left) +-} + +newIndex = do buffer <- newBuffer 0 + withBytes buffer indexItemSize $ \ptr -> ptr `seq` + do poke ptr (IndexItem nullPtr (intToPtr 0) nullPtr nullPtr nullPtr) + return $ Index ptr buffer + +touchIndex (Index _ buffer) = touchBuffer buffer +{- + + +foldIndex node leaf start buffer + = do ptr <- bufferPtr buffer + let loop (-1) = leaf + loop n = do IndexItem size idx left right <- peekElemOff ptr n + restLeft <- loop left + restRight <- loop right + node n size idx restLeft restRight + loop start + +showPrimIndex set + = do let Index buffer = (fromJust (tIndex set)) + let leaf = return [] + node n _size _idx left right + = return $ n:left++right + values <- foldIndex node leaf 0 buffer + free <- readBufferPos buffer + printf " \tSize\tIndex\tLeft\tRight\n" + ptr <- bufferPtr buffer + forM_ [0..free-1] $ \n -> + do IndexItem size idx left right <- peekElemOff ptr n + let isntValue = n `notElem` values + printf "%s%d\t%d\t%d\t%d\t%d\n" (if isntValue then "*" else " ") n size idx left right + +showIndex set + = do let Index buffer = (fromJust (tIndex set)) + elems <- readBufferPos buffer + ptr <- bufferPtr buffer + let leaf = return $ [Node "Leaf" []] + node _ size idx restLeft restRight + = do eIdx <- return idx -- getElement idx -- set =<< extractElemIdx ptr idx + positions <- foldIndex (\_ _ pos left right -> return [Node pos (left++right)]) (return []) idx buffer + return $ [Node (show eIdx ++ ": " ++ show positions) (restLeft++restRight)] + unless (elems==0) + $ do tree <- foldIndex node leaf 0 buffer + putStrLn (drawForest tree) + +{- SPECIALIZE isValid :: DiskSet RawString -> IO () -} +-- isValid :: IO () +isValid set + = do let Index buffer = (fromJust (tIndex set)) + elems <- readBufferPos buffer + ptr <- bufferPtr buffer + let check True _ = return () + check False msg = putStrLn msg >> exitWith (ExitFailure 1) + leaf = return (Nothing, 0) + node n size idx (mbLeft, leftN) (mbRight, rightN) + = do eIdx <- return idx -- getElement idx -- set =<< extractElemIdx ptr idx + check (leftN+rightN+1 == size) $ "Size check failed at " ++ show (n,size,leftN,rightN) + flip (maybe (return ())) mbLeft $ \el -> check (el < eIdx) $ "LT check failed at: " ++ show (idx,el,eIdx) + flip (maybe (return ())) mbRight $ \el -> check (el > eIdx) $ "GT check failed at: " ++ show (idx,el,eIdx) + return (Just eIdx, size) + unless (elems==0) $ foldIndex node leaf 0 buffer >> return () + return () +-} +{- +testSet :: [String] +testSet = [("Hello") + ,("World") + ,("This") + ,("Is") + ,("A Test") + ,("Yay") + ,("I think it works")] + +test :: IO () +test = do idx <- newIndex + forM_ (IntMap.keys testSet) $ \key -> do insert testSet idx key + balanceIndex idx + showIndex testSet idx + isValid testSet idx + putStrLn "Index is valid" +-} + +verify prev !pos | pos == nullPtr = return () +verify prev !pos + = do !top <- extractTop pos + !left <- extractLeft pos + !right <- extractRight pos + !sizeL <- getSize left + !sizeR <- getSize right + size <- getSize pos + unless (size==sizeL+sizeR+1) $ putStrLn $ "Size fail: " ++ show (size,sizeL,sizeR) + unless (top==prev) $ putStrLn "Top fail" + verify pos left + verify pos right + +balanceTree !pos | pos==nullPtr = return () +balanceTree !pos + = do balance pos + !top <- extractTop pos + balanceTree top + +getSize pos | pos == nullPtr = return $! 0 +getSize pos + = extractSize pos + +balance !pos + = do --keyCursor <- extractElemIdx pos + --bs <- peekKeyCursorKey keyCursor + --putStrLn $ "Balancing: " ++ show (decode (Lazy.fromChunks [bs]) :: Integer) + !left <- extractLeft pos + !right <- extractRight pos + !sizeL <- getSize left + !sizeR <- getSize right + putSize pos (sizeL+sizeR+1) + case () of + () | sizeL + sizeR <= 1 -> return () + | sizeR >= delta*sizeL -> rotateL pos left right + | sizeL >= delta*sizeR -> rotateR pos left right + | otherwise -> return () + +rotateL pos left right + = do !sizeLY <- getSize left + !sizeRY <- getSize right + if sizeLY < ratio * sizeRY then singleL pos + else doubleL pos + +rotateR pos left right + = do !sizeLY <- getSize left + !sizeRY <- getSize right + if sizeRY < ratio * sizeLY then singleR pos + else doubleR pos +singleL pos + = do IndexItem kTop kSize kElemIdx p1 k2 <- peek pos + IndexItem k2Top k2Size k2ElemIdx p2 p3 <- peek k2 + --unless (k2Top == pos) $ putStrLn "Assertion failure" + !p2Size <- getSize p2 + let p1Size = ptrToInt kSize-ptrToInt k2Size-1 + poke pos (IndexItem kTop kSize k2ElemIdx k2 p3) -- kSize hasn't changed + poke k2 (IndexItem k2Top (intToPtr $ p2Size+p1Size+1) kElemIdx p1 p2) + putTop p3 pos + putTop p1 k2 +-- recalcSize ptr k2 p1 p2 + +singleR pos + = do IndexItem kTop kSize kElemIdx k2 p3 <- peek pos + IndexItem k2Top k2Size k2ElemIdx p1 p2 <- peek k2 + --unless (k2Top == pos) $ putStrLn "Assertion failure" + !p2Size <- getSize p2 + let p3Size = ptrToInt kSize-ptrToInt k2Size-1 + poke pos (IndexItem kTop kSize k2ElemIdx p1 k2) -- kSize hasn't changed + poke k2 (IndexItem k2Top (intToPtr $ p2Size+p3Size+1) kElemIdx p2 p3) + putTop p1 pos + putTop p3 k2 +-- recalcSize ptr k2 p2 p3 + +doubleL pos + = do IndexItem kTop kSize kElemIdx p1 k2 <- peek pos + IndexItem k2Top k2Size k2ElemIdx k3 p4 <- peek k2 + IndexItem k3Top k3Size k3ElemIdx p2 p3 <- peek k3 + !p2Size <- getSize p2 + !p3Size <- getSize p3 + let p1Size = ptrToInt kSize - ptrToInt k2Size - 1 + p4Size = ptrToInt k2Size - ptrToInt k3Size - 1 + poke pos (IndexItem kTop kSize k3ElemIdx k3 k2) -- kSize hasn't changed + poke k2 (IndexItem k2Top (intToPtr $ p3Size+p4Size+1) k2ElemIdx p3 p4) -- k2ElemIdx and p4 hasn't changed + poke k3 (IndexItem k3Top (intToPtr $ p1Size+p2Size+1) kElemIdx p1 p2) +-- recalcSize ptr k2 p3 p4 +-- recalcSize ptr k3 p1 p2 + +doubleR pos + = do IndexItem kTop kSize kElemIdx k2 p4 <- peek pos + IndexItem k2Top k2Size k2ElemIdx p1 k3 <- peek k2 + IndexItem k3Top k3Size k3ElemIdx p2 p3 <- peek k3 + !p2Size <- getSize p2 + !p3Size <- getSize p3 + let p1Size = ptrToInt k2Size - ptrToInt k3Size - 1 + p4Size = ptrToInt kSize - ptrToInt k2Size - 1 + poke pos (IndexItem kTop kSize k3ElemIdx k3 k2) -- kSize hasn't changed + poke k2 (IndexItem k2Top (intToPtr $ p1Size+p2Size+1) k2ElemIdx p1 p2) -- k2ElemIdx and p1 hasn't changed. + poke k3 (IndexItem k3Top (intToPtr $ p3Size+p4Size+1) kElemIdx p3 p4) +-- recalcSize ptr k2 p1 p2 +-- recalcSize ptr k3 p3 p4 + +{- +recalcSize ptr !pos !left !right + = do !sizeL <- getSize ptr left + !sizeR <- getSize ptr right + putSize ptr pos (sizeL+sizeR+1) +-} + +delta,ratio :: Int +delta = 5 +ratio = 2 + +{- +balance :: k -> a -> Map k a -> Map k a -> Map k a +balance k x l r + | sizeL + sizeR <= 1 = Bin sizeX k x l r + | sizeR >= delta*sizeL = rotateL k x l r + | sizeL >= delta*sizeR = rotateR k x l r + | otherwise = Bin sizeX k x l r + where + sizeL = size l + sizeR = size r + sizeX = sizeL + sizeR + 1 + +-- rotate +rotateL k x l r@(Bin _ _ _ ly ry) + | size ly < ratio*size ry = singleL k x l r + | otherwise = doubleL k x l r + +rotateR k x l@(Bin _ _ _ ly ry) r + | size ry < ratio*size ly = singleR k x l r + | otherwise = doubleR k x l r + +-- basic rotations +singleL k1 x1 t1 (Bin _ k2 x2 t2 t3) = bin k2 x2 (bin k1 x1 t1 t2) t3 +singleR k1 x1 (Bin _ k2 x2 t1 t2) t3 = bin k2 x2 t1 (bin k1 x1 t2 t3) + +doubleL k1 x1 t1 (Bin _ k2 x2 (Bin _ k3 x3 t2 t3) t4) = bin k3 x3 (bin k1 x1 t1 t2) (bin k2 x2 t3 t4) +doubleR k1 x1 (Bin _ k2 x2 t1 (Bin _ k3 x3 t2 t3)) t4 = bin k3 x3 (bin k2 x2 t1 t2) (bin k1 x1 t3 t4) + +-} + diff --git a/src/Data/CompactMap/MemoryMap.hs b/src/Data/CompactMap/MemoryMap.hs new file mode 100644 --- /dev/null +++ b/src/Data/CompactMap/MemoryMap.hs @@ -0,0 +1,104 @@ +{-# LANGUAGE ForeignFunctionInterface #-} +{-# OPTIONS -fasm #-} +module Data.CompactMap.MemoryMap + ( Protection(..) + , Flag(..) + , mmap + , alignSize + , mremap + , munmap +-- , mprotect + , getPageSize + ) where + +import Control.Monad +import GHC.IOBase +import Foreign.C +import Foreign +import Data.Bits +import Numeric +import Data.List +import Data.Char + +foreign import ccall unsafe "mmap" c_mmap :: Ptr a -> CSize -> CInt -> CInt -> CInt -> CInt -> IO (Ptr a) +foreign import ccall unsafe "munmap" c_munmap :: Ptr a -> CSize -> IO CInt +foreign import ccall unsafe "mremap" c_mremap :: Ptr a -> CSize -> CSize -> CInt -> IO (Ptr a) +--foreign import ccall unsafe "mprotect" c_mprotect :: Ptr a -> CSize -> CInt -> IO CInt + +foreign import ccall unsafe "getpagesize" c_getpagesize :: IO CInt + +getPageSize :: IO Int +getPageSize = liftM fromIntegral c_getpagesize + +{- +failWhenNULL :: String -> IO (Ptr a) -> IO (Ptr a) +failWhenNULL name f = do + addr <- f + if addr == nullPtr + then ioError (IOError Nothing ResourceExhausted name + "out of memory" Nothing) + else return addr +-} + +data Protection + = Execute + | Read + | Write + +protToBit Read = 0x1 +protToBit Write = 0x2 +protToBit Execute = 0x4 + +data Flag + = Fixed + | Shared + | Private + | Anonymous + | NoReserve + +flagToBit Fixed = 0x10 +flagToBit Shared = 0x01 +flagToBit Private = 0x02 +flagToBit Anonymous = 0x20 +flagToBit NoReserve = 0x04000 + +errPtr :: Ptr a +errPtr = nullPtr `plusPtr` (-1) +{- +mprotect :: Ptr a -> Int -> [Protection] -> IO () +mprotect ptr size flags + = let cprot = foldr (.|.) 0 (map protToBit flags) + in do throwErrnoIf (== -1) "mprotect" (c_mprotect ptr (fromIntegral size) cprot) + return () +-} +mmap :: Int -> [Protection] -> [Flag] -> IO (Ptr a) +mmap size prot flags + = do let cprot = foldr (.|.) 0 (map protToBit prot) + cflags = foldr (.|.) 0 (map flagToBit flags) + throwErrnoIf (== errPtr) "mmap" (c_mmap nullPtr (fromIntegral size) cprot cflags (-1) 0) + +alignSize :: Int -> IO Int +alignSize size + = do page <- getPageSize + return $ if size <= 0 then page + else (size `div` page) * page + if size `mod` page == 0 then 0 else page + +mremap :: Ptr a -> Int -> Int -> IO (Ptr a) +mremap ptr oldSize newSize + = throwErrnoIf (== errPtr) "mremap" (c_mremap ptr (fromIntegral oldSize) (fromIntegral newSize) 1) + +munmap :: Ptr a -> Int -> IO () +munmap ptr size + = throwErrnoIfMinus1_ "munmap" $ c_munmap ptr (fromIntegral size) +{- +showSize :: Int -> String +showSize n' + = loop sizes (fromIntegral n') + where loop [] n = showFFloat (Just 0) (n::Float) " bytes" + loop ((s,p):xs) n | n >= s = showFFloat (Just 2) (n/s) p + | otherwise = loop xs n + sizes = [ (giga, " GiB") + , (mega, " MiB") + , (kilo, " KiB")] +-} + diff --git a/src/Data/CompactMap/Types.hs b/src/Data/CompactMap/Types.hs new file mode 100644 --- /dev/null +++ b/src/Data/CompactMap/Types.hs @@ -0,0 +1,121 @@ +{-# OPTIONS -fglasgow-exts -fallow-undecidable-instances -fbang-patterns #-} +module Data.CompactMap.Types where + +import Control.Monad +import Foreign +import Foreign.Storable +import Foreign.C + +--import Data.Generics hiding ((:+:),GT) +import Data.Char +import Data.Int +import Data.Binary +import Data.Binary.Put +--import Data.Binary.Get (getInthost) +import Data.ByteString.Internal (memcmp,inlinePerformIO) +import qualified Data.ByteString.Unsafe as B + +--import qualified Data.CompactString as C +--import qualified Data.CompactString.Unsafe as C + +import GHC.IOBase hiding (Buffer) +import GHC.Exts +import GHC.Int + +data Buffer = Buffer + { bufferData :: {-# UNPACK #-} !(IORef (ForeignPtr ())) + , bufferOld :: {-# UNPACK #-} !(IORef [ForeignPtr ()]) + , bufferPos :: {-# UNPACK #-} !FastMutInt + , bufferSize :: {-# UNPACK #-} !FastMutInt + } + + +-- Strict, unboxed IORef +data FastMutInt = FastMutInt (MutableByteArray# RealWorld) +newFastMutInt (I# i) = IO $ \s -> case newByteArray# size s of + (# s, arr #) -> case writeIntArray# arr 0# i s of + s -> (# s, FastMutInt arr #) + where I# size = sizeOf (0::Int) +readFastMutInt (FastMutInt arr) = IO $ \s -> + case readIntArray# arr 0# s of { (# s, i #) -> + (# s, I# i #) } +writeFastMutInt (FastMutInt arr) (I# i) = IO $ \s -> + case writeIntArray# arr 0# i s of { s -> + (# s, () #) } + + +{- +newtype OptInt = OptInt Int deriving (Eq,Ord,Enum,Typeable,Num,Show) +instance Binary OptInt where + {-# INLINE put #-} + put (OptInt i) = putInthost i + {-# INLINE get #-} + get = liftM OptInt getInthost +-} +{- +data RawString = RawString {-# UNPACK #-} !Int {-# UNPACK #-} !(Ptr CChar) +instance Binary RawString where + put (RawString len ptr) = error "put not defined" -- putCString (ptr,len) + get = error "get not defined" {-do bs <- get + return $! RawString (B.length bs) (unsafePerformIO $ B.unsafeUseAsCString bs return) -} +instance Ord RawString where + {-# INLINE compare #-} + compare (RawString len1 ptr1) (RawString len2 ptr2) + = inlinePerformIO $ + do n <- memcmp (castPtr $ ptr1) (castPtr ptr2) (fromIntegral $ min len1 len2) + return $! case n `compare` 0 of + EQ -> compare len1 len2 + x -> x +instance Show RawString where + show (RawString len ptr) = show (unsafePerformIO $ B.unsafePackCStringLen (ptr,len)) +instance Eq RawString where + a == b = compare a b == EQ +-} +{- +instance C.Encoding a => Binary (C.CompactString a) where + {-# INLINE put #-} + put = put . C.toByteString + {-# INLINE get #-} + get = fmap C.unsafeFromByteString get +-} + +data KeyCursor +data DataCursor + +data Index = Index { indexStart :: {-# UNPACK #-} !(Ptr IndexItem) + , indexBuffer :: {-# UNPACK #-} !Buffer } + +data IndexItem = IndexItem {-# UNPACK #-} !(Ptr IndexItem) + {-# UNPACK #-} !(Ptr ()) + {-# UNPACK #-} !(Ptr KeyCursor) + {-# UNPACK #-} !(Ptr IndexItem) + {-# UNPACK #-} !(Ptr IndexItem) -- Top, size, elem idx, left, right + + +type IdxInt = Ptr IndexItem +{-# INLINE extractField #-} +-- Get field 'f' out of the n'th IndexItem. +extractField :: Int -> (Ptr IndexItem) -> IO (Ptr IndexItem) +extractField !f !ptr = do v <- peekByteOff ptr ((sizeOf (undefined::IdxInt) * f)) + return (v::IdxInt) + +{-# INLINE putField #-} +-- Put field 'f' in the n'th IndexItem +putField :: Int -> (Ptr IndexItem) -> Ptr IndexItem -> IO () +putField !f !ptr !v = pokeByteOff ptr ((sizeOf (undefined::IdxInt) * f)) (v :: IdxInt) + + +instance Storable IndexItem where + sizeOf _ = sizeOf (undefined :: IdxInt) * 5 + alignment _ = alignment (undefined :: IdxInt) + {-# INLINE peek #-} + peek ptr = let ptr' = castPtr ptr + get n = (peekElemOff ptr' n :: IO (Ptr a)) + in liftM5 IndexItem (get 0) (get 1) (get 2) (get 3) (get 4) + {-# INLINE poke #-} + poke ptr' (IndexItem a b c d e) + = let ptr = castPtr ptr' + put n v = pokeElemOff ptr n v + in put 0 a >> put 1 b >> put 2 c >> put 3 d >> put 4 e + +