# LANGUAGE GeneralizedNewtypeDeriving #
module Data.Bedrock.Simplify ( simplify, mergeAllocsModule ) where

import           Control.Applicative  (Applicative, pure, (<$>), (<*>))
import           Control.Monad.Reader
import           Control.Monad.State
import           Data.Map             (Map)
import qualified Data.Map             as Map

import           Data.Bedrock
import           Data.Bedrock.Misc

data Env = Env
    { envRenaming  :: Map Variable Variable
    , envAddresses :: Map Variable (Variable, Int)
    , envConstant  :: Map Variable (NodeName, [Variable])
    , envStores    :: Map Variable (NodeName, [Variable])
    , envLoads     :: Map Variable (Variable, Int)
    }
newtype M a = M { unM :: ReaderT Env (State AvailableNamespace) a }
    deriving
        ( Monad, MonadReader Env, Applicative, Functor
        , MonadState AvailableNamespace )

simplify :: Module -> Module
simplify m = runM (simplifyModule m)
  where
    runM action = evalState (runReaderT (unM action) env) st
    st = modNamespace m
    env = Env
        { envRenaming  = Map.empty
        , envAddresses = Map.empty
        , envConstant  = Map.empty
        , envStores    = Map.empty
        , envLoads     = Map.empty }

simplifyModule :: Module -> M Module
simplifyModule m = do
    fns <- mapM simplifyFunction (functions m)
    ns <- get
    return m{ functions = fns, modNamespace = ns }

simplifyFunction :: Function -> M Function
simplifyFunction fn = do
    body <- simplifyBlock (fnBody fn)
    return fn{fnBody = body}

simplifyBlock :: Block -> M Block
simplifyBlock block =
    case block of
        Bind vars (Application fn args) (Return rets) | vars == rets ->
            simplifyBlock $ TailCall fn args
Case scrut _ mbDefaultBranch 
                 [ Alternative ( NodePat node [ ] ) branch ] - > do 
             clone < - cloneVariable scrut 
             bindVariable scrut clone $ 
                 Bind [ clone ] ( Unit ( NodeArg node [ ] ) ) < $ > 
                     simplifyBlock branch
                [Alternative (NodePat node []) branch] -> do
            clone <- cloneVariable scrut
            bindVariable scrut clone $
                Bind [clone] (Unit (NodeArg node [])) <$>
                    simplifyBlock branch-}
        Bind [val] expr@(Load ptr n) rest ->
            bindLoad val ptr n $
            Bind [val]
                <$> simplifyExpression expr
                <*> simplifyBlock rest
        Bind [ptr] expr@(Store name vars) rest ->
            bindStore ptr name vars $
            Bind [ptr]
                <$> simplifyExpression expr
                <*> simplifyBlock rest
        Bind [ptr] expr@(StoreAlloc name _tys) rest ->
            bindStore ptr name [] $
            Bind [ptr]
                <$> simplifyExpression expr
                <*> simplifyBlock rest
        Bind [] expr@(StoreWrite ptr args) rest ->
            bindStoreArgs ptr args $
            Bind []
                <$> simplifyExpression expr
                <*> simplifyBlock rest
        Bind [node] expr@(MkNode name vars) rest ->
            bindConstant node name vars $
            Bind [node]
                <$> simplifyExpression expr
                <*> simplifyBlock rest
        Bind rets (Builtin "apply" [PVariable ptr, PVariable arg]) rest -> do
            mbConst <- lookupStore ptr
            case mbConst of
                Just (FunctionName name 1, args) ->
                    simplifyBlock $
                        Bind rets (Application name (args ++ [arg])) rest
                Just (FunctionName name n, args) ->
                    simplifyBlock $
                        Bind rets (Store (FunctionName name (n-1)) (args ++ [arg])) rest
                _Nothing -> do
                  ptr' <- resolve ptr
                  arg' <- resolve arg
                  Bind rets (Builtin "apply" [PVariable ptr', PVariable arg'])
                    <$> simplifyBlock rest
        Bind [] TypeCast{} rest ->
            simplifyBlock rest
        Bind [] Literal{} rest ->
            simplifyBlock rest
        Bind [dst] (TypeCast src) rest | variableType dst == variableType src -> do
            src' <- resolve src
            bindVariable dst src' (simplifyBlock rest)
        Bind [dst] (TypeCast src) rest | variableType dst /= variableType src -> do
            src' <- resolve src
            mbLoad <- lookupLoad src'
            case mbLoad of
              Nothing -> Bind [dst]
                  <$> simplifyExpression (TypeCast src)
                  <*> simplifyBlock rest
              Just (ptr, idx) -> do
                ptr' <- resolve ptr
                Bind [dst] (Load ptr' idx) <$> simplifyBlock rest
        Bind [dst] (Address src idx) rest -> do
          src' <- resolve src
          bindAddress dst src' idx $
            Bind [dst] <$> lookupIndexed src' idx Address <*> simplifyBlock rest
        Bind [] (Write dst idx src) rest -> do
          dst' <- resolve dst
          src' <- resolve src
          Bind [] <$> lookupIndexed dst' idx (\d i -> Write d i src') <*> simplifyBlock rest
        Bind binds simple rest ->
            Bind binds
                <$> simplifyExpression simple
                <*> simplifyBlock rest
        Recursive binds rest ->
          Recursive binds <$> simplifyBlock rest
        -- FIXME: Sanity check: con == node
        Case scrut (Just defaultBranch) [] -> simplifyBlock defaultBranch
        Case scrut Nothing alts@[Alternative (NodePat _con args) branch] -> do
          mbNode <- lookupConstant scrut
          case mbNode of
            Nothing ->
              Case
                <$> resolve scrut
                <*> pure Nothing
                <*> mapM simplifyAlternative alts
            Just (_node, nodeArgs) ->
              bindVariables (zip args nodeArgs) $ simplifyBlock branch
        Case scrut Nothing alts ->
            Case
                <$> resolve scrut
                <*> pure Nothing
                <*> mapM simplifyAlternative alts
        Case scrut (Just defaultBranch) alts ->
            Case
                <$> resolve scrut
                <*> (Just <$> simplifyBlock defaultBranch)
                <*> mapM simplifyAlternative alts
        Return vars -> Return <$> mapM resolve vars
        TailCall name vars -> TailCall name <$> mapM resolve vars
        Exit -> pure Exit
        Panic msg -> pure $ Panic msg
        Invoke cont vars -> Invoke <$> resolve cont <*> mapM resolve vars
        Raise var -> Raise <$> resolve var

