ContextRef.hs

{-# LANGUAGE GeneralizedNewtypeDeriving, RankNTypes, RecursiveDo #-}
module ContextRefHidden 
         ( ContextRef
         , Context
         , runContext
         , newRef
         , readRef
         , writeRef
         , refKey
         , attempt
         , save
         ) where

import Control.Monad.ST
import Control.Monad.State

import Data.Map
import Data.Maybe
import Data.IORef

import System.Mem.Weak

data ContextRef s a = ContextRef Int (IORef a)
  deriving (Eq)

instance Ord (ContextRef s a) where
  compare (ContextRef a _) (ContextRef b _) = compare a b
  
data ContextData s = ContextData [Int] (Map Int (IO (IO ()))) 

newtype Context s a = Context { unContext :: (StateT (IORef (ContextData s)) (ST s) a)}
 deriving (Monad, MonadFix)

runContext :: (forall s. Context s a) -> a
runContext s = runST $ do 
  initialData <- newSTIORef $ ContextData [0..] empty
  (a, _) <- runStateT (unContext s) initialData
  return a

newRef :: a -> Context s (ContextRef s a)
newRef x = Context $ mdo
  ref <- lift $ newSTIORef x
  dataRef <- get
  key <- lift $ unsafeIOToST $ atomicModifyIORef dataRef $ 
     \(ContextData (nextKey:availableKeys) saveMap) ->  (ContextData availableKeys (insert nextKey (saveIORef weakRef) saveMap), nextKey)
  weakRef <- lift $ mkWeakSTIORef key dataRef ref
  return $ ContextRef key ref
  
readRef :: ContextRef s a -> Context s a
readRef (ContextRef _ ref) = Context $ lift $ readSTIORef ref

writeRef :: ContextRef s a -> a -> Context s ()
writeRef (ContextRef _ ref) x = Context $ lift $ writeSTIORef ref x

refKey :: ContextRef s a -> Int
refKey (ContextRef key _) = key

-- any references made during the attempt will not get altered by the restore
attempt :: Context s (Either a b) -> Context s (Either a b)
attempt m = do
  restore <- save
  result <- m
  when (isLeft result) restore
  return result
  
save :: Context s (Context s ())
save = Context $ do
  dataRef <- get
  ContextData _ saveMap <- lift $ readSTIORef dataRef
  restoreActions <- lift $ unsafeIOToST $ sequence $ elems saveMap
  return ((Context $ lift $ unsafeIOToST $ sequence $ restoreActions) >> return ())



saveIORef :: Weak (IORef a) -> IO (IO ())
saveIORef weakRef = do
  maybeRef <- deRefWeak weakRef
  case maybeRef of 
    Just ref -> do
      x <- readIORef ref
      return (writeIORef ref x)
    Nothing -> return (return ())
    

readSTIORef :: IORef a -> ST s a
readSTIORef x = unsafeIOToST $ readIORef x

newSTIORef :: a -> ST s (IORef a)
newSTIORef x = unsafeIOToST $ newIORef x

writeSTIORef :: IORef a -> a -> ST s ()
writeSTIORef ref x = unsafeIOToST $ writeIORef ref x

mkWeakSTIORef :: Int -> IORef (ContextData s) -> IORef a -> ST s (Weak (IORef a))
mkWeakSTIORef i dataRef ref = unsafeIOToST $ mkWeakIORef ref (atomicModifyIORef dataRef cleanUpRef)
  where
    cleanUpRef (ContextData is saveMap) = (ContextData (i:is) (delete i saveMap), ())


isLeft :: Either a b -> Bool
isLeft (Left _) = True
isLeft _ = False