HsJudy-0.1: Judy/Hash.hs

{-# OPTIONS -fallow-undecidable-instances -fallow-incoherent-instances #-}

{-# INCLUDE "Judy.h" #-}

module Judy.Hash (
    Hash (..),

    -- FIXME: need to move to MapM api
    freeze
) where

import Data.Typeable
import Control.Monad (when)
import Foreign.C.String
-- import Foreign.C.Types
-- import Foreign.ForeignPtr
-- import Foreign.Marshal.Alloc
-- import Foreign.Ptr
-- import Foreign.Storable
import Foreign
import Data.Maybe (fromJust)

import Judy.Private
import qualified Judy.CollectionsM as CM
import Judy.Refeable
import Judy.Freeze
import Judy.Stringable
import qualified Judy.MiniGC as GC

import Prelude hiding (map)

-- FIXME: really necessary/useful restrict types here?
newtype (Stringable k, Refeable a) => Hash k a = Hash { judy :: ForeignPtr JudyHS }
    deriving (Eq, Ord, Typeable)

instance (Stringable k, Refeable a) => CM.MapM (Hash k a) k a IO where
    new = new_
    delete = delete_
    member = member_
    lookup = lookup_
    insert = insert_
    alter = alter_
    fromList = fromList_
    toList = toList_
    elems = elems_
    keys = keys_
    mapToList = mapToList_
    swapMaps = swapMaps_

instance (Stringable k, Refeable a) => Freezable (Hash k a) where
    freeze m = do
        m' <- new_
        swapMaps_ m' m
        return (Frozen m')

instance (Stringable k, Refeable a) => CM.MapF (Frozen (Hash k a)) k a where
    memberF k (Frozen m) = unsafePerformIO $ member_ k m
    lookupF k (Frozen m) = unsafePerformIO $ lookup_ k m
    fromListF l = Frozen $ unsafePerformIO $ fromList_ l
    toListF (Frozen m) = unsafePerformIO $ toList_ m

instance Show (Hash k a) where
    show (Hash j) = "<Hash " ++ show j ++ ">"

foreign import ccall "wrapper" mkFin :: (Ptr JudyHS -> IO ()) -> IO (FunPtr (Ptr JudyHS -> IO ()))

finalize :: Bool -> Ptr JudyHS -> IO ()
finalize need j = do
    when need $ do
        j_ <- newForeignPtr_ j
        es <- rawElems (Hash j_)
        mapM_ GC.freeRef es
    v <- judyHSFreeArray j judyError
    --putStrLn $ "\n(FINALIZER CALLED FOR "++ (show j) ++  ": " ++ (show v) ++ ")\n"
    return ()

rawElems :: Hash k a -> IO [Value]
rawElems = internalMap $ \r _ _ -> peek r

dummy :: Refeable a => Hash k a -> a
dummy = undefined


new_ :: Refeable a => IO (Hash k a)
new_ = do
    fp <- mallocForeignPtr
    withForeignPtr fp $ flip poke nullPtr
    m <- return $ Hash fp

    finalize' <- mkFin $ finalize $ needGC (dummy m)
    addForeignPtrFinalizer finalize' fp
    return m

insert_ :: (Stringable k, Refeable a) => k -> a -> Hash k a -> IO ()
insert_ k v (Hash j) = withForeignPtr j $ \j' -> do
    useAsCSLen k $ \(cp, len) -> do
        -- TODO: maybe there's a better way to convert Int -> Value
        r <- judyHSIns j' cp (fromIntegral len) judyError
        if r == pjerr
            then error "HsJudy: Not enough memory."
            else do
                v' <- toRef v
                poke r v'
                return ()

alter_ :: (Eq a, Stringable k, Refeable a) => (Maybe a -> Maybe a) -> k -> Hash k a -> IO (Maybe a)
alter_ f k m@(Hash j) = do
    j' <- withForeignPtr j peek
    useAsCSLen k $ \(cp, len) -> do
        r <- judyHSGet j' cp (fromIntegral len)
        if r == nullPtr
            then if (f Nothing) == Nothing
                    then return Nothing
                    else insert_ k (fromJust (f Nothing)) m >> return (f Nothing)
            else do
                v' <- peek r
                v <- fromRef v'
                let fv = f (Just v)
                if fv == Nothing
                    then do delete_ k m
                            return Nothing
                    else if v /= (fromJust fv)
                             then do when (needGC (fromJust fv)) $ GC.freeRef v'
                                     x <- toRef (fromJust fv)
                                     poke r x
                                     return fv
                             else return fv

lookup_ :: (Stringable k, Refeable a) => k -> Hash k a -> IO (Maybe a)
lookup_ k (Hash j) = do
    j' <- withForeignPtr j peek
    useAsCSLen k $ \(cp, len) -> do
        r <- judyHSGet j' cp (fromIntegral len)
        if r == nullPtr
            then return Nothing
            else do
                v' <- peek r
                v <- fromRef v'
                return $ Just v

member_ :: Stringable k => k -> Hash k a -> IO Bool
member_ k (Hash j) = do
    j' <- withForeignPtr j peek
    useAsCSLen k $ \(cp, len) -> do
        r <- judyHSGet j' cp (fromIntegral len)
        return $ r /= nullPtr

delete_ :: Stringable k => k -> Hash k a -> IO Bool
delete_ k m@(Hash j) = withForeignPtr j $ \j' -> do
    j'' <- peek j'
    useAsCSLen k $ \(cp, len) -> do
        when (needGC (dummy m)) $ do
            r <- judyHSGet j'' cp (fromIntegral len)
            if r == nullPtr
                then return ()
                else do v' <- peek r
                        GC.freeRef v'
                        return ()
        r <- judyHSDel j' cp (fromIntegral len) judyError
        return $ r /= 0

-- FIXME: may use HashIter type to enforce some safety in its use?
newtype HashIter = HashIter { iter :: ForeignPtr JudyHSIter }
    deriving (Eq, Ord, Typeable)

instance Show HashIter where
    show (HashIter i) = "<Iter "++ show i ++ ">"


newIter :: IO (HashIter)
newIter = do
    fp <- mallocForeignPtr
    addForeignPtrFinalizer judyHSIter_free_ptr fp
    withForeignPtr fp $ flip poke nullPtr
    return $ HashIter fp

fromList_ :: (Stringable k, Refeable a) => [(k,a)] -> IO (Hash k a)
fromList_ xs = do
    m <- new_
    mapM_ (\(k,a) -> insert_ k a m) xs
    return m

internalMap :: (Ptr Value -> Ptr CString -> Ptr Value -> IO b) -> Hash k a -> IO [b]
internalMap f (Hash j) = do
    jj <- withForeignPtr j peek
    (HashIter i) <- newIter
    withForeignPtr i $ \ii -> alloca $ \cp -> alloca $ \len -> do
        poke len 0
        jp_null cp
        let loop act xs = do
            r <- act jj ii cp len judyError
            if r == nullPtr
                then return xs
                else do x <- f r cp len
                        loop judyHSIterNext (x:xs)
        loop judyHSIterFirst []

mapToList_ :: (Stringable k, Refeable a) => (k -> a -> b) -> Hash k a -> IO [b]
mapToList_ f = internalMap $ \r cp len -> do
    l <- peek len
    c <- peek cp
    v <- copyCSLen (c, fromIntegral l)
    d <- peek r
    d' <- fromRef d
    return $ f v d'

toList_ :: (Stringable k, Refeable a) => Hash k a -> IO [(k,a)]
toList_ = mapToList_ $ \k a -> (k, a)

elems_ :: Refeable a => Hash k a -> IO [a]
elems_ = internalMap $ \r _ _ -> do
    d <- peek r
    fromRef d

keys_ :: Stringable k => Hash k a -> IO [k]
keys_ = internalMap $ \_ cp len -> do
    l <- peek len
    c <- peek cp
    v <- copyCSLen (c, fromIntegral l)
    return v


swapMaps_ :: Hash k a -> Hash k a -> IO ()
swapMaps_ (Hash j1) (Hash j2) = do
    withForeignPtr j1 $ \p1 -> withForeignPtr j2 $ \p2 -> do
        v1 <- peek p1
        v2 <- peek p2
        poke p1 v2
        poke p2 v1