{-# LANGUAGE RecordWildCards, ScopedTypeVariables #-}
import Data.Set
import Data.IORef
import Data.Maybe
import Data.Dynamic hiding (cast)
import Data.Function(on)
import Control.Applicative
import System.IO.Unsafe

data Thunk = Thunk
    { result :: IORef (Maybe Dynamic)
    , sub :: IORef (Set Thunk)
    , super :: IORef (Set Thunk)
    , thunk :: IO Dynamic
    , ident :: Int
    }

instance Eq Thunk where
  (==) = (==) `on` ident

instance Ord Thunk where
  compare = compare `on` ident

instance Show Thunk where
  show (Thunk{..}) = "Thunk { ident = " ++ show ident ++ ", result = " ++ (show . unsafePerformIO . readIORef $ result) ++ "}"

{-# NOINLINE counter #-}
counter :: IORef Int
counter = unsafePerformIO $ newIORef 0

athunk :: IO Dynamic -> IO Thunk
athunk thunk = do
    ident <- atomicModifyIORef counter (\c -> (c+1,c))
    result <- newIORef Nothing
    sub <- newIORef mempty
    super <- newIORef mempty
    return Thunk {..}

add :: Thunk -> Thunk -> IO ()
add sp sb = do
    modifyIORef (sub sp) (insert sb)
    modifyIORef (super sb) (insert sp)

remove :: Thunk -> Thunk -> IO ()
remove sp sb = do
    modifyIORef (sub sp) (delete sb)
    modifyIORef (super sp) (delete sp)

compute :: Thunk -> IO Dynamic
compute a@Thunk{..} = do
    r <- readIORef result
    case r of
        Just r -> return r
        Nothing -> do
            mapM_ (remove a) <$> readIORef sub
            v <- thunk
            writeIORef result (Just v)
            return v

dirty :: Thunk -> IO()
dirty (a@Thunk{..}) = do
    writeIORef result Nothing
    mapM_ dirty =<< readIORef super

aref :: Dynamic -> IO Thunk
aref val = do
    ident <- atomicModifyIORef counter (\c -> (c+1,c))
    result <- newIORef (Just val)
    sub <- newIORef mempty
    super <- newIORef mempty
    return Thunk { thunk = fromJust <$> readIORef result, ..}

set :: Thunk -> Maybe Dynamic -> IO ()
set (a@Thunk{..}) val = do
    dirty a
    writeIORef result val

{-# NOINLINE adapting #-}
adapting :: IORef (Maybe Thunk)
adapting = unsafePerformIO $ newIORef Nothing

force :: Thunk -> IO Dynamic
force a@Thunk{..} = do
    prev <- readIORef adapting
    writeIORef adapting (Just a)
    result <- compute a
    writeIORef adapting prev
    case prev of
        Nothing -> return ()
        Just ad -> add ad a
    return result

adapt = athunk

memoize :: (Eq a) => (a -> IO b) -> IO (a -> IO b)
memoize f = do
  s <- newIORef []
  return $ \(x :: a) -> do
    v <- lookup x <$> readIORef s
    case v of
      Just v' -> return v'
      Nothing -> do
        r <- f x
        modifyIORef s ((x,r):)
        return r

amemoize_l :: Eq a => (a -> IO Dynamic) -> IO (a -> IO Thunk)
amemoize_l f = memoize (\x -> adapt (f x))

amemoize :: Eq a => (a -> IO Dynamic) -> IO (a -> IO Dynamic)
amemoize f = do
  f' <- amemoize_l f
  return $ \x -> f' x >>= force

cast :: Typeable a => Dynamic -> a
cast x = let a = fromDyn x (error $ "Bad cast to " ++ show (typeOf a) ++ ": " ++ show x) in a

avar :: IO Dynamic -> IO Thunk
avar expr = aref (toDyn $ adapt expr)
avar_get :: Thunk -> IO Dynamic
avar_get v = do
  thnk <- force v
  val <- cast thnk
  force val

avar_set :: Thunk -> IO Dynamic -> IO ()
avar_set v expr = set v (Just . toDyn $ adapt expr)

n x = do
  putStrLn $ "n " ++ show x
  return . toDyn $ (x :: Integer)

plus a b = toDyn <$> do
  putStrLn $ "plus " ++ show a ++ " " ++ show b
  liftA2 ((+) :: Integer -> Integer -> Integer) (cast <$> avar_get a) (cast <$> avar_get b)

prt :: Dynamic -> String
prt x = show $ (cast x :: Integer)

main = do
  n1 <- avar $ n 1
  n2 <- avar $ n 2
  n3 <- avar $ n 3
  p1 <- avar $ n1 `plus` n2
  p2 <- avar $ p1 `plus` n3
  putStrLn . prt =<< avar_get p1
  putStrLn . prt =<< avar_get p2
  avar_set n1 (n 5)
  putStrLn . prt =<< avar_get p1
  avar_set p2 (n3 `plus` p1)
  putStrLn . prt =<< avar_get p2
  avar_set p1 (n 4)
  putStrLn . prt =<< avar_get p2
  avar_set p1 (n1 `plus` n2)
  putStrLn . prt =<< avar_get p2