good start on * and ? in rewrites

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

tool/playground/T.g

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

tool/playground/TestT.java

@@ -2,7 +2,6 @@ import org.antlr.v4.Tool;
 import org.antlr.v4.automata.ParserATNFactory;
 import org.antlr.v4.runtime.*;
 import org.antlr.v4.runtime.atn.ATN;
-import org.antlr.v4.runtime.tree.Tree;
 import org.antlr.v4.semantics.SemanticPipeline;
 import org.antlr.v4.tool.*;
 
@@ -13,8 +12,8 @@ public class TestT {
 		TLexer t = new TLexer(new ANTLRFileStream(args[0]));
 		CommonTokenStream tokens = new CommonTokenStream(t);
 		TParser p = new TParser(tokens);
-		ParserRuleContext ret = p.a();
+//		ParserRuleContext ret = p.a();
-		System.out.println(((Tree)ret.tree).toStringTree());
+//		System.out.println(((Tree)ret.tree).toStringTree());
 	}
 
 	public static void dump() throws Exception {

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

@@ -389,6 +389,28 @@ TreeRewrite(tr, locals, preamble, ops) ::= <<
 _localctx.tree = _root0;
 >>
 
+RewriteIteratorDecl(d) ::= "Iterator <d.name>;"
+RewriteIteratorInit(i) ::= "<i.decl.name> = <i.decl.listName>.iterator();"
+RewriteIteratorName(elemName,level) ::= "it<level>_<elemName>"
+
+RewriteTreeOptional(o, locals, preamble, ops) ::= <<
+// ?
+<locals; separator="\n">
+<preamble; separator="\n">
+//if ( true ) {
+	<ops; separator="\n">
+//}
+>>
+
+RewriteTreeClosure(c, locals, preamble, ops) ::= <<
+// *
+<locals; separator="\n">
+<preamble; separator="\n">
+//while ( <c.iteratorDecls:{d | <d.name>.hasNext()}; separator=" || "> ) {
+	<ops; separator="\n">
+//}
+>>
+
 RewriteTreeStructure(t, locals, preamble, ops) ::= <<
 {
     <locals; separator="\n">
@@ -397,15 +419,11 @@ RewriteTreeStructure(t, locals, preamble, ops) ::= <<
     _adaptor.addChild(<labelref(t.enclosingBlock.rootDecl)>, <labelref(t.rootDecl)>);
 }
 >>
-RewriteIteratorDecl(d) ::= "Iterator <d.name>;"
-RewriteIteratorInit(i) ::= "<i.decl.name> = <i.decl.listName>.iterator();"
-RewriteIteratorName(elemName,level) ::= "it<level>_<elemName>"
 
-RewriteTokenRef(t) ::= "_adaptor.addChild(<t.enclosingBlock.rootDecl.name>, <t.decl.name>.next());"
+RewriteTokenRef(t) ::= "_adaptor.addChild(<t.rootName>, <t.decl.name>.next());"
-RewriteTokenRefIsRoot(t) ::= <%
+RewriteTokenRefIsRoot(t) ::= <<
-<labelref(t.enclosingBlock.rootDecl)> =
+<t.rootName> = _adaptor.becomeRoot(<t.decl.name>.next(), <t.rootName>);
-	_adaptor.becomeRoot(<t.decl.name>.next(), <labelref(t.enclosingBlock.rootDecl)>);
+>>
-%>
 RewriteRuleRef(r) ::= "_adaptor.addChild(<r.enclosingBlock.rootDecl.name>, <r.decl.name>.next());"
 RewriteRuleRefIsRoot(r) ::= <%
 <labelref(r.enclosingBlock.rootDecl)> =

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

@@ -79,9 +79,13 @@ public abstract class BlankOutputModelFactory implements OutputModelFactory {
 
 	// AST REWRITES
 
-	public TreeRewrite treeRewrite(GrammarAST ast, int rewriteLevel) { return null; }
+	public TreeRewrite treeRewrite(GrammarAST ast) { return null; }
 
-	public RewriteTreeStructure rewrite_tree(GrammarAST root, int rewriteLevel) {	return null; }
+	public RewriteTreeOptional rewrite_optional(GrammarAST ast) { return null; }
+
+	public RewriteTreeClosure rewrite_closure(GrammarAST ast) { return null; }
+
+	public RewriteTreeStructure rewrite_tree(GrammarAST root) { return null; }
 
 	public List<SrcOp> rewrite_ruleRef(GrammarAST ID, boolean isRoot) { return null; }
 

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

@@ -91,6 +91,10 @@ public class CodeGeneratorExtension {
 
 	public TreeRewrite treeRewrite(TreeRewrite r) { return r; }
 
+	public RewriteTreeOptional rewrite_optional(RewriteTreeOptional o) { return o; }
+
+	public RewriteTreeClosure rewrite_closure(RewriteTreeClosure c) { return c; }
+
 	public RewriteTreeStructure rewrite_tree(RewriteTreeStructure t) { return t; }
 
 	public List<SrcOp> rewrite_ruleRef(List<SrcOp> ops) { return ops; }

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

@@ -30,7 +30,7 @@
 package org.antlr.v4.codegen;
 
 import org.antlr.v4.codegen.model.*;
-import org.antlr.v4.codegen.model.decl.CodeBlock;
+import org.antlr.v4.codegen.model.decl.*;
 import org.antlr.v4.tool.*;
 
 import java.util.*;
@@ -93,6 +93,10 @@ public abstract class DefaultOutputModelFactory extends BlankOutputModelFactory
 		return currentBlock;
 	}
 
+	public int getCodeBlockLevel() { return controller.walker.codeBlockLevel; }
+
+	public int getTreeLevel() { return controller.walker.treeLevel; }
+
 	// MISC
 
 	public static List<SrcOp> list(Object... values) {
@@ -105,5 +109,13 @@ public abstract class DefaultOutputModelFactory extends BlankOutputModelFactory
 		}
 		return x;
 	}
+
+	public Decl getCurrentDeclForName(String name) {
+		for (Decl d : getCurrentBlock().locals.elements()) {
+			if ( d.name.equals(name) ) return d;
+		}
+		return null;
+	}
+
 }
 

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

