From 50192782046779e9bb9c6d7eb55c1a86b47ba4ba Mon Sep 17 00:00:00 2001
From: Sam Harwell <sam@tunnelvisionlabs.com>
Date: Tue, 31 Jul 2012 13:32:55 -0500
Subject: [PATCH] Add ATN optimization to collapse multiple alternatives to a
 single SetTransition where possible. Currently disabled for parser grammars
 since the code generator doesn't support set transitions.

---
 .../org/antlr/v4/runtime/atn/ATNState.java    |   4 +
 .../org/antlr/v4/automata/ATNOptimizer.java   | 100 ++++++++++++++++++
 2 files changed, 104 insertions(+)

diff --git a/runtime/Java/src/org/antlr/v4/runtime/atn/ATNState.java b/runtime/Java/src/org/antlr/v4/runtime/atn/ATNState.java
index 54eabb2d8..b13c12e20 100644
--- a/runtime/Java/src/org/antlr/v4/runtime/atn/ATNState.java
+++ b/runtime/Java/src/org/antlr/v4/runtime/atn/ATNState.java
@@ -137,6 +137,10 @@ public class ATNState {
 		transitions.set(i, e);
 	}
 
+	public Transition removeTransition(int index) {
+		return transitions.remove(index);
+	}
+
 	public int getStateType() {
 		return serializationTypes.get(this.getClass());
 	}
diff --git a/tool/src/org/antlr/v4/automata/ATNOptimizer.java b/tool/src/org/antlr/v4/automata/ATNOptimizer.java
index 6f453c3a4..220946300 100644
--- a/tool/src/org/antlr/v4/automata/ATNOptimizer.java
+++ b/tool/src/org/antlr/v4/automata/ATNOptimizer.java
@@ -32,8 +32,19 @@ package org.antlr.v4.automata;
 
 import org.antlr.v4.runtime.atn.ATN;
 import org.antlr.v4.runtime.atn.ATNState;
+import org.antlr.v4.runtime.atn.AtomTransition;
+import org.antlr.v4.runtime.atn.BlockEndState;
+import org.antlr.v4.runtime.atn.DecisionState;
+import org.antlr.v4.runtime.atn.EpsilonTransition;
+import org.antlr.v4.runtime.atn.NotSetTransition;
+import org.antlr.v4.runtime.atn.RangeTransition;
+import org.antlr.v4.runtime.atn.SetTransition;
+import org.antlr.v4.runtime.atn.Transition;
+import org.antlr.v4.runtime.misc.Interval;
+import org.antlr.v4.runtime.misc.IntervalSet;
 import org.antlr.v4.runtime.misc.NotNull;
 import org.antlr.v4.tool.Grammar;
+import org.antlr.v4.tool.Rule;
 
 import java.util.List;
 
@@ -44,9 +55,98 @@ import java.util.List;
 public class ATNOptimizer {
 
 	public static void optimize(@NotNull Grammar g, @NotNull ATN atn) {
+		optimizeSets(g, atn);
 		optimizeStates(atn);
 	}
 
+	private static void optimizeSets(Grammar g, ATN atn) {
+		if (g.isParser()) {
+			// parser codegen doesn't currently support SetTransition
+			return;
+		}
+
+		int removedStates = 0;
+		List<DecisionState> decisions = atn.decisionToState;
+		for (DecisionState decision : decisions) {
+			if (decision.ruleIndex >= 0) {
+				Rule rule = g.getRule(decision.ruleIndex);
+				if (Character.isLowerCase(rule.name.charAt(0))) {
+					// parser codegen doesn't currently support SetTransition
+					continue;
+				}
+			}
+
+			IntervalSet setTransitions = new IntervalSet();
+			for (int i = 0; i < decision.getNumberOfTransitions(); i++) {
+				Transition epsTransition = decision.transition(i);
+				if (!(epsTransition instanceof EpsilonTransition)) {
+					continue;
+				}
+
+				if (epsTransition.target.getNumberOfTransitions() != 1) {
+					continue;
+				}
+
+				Transition transition = epsTransition.target.transition(0);
+				if (!(transition.target instanceof BlockEndState)) {
+					continue;
+				}
+
+				if (transition instanceof NotSetTransition) {
+					// TODO: not yet implemented
+					continue;
+				}
+
+				if (transition instanceof AtomTransition
+					|| transition instanceof RangeTransition
+					|| transition instanceof SetTransition)
+				{
+					setTransitions.add(i);
+				}
+			}
+
+			// due to min alt resolution policies, can only collapse sequential alts
+			for (int i = setTransitions.getIntervals().size() - 1; i >= 0; i--) {
+				Interval interval = setTransitions.getIntervals().get(i);
+				if (interval.length() <= 1) {
+					continue;
+				}
+
+				ATNState blockEndState = decision.transition(interval.a).target.transition(0).target;
+				IntervalSet matchSet = new IntervalSet();
+				for (int j = interval.a; j <= interval.b; j++) {
+					Transition matchTransition = decision.transition(j).target.transition(0);
+					if (matchTransition instanceof NotSetTransition) {
+						throw new UnsupportedOperationException("Not yet implemented.");
+					} else {
+						matchSet.addAll(matchTransition.label());
+					}
+				}
+
+				Transition newTransition;
+				if (matchSet.getIntervals().size() == 1) {
+					if (matchSet.size() == 1) {
+						newTransition = new AtomTransition(blockEndState, matchSet.getMinElement());
+					} else {
+						Interval matchInterval = matchSet.getIntervals().get(0);
+						newTransition = new RangeTransition(blockEndState, matchInterval.a, matchInterval.b);
+					}
+				} else {
+					newTransition = new SetTransition(blockEndState, matchSet);
+				}
+
+				decision.transition(interval.a).target.setTransition(0, newTransition);
+				for (int j = interval.a + 1; j <= interval.b; j++) {
+					Transition removed = decision.removeTransition(interval.a + 1);
+					atn.removeState(removed.target);
+					removedStates++;
+				}
+			}
+		}
+
+		System.out.println("ATN optimizer removed " + removedStates + " states by collapsing sets.");
+	}
+
 	private static void optimizeStates(ATN atn) {
 		List<ATNState> states = atn.states;