got nested trees working:

a : A b C -> ^(A ^(b C))
| B
;

[git-p4: depot-paths = "//depot/code/antlr4/main/": change = 8825]

tool/playground/T.g

@@ -1,9 +1,11 @@
 grammar T;
 options {output=AST;}
 tokens {I;}
-a : A b C -> ^(A b C) ;
+a : A b C -> ^(A ^(b C))
+  | B
+  ;
 
-b : B ;
+b : B | C ;
 
 A : 'a';
 B : 'b';

tool/resources/org/antlr/v4/tool/templates/codegen/Java/Java.stg

@@ -123,9 +123,11 @@ public QStack\<<currentRule.ctxType>\> <currentRule.name>_stk = new QStack\<<cur
 >>
 
 CodeBlockForAlt(c, locals, preamble, ops) ::= <<
+{
 <locals; separator="\n">
 <preamble; separator="\n">
 <ops; separator="\n">
+}
 >>
 
 LL1AltBlock(choice, alts, error) ::= <<
@@ -368,9 +370,9 @@ AddTokenLeaf(a) ::= "_adaptor.addChild(_root0, _adaptor.create(<labelref(a.label
 AddRuleLeaf(a) ::= "_adaptor.addChild(_root0, <labelref(a.label)>.tree);"
 RuleBecomeRoot(r) ::= "_root0 = _adaptor.becomeRoot(<labelref(r.label)>.tree, _root0);"
 TokenBecomeRoot(t) ::= "_root0 = _adaptor.becomeRoot(_adaptor.create(<labelref(t.label)>), _root0);"
-AssignTreeResult(a) ::= <<
+AssignTreeResult(a) ::= "_localctx.tree = _root0;"
-_root0 = _adaptor.rulePostProcessing(_root0);
+RuleASTCleanup(r) ::= <<
-_localctx.tree = _root0;
+_localctx.tree = _adaptor.rulePostProcessing(_localctx.tree);
 _adaptor.setTokenBoundaries(_localctx.tree, _localctx.start, _localctx.stop);
 >>
 
@@ -381,11 +383,19 @@ TrackTokenElement(e) ::= "<e.name>.add(_adaptor.create(<labelref(e.label)>));"
 
 TreeRewrite(tr, locals, preamble, ops) ::= <<
 // rewrite: ...
-<CodeBlockForAlt(tr, locals, preamble, ops)>
+<locals; separator="\n">
+<preamble; separator="\n">
+<ops; separator="\n">
+_localctx.tree = _root0;
 >>
 
-RewriteTreeStructure(t, ops) ::= <<
+RewriteTreeStructure(t, locals, preamble, ops) ::= <<
-<ops; separator="\n">
+{
+    <locals; separator="\n">
+    <preamble; separator="\n">
+    <ops; separator="\n">
+    _adaptor.addChild(<labelref(t.enclosingBlock.rootDecl)>, <labelref(t.rootDecl)>);
+}
 >>
 RewriteIteratorDecl(d) ::= "Iterator <d.name>;"
 RewriteIteratorInit(i) ::= "<i.decl.name> = <i.decl.listName>.iterator();"

tool/src/org/antlr/v4/codegen/BlankOutputModelFactory.java

@@ -30,7 +30,7 @@
 package org.antlr.v4.codegen;
 
 import org.antlr.v4.codegen.model.*;
-import org.antlr.v4.codegen.model.ast.TreeRewrite;
+import org.antlr.v4.codegen.model.ast.*;
 import org.antlr.v4.runtime.misc.IntervalSet;
 import org.antlr.v4.tool.*;
 
