From 580ef3bb4ed523b0027d356344065041bcafa2ce Mon Sep 17 00:00:00 2001
From: Areek Zillur <areek@apache.org>
Date: Thu, 27 Mar 2014 21:00:40 +0000
Subject: [PATCH] LUCENE-5555: Fix SortedInputIterator to correctly
 encode/decode contexts in presence of payload

git-svn-id: https://svn.apache.org/repos/asf/lucene/dev/trunk@1582473 13f79535-47bb-0310-9956-ffa450edef68
---
 lucene/CHANGES.txt                             |  2 ++
 .../search/suggest/SortedInputIterator.java    | 10 +++++-----
 .../search/suggest/TestInputIterator.java      | 18 ++++++++++++++++++
 3 files changed, 25 insertions(+), 5 deletions(-)

diff --git a/lucene/CHANGES.txt b/lucene/CHANGES.txt
index 85408c5d9e5..a943dc9821f 100644
--- a/lucene/CHANGES.txt
+++ b/lucene/CHANGES.txt
@@ -210,6 +210,8 @@ Bug fixes
 
 * LUCENE-5111: Fix WordDelimiterFilter to return offsets in correct order.  (Robert Muir)
 
+* LUCENE-5555: Fix SortedInputIterator to correctly encode/decode contexts in presence of payload (Areek Zillur)
+
 Test Framework
 
 * LUCENE-5449: Rename _TestUtil and _TestHelper to remove the leading _.
diff --git a/lucene/suggest/src/java/org/apache/lucene/search/suggest/SortedInputIterator.java b/lucene/suggest/src/java/org/apache/lucene/search/suggest/SortedInputIterator.java
index e0c875a3c1a..d340134d59b 100644
--- a/lucene/suggest/src/java/org/apache/lucene/search/suggest/SortedInputIterator.java
+++ b/lucene/suggest/src/java/org/apache/lucene/search/suggest/SortedInputIterator.java
@@ -209,7 +209,7 @@ public class SortedInputIterator implements InputIterator {
     }
   }
   
-  /** encodes an entry (bytes+(payload)+(contexts)+weight) to the provided writer */
+  /** encodes an entry (bytes+(contexts)+(payload)+weight) to the provided writer */
   protected void encode(ByteSequencesWriter writer, ByteArrayDataOutput output, byte[] buffer, BytesRef spare, BytesRef payload, Set<BytesRef> contexts, long weight) throws IOException {
     int requiredLength = spare.length + 8 + ((hasPayloads) ? 2 + payload.length : 0);
     if (hasContexts) {
@@ -223,10 +223,6 @@ public class SortedInputIterator implements InputIterator {
     }
     output.reset(buffer);
     output.writeBytes(spare.bytes, spare.offset, spare.length);
-    if (hasPayloads) {
-      output.writeBytes(payload.bytes, payload.offset, payload.length);
-      output.writeShort((short) payload.length);
-    }
     if (hasContexts) {
       for (BytesRef ctx : contexts) {
         output.writeBytes(ctx.bytes, ctx.offset, ctx.length);
@@ -234,6 +230,10 @@ public class SortedInputIterator implements InputIterator {
       }
       output.writeShort((short) contexts.size());
     }
+    if (hasPayloads) {
+      output.writeBytes(payload.bytes, payload.offset, payload.length);
+      output.writeShort((short) payload.length);
+    }
     output.writeLong(weight);
     writer.write(buffer, 0, output.getPosition());
   }
diff --git a/lucene/suggest/src/test/org/apache/lucene/search/suggest/TestInputIterator.java b/lucene/suggest/src/test/org/apache/lucene/search/suggest/TestInputIterator.java
index 2b7cb67d186..3aff3c21d36 100644
--- a/lucene/suggest/src/test/org/apache/lucene/search/suggest/TestInputIterator.java
+++ b/lucene/suggest/src/test/org/apache/lucene/search/suggest/TestInputIterator.java
@@ -48,9 +48,11 @@ public class TestInputIterator extends LuceneTestCase {
     TreeMap<BytesRef, SimpleEntry<Long, BytesRef>> sorted = new TreeMap<>(comparator);
     TreeMap<BytesRef, Long> sortedWithoutPayload = new TreeMap<>(comparator);
     TreeMap<BytesRef, SimpleEntry<Long, Set<BytesRef>>> sortedWithContext = new TreeMap<>(comparator);
+    TreeMap<BytesRef, SimpleEntry<Long, SimpleEntry<BytesRef, Set<BytesRef>>>> sortedWithPayloadAndContext = new TreeMap<>(comparator);
     Input[] unsorted = new Input[num];
     Input[] unsortedWithoutPayload = new Input[num];
     Input[] unsortedWithContexts = new Input[num];
+    Input[] unsortedWithPayloadAndContext = new Input[num];
     Set<BytesRef> ctxs;
     for (int i = 0; i < num; i++) {
       BytesRef key;
@@ -67,9 +69,11 @@ public class TestInputIterator extends LuceneTestCase {
       sortedWithoutPayload.put(key, value);
       sorted.put(key, new SimpleEntry<>(value, payload));
       sortedWithContext.put(key, new SimpleEntry<>(value, ctxs));
+      sortedWithPayloadAndContext.put(key, new SimpleEntry<>(value, new SimpleEntry<>(payload, ctxs)));
       unsorted[i] = new Input(key, value, payload);
       unsortedWithoutPayload[i] = new Input(key, value);
       unsortedWithContexts[i] = new Input(key, value, ctxs);
+      unsortedWithPayloadAndContext[i] = new Input(key, value, payload, ctxs);
     }
     
     // test the sorted iterator wrapper with payloads
@@ -96,6 +100,20 @@ public class TestInputIterator extends LuceneTestCase {
     }
     assertNull(wrapper.next());
     
+    // test the sorted iterator wrapper with contexts and payload
+    wrapper = new SortedInputIterator(new InputArrayIterator(unsortedWithPayloadAndContext), comparator);
+    Iterator<Map.Entry<BytesRef, SimpleEntry<Long, SimpleEntry<BytesRef, Set<BytesRef>>>>> expectedPayloadContextEntries = sortedWithPayloadAndContext.entrySet().iterator();
+    while (expectedPayloadContextEntries.hasNext()) {
+      Map.Entry<BytesRef, SimpleEntry<Long, SimpleEntry<BytesRef, Set<BytesRef>>>> entry = expectedPayloadContextEntries.next();
+      assertEquals(entry.getKey(), wrapper.next());
+      assertEquals(entry.getValue().getKey().longValue(), wrapper.weight());
+      Set<BytesRef> actualCtxs = entry.getValue().getValue().getValue();
+      assertEquals(actualCtxs, wrapper.contexts());
+      BytesRef actualPayload = entry.getValue().getValue().getKey();
+      assertEquals(actualPayload, wrapper.payload());
+    }
+    assertNull(wrapper.next());
+    
     // test the unsorted iterator wrapper with payloads
     wrapper = new UnsortedInputIterator(new InputArrayIterator(unsorted));
     TreeMap<BytesRef, SimpleEntry<Long, BytesRef>> actual = new TreeMap<>();