Minor cleanups to BytesReferenceStreamInput (#62302)

Followup to #61681:

- reuse the current iterator in `reset()` if possible
- simply some integer-overflow-avoidance in `skip()`
- clarify some comments
- address some IntelliJ warnings
This commit is contained in:
David Turner 2020-09-14 17:02:07 +01:00
parent e4275f3749
commit 9acd2fd1fd
1 changed files with 21 additions and 13 deletions

View File

@ -72,7 +72,7 @@ public abstract class AbstractBytesReference implements BytesReference {
return new BytesRefIterator() { return new BytesRefIterator() {
BytesRef ref = length() == 0 ? null : toBytesRef(); BytesRef ref = length() == 0 ? null : toBytesRef();
@Override @Override
public BytesRef next() throws IOException { public BytesRef next() {
BytesRef r = ref; BytesRef r = ref;
ref = null; // only return it once... ref = null; // only return it once...
return r; return r;
@ -189,7 +189,7 @@ public abstract class AbstractBytesReference implements BytesReference {
/** /**
* A StreamInput that reads off a {@link BytesRefIterator}. This is used to provide * A StreamInput that reads off a {@link BytesRefIterator}. This is used to provide
* generic stream access to {@link BytesReference} instances without materializing the * generic stream access to {@link BytesReference} instances without materializing the
* underlying bytes reference. * underlying bytes.
*/ */
private final class BytesReferenceStreamInput extends StreamInput { private final class BytesReferenceStreamInput extends StreamInput {
@ -239,7 +239,8 @@ public abstract class AbstractBytesReference implements BytesReference {
throw new IndexOutOfBoundsException( throw new IndexOutOfBoundsException(
"Cannot read " + len + " bytes from stream with length " + length + " at offset " + offset); "Cannot read " + len + " bytes from stream with length " + length + " at offset " + offset);
} }
read(b, bOffset, len); final int bytesRead = read(b, bOffset, len);
assert bytesRead == len : bytesRead + " vs " + len;
} }
@Override @Override
@ -293,8 +294,11 @@ public abstract class AbstractBytesReference implements BytesReference {
@Override @Override
public long skip(long n) throws IOException { public long skip(long n) throws IOException {
final int skip = (int) Math.min(Integer.MAX_VALUE, n); if (n <= 0L) {
final int numBytesSkipped = Math.min(skip, length() - offset()); return 0L;
}
assert offset() <= length() : offset() + " vs " + length();
final int numBytesSkipped = (int)Math.min(n, length() - offset()); // definitely >= 0 and <= Integer.MAX_VALUE so casting is ok
int remaining = numBytesSkipped; int remaining = numBytesSkipped;
while (remaining > 0) { while (remaining > 0) {
maybeNextSlice(); maybeNextSlice();
@ -308,11 +312,16 @@ public abstract class AbstractBytesReference implements BytesReference {
@Override @Override
public void reset() throws IOException { public void reset() throws IOException {
if (sliceStartOffset <= mark) {
sliceIndex = mark - sliceStartOffset;
} else {
iterator = iterator(); iterator = iterator();
slice = iterator.next(); slice = iterator.next();
sliceStartOffset = 0; sliceStartOffset = 0;
sliceIndex = 0; sliceIndex = 0;
skip(mark); final long skipped = skip(mark);
assert skipped == mark : skipped + " vs " + mark;
}
} }
@Override @Override
@ -322,8 +331,7 @@ public abstract class AbstractBytesReference implements BytesReference {
@Override @Override
public void mark(int readLimit) { public void mark(int readLimit) {
// readLimit is optional it only guarantees that the stream remembers data upto this limit but it can remember more // We ignore readLimit since the data is all in-memory and therefore we can reset the mark no matter how far we advance.
// which we do in our case
this.mark = offset(); this.mark = offset();
} }
} }