@@ -53,6 +53,8 @@ public abstract class BlankOutputModelFactory implements OutputModelFactory {
 
 	public CodeBlockForAlt alternative(Alternative alt) { return null; }
 
+	public CodeBlockForAlt finishAlternative(CodeBlockForAlt blk, List<SrcOp> ops) { return blk; }
+
 	public CodeBlockForAlt epsilon() { return null; }
 
 	public List<SrcOp> ruleRef(GrammarAST ID, GrammarAST label, GrammarAST args) { return null; }
@@ -79,7 +81,7 @@ public abstract class BlankOutputModelFactory implements OutputModelFactory {
 
 	public TreeRewrite treeRewrite(GrammarAST ast, int rewriteLevel) { return null; }
 
-	public List<SrcOp> rewrite_tree(GrammarAST root, List<SrcOp> ops) {	return null; }
+	public RewriteTreeStructure rewrite_tree(GrammarAST root, int rewriteLevel) {	return null; }
 
 	public List<SrcOp> rewrite_ruleRef(GrammarAST ID, boolean isRoot) { return null; }
 

tool/src/org/antlr/v4/codegen/CodeGenerator.java

@@ -41,7 +41,6 @@ import java.util.*;
 /** General controller for code gen.  Can instantiate sub generator(s).
  */
 public class CodeGenerator {
-	public static boolean LAUNCH_ST_INSPECTOR = false;
 	public static final String TEMPLATE_ROOT = "org/antlr/v4/tool/templates/codegen";
 	public static final String VOCAB_FILE_EXTENSION = ".tokens";
 	public final static String vocabFilePattern =
@@ -125,7 +124,7 @@ public class CodeGenerator {
 		OutputModelWalker walker = new OutputModelWalker(g.tool, templates);
 		ST st = walker.walk(outputModel);
 
-		if (CodeGenerator.LAUNCH_ST_INSPECTOR) {
+		if ( g.tool.launch_ST_inspector ) {
 			st.inspect();
 			//if ( templates.isDefined("headerFile") ) headerFileST.inspect();
 		}

tool/src/org/antlr/v4/codegen/CodeGeneratorExtension.java

@@ -31,7 +31,7 @@ package org.antlr.v4.codegen;
 
 
 import org.antlr.v4.codegen.model.*;
-import org.antlr.v4.codegen.model.ast.TreeRewrite;
+import org.antlr.v4.codegen.model.ast.*;
 import org.antlr.v4.tool.GrammarAST;
 
 import java.util.List;
@@ -58,6 +58,8 @@ public class CodeGeneratorExtension {
 
 	public CodeBlockForAlt alternative(CodeBlockForAlt blk) { return blk; }
 
+	public CodeBlockForAlt finishAlternative(CodeBlockForAlt blk) { return blk; }
+
 	public CodeBlockForAlt epsilon(CodeBlockForAlt blk) { return blk; }
 
 	public List<SrcOp> ruleRef(List<SrcOp> ops) { return ops; }
@@ -89,7 +91,7 @@ public class CodeGeneratorExtension {
 
 	public TreeRewrite treeRewrite(TreeRewrite r) { return r; }
 
-	public List<SrcOp> rewrite_tree(List<SrcOp> ops) { return ops; }
+	public RewriteTreeStructure rewrite_tree(RewriteTreeStructure t) { return t; }
 
 	public List<SrcOp> rewrite_ruleRef(List<SrcOp> ops) { return ops; }
 

tool/src/org/antlr/v4/codegen/OutputModelController.java

@@ -31,7 +31,7 @@ package org.antlr.v4.codegen;
 
 import org.antlr.runtime.tree.*;
 import org.antlr.v4.codegen.model.*;
-import org.antlr.v4.codegen.model.ast.TreeRewrite;
+import org.antlr.v4.codegen.model.ast.*;
 import org.antlr.v4.codegen.model.decl.CodeBlock;
 import org.antlr.v4.parse.*;
 import org.antlr.v4.runtime.misc.IntervalSet;
@@ -141,9 +141,15 @@ public class OutputModelController implements OutputModelFactory {
 	public CodeGenerator getGenerator() { return delegate.getGenerator(); }
 
 	public CodeBlockForAlt alternative(Alternative alt) {
-		CodeBlockForAlt code = delegate.alternative(alt);
+		CodeBlockForAlt blk = delegate.alternative(alt);
-		for (CodeGeneratorExtension ext : extensions) code = ext.alternative(code);
+		for (CodeGeneratorExtension ext : extensions) blk = ext.alternative(blk);
-		return code;
+		return blk;
+	}
+
+	public CodeBlockForAlt finishAlternative(CodeBlockForAlt blk, List<SrcOp> ops) {
+		blk = delegate.finishAlternative(blk, ops);
+		for (CodeGeneratorExtension ext : extensions) blk = ext.finishAlternative(blk);
+		return blk;
 	}
 
 	public List<SrcOp> ruleRef(GrammarAST ID, GrammarAST label, GrammarAST args) {
@@ -278,10 +284,10 @@ public class OutputModelController implements OutputModelFactory {
 		return r;
 	}
 
-	public List<SrcOp> rewrite_tree(GrammarAST root, List<SrcOp> ops) {
+	public RewriteTreeStructure rewrite_tree(GrammarAST root, int rewriteLevel) {
-		ops = delegate.rewrite_tree(root, ops);
+		RewriteTreeStructure t = delegate.rewrite_tree(root, rewriteLevel);
-		for (CodeGeneratorExtension ext : extensions) ops = ext.rewrite_tree(ops);
+		for (CodeGeneratorExtension ext : extensions) t = ext.rewrite_tree(t);
-		return ops;
+		return t;
 	}
 
 	public List<SrcOp> rewrite_ruleRef(GrammarAST ID, boolean isRoot) {

tool/src/org/antlr/v4/codegen/OutputModelFactory.java

@@ -30,7 +30,7 @@
 package org.antlr.v4.codegen;
 
 import org.antlr.v4.codegen.model.*;
-import org.antlr.v4.codegen.model.ast.TreeRewrite;
+import org.antlr.v4.codegen.model.ast.*;
 import org.antlr.v4.codegen.model.decl.CodeBlock;
 import org.antlr.v4.runtime.misc.IntervalSet;
 import org.antlr.v4.tool.*;
@@ -58,6 +58,8 @@ public interface OutputModelFactory {
 
 	CodeBlockForAlt alternative(Alternative alt);
 
+	CodeBlockForAlt finishAlternative(CodeBlockForAlt blk, List<SrcOp> ops);
+
 	CodeBlockForAlt epsilon();
 
 	List<SrcOp> ruleRef(GrammarAST ID, GrammarAST label, GrammarAST args);
@@ -98,7 +100,7 @@ public interface OutputModelFactory {
 
 	TreeRewrite treeRewrite(GrammarAST ast, int rewriteLevel);
 
-	List<SrcOp> rewrite_tree(GrammarAST root, List<SrcOp> ops);
+	RewriteTreeStructure rewrite_tree(GrammarAST root, int rewriteLevel);
 
 	List<SrcOp> rewrite_ruleRef(GrammarAST ID, boolean isRoot);
 

tool/src/org/antlr/v4/codegen/OutputModelWalker.java

@@ -106,6 +106,7 @@ public class OutputModelWalker {
 				if ( o instanceof OutputModelObject ) {  // SINGLE MODEL OBJECT?
 					OutputModelObject nestedOmo = (OutputModelObject)o;
 					ST nestedST = walk(nestedOmo);
+//					System.out.println("set ModelElement "+fieldName+"="+nestedST+" in "+templateName);
 					st.add(fieldName, nestedST);
 				}
 				else if ( o instanceof Collection || o instanceof OutputModelObject[] ) {
@@ -117,6 +118,7 @@ public class OutputModelWalker {
 					for (OutputModelObject nestedOmo : nestedOmos) {
 						if ( nestedOmo==null ) continue;
 						ST nestedST = walk(nestedOmo);
+//						System.out.println("set ModelElement "+fieldName+"="+nestedST+" in "+templateName);
 						st.add(fieldName, nestedST);
 					}
 				}
@@ -125,6 +127,7 @@ public class OutputModelWalker {
 					Map<Object, ST> m = new HashMap<Object, ST>();
 					for (Object key : nestedOmoMap.keySet()) {
 						ST nestedST = walk(nestedOmoMap.get(key));
+//						System.out.println("set ModelElement "+fieldName+"="+nestedST+" in "+templateName);
 						m.put(key, nestedST);
 					}
 					st.add(fieldName, m);

tool/src/org/antlr/v4/codegen/ParserASTExtension.java

@@ -49,10 +49,17 @@ public class ParserASTExtension extends CodeGeneratorExtension {
 		return blk;
 	}
 
+	@Override
+	public CodeBlockForAlt finishAlternative(CodeBlockForAlt blk) {
+		Alternative alt = factory.getCurrentAlt();
+		if ( !alt.hasRewrite() ) blk.addOp(new AssignTreeResult(factory));
+		return blk;
+	}
+
 	@Override
 	public List<SrcOp> rulePostamble(List<SrcOp> ops) {
-		AssignTreeResult setReturn = new AssignTreeResult(factory);
+		RuleASTCleanup cleanup = new RuleASTCleanup(factory);
-		return DefaultOutputModelFactory.list(ops, setReturn);
+		return DefaultOutputModelFactory.list(ops, cleanup);
 	}
 
 	@Override
@@ -122,19 +129,4 @@ public class ParserASTExtension extends CodeGeneratorExtension {
 	}
 
 	// REWRITES
-
-	@Override
-	public List<SrcOp> rewrite_ruleRef(List<SrcOp> ops) {
-		return super.rewrite_ruleRef(ops);
-	}
-
-	@Override
-	public List<SrcOp> rewrite_tokenRef(List<SrcOp> ops) {
-		return super.rewrite_tokenRef(ops);
-	}
-
-	@Override
-	public List<SrcOp> rewrite_stringRef(List<SrcOp> ops) {
-		return super.rewrite_stringRef(ops);
-	}
 }

tool/src/org/antlr/v4/codegen/ParserFactory.java

@@ -60,6 +60,12 @@ public class ParserFactory extends DefaultOutputModelFactory {
 
 	public CodeBlockForAlt alternative(Alternative alt) { return new CodeBlockForAlt(this); }
 
+	@Override
+	public CodeBlockForAlt finishAlternative(CodeBlockForAlt blk, List<SrcOp> ops) {
+		blk.ops = ops;
+		return blk;
+	}
+
 	public List<SrcOp> action(GrammarAST ast) { return list(new Action(this, ast)); }
 
 	public List<SrcOp> forcedAction(GrammarAST ast) { return list(new ForcedAction(this, ast)); }
@@ -177,8 +183,8 @@ public class ParserFactory extends DefaultOutputModelFactory {
 	}
 
 	@Override
-	public List<SrcOp> rewrite_tree(GrammarAST root, List<SrcOp> ops) {
+	public RewriteTreeStructure rewrite_tree(GrammarAST root, int rewriteLevel) {
-		return list(new RewriteTreeStructure(this, root, ops));
+		return new RewriteTreeStructure(this, root, rewriteLevel);
 	}
 
 	public List<SrcOp> rewrite_ruleRef(GrammarAST ID, boolean isRoot) {

tool/src/org/antlr/v4/codegen/SourceGenTriggers.g

@@ -10,6 +10,7 @@ package org.antlr.v4.codegen;
 import org.antlr.v4.misc.Utils;
 import org.antlr.v4.codegen.model.*;
 import org.antlr.v4.codegen.model.decl.*;
+import org.antlr.v4.codegen.model.ast.*;
 import org.antlr.v4.tool.*;
 import java.util.Collections;
 import java.util.Map;
@@ -17,7 +18,7 @@ import java.util.HashMap;
 }
 
 @members {
-	public int rewriteLevel = 0;
+	public int codeBlockLevel = 0;
 	public OutputModelFactory factory;
     public SourceGenTriggers(TreeNodeStream input, OutputModelFactory factory) {
     	this(input);
@@ -30,7 +31,12 @@ dummy : block[null, null] ;
 block[GrammarAST label, GrammarAST ebnfRoot] returns [List<? extends SrcOp> omos]
     :	^(	blk=BLOCK (^(OPTIONS .+))?
 			{List<CodeBlockForAlt> alts = new ArrayList<CodeBlockForAlt>();}
-    		( alternative {alts.add($alternative.altCodeBlock);} )+
+    		(	alternative
+    			{
+		    	factory.finishAlternative($alternative.altCodeBlock, $alternative.ops);
+    			alts.add($alternative.altCodeBlock);
+    			}
+    		)+
     	)
     	{
     	if ( alts.size()==1 && ebnfRoot==null) return alts;
@@ -43,28 +49,58 @@ block[GrammarAST label, GrammarAST ebnfRoot] returns [List<? extends SrcOp> omos
     	}
     ;
 
+/*
 alternative returns [CodeBlockForAlt altCodeBlock]
 @init {
+	// set alt if outer ALT only
+	if ( inContext("RULE BLOCK") && ((AltAST)$start).alt!=null ) factory.setCurrentAlt(((AltAST)$start).alt);
+}
+	:	alternative_with_rewrite {$altCodeBlock = $alternative_with_rewrite.altCodeBlock;}
+
+	|	^(ALT EPSILON) {$altCodeBlock = factory.epsilon();}
+
+    |	{
     	List<SrcOp> elems = new ArrayList<SrcOp>();
+		$altCodeBlock = factory.alternative(factory.getCurrentAlt());
+		$ops = elems;
+		factory.setCurrentBlock($altCodeBlock);
+		}
+		^( ALT ( element {if ($element.omos!=null) elems.addAll($element.omos);} )+ )
+	;
+	
+alternative_with_rewrite returns [CodeBlockForAlt altCodeBlock]
+	:	^(ALT_REWRITE
+    		a=alternative
+    		(	rewrite {$a.ops.add($rewrite.code);} // insert at end of alt's code
+    		|
+    		)
+    		{$altCodeBlock=$a.altCodeBlock; $ops=$a.ops;}
+    	 )
+	;
+*/
+	
+alternative returns [CodeBlockForAlt altCodeBlock, List<SrcOp> ops]
+@init {
 	// set alt if outer ALT only
 	if ( inContext("RULE BLOCK") && ((AltAST)$start).alt!=null ) factory.setCurrentAlt(((AltAST)$start).alt);
 }
     :	^(ALT_REWRITE
     		a=alternative
-    		(	rewrite {$a.altCodeBlock.ops.add($rewrite.code);} // insert at end of alt's code
+    		(	rewrite {$a.ops.add($rewrite.code);} // insert at end of alt's code
     		|
     		)
-    		{$altCodeBlock=$a.altCodeBlock;}
+    		{$altCodeBlock=$a.altCodeBlock; $ops=$a.ops;}
     	 )
 
     |	^(ALT EPSILON) {$altCodeBlock = factory.epsilon();}
 
     |	{
+    	List<SrcOp> elems = new ArrayList<SrcOp>();
 		$altCodeBlock = factory.alternative(factory.getCurrentAlt());
+		$ops = elems;
 		factory.setCurrentBlock($altCodeBlock);
 		}
 		^( ALT ( element {if ($element.omos!=null) elems.addAll($element.omos);} )+ )
-    	{$altCodeBlock.ops = elems;}
     ;
 
 element returns [List<? extends SrcOp> omos]
@@ -158,7 +194,7 @@ elementOption
 
 rewrite returns [Rewrite code]
 	:	{
-		$code = factory.treeRewrite($start, rewriteLevel++);
+		$code = factory.treeRewrite($start, codeBlockLevel++);
 		CodeBlock save = factory.getCurrentBlock();
 		factory.setCurrentBlock($code);
 		}
@@ -166,6 +202,7 @@ rewrite returns [Rewrite code]
 		{
 		$code.ops = $nakedRewrite.omos;
 		factory.setCurrentBlock(save);
+		codeBlockLevel--;
 		}
 	;
 
@@ -213,12 +250,22 @@ rewriteTreeEbnf returns [List<SrcOp> omos]
 	;
 	
 rewriteTree returns [List<SrcOp> omos]
-	:	{List<SrcOp> elems = new ArrayList<SrcOp>();}
+	:	{
+		List<SrcOp> elems = new ArrayList<SrcOp>();
+		RewriteTreeStructure t = factory.rewrite_tree($start, codeBlockLevel++);
+		CodeBlock save = factory.getCurrentBlock();
+		factory.setCurrentBlock(t);
+		}
 		^(	TREE_BEGIN
 			rewriteTreeAtom[true] {elems.addAll($rewriteTreeAtom.omos);}
 			( rewriteTreeElement {elems.addAll($rewriteTreeElement.omos);} )*
 		 )
-		{$omos = factory.rewrite_tree($TREE_BEGIN, elems);}
+		{
+		t.ops = elems;
+		$omos = DefaultOutputModelFactory.list(t);
+		factory.setCurrentBlock(save);
+		codeBlockLevel--;
+		}
 	;
 
 rewriteSTAlt returns [List<SrcOp> omos]

tool/src/org/antlr/v4/codegen/model/CaptureNextToken.java

@@ -29,7 +29,12 @@
 
 package org.antlr.v4.codegen.model;
 
+import org.antlr.v4.codegen.OutputModelFactory;
+
 public class CaptureNextToken extends SrcOp {
 	public String varName;
-	public CaptureNextToken(String varName) { this.varName = varName; }
+	public CaptureNextToken(OutputModelFactory factory, String varName) {
+		super(factory);
+		this.varName = varName;
+	}
 }

tool/src/org/antlr/v4/codegen/model/CaptureNextTokenType.java

@@ -29,8 +29,13 @@
 
 package org.antlr.v4.codegen.model;
 
+import org.antlr.v4.codegen.OutputModelFactory;
+
 /** */
 public class CaptureNextTokenType extends SrcOp {
 	public String varName;
-	public CaptureNextTokenType(String varName) { this.varName = varName; }
+	public CaptureNextTokenType(OutputModelFactory factory, String varName) {
+		super(factory);
+		this.varName = varName;
+	}
 }

tool/src/org/antlr/v4/codegen/model/Choice.java

@@ -82,7 +82,7 @@ public abstract class Choice extends RuleElement {
 			TestSetInline e = (TestSetInline)expr;
 			Decl d = new TokenTypeDecl(factory, e.varName);
 			factory.getCurrentRuleFunction().addLocalDecl(d);
-			CaptureNextTokenType nextType = new CaptureNextTokenType(e.varName);
+			CaptureNextTokenType nextType = new CaptureNextTokenType(factory,e.varName);
 			addPreambleOp(nextType);
 		}
 		return expr;

tool/src/org/antlr/v4/codegen/model/CodeBlockForAlt.java

@@ -32,14 +32,7 @@ package org.antlr.v4.codegen.model;
 import org.antlr.v4.codegen.OutputModelFactory;
 import org.antlr.v4.codegen.model.decl.CodeBlock;
 
-import java.util.List;
-
 /** Contains Rewrite block (usually as last op) */
 public class CodeBlockForAlt extends CodeBlock {
 	public CodeBlockForAlt(OutputModelFactory factory) { super(factory); }
-
-	public CodeBlockForAlt(OutputModelFactory factory, List<SrcOp> ops) {
-		super(factory);
-		this.ops = ops;
-	}
 }

tool/src/org/antlr/v4/codegen/model/LL1Loop.java

@@ -57,7 +57,7 @@ public abstract class LL1Loop extends Choice {
 		SrcOp expr = addCodeForLookaheadTempVar(look);
 		if ( expr instanceof TestSetInline ) {
 			TestSetInline e = (TestSetInline)expr;
-			CaptureNextTokenType nextType = new CaptureNextTokenType(e.varName);
+			CaptureNextTokenType nextType = new CaptureNextTokenType(factory, e.varName);
 			addIterationOp(nextType);
 		}
 		return expr;

tool/src/org/antlr/v4/codegen/model/OutputModelObject.java

@@ -39,7 +39,7 @@ public abstract class OutputModelObject {
 
 	public OutputModelObject() {;}
 
-	public OutputModelObject(OutputModelFactory factory) { this.factory = factory; }
+	public OutputModelObject(OutputModelFactory factory) { this(factory, null); }
 
 	public OutputModelObject(OutputModelFactory factory, GrammarAST ast) {
 		this.factory = factory;

tool/src/org/antlr/v4/codegen/model/SrcOp.java

@@ -46,8 +46,7 @@ public abstract class SrcOp extends OutputModelObject {
 	 */
 	public SrcOp enclosingBlock;
 
-	public SrcOp() {;}
+	public SrcOp(OutputModelFactory factory) { this(factory,null); }
-	public SrcOp(OutputModelFactory factory) { super(factory); }
 	public SrcOp(OutputModelFactory factory, GrammarAST ast) {
 		super(factory,ast);
 		//uniqueID = ast.token.getTokenIndex();

tool/src/org/antlr/v4/codegen/model/ast/RewriteTreeStructure.java

@@ -30,17 +30,14 @@
 package org.antlr.v4.codegen.model.ast;
 
 import org.antlr.v4.codegen.OutputModelFactory;
-import org.antlr.v4.codegen.model.*;
 import org.antlr.v4.tool.GrammarAST;
 
-import java.util.List;
-
 /** ^(A B C) */
-public class RewriteTreeStructure extends SrcOp {
+public class RewriteTreeStructure extends TreeRewrite {
-	@ModelElement public List<? extends SrcOp> ops;
+	public RewriteTreeStructure(OutputModelFactory factory,
-
+								GrammarAST ast,
-	public RewriteTreeStructure(OutputModelFactory factory, GrammarAST ast, List<? extends SrcOp> ops) {
+								int level)
-		super(factory, ast);
+	{
-		this.ops = ops;
+		super(factory, level);
 	}
 }

tool/src/org/antlr/v4/codegen/model/ast/RuleASTCleanup.java

@@ -0,0 +1,39 @@
+/*
+ [The "BSD license"]
+ Copyright (c) 2011 Terence Parr
+ All rights reserved.
+
+ Redistribution and use in source and binary forms, with or without
+ modification, are permitted provided that the following conditions
+ are met:
+
+ 1. Redistributions of source code must retain the above copyright
+    notice, this list of conditions and the following disclaimer.
+ 2. Redistributions in binary form must reproduce the above copyright
+    notice, this list of conditions and the following disclaimer in the
+    documentation and/or other materials provided with the distribution.
+ 3. The name of the author may not be used to endorse or promote products
+    derived from this software without specific prior written permission.
+
+ THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR
+ IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
+ OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
+ IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT,
+ INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT
+ NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+ DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+ THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF
+ THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ */
+
+package org.antlr.v4.codegen.model.ast;
+
+import org.antlr.v4.codegen.OutputModelFactory;
+import org.antlr.v4.codegen.model.SrcOp;
+
+public class RuleASTCleanup extends SrcOp {
+	public RuleASTCleanup(OutputModelFactory factory) {
+		super(factory);
+	}
+}

tool/src/org/antlr/v4/codegen/model/decl/CodeBlock.java

@@ -38,7 +38,7 @@ import java.util.*;
 public class CodeBlock extends SrcOp {
 	@ModelElement public OrderedHashSet<Decl> locals;
 	@ModelElement public List<SrcOp> preamble;
-	@ModelElement public List ops; // has to be unchecked so we can add different subclasses of SrcOp :(
+	@ModelElement public List<SrcOp> ops;
 
 	public CodeBlock(OutputModelFactory factory) {
 		super(factory);
@@ -55,4 +55,9 @@ public class CodeBlock extends SrcOp {
 		if ( preamble==null ) preamble = new ArrayList<SrcOp>();
 		preamble.add(op);
 	}
+
+	public void addOp(SrcOp op) {
+		if ( ops==null ) ops = new ArrayList<SrcOp>();
+		ops.add(op);
+	}
 }