@@ -49,6 +49,11 @@ public class OutputModelController implements OutputModelFactory {
 	/** Post-processing CodeGeneratorExtension objects; done in order given. */
 	public List<CodeGeneratorExtension> extensions = new ArrayList<CodeGeneratorExtension>();
 
+	/** While walking code in rules, this is set to the tree walker that
+	 *  triggers actions.
+	 */
+	public SourceGenTriggers walker;
+
 	public OutputModelController(OutputModelFactory factory) {
 		this.delegate = factory;
 	}
@@ -76,9 +81,9 @@ public class OutputModelController implements OutputModelFactory {
 			GrammarASTAdaptor adaptor = new GrammarASTAdaptor(r.ast.token.getInputStream());
 			GrammarAST blk = (GrammarAST)r.ast.getFirstChildWithType(ANTLRParser.BLOCK);
 			CommonTreeNodeStream nodes = new CommonTreeNodeStream(adaptor,blk);
-			SourceGenTriggers genTriggers = new SourceGenTriggers(nodes, this);
+			walker = new SourceGenTriggers(nodes, this);
 			try {
-				function.code = DefaultOutputModelFactory.list(genTriggers.block(null, null)); // walk AST of rule alts/elements
+				function.code = DefaultOutputModelFactory.list(walker.block(null, null)); // walk AST of rule alts/elements
 			}
 			catch (Exception e){
 				e.printStackTrace(System.err);
@@ -278,14 +283,27 @@ public class OutputModelController implements OutputModelFactory {
 
 	// REWRITES
 
-	public TreeRewrite treeRewrite(GrammarAST ast, int rewriteLevel) {
+	public TreeRewrite treeRewrite(GrammarAST ast) {
-		TreeRewrite r = delegate.treeRewrite(ast, rewriteLevel);
+		TreeRewrite r = delegate.treeRewrite(ast);
 		for (CodeGeneratorExtension ext : extensions) r = ext.treeRewrite(r);
 		return r;
 	}
 
-	public RewriteTreeStructure rewrite_tree(GrammarAST root, int rewriteLevel) {
+	public RewriteTreeOptional rewrite_optional(GrammarAST ast) {
-		RewriteTreeStructure t = delegate.rewrite_tree(root, rewriteLevel);
+		RewriteTreeOptional o = delegate.rewrite_optional(ast);
+		for (CodeGeneratorExtension ext : extensions) o = ext.rewrite_optional(o);
+		return o;
+	}
+
+	public RewriteTreeClosure rewrite_closure(GrammarAST ast) {
+		List<GrammarAST> refs = getElementReferencesShallow(ast);
+		RewriteTreeClosure c = delegate.rewrite_closure(ast);
+		for (CodeGeneratorExtension ext : extensions) c = ext.rewrite_closure(c);
+		return c;
+	}
+
+	public RewriteTreeStructure rewrite_tree(GrammarAST root) {
+		RewriteTreeStructure t = delegate.rewrite_tree(root);
 		for (CodeGeneratorExtension ext : extensions) t = ext.rewrite_tree(t);
 		return t;
 	}
@@ -323,4 +341,76 @@ public class OutputModelController implements OutputModelFactory {
 	public void setCurrentBlock(CodeBlock blk) { delegate.setCurrentBlock(blk); }
 
 	public CodeBlock getCurrentBlock() { return delegate.getCurrentBlock(); }
+
+	public int getCodeBlockLevel() { return delegate.getCodeBlockLevel(); }
+
+	public int getTreeLevel() { return delegate.getTreeLevel(); }
+
+	// SUPPORT
+
+	/** Given (('?'|'*') (REWRITE_BLOCK (ALT ...))) return list of element refs at
+	 *  top level of REWRITE_BLOCK.
+	 */
+	public List<GrammarAST> getElementReferencesShallow(GrammarAST ebnfRoot) {
+		if ( ebnfRoot.getType()!=ANTLRParser.CLOSURE &&
+		     ebnfRoot.getType()!=ANTLRParser.OPTIONAL )
+		{
+			return null;
+		}
+		GrammarAST blkAST = (GrammarAST)ebnfRoot.getChild(0);
+		if ( blkAST.getType()!=ANTLRParser.REWRITE_BLOCK ) return null;
+		GrammarAST altAST = (GrammarAST)blkAST.getChild(0);
+		if ( altAST.getType()!=ANTLRParser.ALT ) return null;
+
+		IntervalSet elementTokenTypes = getRewriteElementTokenTypeSet();
+		Alternative alt = getCurrentAlt();
+		List<GrammarAST> elems = new ArrayList<GrammarAST>();
+		for (Object o : altAST.getChildren()) {
+			GrammarAST ref = (GrammarAST)o;
+			if ( elementTokenTypes.member(ref.getType()) ) {
+				boolean imaginary = ref.getType()==ANTLRParser.TOKEN_REF &&
+									!alt.tokenRefs.containsKey(ref.getText());
+				if ( !imaginary ) elems.add(ref);
+			}
+		}
+
+		return elems;
+	}
+
+	/** Given (('?'|'*') (REWRITE_BLOCK (ALT ...))) return list of element refs at
+	 *  or below toplevel REWRITE_BLOCK.
+	 */
+	public List<GrammarAST> getElementReferencesDeep(GrammarAST ebnfRoot) {
+		if ( ebnfRoot.getType()!=ANTLRParser.CLOSURE &&
+		     ebnfRoot.getType()!=ANTLRParser.OPTIONAL )
+		{
+			return null;
+		}
+		GrammarAST blkAST = (GrammarAST)ebnfRoot.getChild(0);
+		if ( blkAST.getType()!=ANTLRParser.REWRITE_BLOCK ) return null;
+		GrammarAST altAST = (GrammarAST)blkAST.getChild(0);
+		if ( altAST.getType()!=ANTLRParser.ALT ) return null;
+
+		List<GrammarAST> elems = new ArrayList<GrammarAST>();
+		Alternative alt = getCurrentAlt();
+		IntervalSet elementTokenTypes = getRewriteElementTokenTypeSet();
+		List<GrammarAST> refs = altAST.getNodesWithType(elementTokenTypes);
+		if ( refs!=null ) {
+			for (GrammarAST ref : refs) {
+				boolean imaginary = ref.getType()==ANTLRParser.TOKEN_REF &&
+									!alt.tokenRefs.containsKey(ref.getText());
+				if ( !imaginary ) elems.add(ref);
+			}
+		}
+		return elems;
+	}
+
+	public IntervalSet getRewriteElementTokenTypeSet() {
+		IntervalSet elementTokenTypes = new IntervalSet();
+		elementTokenTypes.add(ANTLRParser.TOKEN_REF); // might be imaginary
+		elementTokenTypes.add(ANTLRParser.RULE_REF);
+		elementTokenTypes.add(ANTLRParser.STRING_LITERAL);
+		elementTokenTypes.add(ANTLRParser.LABEL);
+		return elementTokenTypes;
+	}
 }

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

@@ -98,9 +98,13 @@ public interface OutputModelFactory {
 	// Though dealing with ASTs, we must deal with here since these are
 	// triggered from elements in ANTLR's internal GrammarAST
 
-	TreeRewrite treeRewrite(GrammarAST ast, int rewriteLevel);
+	TreeRewrite treeRewrite(GrammarAST ast);
 
-	RewriteTreeStructure rewrite_tree(GrammarAST root, int rewriteLevel);
+	RewriteTreeOptional rewrite_optional(GrammarAST ast);
+
+	RewriteTreeClosure rewrite_closure(GrammarAST ast);
+
+	RewriteTreeStructure rewrite_tree(GrammarAST root);
 
 	List<SrcOp> rewrite_ruleRef(GrammarAST ID, boolean isRoot);
 
@@ -110,6 +114,8 @@ public interface OutputModelFactory {
 
 	// CONTEXT MANIPULATION
 
+	// TODO: move to controller?
+
 	OutputModelObject getRoot();
 
 	void setRoot(OutputModelObject root);
@@ -129,4 +135,8 @@ public interface OutputModelFactory {
 	void setCurrentBlock(CodeBlock blk);
 
 	CodeBlock getCurrentBlock();
+
+	int getCodeBlockLevel();
+
+	int getTreeLevel();
 }

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

@@ -45,7 +45,7 @@ public class ParserASTExtension extends CodeGeneratorExtension {
 	@Override
 	public CodeBlockForAlt alternative(CodeBlockForAlt blk) {
 		Alternative alt = factory.getCurrentAlt();
-		if ( !alt.hasRewrite() ) blk.addLocalDecl( new RootDecl(factory, 0) );
+		if ( !alt.hasRewrite() ) blk.addLocalDecl( new RootDecl(factory) );
 		return blk;
 	}
 

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

@@ -169,29 +169,49 @@ public class ParserFactory extends DefaultOutputModelFactory {
 	}
 
 	public boolean needsImplicitLabel(GrammarAST ID, LabeledOp op) {
-		return op.getLabels().size()==0 &&
+		return	op.getLabels().size()==0 &&
-		(getCurrentAlt().tokenRefsInActions.containsKey(ID.getText()) ||
+				(getCurrentAlt().tokenRefsInActions.containsKey(ID.getText()) ||
-		getCurrentAlt().ruleRefsInActions.containsKey(ID.getText()));
+				getCurrentAlt().ruleRefsInActions.containsKey(ID.getText()));
 	}
 
 	// AST REWRITE
 
 
 	@Override
-	public TreeRewrite treeRewrite(GrammarAST ast, int rewriteLevel) {
+	public TreeRewrite treeRewrite(GrammarAST ast) {
-		return new TreeRewrite(this, rewriteLevel);
+		return new TreeRewrite(this, getTreeLevel(), getCodeBlockLevel());
 	}
 
 	@Override
-	public RewriteTreeStructure rewrite_tree(GrammarAST root, int rewriteLevel) {
+	public RewriteTreeOptional rewrite_optional(GrammarAST ast) {
-		return new RewriteTreeStructure(this, root, rewriteLevel);
+		List<GrammarAST> refs = controller.getElementReferencesDeep(ast);
+		return new RewriteTreeOptional(this, ast);
+	}
+
+	@Override
+	public RewriteTreeClosure rewrite_closure(GrammarAST ast) {
+		RewriteTreeClosure c =
+			new RewriteTreeClosure(this, ast, getTreeLevel(), getCodeBlockLevel());
+		List<GrammarAST> refs = controller.getElementReferencesShallow(ast);
+		if ( refs!=null ) {
+			for (GrammarAST ref : refs) {
+				RewriteIteratorDecl d = new RewriteIteratorDecl(this, ref);
+				c.addLocalDecl(d);
+				c.iteratorDecls.add(d);
+				RewriteIteratorInit init = new RewriteIteratorInit(this, d);
+				c.addPreambleOp(init);
+			}
+		}
+		return c;
+	}
+
+	@Override
+	public RewriteTreeStructure rewrite_tree(GrammarAST root) {
+		return new RewriteTreeStructure(this, root, getTreeLevel(), getCodeBlockLevel());
 	}
 
 	public List<SrcOp> rewrite_ruleRef(GrammarAST ID, boolean isRoot) {
-		RewriteIteratorDecl d = new RewriteIteratorDecl(this, ID, 0);
+		RewriteIteratorDecl d = new RewriteIteratorDecl(this, ID);
-		getCurrentBlock().addLocalDecl(d);
-		RewriteIteratorInit init = new RewriteIteratorInit(this, d);
-		getCurrentBlock().addPreambleOp(init);
 		RewriteRuleRef ruleRef;
 		if ( isRoot ) ruleRef = new RewriteRuleRefIsRoot(this, ID, d);
 		else ruleRef = new RewriteRuleRef(this, ID, d);
@@ -199,13 +219,12 @@ public class ParserFactory extends DefaultOutputModelFactory {
 	}
 
 	public List<SrcOp> rewrite_tokenRef(GrammarAST ID, boolean isRoot) {
-		RewriteIteratorDecl d = new RewriteIteratorDecl(this, ID, 0);
+		String itName = gen.target.getRewriteIteratorName(ID, getCodeBlockLevel());
-		getCurrentBlock().addLocalDecl(d);
+		String rootName = gen.target.getRootName(getTreeLevel());
-		RewriteIteratorInit init = new RewriteIteratorInit(this, d);
+		RewriteIteratorDecl d = (RewriteIteratorDecl)getCurrentDeclForName(itName);
-		getCurrentBlock().addPreambleOp(init);
 		RewriteTokenRef tokenRef;
-		if ( isRoot ) tokenRef = new RewriteTokenRefIsRoot(this, ID, d);
+		if ( isRoot ) tokenRef = new RewriteTokenRefIsRoot(this, ID, rootName, d);
-		else tokenRef = new RewriteTokenRef(this, ID, d);
+		else tokenRef = new RewriteTokenRef(this, ID, rootName, d);
 		return list(tokenRef);
 	}
 

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

@@ -18,7 +18,8 @@ import java.util.HashMap;
 }
 
 @members {
-	public int codeBlockLevel = 0;
+	public int codeBlockLevel = -1;
+	public int treeLevel = -1;
 	public OutputModelFactory factory;
     public SourceGenTriggers(TreeNodeStream input, OutputModelFactory factory) {
     	this(input);
@@ -194,7 +195,9 @@ elementOption
 
 rewrite returns [Rewrite code]
 	:	{
-		$code = factory.treeRewrite($start, codeBlockLevel++);
+		treeLevel = 0;
+		codeBlockLevel++;
+		$code = factory.treeRewrite($start);
 		CodeBlock save = factory.getCurrentBlock();
 		factory.setCurrentBlock($code);
 		}
@@ -229,7 +232,7 @@ rewriteTreeAlt returns [List<SrcOp> omos]
 rewriteTreeElement returns [List<SrcOp> omos]
 	:	rewriteTreeAtom[false]				{$omos = $rewriteTreeAtom.omos;}
 	|	rewriteTree							{$omos = $rewriteTree.omos;}
-	|   rewriteTreeEbnf						{$omos = $rewriteTreeEbnf.omos;}
+	|   rewriteTreeEbnf						{$omos = DefaultOutputModelFactory.list($rewriteTreeEbnf.op);}
 	;
 
 rewriteTreeAtom[boolean isRoot] returns [List<SrcOp> omos]
@@ -244,15 +247,32 @@ rewriteTreeAtom[boolean isRoot] returns [List<SrcOp> omos]
 	|	ACTION
 	;
 
-rewriteTreeEbnf returns [List<SrcOp> omos]
+rewriteTreeEbnf returns [CodeBlock op]
-	:	^('?' ^(REWRITE_BLOCK rewriteTreeAlt))
+	:	^(	(a=OPTIONAL|a=CLOSURE)
-	|	^('*' ^(REWRITE_BLOCK rewriteTreeAlt))
+			^(	REWRITE_BLOCK
+				{
+				codeBlockLevel++;
+				if ( $a.getType()==OPTIONAL ) $op = factory.rewrite_optional($start);
+				else $op = factory.rewrite_closure($start);
+				CodeBlock save = factory.getCurrentBlock();
+				factory.setCurrentBlock($op);
+				}
+				alt=rewriteTreeAlt
+			)
+		)
+		{
+		$op.addOps($alt.omos);
+		factory.setCurrentBlock(save);
+		codeBlockLevel--;
+		}
 	;
 	
 rewriteTree returns [List<SrcOp> omos]
 	:	{
+		codeBlockLevel++;
+		treeLevel++;
 		List<SrcOp> elems = new ArrayList<SrcOp>();
-		RewriteTreeStructure t = factory.rewrite_tree($start, codeBlockLevel++);
+		RewriteTreeStructure t = factory.rewrite_tree($start);
 		CodeBlock save = factory.getCurrentBlock();
 		factory.setCurrentBlock(t);
 		}
@@ -264,6 +284,7 @@ rewriteTree returns [List<SrcOp> omos]
 		t.ops = elems;
 		$omos = DefaultOutputModelFactory.list(t);
 		factory.setCurrentBlock(save);
+		treeLevel--;
 		codeBlockLevel--;
 		}
 	;

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

@@ -33,7 +33,7 @@ import org.antlr.v4.codegen.OutputModelFactory;
 import org.antlr.v4.codegen.model.decl.CodeBlock;
 
 public class Rewrite extends CodeBlock {
-	public Rewrite(OutputModelFactory factory) {
+	public Rewrite(OutputModelFactory factory, int treeLevel, int codeBlockLevel) {
-		super(factory);
+		super(factory, treeLevel, codeBlockLevel);
 	}
 }

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

@@ -51,6 +51,7 @@ public class RuleFunction extends OutputModelObject {
 	public ATNState startState;
 	public int index;
 	public Collection<Attribute> args = null;
+	public Rule rule;
 
 	@ModelElement public List<SrcOp> code;
 	@ModelElement public OrderedHashSet<Decl> locals; // TODO: move into ctx?
@@ -59,13 +60,10 @@ public class RuleFunction extends OutputModelObject {
 	@ModelElement public Action finallyAction;
 	@ModelElement public List<SrcOp> postamble;
 
-	public RuleFunction(OutputModelFactory factory) {
-		super(factory);
-	}
-
 	public RuleFunction(OutputModelFactory factory, Rule r) {
 		super(factory);
 		this.name = r.name;
+		this.rule = r;
 		if ( r.modifiers!=null && r.modifiers.size()>0 ) {
 			this.modifiers = new ArrayList<String>();
 			for (GrammarAST t : r.modifiers) modifiers.add(t.getText());

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

@@ -37,8 +37,12 @@ import org.antlr.v4.tool.GrammarAST;
 public class RewriteTokenRef extends SrcOp {
 	/** Which iterator decl are we associated with? */
 	public RewriteIteratorDecl decl;
-	public RewriteTokenRef(OutputModelFactory factory, GrammarAST ast, RewriteIteratorDecl decl) {
+	public String rootName;
+	public RewriteTokenRef(OutputModelFactory factory, GrammarAST ast,
+						   String rootName, RewriteIteratorDecl decl)
+	{
 		super(factory, ast);
+		this.rootName = rootName;
 		this.decl = decl;
 	}
 }

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

@@ -34,7 +34,9 @@ import org.antlr.v4.codegen.model.decl.RewriteIteratorDecl;
 import org.antlr.v4.tool.GrammarAST;
 
 public class RewriteTokenRefIsRoot extends RewriteTokenRef {
-	public RewriteTokenRefIsRoot(OutputModelFactory factory, GrammarAST ast, RewriteIteratorDecl decl) {
+	public RewriteTokenRefIsRoot(OutputModelFactory factory, GrammarAST ast,
-		super(factory, ast, decl);
+								 String rootName, RewriteIteratorDecl decl)
+	{
+		super(factory, ast, rootName, decl);
 	}
 }

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

@@ -0,0 +1,47 @@
+/*
+ [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.decl.*;
+import org.antlr.v4.tool.GrammarAST;
+
+import java.util.*;
+
+public class RewriteTreeClosure extends CodeBlock {
+	public List<Decl> iteratorDecls = new ArrayList<Decl>();
+
+	public RewriteTreeClosure(OutputModelFactory factory, GrammarAST ast,
+							  int treeLevel, int codeBlockLevel)
+	{
+		super(factory, treeLevel, codeBlockLevel);
+		this.ast = ast;
+	}
+}

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

@@ -27,15 +27,19 @@
  THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
  */
 
-package org.antlr.v4.codegen.model.decl;
+package org.antlr.v4.codegen.model.ast;
 
 import org.antlr.v4.codegen.OutputModelFactory;
+import org.antlr.v4.codegen.model.decl.*;
+import org.antlr.v4.tool.GrammarAST;
 
-public class NestedDecl extends Decl {
+import java.util.*;
-	public int level;
 
-	public NestedDecl(OutputModelFactory factory, String name, int level) {
+public class RewriteTreeOptional extends CodeBlock {
-		super(factory, name);
+	public List<Decl> iteratorDecls = new ArrayList<Decl>();
-		this.level = level;
+
+	public RewriteTreeOptional(OutputModelFactory factory, GrammarAST ast) {
+		super(factory);
+		this.ast = ast;
 	}
 }

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

@@ -34,10 +34,14 @@ import org.antlr.v4.tool.GrammarAST;
 
 /** ^(A B C) */
 public class RewriteTreeStructure extends TreeRewrite {
+	public int treeLevel;
+
 	public RewriteTreeStructure(OutputModelFactory factory,
 								GrammarAST ast,
-								int level)
+								int treeLevel,
+								int codeBlockLevel)
 	{
-		super(factory, level);
+		super(factory, treeLevel, codeBlockLevel);
+		this.treeLevel = treeLevel;
 	}
 }

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

@@ -34,13 +34,11 @@ import org.antlr.v4.codegen.model.Rewrite;
 import org.antlr.v4.codegen.model.decl.RootDecl;
 
 public class TreeRewrite extends Rewrite {
-	public int level;
 	public RootDecl rootDecl;
 
-	public TreeRewrite(OutputModelFactory factory, int level) {
+	public TreeRewrite(OutputModelFactory factory, int treeLevel, int codeBlockLevel) {
-		super(factory);
+		super(factory, treeLevel, codeBlockLevel);
-		this.level = level;
+		rootDecl = new RootDecl(factory);
-		rootDecl = new RootDecl(factory, level); // track here too for easy access in template
 		addLocalDecl(rootDecl);
 	}
 }

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

@@ -36,6 +36,9 @@ import org.antlr.v4.runtime.misc.OrderedHashSet;
 import java.util.*;
 
 public class CodeBlock extends SrcOp {
+	public int codeBlockLevel;
+	public int treeLevel;
+
 	@ModelElement public OrderedHashSet<Decl> locals;
 	@ModelElement public List<SrcOp> preamble;
 	@ModelElement public List<SrcOp> ops;
@@ -44,6 +47,12 @@ public class CodeBlock extends SrcOp {
 		super(factory);
 	}
 
+	public CodeBlock(OutputModelFactory factory, int treeLevel, int codeBlockLevel) {
+		super(factory);
+		this.treeLevel = treeLevel;
+		this.codeBlockLevel = codeBlockLevel;
+	}
+
 	/** Add local var decl */
 	public void addLocalDecl(Decl d) {
 		if ( locals==null ) locals = new OrderedHashSet<Decl>();
@@ -60,4 +69,9 @@ public class CodeBlock extends SrcOp {
 		if ( ops==null ) ops = new ArrayList<SrcOp>();
 		ops.add(op);
 	}
+
+	public void addOps(List<SrcOp> ops) {
+		if ( this.ops==null ) this.ops = new ArrayList<SrcOp>();
+		this.ops.addAll(ops);
+	}
 }

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

@@ -32,10 +32,13 @@ package org.antlr.v4.codegen.model.decl;
 import org.antlr.v4.codegen.OutputModelFactory;
 import org.antlr.v4.tool.GrammarAST;
 
-public class RewriteIteratorDecl extends NestedDecl {
+public class RewriteIteratorDecl extends Decl {
 	public String listName;
-	public RewriteIteratorDecl(OutputModelFactory factory, GrammarAST elem, int level) {
+	public RewriteIteratorDecl(OutputModelFactory factory,
-		super(factory, factory.getGenerator().target.getRewriteIteratorName(elem, level), level);
+							   GrammarAST elem)
+	{
+		super(factory, factory.getGenerator().target
+			  				.getRewriteIteratorName(elem, factory.getCodeBlockLevel()));
 		listName = factory.getGenerator().target.getElementListName(elem);
 	}
 }

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

@@ -31,8 +31,9 @@ package org.antlr.v4.codegen.model.decl;
 
 import org.antlr.v4.codegen.OutputModelFactory;
 
-public class RootDecl extends NestedDecl {
+public class RootDecl extends Decl {
-	public RootDecl(OutputModelFactory factory, int level) {
+	public RootDecl(OutputModelFactory factory) {
-		super(factory, factory.getGenerator().target.getRootName(level), level);
+		super(factory,
+			  factory.getGenerator().target.getRootName(factory.getCodeBlockLevel()));
 	}
 }

tool/src/org/antlr/v4/tool/Rule.java

@@ -159,7 +159,6 @@ public class Rule implements AttributeResolver {
         return refs;
     }
 
-	// TODO: called frequently; make it more efficient
     public MultiMap<String, LabelElementPair> getLabelDefs() {
         MultiMap<String, LabelElementPair> defs =
             new MultiMap<String, LabelElementPair>();