module Data.Refcount ( Refcount (..)
, refcount
, refcounted
, insertRef
, deleteRef
, removeRef
, Data.Refcount.fromList
) where
import Control.Applicative ((<$>))
import Data.Foldable as Foldable
import Data.Maybe
import Data.Monoid
import Data.Hashable
import Data.HashMap.Strict as HashMap
import GHC.Generics
import Test.QuickCheck (Arbitrary (..))
newtype Refcount a = Refcount { unRefcount :: HashMap a Int }
deriving (Show, Generic, Eq)
instance (Hashable a, Eq a) => Monoid (Refcount a) where
mempty = Refcount empty
mappend (Refcount a) (Refcount b) = Refcount $ unionWith (+) a b
instance (Arbitrary a, Hashable a, Eq a) => Arbitrary (Refcount a) where
arbitrary = Data.Refcount.fromList <$> arbitrary
refcounted :: Refcount a -> [a]
refcounted = keys . unRefcount
fromList :: (Hashable a, Eq a) => [a] -> Refcount a
fromList = Foldable.foldl' (\rc r -> insertRef r rc) mempty
refcount :: (Hashable a, Eq a) => a -> Refcount a -> Int
refcount a (Refcount h) = fromMaybe 0 $ HashMap.lookup a h
insertRef :: (Hashable a, Eq a) => a -> Refcount a -> Refcount a
insertRef a (Refcount m) = Refcount $ insertWith (+) a 1 m
deleteRef :: (Hashable a, Eq a) => a -> Refcount a -> Maybe (Refcount a)
deleteRef a r = snd <$> removeRef a r
removeRef :: (Hashable a, Eq a) => a -> Refcount a -> Maybe (Bool, Refcount a)
removeRef a (Refcount m) = fmap decRef $ HashMap.lookup a m
where decRef count = if count > 1
then (False, Refcount $ adjust (subtract 1) a m)
else (True, Refcount $ delete a m)