simplifyAlternative :: Alternative -> M Alternative
simplifyAlternative (Alternative pat branch) =
    Alternative pat <$> simplifyBlock branch


simplifyExpression :: Expression -> M Expression
simplifyExpression expr =
    case expr of
        Application fn vars -> Application fn <$> mapM resolve vars
        Fetch ptr -> do
          mbStore <- lookupStore ptr
          case mbStore of
            Nothing -> pure $ Fetch ptr
            Just (node, args) -> pure $ MkNode node args
        Eval var -> do
            mbNode <- lookupStore var
            case mbNode of
                Nothing -> Eval <$> resolve var
                -- FIXME: We should update the heap object with
                --        an indirection to the newly created object.
                Just (FunctionName name 0, args) ->
                    simplifyExpression $
                    Application name args
                Just (_nodeName, _args) ->
                    simplifyExpression $ TypeCast var
        MkNode name vars -> MkNode name <$> mapM resolve vars
        TypeCast var -> TypeCast <$> resolve var
        CCall fn vars -> CCall fn <$> mapM resolve vars
        Literal{} -> pure expr
        Catch fn args handler handlerArgs ->
          Catch fn
            <$> mapM resolve args
            <*> pure handler
            <*> mapM resolve handlerArgs
        InvokeReturn v vs ->
          InvokeReturn
            <$> resolve v
            <*> mapM resolve vs
        Builtin fn params -> Builtin fn <$> mapM simplifyParameter params


simplifyParameter :: Parameter -> M Parameter
simplifyParameter param =
  case param of
    PInt{} -> pure param
    PString{} -> pure param
    PName{} -> pure param
    PNodeName{} -> pure param
    PVariable v -> PVariable <$> resolve v
    PVariables vs -> PVariables <$> mapM resolve vs
    PTypes ts -> PTypes <$> pure ts



  Note : This optimisation is only valid after the CPS transformation .
mergeAllocsModule :: Module -> Module
mergeAllocsModule m =
    m{ functions = [ fn{ fnBody = mergeAllocs (fnBody fn) }
                   | fn <- functions m] }

