-- | A few helpers to work with the AST annotations
module Language.ECMAScript3.Syntax.Annotations where

import Language.ECMAScript3.Syntax
import Data.Traversable
import Control.Applicative
import Control.Arrow
import Control.Monad.State hiding (mapM)
import Prelude hiding (mapM)

-- | Removes annotations from a tree
removeAnnotations :: Traversable t => t a -> t ()
removeAnnotations :: t a -> t ()
removeAnnotations = (a -> ()) -> t a -> t ()
forall (t :: * -> *) a b. Traversable t => (a -> b) -> t a -> t b
reannotate (() -> a -> ()
forall a b. a -> b -> a
const ())

-- | Changes all the labels in the tree to another one, given by a
-- function.
reannotate :: Traversable t => (a -> b) -> t a -> t b
reannotate :: (a -> b) -> t a -> t b
reannotate f :: a -> b
f tree :: t a
tree = (a -> () -> b) -> t a -> () -> t b
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (b -> () -> b
forall (f :: * -> *) a. Applicative f => a -> f a
pure (b -> () -> b) -> (a -> b) -> a -> () -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> b
f) t a
tree ()

-- | add an extra field to the AST labels (the label would look like @
-- (a, b) @)
addExtraAnnotationField :: Traversable t => b -> t a -> t (a, b)
addExtraAnnotationField :: b -> t a -> t (a, b)
addExtraAnnotationField def :: b
def t :: t a
t = (a -> () -> (a, b)) -> t a -> () -> t (a, b)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (\z :: a
z -> (a, b) -> () -> (a, b)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
z, b
def)) t a
t ()

-- | remove an extra field
removeExtraAnnotationField :: Traversable t => t (a, b) -> t a
removeExtraAnnotationField :: t (a, b) -> t a
removeExtraAnnotationField t :: t (a, b)
t = ((a, b) -> () -> a) -> t (a, b) -> () -> t a
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (a -> () -> a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a -> () -> a) -> ((a, b) -> a) -> (a, b) -> () -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a, b) -> a
forall a b. (a, b) -> a
fst) t (a, b)
t ()


-- | Assigns unique numeric (Int) ids to each node in the AST. Returns
-- a pair: the tree annotated with UID's and the last ID that was
-- assigned.
assignUniqueIds :: Traversable t => Int -- ^ starting id
                                 -> t a -- ^ tree root
                                 -> (t (a, Int), Int) 
assignUniqueIds :: Int -> t a -> (t (a, Int), Int)
assignUniqueIds first :: Int
first tree :: t a
tree =
  (t (a, Int) -> t (a, Int)
forall (a :: * -> * -> *) b. Arrow a => a b b
returnA (t (a, Int) -> t (a, Int))
-> (Int -> Int) -> (t (a, Int), Int) -> (t (a, Int), Int)
forall (a :: * -> * -> *) b c b' c'.
Arrow a =>
a b c -> a b' c' -> a (b, b') (c, c')
*** \i :: Int
i -> Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
-1) ((t (a, Int), Int) -> (t (a, Int), Int))
-> (t (a, Int), Int) -> (t (a, Int), Int)
forall a b. (a -> b) -> a -> b
$ State Int (t (a, Int)) -> Int -> (t (a, Int), Int)
forall s a. State s a -> s -> (a, s)
runState ((a -> StateT Int Identity (a, Int))
-> t a -> State Int (t (a, Int))
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM a -> StateT Int Identity (a, Int)
forall a. a -> State Int (a, Int)
f t a
tree) Int
first
  where f :: a -> State Int (a, Int)
        f :: a -> State Int (a, Int)
f a :: a
a = do Int
i <- StateT Int Identity Int
forall s (m :: * -> *). MonadState s m => m s
get
                 Int -> StateT Int Identity ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+1)
                 (a, Int) -> State Int (a, Int)
forall (m :: * -> *) a. Monad m => a -> m a
return (a
a, Int
i)

-- | Things that have annotations -- for example, nodes in a syntax
-- tree
class HasAnnotation a where
  -- | Returns the annotation of the root of the tree
  getAnnotation :: a b -> b
  -- | Sets the annotation of the root of the tree  
  setAnnotation :: b -> a b -> a b

-- | Modify the annotation of the root node of the syntax tree
withAnnotation :: (HasAnnotation a) => (b -> b) -> a b -> a b
withAnnotation :: (b -> b) -> a b -> a b
withAnnotation f :: b -> b
f x :: a b
x = b -> a b -> a b
forall (a :: * -> *) b. HasAnnotation a => b -> a b -> a b
setAnnotation (b -> b
f (b -> b) -> b -> b
forall a b. (a -> b) -> a -> b
$ a b -> b
forall (a :: * -> *) b. HasAnnotation a => a b -> b
getAnnotation a b
x) a b
x

