From f9a63b88100d862ec2c91f2b4a3acb755550538f Mon Sep 17 00:00:00 2001
From: Sam Harwell <sam@tunnelvisionlabs.com>
Date: Mon, 15 Oct 2012 14:57:06 -0500
Subject: [PATCH] Fix bug in PredictionContext.arrayMerge

---
 .../v4/runtime/atn/PredictionContext.java     | 24 +++++++------------
 .../org/antlr/v4/test/TestGraphNodes.java     | 20 ++++++++++++++++
 2 files changed, 28 insertions(+), 16 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 333a54b28..73712e5fc 100644
--- a/runtime/Java/src/org/antlr/v4/runtime/atn/PredictionContext.java
+++ b/runtime/Java/src/org/antlr/v4/runtime/atn/PredictionContext.java
@@ -331,23 +331,15 @@ public abstract class PredictionContext implements Iterable<SingletonPredictionC
 
 		// trim merged if we combined a few that had same stack tops
 		if ( k < mergedParents.length ) { // write index < last position; trim
-			int lastSlot = mergedParents.length - 1;
-			int p = lastSlot; // walk backwards from last index until we find non-null parent
-			while ( p>=0 && mergedParents[p]==null ) { p--; }
-			// p is now last non-null index
-			assert p>=0; // could only happen to be <0 if two arrays with $
-			if ( p < lastSlot ) {
-				int n = p+1; // how many slots we really used in merge
-				if ( n == 1 ) { // for just one merged element, return singleton top
-					PredictionContext a_ =
-						SingletonPredictionContext.create(mergedParents[0],
-														  mergedInvokingStates[0]);
-					if ( mergeCache!=null ) mergeCache.put(a,b,a_);
-					return a_;
-				}
-				mergedParents = Arrays.copyOf(mergedParents, n);
-				mergedInvokingStates = Arrays.copyOf(mergedInvokingStates, n);
+			if ( k == 1 ) { // for just one merged element, return singleton top
+				PredictionContext a_ =
+					SingletonPredictionContext.create(mergedParents[0],
+													  mergedInvokingStates[0]);
+				if ( mergeCache!=null ) mergeCache.put(a,b,a_);
+				return a_;
 			}
+			mergedParents = Arrays.copyOf(mergedParents, k);
+			mergedInvokingStates = Arrays.copyOf(mergedInvokingStates, k);
 		}
 
 		PredictionContext M =
diff --git a/tool/test/org/antlr/v4/test/TestGraphNodes.java b/tool/test/org/antlr/v4/test/TestGraphNodes.java
index c0ec1921a..5354ed91d 100644
--- a/tool/test/org/antlr/v4/test/TestGraphNodes.java
+++ b/tool/test/org/antlr/v4/test/TestGraphNodes.java
@@ -161,6 +161,26 @@ public class TestGraphNodes extends TestCase {
 		assertEquals(expecting, toDOTString(r, rootIsWildcard()));
 	}
 
+	@Test public void test_aa$_a$_$_fullCtx() {
+		PredictionContext empty = PredictionContext.EMPTY;
+		PredictionContext child1 = createSingleton(empty, 8);
+		PredictionContext right = PredictionContext.merge(empty, child1, false, null);
+		PredictionContext left = createSingleton(right, 8);
+		PredictionContext merged = PredictionContext.merge(left, right, false, null);
+		String actual = toDOTString(merged, false);
+		System.out.println(actual);
+		String expecting =
+			"digraph G {\n" +
+			"rankdir=LR;\n" +
+			"  s0[shape=record, label=\"<p0>|<p1>$\"];\n" +
+			"  s1[shape=record, label=\"<p0>|<p1>$\"];\n" +
+			"  s2[label=\"$\"];\n" +
+			"  s0:p0->s1[label=\"8\"];\n" +
+			"  s1:p0->s2[label=\"8\"];\n" +
+			"}\n";
+		assertEquals(expecting, actual);
+	}
+
 	@Test public void test_ax$_a$_fullctx() {
 		PredictionContext x = x();
 		PredictionContext a1 = createSingleton(x, 1);