Improve efficiency of BoundedBreakIteratorScanner fragmentation algorithm (the 2nd)

server/src/main/java/org/elasticsearch/lucene/search/uhighlight/BoundedBreakIteratorScanner.java

@@ -89,16 +89,29 @@ public int preceding(int offset) {
             innerStart = innerEnd;
             innerEnd = windowEnd;
         } else {
-            windowStart = innerStart = mainBreak.preceding(offset);
-            windowEnd = innerEnd = mainBreak.following(offset - 1);
-            // expand to next break until we reach maxLen
-            while (innerEnd - innerStart < maxLen) {
-                int newEnd = mainBreak.following(innerEnd);
-                if (newEnd == DONE || (newEnd - innerStart) > maxLen) {
-                    break;
-                }
-                windowEnd = innerEnd = newEnd;
+            innerStart = Math.max(mainBreak.preceding(offset), 0);
+
+            final long targetEndOffset = (long) offset + Math.max(0, maxLen - (offset - innerStart));
+            final int textEndIndex = getText().getEndIndex();
+
+            if (targetEndOffset + 1 > textEndIndex) {
+                innerEnd = textEndIndex;
+            } else {
+                innerEnd = mainBreak.preceding((int) targetEndOffset + 1);
+            }
+
+            assert innerEnd != DONE && innerEnd >= innerStart;
+
+            // in case no break was found up to maxLen, find one afterwards.
+            if (innerStart == innerEnd) {
+                innerEnd = mainBreak.following((int) targetEndOffset);
+                assert innerEnd - innerStart > maxLen;
+            } else {
+                assert innerEnd - innerStart <= maxLen;
             }
+
+            windowStart = innerStart;
+            windowEnd = innerEnd;
         }
 
         if (innerEnd - innerStart > maxLen) {

server/src/test/java/org/elasticsearch/lucene/search/uhighlight/BoundedBreakIteratorScannerTests.java

@@ -117,4 +117,32 @@ public void testBoundedSentence() {
             testRandomAsciiTextCase(BoundedBreakIteratorScanner.getSentence(Locale.ROOT, maxLen), maxLen);
         }
     }
+
+    public void testTextThatEndsBeforeMaxLen() {
+        BreakIterator bi = BoundedBreakIteratorScanner.getSentence(Locale.ROOT, 1000);
+
+        final String text = "This is the first test sentence. Here is the second one.";
+
+        int offset = text.indexOf("first");
+        bi.setText(text);
+        assertEquals(0, bi.preceding(offset));
+        assertEquals(text.length(), bi.following(offset - 1));
+
+        offset = text.indexOf("second");
+        bi.setText(text);
+        assertEquals(33, bi.preceding(offset));
+        assertEquals(text.length(), bi.following(offset - 1));
+    }
+
+    public void testFragmentSizeThatIsTooBig() {
+        final int fragmentSize = Integer.MAX_VALUE;
+        BreakIterator bi = BoundedBreakIteratorScanner.getSentence(Locale.ROOT, fragmentSize);
+
+        final String text = "Any sentence";
+        final int offset = 0; // find at beggining of text
+
+        bi.setText(text);
+        assertEquals(0, bi.preceding(offset));
+        assertEquals(text.length(), bi.following(offset - 1));
+    }
 }