-- | 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 :: forall (t :: * -> *) a. Traversable t => 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 :: forall (t :: * -> *) a b. Traversable t => (a -> b) -> t a -> t b
reannotate a -> b
f 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)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> t a -> f (t b)
traverse (b -> () -> b
forall a. a -> () -> a
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 :: forall (t :: * -> *) b a. Traversable t => b -> t a -> t (a, b)
addExtraAnnotationField b
def 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)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> t a -> f (t b)
traverse (\a
z -> (a, b) -> () -> (a, b)
forall a. a -> () -> a
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 :: forall (t :: * -> *) a b. Traversable t => t (a, b) -> t a
removeExtraAnnotationField 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)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> t a -> f (t b)
traverse (a -> () -> a
forall a. 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 :: forall (t :: * -> *) a.
Traversable t =>
Int -> t a -> (t (a, Int), Int)
assignUniqueIds Int
first 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 b c b' c'. (b -> c) -> (b' -> c') -> (b, b') -> (c, c')
forall (a :: * -> * -> *) b c b' c'.
Arrow a =>
a b c -> a b' c' -> a (b, b') (c, c')
*** \Int
i -> Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
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)
forall (m :: * -> *) a b. 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 :: forall a. a -> State Int (a, Int)
f 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
+Int
1)
                 (a, Int) -> State Int (a, Int)
forall a. a -> StateT Int Identity a
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 :: forall (a :: * -> *) b. HasAnnotation a => (b -> b) -> a b -> a b
withAnnotation b -> b
f a b
x = b -> a b -> a b
forall b. 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 b. a b -> b
forall (a :: * -> *) b. HasAnnotation a => a b -> b
getAnnotation a b
x) a b
x

instance HasAnnotation Expression where
  getAnnotation :: forall b. Expression b -> b
getAnnotation Expression b
e = case Expression b
e of
   (StringLit b
a String
s)              -> b
a
   (RegexpLit b
a String
s Bool
g Bool
ci)         -> b
a
   (NumLit b
a Double
d)                 -> b
a
   (IntLit b
a Int
i)                 -> b
a
   (BoolLit b
a Bool
b)                -> b
a
   (NullLit b
a)                  -> b
a
   (ArrayLit b
a [Expression b]
exps)            -> b
a
   (ObjectLit b
a [(Prop b, Expression b)]
props)          -> b
a
   (ThisRef b
a)                  -> b
a
   (VarRef b
a Id b
id)                -> b
a
   (DotRef b
a Expression b
exp Id b
id)            -> b
a
   (BracketRef b
a Expression b
container Expression b
key) -> b
a
   (NewExpr b
a Expression b
ctor [Expression b]
params)      -> b
a
   (PrefixExpr b
a PrefixOp
op Expression b
e)          -> b
a
   (UnaryAssignExpr b
a UnaryAssignOp
op LValue b
lv)    -> b
a
   (InfixExpr b
a InfixOp
op Expression b
e1 Expression b
e2)       -> b
a
   (CondExpr b
a Expression b
g Expression b
et Expression b
ef)         -> b
a
   (AssignExpr b
a AssignOp
op LValue b
lv Expression b
e)       -> b
a
   (ListExpr b
a [Expression b]
es)              -> b
a
   (CallExpr b
a Expression b
fn [Expression b]
params)       -> b
a
   (FuncExpr b
a Maybe (Id b)
mid [Id b]
args [Statement b]
s)      -> b
a
  setAnnotation :: forall b. b -> Expression b -> Expression b
