meddling with ast
Published 2 years ago
by
jtog
with tags
ast
package org.example
import static org.codehaus.groovy.control.Phases.INITIALIZATION
import static org.codehaus.groovy.control.Phases.SEMANTIC_ANALYSIS
import java.security.CodeSource
import org.codehaus.groovy.ast.builder.AstBuilder
import org.codehaus.groovy.ast.ModuleNode
import org.codehaus.groovy.ast.expr.ConstantExpression
import org.codehaus.groovy.ast.expr.MethodCallExpression
import org.codehaus.groovy.ast.stmt.ExpressionStatement
import org.codehaus.groovy.control.CompilationUnit
import org.codehaus.groovy.control.CompilerConfiguration
import org.codehaus.groovy.control.SourceUnit
class ExampleASTClassloader extends GroovyClassLoader {
@Override
protected CompilationUnit createCompilationUnit(CompilerConfiguration config, CodeSource source) {
def unit = super.createCompilationUnit(config, source)
def operation = {SourceUnit sourceUnit->
ModuleNode module = sourceUnit.ast
def blocks = [:].withDefault{[:]}
module.statementBlock.statements.eachWithIndex {stmt, idx->
if (stmt.statementLabel) {
def label = stmt.statementLabel
blocks[label].to = idx
}
if (stmt instanceof ExpressionStatement
&& stmt.expression instanceof MethodCallExpression
&& stmt.expression.method.value == 'go_to'
&& stmt.expression.arguments.expressions.size() == 1) {
def label = stmt.expression.arguments.expressions[0].variable
blocks[label].from = idx
}
}
blocks.each {
println "detected jump from $it.value.from to $it.value.to with label $it.key"
}
blocks.each {
module.statementBlock.statements[it.value.from] = createPrintComment()
}
} as CompilationUnit.SourceUnitOperation
unit.addPhaseOperation operation, SEMANTIC_ANALYSIS
unit
}
def createPrintComment() {
new AstBuilder().buildFromSpec {
expression {
methodCall {
variable 'this'
constant 'println'
argumentList {
constant 'we should replace this with a jump'
}
}
}
}[0]
}
}
def sourceScript = '''
def a(){println 'a'}
def b(){println 'b'}
someLabel:
a()
b()
go_to someLabel
a()
b()
'''
new ExampleASTClassloader().parseClass(sourceScript).newInstance().run()
