{-# LANGUAGE Rank2Types, GeneralizedNewtypeDeriving, DeriveFunctor #-}
module STMOracle
    ( Oracle, Ref
    , newRef, readRef, writeRef, modifyRef, needRef
    ) where

import Control.Applicative
import Control.Monad
import Control.Concurrent.STM

instance Applicative STM where
    pure = return
    (<*>) = ap

newtype Ref s a = Ref (TVar (Maybe a))
newtype Oracle s a = Oracle { unOracle :: forall r. (a -> STM r) -> STM r } deriving (Functor)

instance Monad (Oracle s) where
  return x = Oracle (\k -> k x)
  Oracle m >>= f = Oracle (\k -> m (\a -> unOracle (f a) k))

mkOracle m = Oracle (m >>=)

runOracle :: (forall s. Oracle s a) -> IO a
runOracle t = atomically (unOracle t return)

newRef :: a -> Oracle s (Ref s a)
newRef a = mkOracle $ Ref <$> newTVar (Just a)

readRef :: Ref s a -> Oracle s a
readRef (Ref r) = mkOracle $ do
    m <- readTVar r
    maybe retry return m

writeRef :: a -> Ref s a -> Oracle s a
writeRef a (Ref r) = mkOracle $ do
    writeTVar r (Just a)
    return a

modifyRef :: (a -> a) -> Ref s a -> Oracle s a
modifyRef f r = do
    a <- readRef r
    writeRef (f a) r

needRef :: Ref s a -> Oracle s Bool
needRef (Ref slot) = Oracle $ \k ->
             (writeTVar slot Nothing >> k False)
    `orElse` k True

-- test case:                                                                                                                                
refMaybe b dflt ref = if b then readRef ref else return dflt
refIgnore ref = return "blablabla"
refFst ref = fst `fmap` readRef ref
test = do
     a <- newRef "x"
     b <- newRef 1
     c <- newRef ('z', Just 0)
     -- no performLocalGC required                                                                                                           
     x <- needRef a
     y <- needRef b
     z <- needRef c
     u <- refMaybe y "t" a -- note that it wouldn't actually read "a",                                                                       
                           -- but it won't be known until runtime.                                                                           
     w <- refIgnore b
     v <- refFst c
     return (x, y, z)