lucene/core/src/java/org/apache/lucene/util/BytesRefArray.java - lucene-solr - Git at Google

 /*
  * Licensed to the Apache Software Foundation (ASF) under one or more
  * contributor license agreements.  See the NOTICE file distributed with
  * this work for additional information regarding copyright ownership.
  * The ASF licenses this file to You under the Apache License, Version 2.0
  * (the "License"); you may not use this file except in compliance with
  * the License.  You may obtain a copy of the License at
  *
  *     http://www.apache.org/licenses/LICENSE-2.0
  *
  * Unless required by applicable law or agreed to in writing, software
  * distributed under the License is distributed on an "AS IS" BASIS,
  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
 package org.apache.lucene.util;

 import java.util.Arrays;
 import java.util.Comparator;
 import java.util.function.IntBinaryOperator;

 /**
  * A simple append only random-access {@link BytesRef} array that stores full
  * copies of the appended bytes in a {@link ByteBlockPool}.
  *
  *
  * <b>Note: This class is not Thread-Safe!</b>
  *
  * @lucene.internal
  * @lucene.experimental
  */
 public final class BytesRefArray implements SortableBytesRefArray {
   private final ByteBlockPool pool;
   private int[] offsets = new int[1];
   private int lastElement = 0;
   private int currentOffset = 0;
   private final Counter bytesUsed;

   /**
    * Creates a new {@link BytesRefArray} with a counter to track allocated bytes
    */
   public BytesRefArray(Counter bytesUsed) {
     this.pool = new ByteBlockPool(new ByteBlockPool.DirectTrackingAllocator(
         bytesUsed));
     pool.nextBuffer();
     bytesUsed.addAndGet(RamUsageEstimator.NUM_BYTES_ARRAY_HEADER * Integer.BYTES);
     this.bytesUsed = bytesUsed;
   }

   /**
    * Clears this {@link BytesRefArray}
    */
   @Override
   public void clear() {
     lastElement = 0;
     currentOffset = 0;
     // TODO: it's trappy that this does not return storage held by int[] offsets array!
     Arrays.fill(offsets, 0);
     pool.reset(false, true); // no need to 0 fill the buffers we control the allocator
   }

   /**
    * Appends a copy of the given {@link BytesRef} to this {@link BytesRefArray}.
    * @param bytes the bytes to append
    * @return the index of the appended bytes
    */
   @Override
   public int append(BytesRef bytes) {
     if (lastElement >= offsets.length) {
       int oldLen = offsets.length;
       offsets = ArrayUtil.grow(offsets, offsets.length + 1);
       bytesUsed.addAndGet((offsets.length - oldLen) * Integer.BYTES);
     }
     pool.append(bytes);
     offsets[lastElement++] = currentOffset;
     currentOffset += bytes.length;
     return lastElement-1;
   }

   /**
    * Returns the current size of this {@link BytesRefArray}
    * @return the current size of this {@link BytesRefArray}
    */
   @Override
   public int size() {
     return lastElement;
   }

   /**
    * Returns the <i>n'th</i> element of this {@link BytesRefArray}
    * @param spare a spare {@link BytesRef} instance
    * @param index the elements index to retrieve
    * @return the <i>n'th</i> element of this {@link BytesRefArray}
    */
   public BytesRef get(BytesRefBuilder spare, int index) {
     FutureObjects.checkIndex(index, lastElement);
     int offset = offsets[index];
     int length = index == lastElement - 1 ? currentOffset - offset
         : offsets[index + 1] - offset;
     spare.grow(length);
     spare.setLength(length);
     pool.readBytes(offset, spare.bytes(), 0, spare.length());
     return spare.get();
   }

   /** Used only by sort below, to set a {@link BytesRef} with the specified slice, avoiding copying bytes in the common case when the slice
    *  is contained in a single block in the byte block pool. */
   private void setBytesRef(BytesRefBuilder spare, BytesRef result, int index) {
     FutureObjects.checkIndex(index, lastElement);
     int offset = offsets[index];
     int length;
     if (index == lastElement - 1) {
       length = currentOffset - offset;
     } else {
       length = offsets[index + 1] - offset;
     }
     pool.setBytesRef(spare, result, offset, length);
   }


   /**
    * Returns a {@link SortState} representing the order of elements in this array. This is a non-destructive operation.
    */
   public SortState sort(final Comparator<BytesRef> comp, final IntBinaryOperator tieComparator) {
     final int[] orderedEntries = new int[size()];
     for (int i = 0; i < orderedEntries.length; i++) {
       orderedEntries[i] = i;
     }
     new IntroSorter() {
       @Override
       protected void swap(int i, int j) {
         final int o = orderedEntries[i];
         orderedEntries[i] = orderedEntries[j];
         orderedEntries[j] = o;
       }

       @Override
       protected int compare(int i, int j) {
         final int idx1 = orderedEntries[i], idx2 = orderedEntries[j];
         setBytesRef(scratch1, scratchBytes1, idx1);
         setBytesRef(scratch2, scratchBytes2, idx2);
         return compare(idx1, scratchBytes1, idx2, scratchBytes2);
       }

       @Override
       protected void setPivot(int i) {
         pivotIndex = orderedEntries[i];
         setBytesRef(pivotBuilder, pivot, pivotIndex);
       }

       @Override
       protected int comparePivot(int j) {
         final int index = orderedEntries[j];
         setBytesRef(scratch2, scratchBytes2, index);
         return compare(pivotIndex, pivot, index, scratchBytes2);
       }

       private int compare(int i1, BytesRef b1, int i2, BytesRef b2) {
         int res = comp.compare(b1, b2);
         return res == 0 ? tieComparator.applyAsInt(i1, i2) : res;
       }

       private int pivotIndex;
       private final BytesRef pivot = new BytesRef();
       private final BytesRef scratchBytes1 = new BytesRef();
       private final BytesRef scratchBytes2 = new BytesRef();
       private final BytesRefBuilder pivotBuilder = new BytesRefBuilder();
       private final BytesRefBuilder scratch1 = new BytesRefBuilder();
       private final BytesRefBuilder scratch2 = new BytesRefBuilder();
     }.sort(0, size());
     return new SortState(orderedEntries);
   }

   /**
    * sugar for {@link #iterator(Comparator)} with a <code>null</code> comparator
    */
   public BytesRefIterator iterator() {
     return iterator((SortState) null);
   }

   /**
    * <p>
    * Returns a {@link BytesRefIterator} with point in time semantics. The
    * iterator provides access to all so far appended {@link BytesRef} instances.
    * </p>
    * <p>
    * If a non <code>null</code> {@link Comparator} is provided the iterator will
    * iterate the byte values in the order specified by the comparator. Otherwise
    * the order is the same as the values were appended.
    * </p>
    * <p>
    * This is a non-destructive operation.
    * </p>
    */
   @Override
   public BytesRefIterator iterator(final Comparator<BytesRef> comp) {
     return iterator(sort(comp, (i, j) -> 0));
   }

   /**
    * Returns an {@link IndexedBytesRefIterator} with point in time semantics. The iterator provides access to all
    * so far appended {@link BytesRef} instances. If a non-null sortState is specified then the iterator will iterate
    * the byte values in the order of the sortState; otherwise, the order is the same as the values were appended.
    */
   public IndexedBytesRefIterator iterator(final SortState sortState) {
     final int size = size();
     final int[] indices = sortState == null ? null : sortState.indices;
     assert indices == null || indices.length == size : indices.length + " != " + size;
     final BytesRefBuilder spare = new BytesRefBuilder();
     final BytesRef result = new BytesRef();

     return new IndexedBytesRefIterator() {
       int pos = -1;
       int ord = 0;
       @Override
       public BytesRef next() {
         ++pos;
         if (pos < size) {
           ord = indices == null ? pos : indices[pos];
           setBytesRef(spare, result, ord);
           return result;
         }
         return null;
       }

       @Override
       public int ord() {
         return ord;
       }
     };
   }

   /**
    * Used to iterate the elements of an array in a given order.
    */
   public final static class SortState implements Accountable {
     private final int[] indices;

     private SortState(int[] indices) {
       this.indices = indices;
     }

     @Override
     public long ramBytesUsed() {
       return RamUsageEstimator.NUM_BYTES_ARRAY_HEADER + indices.length * Integer.BYTES;
     }
   }

   /**
    * An extension of {@link BytesRefIterator} that allows retrieving the index of the current element
    */
   public interface IndexedBytesRefIterator extends BytesRefIterator {
     /**
      * Returns the ordinal position of the element that was returned in the latest call of {@link #next()}.
      * Do not call this method if {@link #next()} is not called yet or the last call returned a null value.
      */
     int ord();
   }
 }
	/*
	* Licensed to the Apache Software Foundation (ASF) under one or more
	* contributor license agreements. See the NOTICE file distributed with
	* this work for additional information regarding copyright ownership.
	* The ASF licenses this file to You under the Apache License, Version 2.0
	* (the "License"); you may not use this file except in compliance with
	* the License. You may obtain a copy of the License at
	*
	* http://www.apache.org/licenses/LICENSE-2.0
	*
	* Unless required by applicable law or agreed to in writing, software
	* distributed under the License is distributed on an "AS IS" BASIS,
	* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	* See the License for the specific language governing permissions and
	* limitations under the License.
	*/
	package org.apache.lucene.util;

	import java.util.Arrays;
	import java.util.Comparator;
	import java.util.function.IntBinaryOperator;

	/**
	* A simple append only random-access {@link BytesRef} array that stores full
	* copies of the appended bytes in a {@link ByteBlockPool}.
	*
	*
	* <b>Note: This class is not Thread-Safe!</b>
	*
	* @lucene.internal
	* @lucene.experimental
	*/
	public final class BytesRefArray implements SortableBytesRefArray {
	private final ByteBlockPool pool;
	private int[] offsets = new int[1];
	private int lastElement = 0;
	private int currentOffset = 0;
	private final Counter bytesUsed;

	/**
	* Creates a new {@link BytesRefArray} with a counter to track allocated bytes
	*/
	public BytesRefArray(Counter bytesUsed) {
	this.pool = new ByteBlockPool(new ByteBlockPool.DirectTrackingAllocator(
	bytesUsed));
	pool.nextBuffer();
	bytesUsed.addAndGet(RamUsageEstimator.NUM_BYTES_ARRAY_HEADER * Integer.BYTES);
	this.bytesUsed = bytesUsed;
	}

	/**
	* Clears this {@link BytesRefArray}
	*/
	@Override
	public void clear() {
	lastElement = 0;
	currentOffset = 0;
	// TODO: it's trappy that this does not return storage held by int[] offsets array!
	Arrays.fill(offsets, 0);
	pool.reset(false, true); // no need to 0 fill the buffers we control the allocator
	}

	/**
	* Appends a copy of the given {@link BytesRef} to this {@link BytesRefArray}.
	* @param bytes the bytes to append
	* @return the index of the appended bytes
	*/
	@Override
	public int append(BytesRef bytes) {
	if (lastElement >= offsets.length) {
	int oldLen = offsets.length;
	offsets = ArrayUtil.grow(offsets, offsets.length + 1);
	bytesUsed.addAndGet((offsets.length - oldLen) * Integer.BYTES);
	}
	pool.append(bytes);
	offsets[lastElement++] = currentOffset;
	currentOffset += bytes.length;
	return lastElement-1;
	}

	/**
	* Returns the current size of this {@link BytesRefArray}
	* @return the current size of this {@link BytesRefArray}
	*/
	@Override
	public int size() {
	return lastElement;
	}

	/**
	* Returns the <i>n'th</i> element of this {@link BytesRefArray}
	* @param spare a spare {@link BytesRef} instance
	* @param index the elements index to retrieve
	* @return the <i>n'th</i> element of this {@link BytesRefArray}
	*/
	public BytesRef get(BytesRefBuilder spare, int index) {
	FutureObjects.checkIndex(index, lastElement);
	int offset = offsets[index];
	int length = index == lastElement - 1 ? currentOffset - offset
	: offsets[index + 1] - offset;
	spare.grow(length);
	spare.setLength(length);
	pool.readBytes(offset, spare.bytes(), 0, spare.length());
	return spare.get();
	}

	/** Used only by sort below, to set a {@link BytesRef} with the specified slice, avoiding copying bytes in the common case when the slice
	* is contained in a single block in the byte block pool. */
	private void setBytesRef(BytesRefBuilder spare, BytesRef result, int index) {
	FutureObjects.checkIndex(index, lastElement);
	int offset = offsets[index];
	int length;
	if (index == lastElement - 1) {
	length = currentOffset - offset;
	} else {
	length = offsets[index + 1] - offset;
	}
	pool.setBytesRef(spare, result, offset, length);
	}


	/**
	* Returns a {@link SortState} representing the order of elements in this array. This is a non-destructive operation.
	*/
	public SortState sort(final Comparator<BytesRef> comp, final IntBinaryOperator tieComparator) {
	final int[] orderedEntries = new int[size()];
	for (int i = 0; i < orderedEntries.length; i++) {
	orderedEntries[i] = i;
	}
	new IntroSorter() {
	@Override
	protected void swap(int i, int j) {
	final int o = orderedEntries[i];
	orderedEntries[i] = orderedEntries[j];
	orderedEntries[j] = o;
	}

	@Override
	protected int compare(int i, int j) {
	final int idx1 = orderedEntries[i], idx2 = orderedEntries[j];
	setBytesRef(scratch1, scratchBytes1, idx1);
	setBytesRef(scratch2, scratchBytes2, idx2);
	return compare(idx1, scratchBytes1, idx2, scratchBytes2);
	}

	@Override
	protected void setPivot(int i) {
	pivotIndex = orderedEntries[i];
	setBytesRef(pivotBuilder, pivot, pivotIndex);
	}

	@Override
	protected int comparePivot(int j) {
	final int index = orderedEntries[j];
	setBytesRef(scratch2, scratchBytes2, index);
	return compare(pivotIndex, pivot, index, scratchBytes2);
	}

	private int compare(int i1, BytesRef b1, int i2, BytesRef b2) {
	int res = comp.compare(b1, b2);
	return res == 0 ? tieComparator.applyAsInt(i1, i2) : res;
	}

	private int pivotIndex;
	private final BytesRef pivot = new BytesRef();
	private final BytesRef scratchBytes1 = new BytesRef();
	private final BytesRef scratchBytes2 = new BytesRef();
	private final BytesRefBuilder pivotBuilder = new BytesRefBuilder();
	private final BytesRefBuilder scratch1 = new BytesRefBuilder();
	private final BytesRefBuilder scratch2 = new BytesRefBuilder();
	}.sort(0, size());
	return new SortState(orderedEntries);
	}

	/**
	* sugar for {@link #iterator(Comparator)} with a <code>null</code> comparator
	*/
	public BytesRefIterator iterator() {
	return iterator((SortState) null);
	}

	/**
	* <p>
	* Returns a {@link BytesRefIterator} with point in time semantics. The
	* iterator provides access to all so far appended {@link BytesRef} instances.
	* </p>
	* <p>
	* If a non <code>null</code> {@link Comparator} is provided the iterator will
	* iterate the byte values in the order specified by the comparator. Otherwise
	* the order is the same as the values were appended.
	* </p>
	* <p>
	* This is a non-destructive operation.
	* </p>
	*/
	@Override
	public BytesRefIterator iterator(final Comparator<BytesRef> comp) {
	return iterator(sort(comp, (i, j) -> 0));
	}

	/**
	* Returns an {@link IndexedBytesRefIterator} with point in time semantics. The iterator provides access to all
	* so far appended {@link BytesRef} instances. If a non-null sortState is specified then the iterator will iterate
	* the byte values in the order of the sortState; otherwise, the order is the same as the values were appended.
	*/
	public IndexedBytesRefIterator iterator(final SortState sortState) {
	final int size = size();
	final int[] indices = sortState == null ? null : sortState.indices;
	assert indices == null \|\| indices.length == size : indices.length + " != " + size;
	final BytesRefBuilder spare = new BytesRefBuilder();
	final BytesRef result = new BytesRef();

	return new IndexedBytesRefIterator() {
	int pos = -1;
	int ord = 0;
	@Override
	public BytesRef next() {
	++pos;
	if (pos < size) {
	ord = indices == null ? pos : indices[pos];
	setBytesRef(spare, result, ord);
	return result;
	}
	return null;
	}

	@Override
	public int ord() {
	return ord;
	}
	};
	}

	/**
	* Used to iterate the elements of an array in a given order.
	*/
	public final static class SortState implements Accountable {
	private final int[] indices;

	private SortState(int[] indices) {
	this.indices = indices;
	}

	@Override
	public long ramBytesUsed() {
	return RamUsageEstimator.NUM_BYTES_ARRAY_HEADER + indices.length * Integer.BYTES;
	}
	}

	/**
	* An extension of {@link BytesRefIterator} that allows retrieving the index of the current element
	*/
	public interface IndexedBytesRefIterator extends BytesRefIterator {
	/**
	* Returns the ordinal position of the element that was returned in the latest call of {@link #next()}.
	* Do not call this method if {@link #next()} is not called yet or the last call returned a null value.
	*/
	int ord();
	}
	}