instance HasAnnotation Expression where
  getAnnotation :: Expression b -> b
getAnnotation e :: Expression b
e = case Expression b
e of
   (StringLit a :: b
a s :: String
s)              -> b
a
   (RegexpLit a :: b
a s :: String
s g :: Bool
g ci :: Bool
ci)         -> b
a
   (NumLit a :: b
a d :: Double
d)                 -> b
a
   (IntLit a :: b
a i :: Int
i)                 -> b
a
   (BoolLit a :: b
a b :: Bool
b)                -> b
a
   (NullLit a :: b
a)                  -> b
a
   (ArrayLit a :: b
a exps :: [Expression b]
exps)            -> b
a
   (ObjectLit a :: b
a props :: [(Prop b, Expression b)]
props)          -> b
a
   (ThisRef a :: b
a)                  -> b
a
   (VarRef a :: b
a id :: Id b
id)                -> b
a
   (DotRef a :: b
a exp :: Expression b
exp id :: Id b
id)            -> b
a
   (BracketRef a :: b
a container :: Expression b
container key :: Expression b
key) -> b
a
   (NewExpr a :: b
a ctor :: Expression b
ctor params :: [Expression b]
params)      -> b
a
   (PrefixExpr a :: b
a op :: PrefixOp
op e :: Expression b
e)          -> b
a
   (UnaryAssignExpr a :: b
a op :: UnaryAssignOp
op lv :: LValue b
lv)    -> b
a
   (InfixExpr a :: b
a op :: InfixOp
op e1 :: Expression b
e1 e2 :: Expression b
e2)       -> b
a
   (CondExpr a :: b
a g :: Expression b
g et :: Expression b
et ef :: Expression b
ef)         -> b
a
   (AssignExpr a :: b
a op :: AssignOp
op lv :: LValue b
lv e :: Expression b
e)       -> b
a
   (ListExpr a :: b
a es :: [Expression b]
es)              -> b
a
   (CallExpr a :: b
a fn :: Expression b
fn params :: [Expression b]
params)       -> b
a
   (FuncExpr a :: b
a mid :: Maybe (Id b)
mid args :: [Id b]
args s :: [Statement b]
s)      -> b
a
  setAnnotation :: b -> Expression b -> Expression b
