-- Intended to be imported qualified.

module GCStore(Store, Id, empty, store, retrieve, update, refcount, ref, deref) where

import qualified Data.Map.Strict as Map
import qualified Data.Set as Set

-- import Debug.Trace


type Id = Int

data Store a = Store (Map.Map Id (a, Int)) (Set.Set Id) (Set.Set Id)
  deriving (Show)

empty :: Store a
empty = Store Map.empty Set.empty Set.empty

store :: Show a => Store a -> a -> (Id, Store a)
store (Store vm fs us) a =
    let (i, fs') = if Set.null fs
                       then (maybe 1 succ (Set.lookupMax us), fs)
                       else Set.deleteFindMin fs
    in (i, Store (Map.insert i (a, 1) vm) fs' (Set.insert i us))

retrieve :: Store a -> Id -> Maybe a
retrieve (Store vm _ _) i = fst <$> Map.lookup i vm

update :: Show a => Store a -> Id -> a -> Store a
update (Store vm fs us) i a =
    maybe (error $ "GCStore.update on nonexistent id " ++ show i)
          (\(_, rc) -> Store (Map.insert i (a, rc) vm) fs us)
          (Map.lookup i vm)

refcount :: Store a -> Id -> Int
refcount (Store vm _ _) i = maybe 0 snd $ Map.lookup i vm

ref :: Store a -> Id -> Store a
ref (Store vm fs us) i = Store (Map.alter updf i vm) fs us
  where updf Nothing = error $ "GCStore.ref on nonexistent id " ++ show i
        updf (Just (a, rc)) = Just (a, rc + 1)

deref :: Store a -> Id -> Store a
deref (Store vm fs us) i =
    maybe (error $ "GCStore.deref on nonexistent id " ++ show i)
          (\(a, rc) -> if rc == 1
              then Store (Map.delete i vm) (Set.insert i fs) (Set.delete i us)
              else Store (Map.insert i (a, rc - 1) vm) fs us)
          (Map.lookup i vm)