package org.apache.lucene.util.bkd;
import org.apache.lucene.util.BytesRef;
import org.apache.lucene.util.FutureArrays;
public final class HeapPointWriter implements PointWriter {
public final byte[] block;
final int size;
final int packedBytesLength;
final int packedBytesDocIDLength;
private final byte[] scratch;
private int nextWrite;
private boolean closed;
private HeapPointReader.HeapPointValue pointValue;
public HeapPointWriter(int size, int packedBytesLength) {
this.packedBytesDocIDLength = packedBytesLength + Integer.BYTES;
this.packedBytesLength = packedBytesLength;
this.block = new byte[packedBytesDocIDLength * size];
this.size = size;
this.scratch = new byte[packedBytesDocIDLength];
if (size > 0) {
pointValue = new HeapPointReader.HeapPointValue(block, packedBytesLength);
} else {
pointValue = null;
}
}
public PointValue getPackedValueSlice(int index) {
assert index < nextWrite : "nextWrite=" + (nextWrite) + " vs index=" + index;
pointValue.setOffset(index * packedBytesDocIDLength);
return pointValue;
}
@Override
public void append(byte[] packedValue, int docID) {
assert closed == false : "point writer is already closed";
assert packedValue.length == packedBytesLength : "[packedValue] must have length [" + packedBytesLength + "] but was [" + packedValue.length + "]";
assert nextWrite < size : "nextWrite=" + (nextWrite + 1) + " vs size=" + size;
System.arraycopy(packedValue, 0, block, nextWrite * packedBytesDocIDLength, packedBytesLength);
int position = nextWrite * packedBytesDocIDLength + packedBytesLength;
block[position] = (byte) (docID >> 24);
block[++position] = (byte) (docID >> 16);
block[++position] = (byte) (docID >> 8);
block[++position] = (byte) (docID >> 0);
nextWrite++;
}
@Override
public void append(PointValue pointValue) {
assert closed == false : "point writer is already closed";
assert nextWrite < size : "nextWrite=" + (nextWrite + 1) + " vs size=" + size;
BytesRef packedValueDocID = pointValue.packedValueDocIDBytes();
assert packedValueDocID.length == packedBytesDocIDLength : "[packedValue] must have length [" + (packedBytesDocIDLength) + "] but was [" + packedValueDocID.length + "]";
System.arraycopy(packedValueDocID.bytes, packedValueDocID.offset, block, nextWrite * packedBytesDocIDLength, packedBytesDocIDLength);
nextWrite++;
}
public void swap(int i, int j) {
int indexI = i * packedBytesDocIDLength;
int indexJ = j * packedBytesDocIDLength;
System.arraycopy(block, indexI, scratch, 0, packedBytesDocIDLength);
System.arraycopy(block, indexJ, block, indexI, packedBytesDocIDLength);
System.arraycopy(scratch, 0, block, indexJ, packedBytesDocIDLength);
}
public int computeCardinality(int from, int to, int numDataDims, int bytesPerDim, int[] commonPrefixLengths) {
assert packedBytesLength == numDataDims * bytesPerDim;
int leafCardinality = 1;
for (int i = from + 1; i < to; i++) {
for (int dim = 0; dim < numDataDims; dim++) {
final int start = dim * bytesPerDim + commonPrefixLengths[dim];
final int end = dim * bytesPerDim + bytesPerDim;
if (FutureArrays.mismatch(block, i * packedBytesDocIDLength + start, i * packedBytesDocIDLength + end,
block, (i - 1) * packedBytesDocIDLength + start, (i - 1) * packedBytesDocIDLength + end) != -1) {
leafCardinality++;
break;
}
}
}
return leafCardinality;
}
@Override
public long count() {
return nextWrite;
}
@Override
public PointReader getReader(long start, long length) {
assert closed : "point writer is still open and trying to get a reader";
assert start + length <= size: "start=" + start + " length=" + length + " docIDs.length=" + size;
assert start + length <= nextWrite: "start=" + start + " length=" + length + " nextWrite=" + nextWrite;
return new HeapPointReader(block, packedBytesLength, (int) start, Math.toIntExact(start+length));
}
@Override
public void close() {
closed = true;
}
@Override
public void destroy() {
}
@Override
public String toString() {
return "HeapPointWriter(count=" + nextWrite + " size=" + size + ")";
}
}