setAnnotation a :: b
a e :: Expression b
e = case Expression b
e of
    (StringLit _ s :: String
s)              -> (b -> String -> Expression b
forall a. a -> String -> Expression a
StringLit b
a String
s)
    (RegexpLit _ s :: String
s g :: Bool
g ci :: Bool
ci)         -> (b -> String -> Bool -> Bool -> Expression b
forall a. a -> String -> Bool -> Bool -> Expression a
RegexpLit b
a String
s Bool
g Bool
ci)
    (NumLit _ d :: Double
d)                 -> (b -> Double -> Expression b
forall a. a -> Double -> Expression a
NumLit b
a Double
d)
    (IntLit _ i :: Int
i)                 -> (b -> Int -> Expression b
forall a. a -> Int -> Expression a
IntLit b
a Int
i)
    (BoolLit _ b :: Bool
b)                -> (b -> Bool -> Expression b
forall a. a -> Bool -> Expression a
BoolLit b
a Bool
b)
    (NullLit _)                  -> (b -> Expression b
forall a. a -> Expression a
NullLit b
a)
    (ArrayLit _ exps :: [Expression b]
exps)            -> (b -> [Expression b] -> Expression b
forall a. a -> [Expression a] -> Expression a
ArrayLit b
a [Expression b]
exps)
    (ObjectLit _ props :: [(Prop b, Expression b)]
props)          -> (b -> [(Prop b, Expression b)] -> Expression b
forall a. a -> [(Prop a, Expression a)] -> Expression a
ObjectLit b
a [(Prop b, Expression b)]
props)
    (ThisRef _)                  -> (b -> Expression b
forall a. a -> Expression a
ThisRef b
a)
    (VarRef _ id :: Id b
id)                -> (b -> Id b -> Expression b
forall a. a -> Id a -> Expression a
VarRef b
a Id b
id)
    (DotRef _ exp :: Expression b
exp id :: Id b
id)            -> (b -> Expression b -> Id b -> Expression b
forall a. a -> Expression a -> Id a -> Expression a
DotRef b
a Expression b
exp Id b
id)
    (BracketRef _ container :: Expression b
container key :: Expression b
key) -> (b -> Expression b -> Expression b -> Expression b
forall a. a -> Expression a -> Expression a -> Expression a
BracketRef b
a Expression b
container Expression b
key)
    (NewExpr _ ctor :: Expression b
ctor params :: [Expression b]
params)      -> (b -> Expression b -> [Expression b] -> Expression b
forall a. a -> Expression a -> [Expression a] -> Expression a
NewExpr b
a Expression b
ctor [Expression b]
params)
    (PrefixExpr _ op :: PrefixOp
op e :: Expression b
e)          -> (b -> PrefixOp -> Expression b -> Expression b
forall a. a -> PrefixOp -> Expression a -> Expression a
PrefixExpr b
a PrefixOp
op Expression b
e)
    (UnaryAssignExpr _ op :: UnaryAssignOp
op lv :: LValue b
lv)    -> (b -> UnaryAssignOp -> LValue b -> Expression b
forall a. a -> UnaryAssignOp -> LValue a -> Expression a
UnaryAssignExpr b
a UnaryAssignOp
op LValue b
lv)
    (InfixExpr _ op :: InfixOp
op e1 :: Expression b
e1 e2 :: Expression b
e2)       -> (b -> InfixOp -> Expression b -> Expression b -> Expression b
forall a.
a -> InfixOp -> Expression a -> Expression a -> Expression a
InfixExpr b
a InfixOp
op Expression b
e1 Expression b
e2)
    (CondExpr _ g :: Expression b
g et :: Expression b
et ef :: Expression b
ef)         -> (b -> Expression b -> Expression b -> Expression b -> Expression b
forall a.
a -> Expression a -> Expression a -> Expression a -> Expression a
CondExpr b
a Expression b
g Expression b
et Expression b
ef)
    (AssignExpr _ op :: AssignOp
op lv :: LValue b
lv e :: Expression b
e)       -> (b -> AssignOp -> LValue b -> Expression b -> Expression b
forall a. a -> AssignOp -> LValue a -> Expression a -> Expression a
AssignExpr b
a AssignOp
op LValue b
lv Expression b
e)
    (ListExpr _ es :: [Expression b]
es)              -> (b -> [Expression b] -> Expression b
forall a. a -> [Expression a] -> Expression a
ListExpr b
a [Expression b]
es)
    (CallExpr _ fn :: Expression b
fn params :: [Expression b]
params)       -> (b -> Expression b -> [Expression b] -> Expression b
forall a. a -> Expression a -> [Expression a] -> Expression a
CallExpr b
a Expression b
fn [Expression b]
params)
    (FuncExpr _ mid :: Maybe (Id b)
mid args :: [Id b]
args s :: [Statement b]
s)      -> (b -> Maybe (Id b) -> [Id b] -> [Statement b] -> Expression b
forall a.
a -> Maybe (Id a) -> [Id a] -> [Statement a] -> Expression a
FuncExpr b
a Maybe (Id b)
mid [Id b]
args [Statement b]
s)   

instance HasAnnotation Statement where
  getAnnotation :: Statement b -> b
getAnnotation s :: Statement b
s = case Statement b
s of
    BlockStmt a :: b
a _        -> b
a
    EmptyStmt a :: b
a          -> b
a
    ExprStmt a :: b
a _         -> b
a
    IfStmt a :: b
a _ _ _       -> b
a
    IfSingleStmt a :: b
a _ _   -> b
a
    SwitchStmt a :: b
a _ _     -> b
a
    WhileStmt a :: b
a _ _      -> b
a
    DoWhileStmt a :: b
a _ _    -> b
a
    BreakStmt a :: b
a _        -> b
a
    ContinueStmt a :: b
a _     -> b
a
    LabelledStmt a :: b
a _ _   -> b
a
    ForInStmt a :: b
a _ _ _    -> b
a
    ForStmt a :: b
a _ _ _ _    -> b
a
    TryStmt a :: b
a _ _ _      -> b
a
    ThrowStmt a :: b
a _        -> b
a
    ReturnStmt a :: b
a _       -> b
a
    WithStmt a :: b
a _ _       -> b
a
    VarDeclStmt a :: b
a _      -> b
a
    FunctionStmt a :: b
a _ _ _ -> b
a
  setAnnotation :: b -> Statement b -> Statement b
setAnnotation a :: b
a s :: Statement b
s = case Statement b
s of
    BlockStmt _ ss :: [Statement b]
