module WebBits.JavaScript.Environment
( staticEnvironment
, Ann
, LabelledStatement
, LabelledExpression
, Env
) where
import Data.Generics hiding (GT)
import Data.Maybe (fromJust)
import qualified Data.Set as S
import WebBits.Data.Zipper (ZipperT,ZipperT)
import qualified Data.Map as M
import qualified WebBits.Data.Zipper as Z
import Control.Monad.State
import qualified Data.Foldable as F
import qualified Data.List as L
import WebBits.JavaScript.JavaScript
import Text.ParserCombinators.Parsec.Pos (SourcePos,initialPos)
thisStmt :: Statement SourcePos -> Statement SourcePos
thisStmt (FunctionStmt loc id args body) =
FunctionStmt loc id ((Id loc "this"):(Id loc "arguments"):args) body
thisStmt s = s
thisExpr :: Expression SourcePos -> Expression SourcePos
thisExpr (FuncExpr l args s) =
FuncExpr l ((Id l "this"):(Id l "arguments"):args) s
thisExpr e = e
removeParens :: Expression SourcePos -> Expression SourcePos
removeParens (ParenExpr _ e) = e
removeParens e = e
removeSingletons :: Expression SourcePos -> Expression SourcePos
removeSingletons (ListExpr _ [e]) = e
removeSingletons e = e
explicitThis :: [Statement SourcePos] -> [Statement SourcePos]
explicitThis = everywhere $ (mkT removeSingletons) . (mkT removeParens)
. (mkT thisExpr) . (mkT thisStmt)
type PartialEnv = (Env,S.Set String)
type Env = M.Map String Int
emptyPartialEnv :: PartialEnv
emptyPartialEnv = (M.empty,S.empty)
type RefM a = Z.ZipperT PartialEnv (State Int) a
nextLabel :: (MonadTrans t, Monad (t (State Int))) => t (State Int) Int
nextLabel = do
l <- lift get
lift $ modify (+1)
return l
bind :: Id SourcePos -> RefM ()
bind (Id _ s) = do
lbl <- nextLabel
(env,freeIds) <- Z.getNode
Z.setNode (M.insert s lbl env, S.delete s freeIds)
use :: Id SourcePos -> RefM ()
use (Id _ s) = do
(env,freeIds) <- Z.getNode
case M.lookup s env of
Just _ -> return ()
Nothing -> Z.setNode (env, S.insert s freeIds)
buildExpr :: Expression SourcePos -> RefM (Expression SourcePos)
buildExpr (FuncExpr loc args stmt) = do
stmt' <- Z.nest emptyPartialEnv (mapM_ bind args >> buildStmt stmt)
return (FuncExpr loc args stmt')
buildExpr e@(VarRef _ id) = do
use id
return e
buildExpr (AssignExpr loc op (VarRef loc' id) e) = do
use id
e' <- buildExpr e
return (AssignExpr loc op (VarRef loc' id) e')
buildExpr e =
gmapM buildAny e
buildCatchClause :: CatchClause SourcePos -> RefM (CatchClause SourcePos)
buildCatchClause (CatchClause loc id stmt) = do
bind id
stmt' <- buildStmt stmt
return (CatchClause loc id stmt')
buildVarDecl :: VarDecl SourcePos -> RefM (VarDecl SourcePos)
buildVarDecl (VarDecl loc id ye) = do
bind id
ye' <- gmapM buildAny ye
return (VarDecl loc id ye')
buildForInInit :: ForInInit SourcePos -> RefM (ForInInit SourcePos)
buildForInInit e@(ForInVar id) = do
bind id
return e
buildForInInit e@(ForInNoVar id) = do
use id
return e
buildStmt :: Statement SourcePos -> RefM (Statement SourcePos)
buildStmt (FunctionStmt loc id args stmt) = do
stmt' <- Z.nest emptyPartialEnv (mapM_ bind args >> buildStmt stmt)
return (FunctionStmt loc id args stmt')
buildStmt s = gmapM buildAny s
buildAny' :: (Data a, Typeable a) => a -> RefM a
buildAny' v = gmapM buildAny v
buildAny :: GenericM RefM
buildAny =
buildAny' `extM` buildExpr `extM` buildCatchClause `extM`
buildVarDecl `extM` buildForInInit `extM` buildStmt
resolveFreeId :: Env -> String -> StateT Env (State Int) Env
resolveFreeId env freeId = case M.lookup freeId env of
Just lbl -> return env
Nothing -> do
globals <- get
case M.lookup freeId globals of
Just lbl -> return (M.insert freeId lbl env)
Nothing -> do
lbl <- nextLabel
put (M.insert freeId lbl globals)
return (M.insert freeId lbl env)
completeEnvM' :: Z.Tree PartialEnv -> StateT Env (State Int) (Z.Tree Env)
completeEnvM' pt = Z.dfsFoldM f M.empty pt where
f enclosingEnv (env,freeIds) =
foldM resolveFreeId (M.union env enclosingEnv) (S.elems freeIds)
completeEnvM :: Z.Tree PartialEnv -> StateT Env (State Int) (Z.Tree Env)
completeEnvM pt = do
tree <- completeEnvM' pt
globals <- get
return (fmap (\lexicals -> M.union lexicals globals) tree)
type Ann = (Env,Int,SourcePos)
type LabelledStatement = Statement Ann
type LabelledExpression = Expression Ann
insertEmptyAnn :: (Functor f) => f SourcePos -> f Ann
insertEmptyAnn = fmap (\loc -> (M.empty,0,loc))
locOf :: F.Foldable t => t SourcePos -> SourcePos
locOf = fromJust . (F.find $ const True)
labelEnv :: Env -> Z.ZipperT Env (State Int) Env
labelEnv env = Z.getNode
labelId :: Id Ann -> Z.ZipperT Env (State Int) (Id Ann)
labelId id@(Id (_,_,loc) s) = do
env <- Z.getNode
case M.lookup s env of
Nothing -> fail $ "BUG: unbound identifier while labelling" ++ show id
Just lbl -> return (Id (env,lbl,loc) s)
labelIdNoVar :: Id Ann -> Z.ZipperT Env (State Int) (Id Ann)
labelIdNoVar (Id (_,_,loc) s) = do
env <- Z.getNode
lbl <- nextLabel
return (Id (env,lbl,loc) s)
labelProp :: Prop Ann -> Z.ZipperT Env (State Int) (Prop Ann)
labelProp (PropId (_,_,loc) id) = do
env <- Z.getNode
lbl <- nextLabel
id' <- labelIdNoVar id
return (PropId (env,lbl,loc) id')
labelProp e = gmapM labelAny e
labelExpr :: Expression Ann
-> Z.ZipperT Env (State Int) (Expression Ann)
labelExpr (ThisRef (_,_,loc)) = do
env <- Z.getNode
lbl <- M.lookup "this" env
return (ThisRef (env,lbl,loc))
labelExpr (DotRef (_,_,loc) expr id) = do
env <- Z.getNode
lbl <- nextLabel
id' <- labelIdNoVar id
expr' <- labelExpr expr
return (DotRef (env,lbl,loc) expr' id')
labelExpr (FuncExpr (_,_,loc) args stmt) = do
env <- Z.getNode
lbl <- nextLabel
args' <- Z.withCurrentChild (mapM labelId args)
stmt' <- Z.withCurrentChild (labelStmt stmt)
Z.shiftRight'
return (FuncExpr (env,lbl,loc) args' stmt')
labelExpr e = gmapM labelAny e
labelStmt :: Statement Ann
-> Z.ZipperT Env (State Int) (Statement Ann)
labelStmt (BreakStmt (_,_,loc) (Just id)) = do
env <- Z.getNode
lbl <- nextLabel
id' <- labelIdNoVar id
return (BreakStmt (env,lbl,loc) (Just id'))
labelStmt (ContinueStmt (_,_,loc) (Just id)) = do
env <- Z.getNode
lbl <- nextLabel
id' <- labelIdNoVar id
return (ContinueStmt (env,lbl,loc) (Just id'))
labelStmt (LabelledStmt (_,_,loc) id stmt) = do
env <- Z.getNode
lbl <- nextLabel
id' <- labelIdNoVar id
stmt' <- labelStmt stmt
return (LabelledStmt (env,lbl,loc) id' stmt')
labelStmt (FunctionStmt (_,_,loc) id args stmt) = do
env <- Z.getNode
lbl <- nextLabel
args' <- Z.withCurrentChild (mapM labelId args)
stmt' <- Z.withCurrentChild (labelStmt stmt)
Z.shiftRight'
return (FunctionStmt (env,lbl,loc) id args' stmt')
labelStmt e = gmapM labelAny e
labelAny' :: (Data a, Typeable a) => a -> Z.ZipperT Env (State Int) a
labelAny' a = gmapM labelAny a
labelAny :: GenericM (Z.ZipperT Env (State Int))
labelAny a = (labelAny' `extM` labelEnv `extM` labelId `extM` labelProp `extM`
labelExpr `extM` labelStmt) a
staticEnvironment :: [Statement SourcePos]
-> ([Statement Ann],Env,Int)
staticEnvironment stmts =
let stmts' = explicitThis stmts
labelM = do
partialEnvTree <- Z.execZipperT (mapM buildStmt stmts')
(Z.toLocation (Z.Node emptyPartialEnv []))
(envTree,globals) <- runStateT (completeEnvM partialEnvTree) M.empty
let stmts'' = map insertEmptyAnn stmts'
labelledStmts <- Z.evalZipperT (mapM labelStmt stmts'')
(Z.toLocation envTree)
return (labelledStmts,globals)
((labelledStmts,globals),nextLabel) = runState labelM 0
in (labelledStmts,globals,nextLabel)