Mini adapton in Haskell

Miniadapton.hs

{-# LANGUAGE RecordWildCards, FlexibleContexts, ScopedTypeVariables, FlexibleInstances, OverlappingInstances, GADTs, StandaloneDeriving #-}
import Data.Set as Set
import Data.IORef
import Data.Maybe
import Data.Dynamic hiding (cast)
import Data.Function(on)
import Control.Applicative
import Control.Monad.Fix
import System.IO.Unsafe

data AThunk = forall a. AThunk (Thunk a)

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

instance (Show a, Typeable a) => Show (Thunk a) where
  show Thunk{..} = "Thunk {" ++ show ident
      ++ (show . toList . unsafePerformIO . readIORef $ sub)
      ++ (show . toList . unsafePerformIO . readIORef $ super)
      ++ show thunk
      ++ " = " ++ (show . unsafePerformIO . readIORef $ result) ++ "}"

instance Show AThunk where
  show (AThunk x) = show (ident x)

instance Eq AThunk where
  AThunk a == AThunk b = ident a == ident b

instance Eq (Thunk a) where
  a == b = ident a == ident b

instance Ord AThunk where
  AThunk a `compare` AThunk b = ident a `compare` ident b

instance Typeable a => Show (IO a) where
  show x = "<" ++ (show . typeOf . unsafePerformIO $ x) ++ ">"

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

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

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

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

compute a@Thunk{..} = do
    r <- readIORef result
    case r of
        Just r -> do
          putStrLn $ "read " ++ show ident ++ ": " ++ show a
          return r
        Nothing -> do
            mapM_ (\(AThunk x) -> remove a x) <$> readIORef sub
            v <- mfix (\v -> do
              writeIORef result (Just v)
              thunk)
            vf <- readIORef result
            putStrLn $ "eval " ++ show ident ++ ": " ++ show a ++ " -> " ++ show v
            return v

dirty :: Thunk a -> IO ()
dirty (a@Thunk{..}) = do
  x <- readIORef result
  case x of
    Nothing -> return ()
    Just r -> do
      writeIORef result Nothing
      putStrLn $ "dirty " ++ show ident
      mapM_ (\(AThunk x) -> dirty x) =<< readIORef super

aref :: a -> IO (Thunk a)
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 (a@Thunk{..}) val = do
    x <- readIORef result
    case x of
      v | v == val -> return ()
      _ -> do
        dirty a
        writeIORef result val

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

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

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 b) -> IO (a -> IO (Thunk b))
amemoize_l f = memoize (\x -> athunk (f x))

amemoize f = do
  f' <- amemoize_l f
  return $ \x -> f' x >>= force

type AVar b = Thunk (Thunk b)

avar expr = aref =<< athunk expr
avar_get v = do
  thnk <- force v
  force thnk

avar_set v expr = do
  e <- athunk expr
  set v (Just e)

plus :: (Num a, Show a, Typeable a) => AVar a -> AVar a -> IO a
plus a b = liftA2 (+) (avar_get a) (avar_get b)

n = return :: Integer -> IO Integer

spreadsheet = do
  n1 <- avar $ n 1
  n2 <- avar $ n 2
  n3 <- avar $ n 3
  p1 <- avar $ n1 `plus` n2
  p2 <- avar $ p1 `plus` n3
  print =<< avar_get p1
  print =<< avar_get p2
  avar_set n1 (n 5)
  print =<< avar_get p1
  print =<< avar_get p2
  avar_set p2 (n3 `plus` p1)
  print =<< avar_get p2
  avar_set p1 (n 4)
  print =<< avar_get p1
  print =<< avar_get p2
  avar_set p1 (n1 `plus` n2)
  print =<< avar_get p1
  print =<< avar_get p2

  
violation = do
  nondet <- newIORef False
  ref <- avar $ n 5
  r2 <- avar $ do
    a <- avar_get ref
    n <- readIORef nondet
    let y = if n then a else a*a
    avar_set ref $ return y
    return y
  return (ref, r2)

I'm not sure if it's actually doing incremental computation.

Fixed now. It has (my) desired behavior, e.g. for violation:

(r1, r2) <- violation
avar_get r1
avar_get r2
avar_get r1

At this point r1 = 25 is clean; calling avar_get r2 will square r1 again.