ss       -> b -> [Statement b] -> Statement b
forall a. a -> [Statement a] -> Statement a
BlockStmt b
a [Statement b]
ss
    EmptyStmt _          -> b -> Statement b
forall a. a -> Statement a
EmptyStmt b
a
    ExprStmt _ e :: Expression b
e         -> b -> Expression b -> Statement b
forall a. a -> Expression a -> Statement a
ExprStmt b
a Expression b
e
    IfStmt _ g :: Expression b
g t :: Statement b
t e :: Statement b
e       -> b -> Expression b -> Statement b -> Statement b -> Statement b
forall a.
a -> Expression a -> Statement a -> Statement a -> Statement a
IfStmt b
a Expression b
g Statement b
t Statement b
e
    IfSingleStmt _ g :: Expression b
g t :: Statement b
t   -> b -> Expression b -> Statement b -> Statement b
forall a. a -> Expression a -> Statement a -> Statement a
IfSingleStmt b
a Expression b
g Statement b
t
    SwitchStmt _ g :: Expression b
g cs :: [CaseClause b]
cs    -> b -> Expression b -> [CaseClause b] -> Statement b
forall a. a -> Expression a -> [CaseClause a] -> Statement a
SwitchStmt b
a Expression b
g [CaseClause b]
cs
    WhileStmt _ g :: Expression b
g ss :: Statement b
ss     -> b -> Expression b -> Statement b -> Statement b
forall a. a -> Expression a -> Statement a -> Statement a
WhileStmt b
a Expression b
g Statement b
ss
    DoWhileStmt _ ss :: Statement b
ss g :: Expression b
g   -> b -> Statement b -> Expression b -> Statement b
forall a. a -> Statement a -> Expression a -> Statement a
DoWhileStmt b
a Statement b
ss Expression b
g
    BreakStmt _ l :: Maybe (Id b)
l        -> b -> Maybe (Id b) -> Statement b
forall a. a -> Maybe (Id a) -> Statement a
BreakStmt b
a Maybe (Id b)
l
    ContinueStmt _ l :: Maybe (Id b)
l     -> b -> Maybe (Id b) -> Statement b
forall a. a -> Maybe (Id a) -> Statement a
ContinueStmt b
a Maybe (Id b)
l
    LabelledStmt _ l :: Id b
l s :: Statement b
s   -> b -> Id b -> Statement b -> Statement b
forall a. a -> Id a -> Statement a -> Statement a
LabelledStmt b
a Id b
l Statement b
s
    ForInStmt _ i :: ForInInit b
i o :: Expression b
o ss :: Statement b
ss   -> b -> ForInInit b -> Expression b -> Statement b -> Statement b
forall a.
a -> ForInInit a -> Expression a -> Statement a -> Statement a
ForInStmt b
a ForInInit b
i Expression b
o Statement b
ss
    ForStmt _ i :: ForInit b
i t :: Maybe (Expression b)
t inc :: Maybe (Expression b)
inc ss :: Statement b
ss -> b
-> ForInit b
-> Maybe (Expression b)
-> Maybe (Expression b)
-> Statement b
-> Statement b
forall a.
a
-> ForInit a
-> Maybe (Expression a)
-> Maybe (Expression a)
-> Statement a
-> Statement a
ForStmt b
a ForInit b
i Maybe (Expression b)
t Maybe (Expression b)
inc Statement b
ss
    TryStmt _ tb :: Statement b
tb mcb :: Maybe (CatchClause b)
mcb mfb :: Maybe (Statement b)
mfb -> b
-> Statement b
-> Maybe (CatchClause b)
-> Maybe (Statement b)
-> Statement b
forall a.
a
-> Statement a
-> Maybe (CatchClause a)
-> Maybe (Statement a)
-> Statement a
TryStmt b
a Statement b
tb Maybe (CatchClause b)
mcb Maybe (Statement b)
mfb
    ThrowStmt _ e :: Expression b
e        -> b -> Expression b -> Statement b
forall a. a -> Expression a -> Statement a
ThrowStmt b
a Expression b
e
    ReturnStmt _ e :: Maybe (Expression b)
e       -> b -> Maybe (Expression b) -> Statement b
forall a. a -> Maybe (Expression a) -> Statement a
ReturnStmt b
a Maybe (Expression b)
e
    WithStmt _ o :: Expression b
o b :: Statement b
b       -> b -> Expression b -> Statement b -> Statement b
forall a. a -> Expression a -> Statement a -> Statement a
WithStmt b
a Expression b
o Statement b
b
    VarDeclStmt _ vds :: [VarDecl b]
