Spreadsheet.hs

{-# LANGUAGE GADTs #-}

module Spreadsheet where

import Control.Applicative
import Control.Monad (forM_)
import Data.IORef
import Data.List (union)
import Data.Unique

data Result a = Result a [ECell]

data Exp a = Exp { runExp :: IO (Result a) }

data Cell a = Cell {
    _code      :: IORef (Exp a )
  , _value     :: IORef (Maybe a)
  , _reads     :: IORef [ECell]
  , _observers :: IORef [ECell]
  , _id        :: Unique
}

data ECell where
    Pack :: Cell a -> ECell

instance Eq (ECell) where
    Pack a == Pack b = _id a == _id b

instance Functor Exp where
    fmap f exp = Exp $ do
        Result a bs <- runExp exp
        return $ Result (f a) bs

instance Applicative Exp where
    pure a = Exp $ return $ Result a []
    expf <*> expa = Exp $ do
        Result f cs <- runExp expf
        Result a ds <- runExp expa
        return $ Result (f a) (cs `union` ds)

instance Monad Exp where
    return a = Exp $ return $ Result a []
    exp >>= f = Exp $ do
        Result a cs <- runExp $ exp
        Result b ds <- runExp $ f a
        return $ Result b (cs `union` ds)

cell :: Exp a -> IO (Cell a)
cell exp =
    Cell <$> newIORef exp
         <*> newIORef Nothing
         <*> newIORef []
         <*> newIORef []
         <*> newUnique

get :: Cell a -> IO (Result a)
get c = do
    v <- readIORef (_value c)
    case v of
        Just a  -> return $ Result a [Pack c]
        Nothing -> do
            Result a ds <- runExp =<< readIORef (_code c)
            writeIORef (_value c) (Just a)
            writeIORef (_reads c) ds
            forM_ ds $ \(Pack d) -> modifyIORef (_observers d) (Pack c:)
            return $ Result a [Pack c]

removeObserver :: ECell -> ECell -> IO ()
removeObserver o (Pack c) =
    modifyIORef (_observers c) (filter (/= o))

invalidate :: ECell -> IO ()
invalidate (Pack c) = do
    os <- readIORef (_observers c)
    rs <- readIORef (_reads c)
    writeIORef (_observers c) []
    writeIORef (_reads c)     []
    writeIORef (_value c)     Nothing
    forM_ rs $ removeObserver (Pack c)
    forM_ os $ invalidate
    
set :: Cell a -> Exp a -> IO ()
set c exp = do
    writeIORef (_code c) exp
    invalidate (Pack c)

run :: Exp a -> IO a
run exp = runExp exp >>= \(Result a _ ) -> return a


runCell c = readIORef (_code c) >>= run


------ Example -------------

expr :: IO a -> Exp a
expr io = Exp $ do
    a <- io
    return $ Result a []

unop :: (a -> b) -> Cell a -> Exp b
unop f c = Exp $ do
    Result a bs <- get c
    return $ Result (f a) bs

binop :: (a -> b -> c) -> Cell a -> Cell b -> Exp c
binop f ca cb = Exp $ do
    Result a cs <- get ca
    Result b ds <- get cb
    return $ Result (f a b) (cs `union` ds)

foldCells :: (a -> b -> b) -> Exp b -> [Cell a] -> Exp b
foldCells f exp []       = exp
foldCells f exp (ca:cas) = Exp $ do
    Result a cs <- get ca
    Result b ds <- runExp (foldCells f exp cas)
    return $ Result (f a b) (cs `union` ds)

sumCells :: Num a => [Cell a] -> Exp a
sumCells cs = foldCells (+) (return 0) cs

add :: Num a => Cell a -> Cell a -> Exp a
add = binop (+)

sub :: Num a => Cell a -> Cell a -> Exp a
sub = binop (-)

mul :: Num a => Cell a -> Cell a -> Exp a
mul = binop (*)

ratio :: Fractional a => Cell a -> Cell a -> Exp a
ratio = binop (/) 

exponential :: Floating a => Cell a -> Exp a
exponential = unop exp

logarithm :: Floating a => Cell a -> Exp a
logarithm = unop exp