setAnnotation b
a Expression b
e = case Expression b
e of
    (StringLit b
_ String
s)              -> (b -> String -> Expression b
forall a. a -> String -> Expression a
StringLit b
a String
s)
    (RegexpLit b
_ String
s Bool
g 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 b
_ Double
d)                 -> (b -> Double -> Expression b
forall a. a -> Double -> Expression a
NumLit b
a Double
d)
    (IntLit b
_ 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
_)                  -> (b -> Expression b
forall a. a -> Expression a
NullLit b
a)
    (ArrayLit b
_ [Expression b]
exps)            -> (b -> [Expression b] -> Expression b
forall a. a -> [Expression a] -> Expression a
ArrayLit b
a [Expression b]
exps)
    (ObjectLit b
_ [(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
_)                  -> (b -> Expression b
forall a. a -> Expression a
ThisRef b
a)
    (VarRef b
_ Id b
id)                -> (b -> Id b -> Expression b
forall a. a -> Id a -> Expression a
VarRef b
a Id b
id)
    (DotRef b
_ Expression b
exp 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 b
_ Expression b
container 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 b
_ Expression b
ctor [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 b
_ PrefixOp
op 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 b
_ UnaryAssignOp
op 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 b
_ InfixOp
op Expression b
e1 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 b
_ Expression b
g Expression b
et 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 b
_ AssignOp
op LValue b
lv 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 b
_ [Expression b]
es)              -> (b -> [Expression b] -> Expression b
forall a. a -> [Expression a] -> Expression a
ListExpr b
a [Expression b]
es)
    (CallExpr b
_ Expression b
fn [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 b
_ Maybe (Id b)
mid [Id b]
args [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 :: forall b. Statement b -> b
getAnnotation Statement b
s = case Statement b
s of
    BlockStmt b
a [Statement b]
_        -> b
a
    EmptyStmt b
a          -> b
a
    ExprStmt b
a Expression b
_         -> b
a
    IfStmt b
a Expression b
_ Statement b
_ Statement b
_       -> b
a
    IfSingleStmt b
a Expression b
_ Statement b
_   -> b
a
    SwitchStmt b
a Expression b
_ [CaseClause b]
_     -> b
a
    WhileStmt b
a Expression b
_ Statement b
_      -> b
a
    DoWhileStmt b
a Statement b
_ Expression b
_    -> b
a
    BreakStmt b
a Maybe (Id b)
_        -> b
a
    ContinueStmt b
a Maybe (Id b)
_     -> b
a
    LabelledStmt b
a Id b
_ Statement b
_   -> b
a
    ForInStmt b
a ForInInit b
_ Expression b
_ Statement b
_    -> b
a
    ForStmt b
a ForInit b
_ Maybe (Expression b)
_ Maybe (Expression b)
_ Statement b
_    -> b
a
    TryStmt b
a Statement b
_ Maybe (CatchClause b)
_ Maybe (Statement b)
_      -> b
a
    ThrowStmt b
a Expression b
_        -> b
a
    ReturnStmt b
a Maybe (Expression b)
_       -> b
a
    WithStmt b
a Expression b
_ Statement b
_       -> b
a
    VarDeclStmt b
a [VarDecl b]
_      -> b
a
    FunctionStmt b
a Id b
_ [Id b]
_ [Statement b]
_ -> b
a
  setAnnotation :: forall b. b -> Statement b -> Statement b
setAnnotation b
a Statement b
s = case Statement b
s of
    BlockStmt b
_ [Statement b]
ss       -> b -> [Statement b] -> Statement b
forall a. a -> [Statement a] -> Statement a
BlockStmt b
a [Statement b]
ss
    EmptyStmt b
_          -> b -> Statement b
forall a. a -> Statement a
EmptyStmt b
a
    ExprStmt b
_ Expression b
e         -> b -> Expression b -> Statement b
forall a. a -> Expression a -> Statement a
ExprStmt b
a Expression b
e
    IfStmt b
_ Expression b
g Statement b
t 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 b
_ Expression b
g 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 b
_ Expression b
g [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 b
_ Expression b
g 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 b
_ Statement b
ss 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 b
_ 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 b
_ 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 b
_ Id b
l 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 b
_ ForInInit b
i Expression b
o 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 b
_ ForInit b
i Maybe (Expression b)
t Maybe (Expression b)
inc 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 b
_ Statement b
tb Maybe (CatchClause b)
mcb 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 b
_ Expression b
e        -> b -> Expression b -> Statement b
forall a. a -> Expression a -> Statement a
ThrowStmt b
a Expression b
e
    ReturnStmt b
_ 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 b
_ Expression b
o 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 b
_ [VarDecl b]
vds    -> b -> [VarDecl b] -> Statement b
forall a. a -> [VarDecl a] -> Statement a
VarDeclStmt b
a [VarDecl b]
vds
    FunctionStmt b
_ Id b
n [Id b]
as [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 :: forall b. LValue b -> b
getAnnotation LValue b
lv = case LValue b
lv of
    LVar b
a String
_ -> b
a
    LDot b
a Expression b
_ String
_ -> b
a
    LBracket b
a Expression b
_ Expression b
_ -> b
a
  setAnnotation :: forall b. b -> LValue b -> LValue b
setAnnotation b
a LValue b
lv = case LValue b
lv of
    LVar b
_ String
n -> b -> String -> LValue b
forall a. a -> String -> LValue a
LVar b
a String
n
    LDot b
_ Expression b
o 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 b
a Expression b
o 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 :: forall b. VarDecl b -> b
getAnnotation (VarDecl b
a Id b
_ Maybe (Expression b)
_) = b
a
  setAnnotation :: forall b. b -> VarDecl b -> VarDecl b
setAnnotation b
a (VarDecl b
_ Id b
vn 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 :: forall b. Prop b -> b
getAnnotation Prop b
p = case Prop b
p of
    PropId b
a Id b
_ -> b
a
    PropString b
a String
_ -> b
a
    PropNum b
a Integer
_ -> b
a
  setAnnotation :: forall b. b -> Prop b -> Prop b
setAnnotation b
a Prop b
p = case Prop b
p of
    PropId b
_ Id b
id -> b -> Id b -> Prop b
forall a. a -> Id a -> Prop a
PropId b
a Id b
id
    PropString b
_ String
s -> b -> String -> Prop b
forall a. a -> String -> Prop a
PropString b
a String
s
    PropNum b
_ Integer
n -> b -> Integer -> Prop b
forall a. a -> Integer -> Prop a
PropNum b
a Integer
n    
  
instance HasAnnotation CaseClause where
  getAnnotation :: forall b. CaseClause b -> b
getAnnotation CaseClause b
c = case CaseClause b
c of
    CaseClause b
a Expression b
_ [Statement b]
_ -> b
a
    CaseDefault b
a [Statement b]
_ -> b
a
  setAnnotation :: forall b. b -> CaseClause b -> CaseClause b
setAnnotation b
a CaseClause b
c = case CaseClause b
c of
    CaseClause b
_ Expression b
e [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 :: forall b. CatchClause b -> b
getAnnotation (CatchClause b
a Id b
_ Statement b
_) = b
a
  setAnnotation :: forall b. b -> CatchClause b -> CatchClause b
setAnnotation b
a (CatchClause b
_ Id b
id 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 :: forall b. Id b -> b
getAnnotation (Id b
a String
_) = b
a
  setAnnotation :: forall b. b -> Id b -> Id b
setAnnotation b
a (Id b
_ String
s) = b -> String -> Id b
forall a. a -> String -> Id a
Id b
a String
s