{-# LANGUAGE ExistentialQuantification, DeriveDataTypeable, RecursiveDo #-}
{-# LANGUAGE Trustworthy #-}

module Data.WeakDict (Dyn, value, makeDyn, WeakDict, emptyWeak, sizeWeak, filterWeak, insertWeak, fastDeleteWeak, deleteWeak, lookupWeak, fastDeleteWeakDyn) where

import GHC.Weak
import Control.Concurrent.MVar
import Control.Monad
import Data.Dynamic
import Data.IORef

-- | Dynamic values equipped with the ability to compare to other values.
data Dyn = forall t. (Eq t, Typeable t) => Dyn t (t -> Bool)

value (Dyn x _) = cast x

makeDyn x = Dyn x (x==)

instance Eq Dyn where
	Dyn _ eq == Dyn x _ = maybe False eq (cast x)

-----------------------------------

data WeakDict0 t u = Empty | Cons !(Weak t) u !(WeakDict t u)

data WeakDict t u = WeakDict !(MVar (WeakDict0 t u)) deriving Typeable

-- | An empty WeakDict.
emptyWeak = liftM WeakDict $ newMVar Empty

-- | Size of a WeakDict.
sizeWeak :: WeakDict t u -> IO Int
sizeWeak dict = size dict 0 where
	size (WeakDict r) n = do
		dict <- readMVar r
		case dict of
			Cons _ _ dict -> size dict (n + 1)
			Empty -> return n

-- | Filtering contents by a predicate.
filterWeak f (WeakDict r) = modifyMVar r $ \dict -> liftM ((,) dict) $ case dict of
	Cons w x dict -> deRefWeak w >>= maybe
		(filterWeak f dict)
		(\k -> if f k x then
				-- Entries before the entry to be deleted have to be reconstructed with 'insertWeak'.
				filterWeak f dict >>= insertWeak k x
			else
				filterWeak f dict)
	Empty -> return (WeakDict r)

-- | Insert a key-value pair into the dictionary, in such a way as that the pair does not keep the key (or value) alive.
insertWeak k x dict = mdo
	weak1 <- mkWeak k k $ Just $ deRefWeak weak >>= maybe (return ()) (fastDeleteWeak k)
	dict1 <- liftM WeakDict
		$ newMVar
		$! Cons weak1 x dict
	weak <- mkWeak dict1 dict1 Nothing
	return dict1

-- | Deletes a key-value pair, updating the original structure.
fastDeleteWeak k (WeakDict r) = modifyMVar_ r $ \dict -> case dict of
	Cons w2 _ (WeakDict r2) -> deRefWeak w2 >>= maybe
		(do
			fastDeleteWeak k (WeakDict r2)
			return dict)
		(\k2 -> if k == k2 then
				readMVar r2
			else do
				fastDeleteWeak k (WeakDict r2)
				return dict)
	Empty -> return dict

-- | Deletes a key-value pair, so that the original structure is not affected.
deleteWeak k (WeakDict r) = modifyMVar r $ \dict -> liftM ((,) dict) $ case dict of
	Cons w2 x dict -> deRefWeak w2 >>= maybe
		(deleteWeak k dict)
		(\k2 -> if k == k2 then
				return dict
			else
				deleteWeak k dict >>= insertWeak k2 x)

-- | Look up a value based on a key.
lookupWeak k (WeakDict r) = modifyMVar r $ \dict -> liftM ((,) dict) $ case dict of
	Cons w2 x dict -> deRefWeak w2 >>= maybe
		(lookupWeak k dict)
		(\k2 -> if k == k2 then
				return $! Just x
			else
				lookupWeak k dict)
	Empty -> return Nothing

fastDeleteWeakDyn k (WeakDict r) = modifyMVar_ r $ \dict -> case dict of
	Cons w2 _ (WeakDict r2) -> deRefWeak w2 >>= maybe
		(do
			fastDeleteWeakDyn k (WeakDict r2)
			return dict)
		(\k2 -> if maybe False (==k) $ fromDynamic k2 then
				readMVar r2
			else do
				fastDeleteWeakDyn k (WeakDict r2)
				return dict)
	Empty -> return dict