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

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

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

import System.Exit

-- | 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 {-# NOUNPACK #-} !(IORef (WeakDict0 t u)) deriving Typeable

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

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

assocsWeak :: WeakDict t u -> IO [(t, u)]
assocsWeak (WeakDict r) = do
	dict <- readIORef r
	case dict of
		Cons w x dict -> do
			y <- deRefWeak w
			case y of
				Just y -> liftM ((y, x):) $ assocsWeak dict
				Nothing -> assocsWeak dict
		Empty -> return []

-- | Filtering contents by a predicate.
filterWeak f (WeakDict r) = do
	dict <- readIORef r
	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
		$ newIORef
		$! Cons weak1 x dict
	weak <- mkWeak dict1 dict1 Nothing
	return dict1

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

-- | Deletes a key-value pair, so that the original structure is not affected.
deleteWeak k (WeakDict r) = do
	dict <- readIORef r
	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) = do
	dict <- readIORef r
	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) = do
	dict <- readIORef r
	case dict of
		Cons w2 _ (WeakDict r2) -> deRefWeak w2 >>= maybe
			(do
				fastDeleteWeakDyn k (WeakDict r2)
				readIORef r2 >>= writeIORef r)
			(\k2 -> if maybe False (==k) $ fromDynamic k2 then
				readIORef r2 >>= writeIORef r
			else
				fastDeleteWeakDyn k (WeakDict r2))
		Empty -> return ()