mergeAllocs :: Block -> Block
mergeAllocs block0 = applyAlloc 0 nFinal block'
  where
    applyAlloc slob allocs block =
        case slob `compare` allocs of
  XXX : No need to decrement the heap pointer since  only
            --      make sure there's available space but don't actually
            --      increment the pointer.
            GT -> block -- Bind [] (BumpHeapPtr (negate (slob-allocs))) block
            EQ -> block
            LT -> Bind [] (Alloc (allocs-slob)) block
    (nFinal, block') = worker 0 block0
    worker n block =
        case block of
  Hoist  out of case alternatives iff :
   1 . There are  before the Case statement , or
   2 . There are  in all of the Case branches .
            -- Otherwise treat the branches on their own.
            Case scrut defaultBranch alts ->
                let hoist = n > 0 || all (>0) branches
                    slob = blockAlloc - n
                    maxBranchAlloc = maximum (0:branches)
                    branches = maybe id (:) (fmap fst def) (map fst alts')
                    def = fmap (worker 0) defaultBranch
                    alts' = [ (branchAlloc, Alternative pat branch'')
                            | Alternative pat branch <- alts
                            , let (branchAlloc, branch') = worker 0 branch
                            , let branch'' = applyAlloc slob branchAlloc branch' ]
                    blockAlloc
                        | hoist = n + maxBranchAlloc
                        | otherwise = n
                in (blockAlloc, Case scrut (fmap snd def) (map snd alts'))
            Bind [] (Alloc size) rest ->
                worker (n+size) rest
            Bind binds expr rest ->
                let (n', rest') = worker n rest
                in (n', Bind binds expr rest')
            _ -> (n, block)




--------------------------------------------------------
  Utils

_cloneVariable :: Variable -> M Variable
_cloneVariable var = do
    ns <- get
    let (idNum, ns') = newIDByType ns (variableType var)
        name = (variableName var){ nameUnique = idNum }
    put ns'
    return var{ variableName = name }

resolve :: Variable -> M Variable
resolve var = do
    m <- asks envRenaming
    return $ Map.findWithDefault var var m

bindVariable :: Variable -> Variable -> M a -> M a
bindVariable old new fn = do
    new' <- resolve new
    local
        (\env -> env{
            envRenaming = Map.insert old new' (envRenaming env),
            envStores = case Map.lookup old (envStores env) of
              Nothing   -> envStores env
              Just node -> Map.insert new' node (envStores env),
            envConstant = case Map.lookup old (envConstant env) of
                Nothing   -> envConstant env
                Just node -> Map.insert new' node (envConstant env)
          })
        fn

bindVariables :: [(Variable, Variable)] -> M a -> M a
bindVariables [] = id
bindVariables ((x,y):xs) = bindVariables xs . bindVariable x y

lookupConstant :: Variable -> M (Maybe (NodeName, [Variable]))
lookupConstant var = do
    var' <- resolve var
    m <- asks envConstant
    return $ Map.lookup var' m

bindConstant :: Variable -> NodeName -> [Variable] -> M a -> M a
bindConstant var node vars = local $ \env ->
    env{ envConstant = Map.insert var (node, vars) (envConstant env) }

bindAddress :: Variable -> Variable -> Int  -> M a -> M a
bindAddress dst src idx = local $ \env ->
  env{ envAddresses = Map.insert dst (src,idx) (envAddresses env) }

lookupIndexed :: Variable -> Int -> (Variable -> Int -> Expression) -> M Expression
lookupIndexed src idx fn = do
  addrs <- asks envAddresses
  case Map.lookup src addrs of
    Just (orig, prevIdx) -> fn <$> resolve orig <*> pure (prevIdx+idx)
    Nothing              -> fn <$> resolve src <*> pure idx

lookupStore :: Variable -> M (Maybe (NodeName, [Variable]))
lookupStore var = do
    var' <- resolve var
    m <- asks envStores
    return $ Map.lookup var' m

lookupLoad :: Variable -> M (Maybe (Variable, Int))
lookupLoad var = do
    var' <- resolve var
    m <- asks envLoads
    return $ Map.lookup var' m

bindStore :: Variable -> NodeName -> [Variable] -> M a -> M a
bindStore var node vars = local $ \env ->
    env{ envStores = Map.insert var (node, vars) (envStores env) }

bindStoreArgs :: Variable -> [Variable] -> M a -> M a
bindStoreArgs var vars = local $ \env ->
    env{ envStores = Map.adjust (\(node, _) -> (node, vars)) var (envStores env) }

bindLoad :: Variable -> Variable -> Int -> M a -> M a
bindLoad var ptr n = local $ \env ->
    env{ envLoads = Map.insert var (ptr, n) (envLoads env) }

 FIXME: Sanity check: con == node
 FIXME: We should update the heap object with
        an indirection to the newly created object.
      make sure there's available space but don't actually
      increment the pointer.
 Bind [] (BumpHeapPtr (negate (slob-allocs))) block
 Otherwise treat the branches on their own.
------------------------------------------------------
# LANGUAGE GeneralizedNewtypeDeriving #
module Data.Bedrock.Simplify ( simplify, mergeAllocsModule ) where

import           Control.Applicative  (Applicative, pure, (<$>), (<*>))
import           Control.Monad.Reader
import           Control.Monad.State
import           Data.Map             (Map)
import qualified Data.Map             as Map

import           Data.Bedrock
import           Data.Bedrock.Misc

data Env = Env
    { envRenaming  :: Map Variable Variable
    , envAddresses :: Map Variable (Variable, Int)
    , envConstant  :: Map Variable (NodeName, [Variable])
    , envStores    :: Map Variable (NodeName, [Variable])
    , envLoads     :: Map Variable (Variable, Int)
    }
newtype M a = M { unM :: ReaderT Env (State AvailableNamespace) a }
    deriving
        ( Monad, MonadReader Env, Applicative, Functor
        , MonadState AvailableNamespace )

simplify :: Module -> Module
simplify m = runM (simplifyModule m)
  where
    runM action = evalState (runReaderT (unM action) env) st
    st = modNamespace m
    env = Env
        { envRenaming  = Map.empty
        , envAddresses = Map.empty
        , envConstant  = Map.empty
        , envStores    = Map.empty
        , envLoads     = Map.empty }

simplifyModule :: Module -> M Module
simplifyModule m = do
    fns <- mapM simplifyFunction (functions m)
    ns <- get
    return m{ functions = fns, modNamespace = ns }

simplifyFunction :: Function -> M Function
simplifyFunction fn = do
    body <- simplifyBlock (fnBody fn)
    return fn{fnBody = body}

simplifyBlock :: Block -> M Block
simplifyBlock block =
    case block of
        Bind vars (Application fn args) (Return rets) | vars == rets ->
            simplifyBlock $ TailCall fn args
Case scrut _ mbDefaultBranch 
                 [ Alternative ( NodePat node [ ] ) branch ] - > do 
             clone < - cloneVariable scrut 
             bindVariable scrut clone $ 
                 Bind [ clone ] ( Unit ( NodeArg node [ ] ) ) < $ > 
                     simplifyBlock branch
                [Alternative (NodePat node []) branch] -> do
            clone <- cloneVariable scrut
            bindVariable scrut clone $
                Bind [clone] (Unit (NodeArg node [])) <$>
                    simplifyBlock branch-}
        Bind [val] expr@(Load ptr n) rest ->
            bindLoad val ptr n $
            Bind [val]
                <$> simplifyExpression expr
                <*> simplifyBlock rest
        Bind [ptr] expr@(Store name vars) rest ->
            bindStore ptr name vars $
            Bind [ptr]
                <$> simplifyExpression expr
                <*> simplifyBlock rest
        Bind [ptr] expr@(StoreAlloc name _tys) rest ->
            bindStore ptr name [] $
            Bind [ptr]
                <$> simplifyExpression expr
                <*> simplifyBlock rest
        Bind [] expr@(StoreWrite ptr args) rest ->
            bindStoreArgs ptr args $
            Bind []
                <$> simplifyExpression expr
                <*> simplifyBlock rest
        Bind [node] expr@(MkNode name vars) rest ->
            bindConstant node name vars $
            Bind [node]
                <$> simplifyExpression expr
                <*> simplifyBlock rest
        Bind rets (Builtin "apply" [PVariable ptr, PVariable arg]) rest -> do
            mbConst <- lookupStore ptr
            case mbConst of
                Just (FunctionName name 1, args) ->
                    simplifyBlock $
                        Bind rets (Application name (args ++ [arg])) rest
                Just (FunctionName name n, args) ->
                    simplifyBlock $
                        Bind rets (Store (FunctionName name (n-1)) (args ++ [arg])) rest
                _Nothing -> do
                  ptr' <- resolve ptr
                  arg' <- resolve arg
                  Bind rets (Builtin "apply" [PVariable ptr', PVariable arg'])
                    <$> simplifyBlock rest
        Bind [] TypeCast{} rest ->
            simplifyBlock rest
        Bind [] Literal{} rest ->
            simplifyBlock rest
        Bind [dst] (TypeCast src) rest | variableType dst == variableType src -> do
            src' <- resolve src
            bindVariable dst src' (simplifyBlock rest)
        Bind [dst] (TypeCast src) rest | variableType dst /= variableType src -> do
            src' <- resolve src
            mbLoad <- lookupLoad src'
            case mbLoad of
              Nothing -> Bind [dst]
                  <$> simplifyExpression (TypeCast src)
                  <*> simplifyBlock rest
              Just (ptr, idx) -> do
                ptr' <- resolve ptr
                Bind [dst] (Load ptr' idx) <$> simplifyBlock rest
        Bind [dst] (Address src idx) rest -> do
          src' <- resolve src
          bindAddress dst src' idx $
            Bind [dst] <$> lookupIndexed src' idx Address <*> simplifyBlock rest
        Bind [] (Write dst idx src) rest -> do
          dst' <- resolve dst
          src' <- resolve src
          Bind [] <$> lookupIndexed dst' idx (\d i -> Write d i src') <*> simplifyBlock rest
        Bind binds simple rest ->
            Bind binds
                <$> simplifyExpression simple
                <*> simplifyBlock rest
        Recursive binds rest ->
          Recursive binds <$> simplifyBlock rest
        Case scrut (Just defaultBranch) [] -> simplifyBlock defaultBranch
        Case scrut Nothing alts@[Alternative (NodePat _con args) branch] -> do
          mbNode <- lookupConstant scrut
          case mbNode of
            Nothing ->
              Case
                <$> resolve scrut
                <*> pure Nothing
                <*> mapM simplifyAlternative alts
            Just (_node, nodeArgs) ->
              bindVariables (zip args nodeArgs) $ simplifyBlock branch
        Case scrut Nothing alts ->
            Case
                <$> resolve scrut
                <*> pure Nothing
                <*> mapM simplifyAlternative alts
        Case scrut (Just defaultBranch) alts ->
            Case
                <$> resolve scrut
                <*> (Just <$> simplifyBlock defaultBranch)
                <*> mapM simplifyAlternative alts
        Return vars -> Return <$> mapM resolve vars
        TailCall name vars -> TailCall name <$> mapM resolve vars
        Exit -> pure Exit
        Panic msg -> pure $ Panic msg
        Invoke cont vars -> Invoke <$> resolve cont <*> mapM resolve vars
        Raise var -> Raise <$> resolve var

simplifyAlternative :: Alternative -> M Alternative
simplifyAlternative (Alternative pat branch) =
    Alternative pat <$> simplifyBlock branch


simplifyExpression :: Expression -> M Expression
simplifyExpression expr =
    case expr of
        Application fn vars -> Application fn <$> mapM resolve vars
        Fetch ptr -> do
          mbStore <- lookupStore ptr
          case mbStore of
            Nothing -> pure $ Fetch ptr
            Just (node, args) -> pure $ MkNode node args
        Eval var -> do
            mbNode <- lookupStore var
            case mbNode of
                Nothing -> Eval <$> resolve var
                Just (FunctionName name 0, args) ->
                    simplifyExpression $
                    Application name args
                Just (_nodeName, _args) ->
                    simplifyExpression $ TypeCast var
        MkNode name vars -> MkNode name <$> mapM resolve vars
        TypeCast var -> TypeCast <$> resolve var
        CCall fn vars -> CCall fn <$> mapM resolve vars
        Literal{} -> pure expr
        Catch fn args handler handlerArgs ->
          Catch fn
            <$> mapM resolve args
            <*> pure handler
            <*> mapM resolve handlerArgs
        InvokeReturn v vs ->
          InvokeReturn
            <$> resolve v
            <*> mapM resolve vs
        Builtin fn params -> Builtin fn <$> mapM simplifyParameter params


simplifyParameter :: Parameter -> M Parameter
simplifyParameter param =
  case param of
    PInt{} -> pure param
    PString{} -> pure param
    PName{} -> pure param
    PNodeName{} -> pure param
    PVariable v -> PVariable <$> resolve v
    PVariables vs -> PVariables <$> mapM resolve vs
    PTypes ts -> PTypes <$> pure ts



  Note : This optimisation is only valid after the CPS transformation .
mergeAllocsModule :: Module -> Module
mergeAllocsModule m =
    m{ functions = [ fn{ fnBody = mergeAllocs (fnBody fn) }
                   | fn <- functions m] }

mergeAllocs :: Block -> Block
mergeAllocs block0 = applyAlloc 0 nFinal block'
  where
    applyAlloc slob allocs block =
        case slob `compare` allocs of
  XXX : No need to decrement the heap pointer since  only
            EQ -> block
            LT -> Bind [] (Alloc (allocs-slob)) block
    (nFinal, block') = worker 0 block0
    worker n block =
        case block of
  Hoist  out of case alternatives iff :
   1 . There are  before the Case statement , or
   2 . There are  in all of the Case branches .
            Case scrut defaultBranch alts ->
                let hoist = n > 0 || all (>0) branches
                    slob = blockAlloc - n
                    maxBranchAlloc = maximum (0:branches)
                    branches = maybe id (:) (fmap fst def) (map fst alts')
                    def = fmap (worker 0) defaultBranch
                    alts' = [ (branchAlloc, Alternative pat branch'')
                            | Alternative pat branch <- alts
                            , let (branchAlloc, branch') = worker 0 branch
                            , let branch'' = applyAlloc slob branchAlloc branch' ]
                    blockAlloc
                        | hoist = n + maxBranchAlloc
                        | otherwise = n
                in (blockAlloc, Case scrut (fmap snd def) (map snd alts'))
            Bind [] (Alloc size) rest ->
                worker (n+size) rest
            Bind binds expr rest ->
                let (n', rest') = worker n rest
                in (n', Bind binds expr rest')
            _ -> (n, block)




  Utils

_cloneVariable :: Variable -> M Variable
_cloneVariable var = do
    ns <- get
    let (idNum, ns') = newIDByType ns (variableType var)
        name = (variableName var){ nameUnique = idNum }
    put ns'
    return var{ variableName = name }

resolve :: Variable -> M Variable
resolve var = do
    m <- asks envRenaming
    return $ Map.findWithDefault var var m

bindVariable :: Variable -> Variable -> M a -> M a
bindVariable old new fn = do
    new' <- resolve new
    local
        (\env -> env{
            envRenaming = Map.insert old new' (envRenaming env),
            envStores = case Map.lookup old (envStores env) of
              Nothing   -> envStores env
              Just node -> Map.insert new' node (envStores env),
            envConstant = case Map.lookup old (envConstant env) of
                Nothing   -> envConstant env
                Just node -> Map.insert new' node (envConstant env)
          })
        fn

bindVariables :: [(Variable, Variable)] -> M a -> M a
bindVariables [] = id
bindVariables ((x,y):xs) = bindVariables xs . bindVariable x y

lookupConstant :: Variable -> M (Maybe (NodeName, [Variable]))
lookupConstant var = do
    var' <- resolve var
    m <- asks envConstant
    return $ Map.lookup var' m

bindConstant :: Variable -> NodeName -> [Variable] -> M a -> M a
bindConstant var node vars = local $ \env ->
    env{ envConstant = Map.insert var (node, vars) (envConstant env) }

bindAddress :: Variable -> Variable -> Int  -> M a -> M a
bindAddress dst src idx = local $ \env ->
  env{ envAddresses = Map.insert dst (src,idx) (envAddresses env) }

lookupIndexed :: Variable -> Int -> (Variable -> Int -> Expression) -> M Expression
lookupIndexed src idx fn = do
  addrs <- asks envAddresses
  case Map.lookup src addrs of
    Just (orig, prevIdx) -> fn <$> resolve orig <*> pure (prevIdx+idx)
    Nothing              -> fn <$> resolve src <*> pure idx

lookupStore :: Variable -> M (Maybe (NodeName, [Variable]))
lookupStore var = do
    var' <- resolve var
    m <- asks envStores
    return $ Map.lookup var' m

lookupLoad :: Variable -> M (Maybe (Variable, Int))
lookupLoad var = do
    var' <- resolve var
    m <- asks envLoads
    return $ Map.lookup var' m

bindStore :: Variable -> NodeName -> [Variable] -> M a -> M a
bindStore var node vars = local $ \env ->
    env{ envStores = Map.insert var (node, vars) (envStores env) }

bindStoreArgs :: Variable -> [Variable] -> M a -> M a
bindStoreArgs var vars = local $ \env ->
    env{ envStores = Map.adjust (\(node, _) -> (node, vars)) var (envStores env) }

bindLoad :: Variable -> Variable -> Int -> M a -> M a
bindLoad var ptr n = local $ \env ->
    env{ envLoads = Map.insert var (ptr, n) (envLoads env) }