vds    -> b -> [VarDecl b] -> Statement b
forall a. a -> [VarDecl a] -> Statement a
VarDeclStmt b
a [VarDecl b]
vds
    FunctionStmt _ n :: Id b
n as :: [Id b]
as b :: [Statement b]
b-> b -> Id b -> [Id b] -> [Statement b] -> Statement b
forall a. a -> Id a -> [Id a] -> [Statement a] -> Statement a
FunctionStmt b
a Id b
n [Id b]
as [Statement b]
b    
    
instance HasAnnotation LValue where
  getAnnotation :: LValue b -> b
getAnnotation lv :: LValue b
lv = case LValue b
lv of
    LVar a :: b
a _ -> b
a
    LDot a :: b
a _ _ -> b
a
    LBracket a :: b
a _ _ -> b
a
  setAnnotation :: b -> LValue b -> LValue b
setAnnotation a :: b
a lv :: LValue b
lv = case LValue b
lv of
    LVar _ n :: String
n -> b -> String -> LValue b
forall a. a -> String -> LValue a
LVar b
a String
n
    LDot _ o :: Expression b
o f :: String
f -> b -> Expression b -> String -> LValue b
forall a. a -> Expression a -> String -> LValue a
LDot b
a Expression b
o String
f
    LBracket a :: b
a o :: Expression b
o fe :: Expression b
fe -> b -> Expression b -> Expression b -> LValue b
forall a. a -> Expression a -> Expression a -> LValue a
LBracket b
a Expression b
o Expression b
fe    
  
instance HasAnnotation VarDecl where
  getAnnotation :: VarDecl b -> b
getAnnotation (VarDecl a :: b
a _ _) = b
a
  setAnnotation :: b -> VarDecl b -> VarDecl b
setAnnotation a :: b
a (VarDecl _ vn :: Id b
vn e :: Maybe (Expression b)
e) = b -> Id b -> Maybe (Expression b) -> VarDecl b
forall a. a -> Id a -> Maybe (Expression a) -> VarDecl a
VarDecl b
a Id b
vn Maybe (Expression b)
e  

instance HasAnnotation Prop  where
  getAnnotation :: Prop b -> b
getAnnotation p :: Prop b
p = case Prop b
p of
    PropId a :: b
a _ -> b
a
    PropString a :: b
a _ -> b
a
    PropNum a :: b
a _ -> b
a
  setAnnotation :: b -> Prop b -> Prop b
setAnnotation a :: b
a p :: Prop b
p = case Prop b
p of
    PropId _ id :: Id b
id -> b -> Id b -> Prop b
forall a. a -> Id a -> Prop a
PropId b
a Id b
id
    PropString _ s :: String
s -> b -> String -> Prop b
forall a. a -> String -> Prop a
PropString b
a String
s
    PropNum _ n :: Integer
n -> b -> Integer -> Prop b
forall a. a -> Integer -> Prop a
PropNum b
a Integer
n    
  
instance HasAnnotation CaseClause where
  getAnnotation :: CaseClause b -> b
getAnnotation c :: CaseClause b
c = case CaseClause b
c of
    CaseClause a :: b
a _ _ -> b
a
    CaseDefault a :: b
a _ -> b
a
  setAnnotation :: b -> CaseClause b -> CaseClause b
setAnnotation a :: b
a c :: CaseClause b
c = case CaseClause b
c of
    CaseClause _ e :: Expression b
e b :: [Statement b]
b -> b -> Expression b -> [Statement b] -> CaseClause b
forall a. a -> Expression a -> [Statement a] -> CaseClause a
CaseClause b
a Expression b
e [Statement b]
b
    CaseDefault _ b :: [Statement b]
b  -> b -> [Statement b] -> CaseClause b
forall a. a -> [Statement a] -> CaseClause a
CaseDefault b
a [Statement b]
b    
    
instance HasAnnotation CatchClause where
  getAnnotation :: CatchClause b -> b
getAnnotation (CatchClause a :: b
a _ _) = b
a
  setAnnotation :: b -> CatchClause b -> CatchClause b
setAnnotation a :: b
a (CatchClause _ id :: Id b
id b :: Statement b
b) = b -> Id b -> Statement b -> CatchClause b
forall a. a -> Id a -> Statement a -> CatchClause a
CatchClause b
a Id b
id Statement b
b

instance HasAnnotation Id where
  getAnnotation :: Id b -> b
getAnnotation (Id a :: b
a _) = b
a
  setAnnotation :: b -> Id b -> Id b
setAnnotation a :: b
a (Id _ s :: String
s) = b -> String -> Id b
forall a. a -> String -> Id a
Id b
a String
s