BUBS.hs

{-# LANGUAGE PatternGuards, FlexibleInstances, MultiParamTypeClasses #-}

-- An implementation of bottom-up beta substitution.  From the paper: 
-- Bottom-Up β-Substitution: Uplinks and λ-DAGs
-- By Olin Shivers & Mitchell Wand.  2004.

module BUBS 
    ( Term, eval )
where

import qualified HOAS
import Data.IORef
import Control.Monad (forM_, (<=<), when)
import Control.Applicative
import Control.Monad.IO.Class
import Control.Monad.Trans.Class
import Control.Monad.Trans.State
import Control.Monad.Trans.Writer
import System.Process (system)
import Data.Maybe (fromJust, catMaybes)
import Data.List (delete)

data UplinkType = UplinkAppL | UplinkAppR | UplinkLambda | UplinkVar
    deriving (Eq)

type Uplink a = (UplinkType, NodeRef a)

data NodeData a
    = AppNode (NodeRef a) (NodeRef a)
    | LambdaNode (NodeRef a) (NodeRef a)
    | VarNode
    | PrimNode a

data Node a = Node {
    nodeCache :: Maybe (NodeRef a),
    nodeUplinks :: [Uplink a],
    nodeData :: NodeData a
  }

type NodeRef a = IORef (Node a)


upcopy :: NodeRef a -> NodeRef a -> Uplink a -> IO ()
upcopy stop newchild (uplinkType, intoref) | intoref == stop = return ()
                                           | otherwise = do
    into <- readIORef intoref

    let traverse newnode = mapM_ (upcopy stop newnode) (nodeUplinks into)
    
    case nodeData into of
        AppNode left right -> do
            case nodeCache into of
                Nothing -> do
                    let dat' | UplinkAppL <- uplinkType = AppNode newchild right
                             | UplinkAppR <- uplinkType = AppNode left newchild
                    newnode <- newNodeRef dat'
                    setCache intoref (Just newnode)
                    traverse newnode
                Just cache -> do
                    case uplinkType of
                        UplinkAppL -> replaceLeft newchild cache
                        UplinkAppR -> replaceRight newchild cache
        LambdaNode var body -> do
            var' <- newNodeRef VarNode
            lambda' <- newNodeRef (LambdaNode var' newchild)
            setCache intoref (Just lambda')
            upcopy lambda' var' (UplinkVar, var)
            traverse lambda'
        VarNode -> do
            setCache intoref (Just newchild)
            traverse newchild

setCache :: NodeRef a -> Maybe (NodeRef a) -> IO ()
setCache ref newcache = modifyIORef ref (\n -> n { nodeCache = newcache })

newNodeRef :: NodeData a -> IO (NodeRef a)
newNodeRef dat = newIORef $ Node { nodeCache = Nothing, nodeUplinks = [], nodeData = dat }

replaceLeft :: NodeRef a -> NodeRef a -> IO ()
replaceLeft newchild node = modifyIORef node $ \n -> n { nodeData = go (nodeData n) }
    where
    go (AppNode l r) = AppNode newchild r

replaceRight :: NodeRef a -> NodeRef a -> IO ()
replaceRight newchild node = modifyIORef node $ \n -> n { nodeData = go (nodeData n) }
    where
    go (AppNode l r) = AppNode l newchild

replaceBody :: NodeRef a -> NodeRef a -> IO ()
replaceBody newchild node = modifyIORef node $ \n -> n { nodeData = go (nodeData n) }
    where
    go (LambdaNode v b) = LambdaNode v newchild

addUplink :: Uplink a -> NodeRef a -> IO ()
addUplink uplink node = modifyIORef node $ \n -> n { nodeUplinks = uplink : nodeUplinks n }

deleteUplink :: Uplink a -> NodeRef a -> IO ()
deleteUplink uplink node = modifyIORef node $ \n -> n { nodeUplinks = delete uplink (nodeUplinks n) }

getLeft :: NodeRef a -> IO (NodeRef a)
getLeft ref = (\(AppNode l r) -> l) . nodeData <$> readIORef ref

getRight :: NodeRef a -> IO (NodeRef a)
getRight ref = (\(AppNode l r) -> r) . nodeData <$> readIORef ref

getBody :: NodeRef a -> IO (NodeRef a)
getBody ref = (\(LambdaNode _ b) -> b) . nodeData <$> readIORef ref

clear :: NodeRef a -> IO ()
clear noderef = do
    node <- readIORef noderef
    forM_ (nodeUplinks node) $ \(uplinkType, uplinkRef) -> do
        upnode <- readIORef uplinkRef
        case nodeCache upnode of
            Nothing -> return ()
            Just cache -> do
                case nodeData upnode of
                    AppNode _ _ -> do
                        addUplink (UplinkAppL, cache) =<< getLeft cache
                        addUplink (UplinkAppR, cache) =<< getRight cache
                        setCache uplinkRef Nothing
                    LambdaNode var _ -> do
                        addUplink (UplinkLambda, cache) =<< getBody cache
                        setCache uplinkRef Nothing
                        clear var
                clear uplinkRef
    setCache noderef Nothing

cleanup :: NodeRef a -> IO ()
cleanup noderef = do
    node <- readIORef noderef
    when (null (nodeUplinks node)) $ case nodeData node of
        AppNode left right -> do
            deleteUplink (UplinkAppL, noderef) left
            cleanup left
            deleteUplink (UplinkAppR, noderef) right
            cleanup right
        LambdaNode var body -> do
            deleteUplink (UplinkLambda, noderef) body
            cleanup body
        _ -> return ()

upreplace :: NodeRef a -> Uplink a -> IO ()
upreplace newchild (uplinkType, intoref) = do
    into <- readIORef intoref
    case (nodeData into, uplinkType) of
        (AppNode left right, UplinkAppL) -> do
            deleteUplink (UplinkAppL, intoref) left
            replaceLeft newchild intoref
            addUplink (UplinkAppL, intoref) newchild
            cleanup left
        (AppNode left right, UplinkAppR) -> do
            deleteUplink (UplinkAppR, intoref) right
            replaceRight newchild intoref
            addUplink (UplinkAppR, intoref) newchild
            cleanup right
        (LambdaNode var body, UplinkLambda) -> do
            deleteUplink (UplinkLambda, intoref) body
            replaceBody newchild intoref
            addUplink (UplinkLambda, intoref) newchild
            cleanup body
            

betaReduce :: NodeRef a -> IO (NodeRef a)
betaReduce appref = do
    app <- readIORef appref
    let AppNode leftref rightref = nodeData app
    left <- readIORef leftref
    let LambdaNode varref bodyref = nodeData left
    var <- readIORef varref
    result <- case nodeUplinks var of
        [] -> return bodyref
        _ -> do
            upcopy leftref rightref (UplinkVar, varref)
            result <- fromJust . nodeCache <$> (readIORef =<< getBody leftref)
            setCache leftref Nothing
            clear varref
            return result
    mapM_ (upreplace result) =<< nodeUplinks <$> readIORef appref
    return result

hnfReduce :: (HOAS.Primitive a) => NodeRef a -> IO ()
hnfReduce noderef = do
    node <- readIORef noderef
    case nodeData node of
        LambdaNode var body -> hnfReduce body 
        AppNode left right -> do
            hnfReduce left
            left' <- readIORef =<< getLeft noderef 
            case nodeData left' of
                LambdaNode {} -> hnfReduce =<< betaReduce noderef
                PrimNode p -> do
                    hnfReduce right
                    right' <- readIORef =<< getRight noderef
                    case nodeData right' of
                        PrimNode p' -> do
                            result <- newNodeRef $ PrimNode (p `HOAS.apply` p')
                            mapM_ (upreplace result) =<< nodeUplinks <$> readIORef noderef
                        _ -> return ()
                _ -> return ()
        _ -> return ()

graphviz :: (HOAS.Primitive a) => NodeRef a -> IO String
graphviz noderef_ = do
    output <- evalStateT (execWriterT (go noderef_)) ([], 0)
    return $ "digraph Lambda {\n" ++ output ++ "}\n"
    where
    go noderef = do
        cached <- lift $ gets (lookup noderef . fst)
        case cached of
            Just ident -> return ident
            Nothing -> do
                (seen, ident) <- lift get
                lift $ put ((noderef,ident):seen, ident+1)
                node <- liftIO $ readIORef noderef
                let color | noderef == noderef_ = "color=orange,style=filled"
                          | otherwise           = "colorblack"
                forM_ (nodeUplinks node) $ \(ty, uplink) -> do
                    uplinkid <- go uplink
                    tell $ "p" ++ show ident ++ " -> p" ++ show uplinkid ++ " [weight=1,color=red];\n"
                case nodeData node of
                    AppNode left right -> do
                        tell $ "p" ++ show ident ++ " [label=\"*\"," ++ color ++ "];\n"
                        leftid <- go left
                        tell $ "p" ++ show ident ++ " -> p" ++ show leftid ++ " [weight=1,color=\"#007f00\",label=\"fv\"];\n"
                        rightid <- go right
                        tell $ "p" ++ show ident ++ " -> p" ++ show rightid ++ " [weight=1,label=\"av\"];\n"
                    LambdaNode var body -> do
                        tell $ "p" ++ show ident ++ " [label=\"\\\\\"," ++ color ++ "];\n"
                        bodyid <- go body
                        tell $ "p" ++ show ident ++ " -> p" ++ show bodyid ++ " [weight=1];\n"
                        varid <- go var
                        tell $ "p" ++ show ident ++ " -> p" ++ show varid ++ " [weight=0,color=blue];\n"
                    VarNode -> do
                        tell $ "p" ++ show ident ++ " [label=\"x\"," ++ color ++ "];\n"
                    PrimNode x -> do
                        tell $ "p" ++ show ident ++ " [label=\"" ++ show x ++ "\"];\n"
                case nodeCache node of
                    Just cacher -> do
                        cacheid <- go cacher
                        tell $ "p" ++ show ident ++ " -> p" ++ show cacheid ++ " [weight=0,style=dotted];\n"
                    Nothing -> return ()
                return ident

runGraphviz :: (HOAS.Primitive a) => NodeRef a -> IO ()
runGraphviz node = do
    writeFile "graph.dot" =<< graphviz node
    system "dot -T png -o graph.png graph.dot"
    system "eog graph.png"
    return ()

eval :: (HOAS.Primitive a) => Term a -> IO a
eval t = do
    noderef <- getTerm $ fun (\z -> t)
    hnfReduce noderef
    dat <- nodeData <$> (readIORef =<< getBody noderef)
    case dat of
        PrimNode p -> return p
        _ -> fail "Not a primitive!"
    

newtype Term a = Term { getTerm :: IO (NodeRef a) }

infixl 9 %
(%) :: Term a -> Term a -> Term a
Term left % Term right = Term $ do
    left' <- left
    right' <- right
    newref <- newNodeRef $ AppNode left' right'
    addUplink (UplinkAppL, newref) left'
    addUplink (UplinkAppR, newref) right'
    return newref

fun :: (Term a -> Term a) -> Term a
fun bodyf = Term $ do
    var <- newNodeRef $ VarNode
    body <- getTerm . bodyf . Term $ return var
    newref <- newNodeRef $ LambdaNode var body
    addUplink (UplinkLambda, newref) body
    return newref

prim :: a -> Term a
prim x = Term $ do
    newref <- newNodeRef $ PrimNode x
    return newref

instance HOAS.Term (Term a) where
    (%) = (%)
    fun = fun

instance HOAS.PrimTerm a (Term a) where
    prim = prim