Groovy web console

subscribe to the feed Subscribe
to this
site
meddling with ast (via #groovywebconsole)
tweet this snippet Tweet
this
script

meddling with ast

Published 3 years ago by jtog with tags ast
Actions  ➤ Edit in console Back to console Show/hide line numbers View recent scripts
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()