From d8a9207041fc12f358f96369a6afd27a8cdb3013 Mon Sep 17 00:00:00 2001
From: Terence Parr <parrt@antlr.org>
Date: Sun, 22 Jul 2012 20:07:46 -0700
Subject: [PATCH] improve combine common parents, return a or b in merge array
 of merged is a or b; new unit test.

---
 .../v4/runtime/atn/PredictionContext.java     | 25 +++++++++----
 .../org/antlr/v4/test/TestGraphNodes.java     | 37 ++++++++++++++++---
 2 files changed, 48 insertions(+), 14 deletions(-)

diff --git a/runtime/Java/src/org/antlr/v4/runtime/atn/PredictionContext.java b/runtime/Java/src/org/antlr/v4/runtime/atn/PredictionContext.java
index 08ae11c7e..269e96746 100644
--- a/runtime/Java/src/org/antlr/v4/runtime/atn/PredictionContext.java
+++ b/runtime/Java/src/org/antlr/v4/runtime/atn/PredictionContext.java
@@ -8,6 +8,7 @@ import java.util.ArrayDeque;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Deque;
+import java.util.HashMap;
 import java.util.IdentityHashMap;
 import java.util.Iterator;
 import java.util.List;
@@ -203,7 +204,7 @@ public abstract class PredictionContext implements Iterable<SingletonPredictionC
 	public static PredictionContext merge(PredictionContext a, PredictionContext b,
 										  boolean rootIsWildcard)
 	{
-		if ( a.equals(b) ) return a; // share same graph if both same
+		if ( (a==null&&b==null) || a.equals(b) ) return a; // share same graph if both same
 
 		if ( a instanceof SingletonPredictionContext && b instanceof SingletonPredictionContext) {
 			return mergeSingletons((SingletonPredictionContext)a,
@@ -401,7 +402,10 @@ public abstract class PredictionContext implements Iterable<SingletonPredictionC
 		ArrayPredictionContext M =
 			new ArrayPredictionContext(mergedParents, mergedInvokingStates);
 
-		// TODO: if we created same array as a or b, return that instead
+		// if we created same array as a or b, return that instead
+		// TODO: track whether this is possible above during merge sort for speed
+		if ( M.equals(a) ) return a;
+		if ( M.equals(b) ) return b;
 
 		combineCommonParents(mergedParents);
 
@@ -410,14 +414,19 @@ public abstract class PredictionContext implements Iterable<SingletonPredictionC
 
 	/** make pass over all M parents; merge any equals() ones */
 	protected static void combineCommonParents(PredictionContext[] parents) {
-		// compare pairwise
-		for (int i = 0; i < parents.length; i++) {
-			for (int j = i+1; j < parents.length; j++) {
-				if ( parents[i].equals(parents[j]) ) {
-					parents[j] = parents[i]; // use left one if same
-				}
+		Map<PredictionContext, PredictionContext> uniqueParents =
+			new HashMap<PredictionContext, PredictionContext>();
+
+		for (int p = 0; p < parents.length; p++) {
+			PredictionContext parent = parents[p];
+			if ( !uniqueParents.containsKey(parent) ) { // don't replace
+				uniqueParents.put(parent, parent);
 			}
 		}
+
+		for (int p = 0; p < parents.length; p++) {
+			parents[p] = uniqueParents.get(parents[p]);
+		}
 	}
 
 	public static String toDotString(PredictionContext context) {
diff --git a/tool/test/org/antlr/v4/test/TestGraphNodes.java b/tool/test/org/antlr/v4/test/TestGraphNodes.java
index 597714e55..0015b95f3 100644
--- a/tool/test/org/antlr/v4/test/TestGraphNodes.java
+++ b/tool/test/org/antlr/v4/test/TestGraphNodes.java
@@ -491,10 +491,10 @@ public class TestGraphNodes extends TestCase {
 		String expecting =
 			"digraph G {\n" +
 			"rankdir=LR;\n" +
-			"  s6 [shape=box, label=\"[1, 2]\"];\n" +
+			"  s3 [shape=box, label=\"[1, 2]\"];\n" +
 			"  s0 [label=\"$\"];\n" +
-			"  s6->s0 [label=\"parent[0]\"];\n" +
-			"  s6->s0 [label=\"parent[1]\"];\n" +
+			"  s3->s0 [label=\"parent[0]\"];\n" +
+			"  s3->s0 [label=\"parent[1]\"];\n" +
 			"}\n";
 		assertEquals(expecting, PredictionContext.toDotString(r));
 	}
@@ -507,10 +507,10 @@ public class TestGraphNodes extends TestCase {
 		String expecting =
 			"digraph G {\n" +
 			"rankdir=LR;\n" +
-			"  s6 [shape=box, label=\"[1, 2]\"];\n" +
+			"  s3 [shape=box, label=\"[1, 2]\"];\n" +
 			"  s0 [label=\"$\"];\n" +
-			"  s6->s0 [label=\"parent[0]\"];\n" +
-			"  s6->s0 [label=\"parent[1]\"];\n" +
+			"  s3->s0 [label=\"parent[0]\"];\n" +
+			"  s3->s0 [label=\"parent[1]\"];\n" +
 			"}\n";
 		assertEquals(expecting, PredictionContext.toDotString(r));
 	}
@@ -689,6 +689,31 @@ public class TestGraphNodes extends TestCase {
 		assertEquals(expecting, PredictionContext.toDotString(r));
 	}
 
+	@Test public void test_Aaubu_Acudu() { // au,bu + cu,du -> [a,b,c,d]->[u,u,u,u]
+		SingletonPredictionContext a = createSingleton(u(), 1);
+		SingletonPredictionContext b = createSingleton(u(), 2);
+		SingletonPredictionContext c = createSingleton(u(), 3);
+		SingletonPredictionContext d = createSingleton(u(), 4);
+		ArrayPredictionContext A1 = array(a, b);
+		ArrayPredictionContext A2 = array(c, d);
+		PredictionContext r = PredictionContext.merge(A1, A2, rootIsWildcard());
+		System.out.println(PredictionContext.toDotString(r));
+		String expecting =
+			"digraph G {\n" +
+			"rankdir=LR;\n" +
+			"  s11 [shape=box, label=\"[1, 2, 3, 4]\"];\n" +
+			"  s1 [label=\"6\"];\n" +
+			"  s0 [label=\"$\"];\n" +
+			"  s11->s1 [label=\"parent[0]\"];\n" +
+			"  s11->s1 [label=\"parent[1]\"];\n" +
+			"  s11->s1 [label=\"parent[2]\"];\n" +
+			"  s11->s1 [label=\"parent[3]\"];\n" +
+			"  s1->s0;\n" +
+			"}\n";
+		assertEquals(expecting, PredictionContext.toDotString(r));
+	}
+
+
 	// ------------ SUPPORT -------------------------
 
 	protected SingletonPredictionContext a() {