New API, BufferedSource.indexOf(ByteString, fromIndex, toIndex) (#1626)

* New API, BufferedSource.indexOf(ByteString, fromIndex, toIndex)

This is surprisingly interesting. To minimize unnecessary reads for
toIndex it is necessary to check whether a prefix of the query matches
a suffix of the currently-loaded data.

This read-avoidance is useful in practice. When doing HTTP multipart
decoding the caller may scan for a boundary separator with a bounded
range, and we don't want to block reading when doing so won't impact
the result of the call.

* apiDump

* Check both maximum and minimum prefix sizes

---------

Co-authored-by: Jesse Wilson <[email protected]>

Author: swankjesse

okio/api/okio.api

@@ -93,6 +93,7 @@ public final class okio/Buffer : java/lang/Cloneable, java/nio/channels/ByteChan
 	public fun indexOf (BJJ)J
 	public fun indexOf (Lokio/ByteString;)J
 	public fun indexOf (Lokio/ByteString;J)J
+	public fun indexOf (Lokio/ByteString;JJ)J
 	public fun indexOfElement (Lokio/ByteString;)J
 	public fun indexOfElement (Lokio/ByteString;J)J
 	public fun inputStream ()Ljava/io/InputStream;
@@ -250,6 +251,7 @@ public abstract interface class okio/BufferedSource : java/nio/channels/Readable
 	public abstract fun indexOf (BJJ)J
 	public abstract fun indexOf (Lokio/ByteString;)J
 	public abstract fun indexOf (Lokio/ByteString;J)J
+	public abstract fun indexOf (Lokio/ByteString;JJ)J
 	public abstract fun indexOfElement (Lokio/ByteString;)J
 	public abstract fun indexOfElement (Lokio/ByteString;J)J
 	public abstract fun inputStream ()Ljava/io/InputStream;

okio/src/commonMain/kotlin/okio/Buffer.kt

@@ -118,6 +118,7 @@ expect class Buffer() : BufferedSource, BufferedSink {
   override fun indexOf(b: Byte, fromIndex: Long, toIndex: Long): Long
   override fun indexOf(bytes: ByteString): Long
   override fun indexOf(bytes: ByteString, fromIndex: Long): Long
+  override fun indexOf(bytes: ByteString, fromIndex: Long, toIndex: Long): Long
   override fun indexOfElement(targetBytes: ByteString): Long
   override fun indexOfElement(targetBytes: ByteString, fromIndex: Long): Long
   override fun peek(): BufferedSource

okio/src/commonMain/kotlin/okio/BufferedSource.kt

@@ -514,6 +514,28 @@ expect sealed interface BufferedSource : Source {
    */
   fun indexOf(bytes: ByteString, fromIndex: Long): Long
 
+  /**
+   * Returns the index of the first match for `bytes` in the buffer that is at or after `fromIndex`,
+   * and that is also less than `toIndex`. Returns -1 if `bytes` isn't found in that range. If
+   * `fromIndex == toIndex` then search range is empty and -1 is returned.
+   *
+   * This may attempt to expand the buffer. It won't attempt to expand the buffer if doing so
+   * couldn't change the result.
+   *
+   * ```java
+   * ByteString MOVE = ByteString.encodeUtf8("move");
+   *
+   * Buffer buffer = new Buffer();
+   * buffer.writeUtf8("Don't move! He can't see us if we don't move.");
+   *
+   * assertEquals( 6, buffer.indexOf(MOVE, 0,  40));
+   * assertEquals( 6, buffer.indexOf(MOVE, 0, 100));
+   * assertEquals(-1, buffer.indexOf(MOVE, 7,  40));
+   * assertEquals(40, buffer.indexOf(MOVE, 7, 100));
+   * ```
+   */
+  fun indexOf(bytes: ByteString, fromIndex: Long, toIndex: Long): Long
+
   /** Equivalent to [indexOfElement(targetBytes, 0)][indexOfElement]. */
   fun indexOfElement(targetBytes: ByteString): Long
 

okio/src/commonMain/kotlin/okio/RealBufferedSource.kt

@@ -30,6 +30,7 @@ internal expect class RealBufferedSource(
   override fun indexOf(b: Byte, fromIndex: Long, toIndex: Long): Long
   override fun indexOf(bytes: ByteString): Long
   override fun indexOf(bytes: ByteString, fromIndex: Long): Long
+  override fun indexOf(bytes: ByteString, fromIndex: Long, toIndex: Long): Long
   override fun indexOfElement(targetBytes: ByteString): Long
   override fun indexOfElement(targetBytes: ByteString, fromIndex: Long): Long
   override fun peek(): BufferedSource

okio/src/commonMain/kotlin/okio/internal/Buffer.kt

@@ -1285,10 +1285,19 @@ internal inline fun Buffer.commonIndexOf(b: Byte, fromIndex: Long, toIndex: Long
   }
 }
 
-internal inline fun Buffer.commonIndexOf(bytes: ByteString, fromIndex: Long): Long {
-  var fromIndex = fromIndex
+internal inline fun Buffer.commonIndexOf(
+  bytes: ByteString,
+  fromIndex: Long,
+  toIndex: Long = Long.MAX_VALUE,
+): Long {
   require(bytes.size > 0) { "bytes is empty" }
-  require(fromIndex >= 0L) { "fromIndex < 0: $fromIndex" }
+  require(fromIndex >= 0) { "fromIndex < 0: $fromIndex" }
+  require(fromIndex <= toIndex) { "fromIndex > toIndex: $fromIndex > $toIndex" }
+
+  var fromIndex = fromIndex
+  var toIndex = toIndex
+  if (toIndex > size) toIndex = size
+  if (fromIndex == toIndex) return -1L
 
   seek(fromIndex) { s, offset ->
     var s = s ?: return -1L
@@ -1299,11 +1308,11 @@ internal inline fun Buffer.commonIndexOf(bytes: ByteString, fromIndex: Long): Lo
     val targetByteArray = bytes.internalArray()
     val b0 = targetByteArray[0]
     val bytesSize = bytes.size
-    val resultLimit = size - bytesSize + 1L
+    val resultLimit = minOf(toIndex, size - bytesSize + 1L)
     while (offset < resultLimit) {
       // Scan through the current segment.
       val data = s.data
-      val segmentLimit = okio.minOf(s.limit, s.pos + resultLimit - offset).toInt()
+      val segmentLimit = minOf(s.limit, s.pos + resultLimit - offset).toInt()
       for (pos in (s.pos + fromIndex - offset).toInt() until segmentLimit) {
         if (data[pos] == b0 && rangeEquals(s, pos + 1, targetByteArray, 1, bytesSize)) {
           return pos - s.pos + offset
@@ -1624,7 +1633,7 @@ internal inline fun UnsafeCursor.commonResizeBuffer(newSize: Long): Long {
       val tailSize = tail!!.limit - tail.pos
       if (tailSize <= bytesToSubtract) {
         buffer.head = tail.pop()
-        okio.SegmentPool.recycle(tail)
+        SegmentPool.recycle(tail)
         bytesToSubtract -= tailSize.toLong()
       } else {
         tail.limit -= bytesToSubtract.toInt()

okio/src/commonMain/kotlin/okio/internal/RealBufferedSource.kt

@@ -33,6 +33,7 @@ import okio.Segment
 import okio.Sink
 import okio.buffer
 import okio.checkOffsetAndCount
+import okio.minOf
 
 internal inline fun RealBufferedSource.commonRead(sink: Buffer, byteCount: Long): Long {
   require(byteCount >= 0L) { "byteCount < 0: $byteCount" }
@@ -332,22 +333,67 @@ internal inline fun RealBufferedSource.commonIndexOf(b: Byte, fromIndex: Long, t
   return -1L
 }
 
-internal inline fun RealBufferedSource.commonIndexOf(bytes: ByteString, fromIndex: Long): Long {
+internal inline fun RealBufferedSource.commonIndexOf(
+  bytes: ByteString,
+  fromIndex: Long,
+  toIndex: Long = Long.MAX_VALUE,
+): Long {
   var fromIndex = fromIndex
   check(!closed) { "closed" }
 
   while (true) {
-    val result = buffer.indexOf(bytes, fromIndex)
+    val result = buffer.indexOf(bytes, fromIndex, toIndex)
     if (result != -1L) return result
 
     val lastBufferSize = buffer.size
+    val nextFromIndex = lastBufferSize - bytes.size + 1
+    if (nextFromIndex >= toIndex) return -1L
+
+    if (!matchPossibleByExpandingBuffer(buffer, bytes, fromIndex, toIndex)) return -1L
     if (source.read(buffer, Segment.SIZE.toLong()) == -1L) return -1L
 
     // Keep searching, picking up from where we left off.
-    fromIndex = maxOf(fromIndex, lastBufferSize - bytes.size + 1)
+    fromIndex = maxOf(fromIndex, nextFromIndex)
   }
 }
 
+/**
+ * Returns true if loading more data could result in an `indexOf` match.
+ *
+ * This function's utility is avoiding potentially-slow `read` calls that cannot impact the result
+ * of an `indexOf` call. For example, consider this situation:
+ *
+ * ```
+ * val source = ...
+ * source.indexOf("hello", fromIndex = 0, toIndex = 4)
+ * ```
+ *
+ * If the source's loaded content is the string "shell", it is necessary to load more data because
+ * if the next loaded byte is 'o' then the result will be 1. But if the source's loaded content is
+ * 'look', we know the result is -1 without loading more data.
+ */
+private fun matchPossibleByExpandingBuffer(
+  buffer: Buffer,
+  bytes: ByteString,
+  fromIndex: Long,
+  toIndex: Long,
+): Boolean {
+  // Load new data if the match could come entirely in that new data.
+  if (buffer.size < toIndex) return true
+
+  // Load new data if a prefix of 'bytes' matches a suffix of 'buffer'.
+  val begin = maxOf(1, buffer.size - toIndex + 1).toInt()
+  val limit = minOf(bytes.size, buffer.size - fromIndex + 1).toInt()
+  for (i in limit - 1 downTo begin) {
+    if (buffer.rangeEquals(buffer.size - i, bytes, 0, i)) {
+      return true
+    }
+  }
+
+  // No matter what we load, we won't find a match.
+  return false
+}
+
 internal inline fun RealBufferedSource.commonIndexOfElement(targetBytes: ByteString, fromIndex: Long): Long {
   var fromIndex = fromIndex
   check(!closed) { "closed" }

okio/src/jvmMain/kotlin/okio/Buffer.kt

@@ -485,6 +485,10 @@ actual class Buffer : BufferedSource, BufferedSink, Cloneable, ByteChannel {
   @Throws(IOException::class)
   actual override fun indexOf(bytes: ByteString, fromIndex: Long): Long = commonIndexOf(bytes, fromIndex)
 
+  @Throws(IOException::class)
+  actual override fun indexOf(bytes: ByteString, fromIndex: Long, toIndex: Long): Long =
+    commonIndexOf(bytes, fromIndex, toIndex)
+
   actual override fun indexOfElement(targetBytes: ByteString) = indexOfElement(targetBytes, 0L)
 
   actual override fun indexOfElement(targetBytes: ByteString, fromIndex: Long): Long =

okio/src/jvmMain/kotlin/okio/BufferedSource.kt

@@ -147,6 +147,9 @@ actual sealed interface BufferedSource : Source, ReadableByteChannel {
   @Throws(IOException::class)
   actual fun indexOf(bytes: ByteString, fromIndex: Long): Long
 
+  @Throws(IOException::class)
+  actual fun indexOf(bytes: ByteString, fromIndex: Long, toIndex: Long): Long
+
   @Throws(IOException::class)
   actual fun indexOfElement(targetBytes: ByteString): Long
 

okio/src/jvmMain/kotlin/okio/RealBufferedSource.kt

@@ -127,6 +127,8 @@ internal actual class RealBufferedSource actual constructor(
   actual override fun indexOf(bytes: ByteString): Long = indexOf(bytes, 0L)
   actual override fun indexOf(bytes: ByteString, fromIndex: Long): Long =
     commonIndexOf(bytes, fromIndex)
+  actual override fun indexOf(bytes: ByteString, fromIndex: Long, toIndex: Long): Long =
+    commonIndexOf(bytes, fromIndex, toIndex)
   actual override fun indexOfElement(targetBytes: ByteString): Long =
     indexOfElement(targetBytes, 0L)
   actual override fun indexOfElement(targetBytes: ByteString, fromIndex: Long): Long =