rewrite.go

package main

import (
	"fmt"
	"log"

	"github.com/robertkrimen/otto/ast"
	"github.com/robertkrimen/otto/file"
	"github.com/robertkrimen/otto/parser"
)

// Visitor Visit method is invoked for each node encountered by Walk.
// If the result visitor w is not nil, Walk visits each of the children
// of node with the visitor w, followed by a call of w.Visit(nil).
type Visitor interface {
	Visit(n ast.Node) (w Visitor)
}

type VisitorFunc func(v Visitor, n ast.Node) (w Visitor)

func (vf VisitorFunc) Visit(n ast.Node) (w Visitor) {
	return vf(vf, n)
}

// Walk traverses an AST in depth-first order: It starts by calling
// v.Visit(node); node must not be nil. If the visitor w returned by
// v.Visit(node) is not nil, Walk is invoked recursively with visitor
// w for each of the non-nil children of node, followed by a call
// of w.Visit(nil).
func Walk(v Visitor, n ast.Node) {
	if n == nil {
		return
	}
	if v = v.Visit(n); v == nil {
		return
	}

	switch n := n.(type) {
	case *ast.ArrayLiteral:
		for _, ex := range n.Value {
			Walk(v, ex)
		}
	case *ast.AssignExpression:
		Walk(v, n.Left)
		Walk(v, n.Right)
	case *ast.BadExpression:
	case *ast.BinaryExpression:
		Walk(v, n.Left)
		Walk(v, n.Right)
	case *ast.BlockStatement:
		for _, s := range n.List {
			Walk(v, s)
		}
	case *ast.BooleanLiteral:
	case *ast.BracketExpression:
		Walk(v, n.Left)
		Walk(v, n.Member)
	case *ast.BranchStatement:
		Walk(v, n.Label)
	case *ast.CallExpression:
		Walk(v, n.Callee)
		for _, a := range n.ArgumentList {
			Walk(v, a)
		}
	case *ast.CaseStatement:
		Walk(v, n.Test)
		for _, c := range n.Consequent {
			Walk(v, c)
		}
	case *ast.CatchStatement:
		Walk(v, n.Parameter)
		Walk(v, n.Body)
	case *ast.ConditionalExpression:
		Walk(v, n.Test)
		Walk(v, n.Consequent)
		Walk(v, n.Alternate)
	case *ast.DebuggerStatement:
	case *ast.DoWhileStatement:
		Walk(v, n.Test)
		Walk(v, n.Body)
	case *ast.DotExpression:
		Walk(v, n.Left)
	case *ast.EmptyExpression:
	case *ast.EmptyStatement:
	case *ast.ExpressionStatement:
		Walk(v, n.Expression)
	case *ast.ForInStatement:
		Walk(v, n.Into)
		Walk(v, n.Source)
		Walk(v, n.Body)
	case *ast.ForStatement:
		Walk(v, n.Initializer)
		Walk(v, n.Update)
		Walk(v, n.Test)
		Walk(v, n.Body)
	case *ast.FunctionLiteral:
		Walk(v, n.Name)
		for _, p := range n.ParameterList.List {
			Walk(v, p)
		}
		Walk(v, n.Body)
	case *ast.FunctionStatement:
		Walk(v, n.Function)
	case *ast.Identifier:
	case *ast.IfStatement:
		Walk(v, n.Test)
		Walk(v, n.Consequent)
		Walk(v, n.Alternate)
	case *ast.LabelledStatement:
		Walk(v, n.Statement)
	case *ast.NewExpression:
		Walk(v, n.Callee)
		for _, a := range n.ArgumentList {
			Walk(v, a)
		}
	case *ast.NullLiteral:
	case *ast.NumberLiteral:
	case *ast.ObjectLiteral:
	case *ast.Program:
		for _, b := range n.Body {
			Walk(v, b)
		}
	case *ast.RegExpLiteral:
	case *ast.ReturnStatement:
		Walk(v, n.Argument)
	case *ast.SequenceExpression:
		for _, e := range n.Sequence {
			Walk(v, e)
		}
	case *ast.StringLiteral:
	case *ast.SwitchStatement:
		Walk(v, n.Discriminant)
		for _, c := range n.Body {
			Walk(v, c)
		}
	case *ast.ThisExpression:
	case *ast.ThrowStatement:
		Walk(v, n.Argument)
	case *ast.TryStatement:
		Walk(v, n.Body)
		Walk(v, n.Catch)
		Walk(v, n.Finally)
	case *ast.UnaryExpression:
		Walk(v, n.Operand)
	case *ast.VariableExpression:
		Walk(v, n.Initializer)
	case *ast.VariableStatement:
		for _, e := range n.List {
			Walk(v, e)
		}
	case *ast.WhileStatement:
		Walk(v, n.Test)
		Walk(v, n.Body)
	case *ast.WithStatement:
		Walk(v, n.Object)
		Walk(v, n.Body)
	default:
		panic(fmt.Sprintf("Walk: unexpected node type %T", n))
	}

	Walk(v, nil)
}

func main() {
	source := `var b = function() {test(); var test = "test(); var test = 1"} // test`
	program, err := parser.ParseFile(nil, "", source, 0)
	if err != nil {
		log.Fatal(err)
	}

	var shift file.Idx

	Walk(VisitorFunc(func(v Visitor, n ast.Node) Visitor {
		if n == nil {
			return v
		}
		if id, ok := n.(*ast.Identifier); ok && id != nil {
			idx := n.Idx0() + shift - 1
			s := source[:idx] + "proxy_" + source[idx:]
			source = s
			shift += 6
		}
		if v, ok := n.(*ast.VariableExpression); ok && v != nil {
			idx := n.Idx0() + shift - 1
			s := source[:idx] + "varproxy_" + source[idx:]
			source = s
			shift += 9
		}

		return v
	}), program)

	println(source)
	// Output: var varproxy_b = function() {proxy_test(); var varproxy_test = "test(); var test = 1"} // test
}