Haskell mutable map/tree

I am looking for a mutable (balanced) tree/map/hash table in Haskell or a way how to simulate it inside a function. I.e. when I call the same function several times, the structure is preserved. So far I have tried Data.HashTable (which is OK, but somewhat slow) and tried Data.Array.Judy but I w开发者_如何学编程as unable to make it work with GHC 6.10.4. Are there any other options?

If you want mutable state, you can have it. Just keep passing the updated map around, or keep it in a state monad (which turns out to be the same thing).

import qualified Data.Map as Map
import Control.Monad.ST
import Data.STRef
memoize :: Ord k => (k -> ST s a) -> ST s (k -> ST s a)
memoize f = do
    mc <- newSTRef Map.empty
    return $ \k -> do
        c <- readSTRef mc
        case Map.lookup k c of
            Just a -> return a
            Nothing -> do a <- f k
                          writeSTRef mc (Map.insert k a c) >> return a

You can use this like so. (In practice, you might want to add a way to clear items from the cache, too.)

import Control.Monad
main :: IO ()
main = do
    fib <- stToIO $ fixST $ \fib -> memoize $ \n ->
        if n < 2 then return n else liftM2 (+) (fib (n-1)) (fib (n-2))
    mapM_ (print <=< stToIO . fib) [1..10000]

At your own risk, you can unsafely escape from the requirement of threading state through everything that needs it.

import System.IO.Unsafe
unsafeMemoize :: Ord k => (k -> a) -> k -> a
unsafeMemoize f = unsafePerformIO $ do
    f' <- stToIO $ memoize $ return . f
    return $ unsafePerformIO . stToIO . f'

fib :: Integer -> Integer
fib = unsafeMemoize $ \n -> if n < 2 then n else fib (n-1) + fib (n-2)

main :: IO ()
main = mapM_ (print . fib) [1..1000]

Building on @Ramsey's answer, I also suggest you reconceive your function to take a map and return a modified one. Then code using good ol' Data.Map, which is pretty efficient at modifications. Here is a pattern:

import qualified Data.Map as Map

-- | takes input and a map, and returns a result and a modified map
myFunc :: a -> Map.Map k v -> (r, Map.Map k v)
myFunc a m = … -- put your function here

-- | run myFunc over a list of inputs, gathering the outputs
mapFuncWithMap :: [a] -> Map.Map k v -> ([r], Map.Map k v)
mapFuncWithMap as m0 = foldr step ([], m0) as
    where step a (rs, m) = let (r, m') = myFunc a m in (r:rs, m')
    -- this starts with an initial map, uses successive versions of the map
    -- on each iteration, and returns a tuple of the results, and the final map

-- | run myFunc over a list of inputs, gathering the outputs
mapFunc :: [a] -> [r]
mapFunc as = fst $ mapFuncWithMap as Map.empty
    -- same as above, but starts with an empty map, and ignores the final map

It is easy to abstract this pattern and make mapFuncWithMap generic over functions that use maps in this way.

Although you ask for a mutable type, let me suggest that you use an immutable data structure and that you pass successive versions to your functions as an argument.

Regarding which data structure to use,

• There is an implementation of red-black trees at Kent

• If you have integer keys, Data.IntMap is extremely efficient.

• If you have string keys, the bytestring-trie package from Hackage looks very good.

The problem is that I cannot use (or I don't know how to) use a non-mutable type.

If you're lucky, you can pass your table data structure as an extra parameter to every function that needs it. If, however, your table needs to be widely distributed, you may wish to use a state monad where the state is the contents of your table.

If you are trying to memoize, you can try some of the lazy memoization tricks from Conal Elliott's blog, but as soon as you go beyond integer arguments, lazy memoization becomes very murky—not something I would recommend you try as a beginner. Maybe you can post a question about the broader problem you are trying to solve? Often with Haskell and mutability the issue is how to contain the mutation or updates within some kind of scope.

It's not so easy learning to program without any global mutable variables.

If I read your comments right, then you have a structure with possibly ~500k total values to compute. The computations are expensive, so you want them done only once, and on subsequent accesses, you just want the value without recomputation.

In this case, use Haskell's laziness to your advantage! ~500k is not so big: Just build a map of all the answers, and then fetch as needed. The first fetch will force computation, subsequent fetches of the same answer will reuse the same result, and if you never fetch a particular computation - it never happens!

You can find a small implementation of this idea using 3D point distances as the computation in the file PointCloud.hs. That file uses Debug.Trace to log when the computation actually gets done:

> ghc --make PointCloud.hs 
[1 of 1] Compiling Main             ( PointCloud.hs, PointCloud.o )
Linking PointCloud ...

> ./PointCloud 
(1,2)
(<calc (1,2)>)
Just 1.0
(1,2)
Just 1.0
(1,5)
(<calc (1,5)>)
Just 1.0
(1,2)
Just 1.0

Are there any other options?

A mutable reference to a purely functional dictionary like Data.Map.