Allow FST builder to use different writer

Author: dungba88

lucene/core/src/java/org/apache/lucene/util/fst/BytesStore.java

@@ -26,7 +26,8 @@
 // TODO: merge with PagedBytes, except PagedBytes doesn't
 // let you read while writing which FST needs
 
-class BytesStore extends DataOutput implements FSTReader {
+// TODO: Separate the scratch writer and reader functionality
+class BytesStore extends DataOutput implements FSTReader, Freezeable {
 
   private static final long BASE_RAM_BYTES_USED =
       RamUsageEstimator.shallowSizeOfInstance(BytesStore.class)
@@ -40,6 +41,7 @@ class BytesStore extends DataOutput implements FSTReader {
 
   private byte[] current;
   private int nextWrite;
+  private boolean frozen;
 
   public BytesStore(int blockBits) {
     this.blockBits = blockBits;
@@ -50,13 +52,15 @@ public BytesStore(int blockBits) {
 
   /** Absolute write byte; you must ensure dest is < max position written so far. */
   public void writeByte(long dest, byte b) {
+    assert frozen == false;
     int blockIndex = (int) (dest >> blockBits);
     byte[] block = blocks.get(blockIndex);
     block[(int) (dest & blockMask)] = b;
   }
 
   @Override
   public void writeByte(byte b) {
+    assert frozen == false;
     if (nextWrite == blockSize) {
       current = new byte[blockSize];
       blocks.add(current);
@@ -67,6 +71,7 @@ public void writeByte(byte b) {
 
   @Override
   public void writeBytes(byte[] b, int offset, int len) {
+    assert frozen == false;
     while (len > 0) {
       int chunk = blockSize - nextWrite;
       if (len <= chunk) {
@@ -97,6 +102,7 @@ int getBlockBits() {
    * so you must only call it on already written parts.
    */
   void writeBytes(long dest, byte[] b, int offset, int len) {
+    assert frozen == false;
     // System.out.println("  BS.writeBytes dest=" + dest + " offset=" + offset + " len=" + len);
     assert dest + len <= getPosition() : "dest=" + dest + " pos=" + getPosition() + " len=" + len;
 
@@ -156,6 +162,7 @@ void writeBytes(long dest, byte[] b, int offset, int len) {
 
   @Override
   public void copyBytes(DataInput input, long numBytes) throws IOException {
+    assert frozen == false;
     assert numBytes >= 0 : "numBytes=" + numBytes;
     assert input != null;
     long len = numBytes;
@@ -180,6 +187,7 @@ public void copyBytes(DataInput input, long numBytes) throws IOException {
    * bytes, so must only call it on already written parts.
    */
   public void copyBytes(long src, long dest, int len) {
+    assert frozen == false;
     // System.out.println("BS.copyBytes src=" + src + " dest=" + dest + " len=" + len);
     assert src < dest;
 
@@ -238,6 +246,7 @@ public void copyBytes(long src, long dest, int len) {
 
   /** Copies bytes from this store to a target byte array. */
   public void copyBytes(long src, byte[] dest, int offset, int len) {
+    assert frozen == false;
     int blockIndex = (int) (src >> blockBits);
     int upto = (int) (src & blockMask);
     byte[] block = blocks.get(blockIndex);
@@ -259,6 +268,7 @@ public void copyBytes(long src, byte[] dest, int offset, int len) {
 
   /** Writes an int at the absolute position without changing the current pointer. */
   public void writeInt(long pos, int value) {
+    assert frozen == false;
     int blockIndex = (int) (pos >> blockBits);
     int upto = (int) (pos & blockMask);
     byte[] block = blocks.get(blockIndex);
@@ -276,6 +286,7 @@ public void writeInt(long pos, int value) {
 
   /** Reverse from srcPos, inclusive, to destPos, inclusive. */
   public void reverse(long srcPos, long destPos) {
+    assert frozen == false;
     assert srcPos < destPos;
     assert destPos < getPosition();
     // System.out.println("reverse src=" + srcPos + " dest=" + destPos);
@@ -314,6 +325,7 @@ public void reverse(long srcPos, long destPos) {
   }
 
   public void skipBytes(int len) {
+    assert frozen == false;
     while (len > 0) {
       int chunk = blockSize - nextWrite;
       if (len <= chunk) {
@@ -337,11 +349,23 @@ public long size() {
     return getPosition();
   }
 
+  /** Similar to {@link #truncate(long)} with newLen=0 but keep the first block to reduce GC. */
+  public void reset() {
+    assert frozen == false;
+    if (blocks.isEmpty()) {
+      return;
+    }
+    nextWrite = 0;
+    current = blocks.get(0);
+    blocks.subList(1, blocks.size()).clear();
+  }
+
   /**
    * Pos must be less than the max position written so far! Ie, you cannot "grow" the file with
    * this!
    */
   public void truncate(long newLen) {
+    assert frozen == false;
     assert newLen <= getPosition();
     assert newLen >= 0;
     int blockIndex = (int) (newLen >> blockBits);
@@ -359,7 +383,9 @@ public void truncate(long newLen) {
     assert newLen == getPosition();
   }
 
-  public void finish() {
+  @Override
+  public void freeze() {
+    this.frozen = true;
     if (current != null) {
       byte[] lastBuffer = new byte[nextWrite];
       System.arraycopy(current, 0, lastBuffer, 0, nextWrite);
@@ -368,9 +394,22 @@ public void finish() {
     }
   }
 
+  /** Writes all of our bytes to the target {@link FSTDataOutputWriter}. */
+  public void writeTo(FSTDataOutputWriter out) throws IOException {
+    reverse(0, getPosition() - 1);
+    for (byte[] block : blocks) {
+      if (block == current) { // last block
+        out.writeBytes(block, 0, nextWrite);
+      } else {
+        out.writeBytes(block, 0, block.length);
+      }
+    }
+  }
+
   /** Writes all of our bytes to the target {@link DataOutput}. */
   @Override
   public void writeTo(DataOutput out) throws IOException {
+    assert frozen;
     for (byte[] block : blocks) {
       out.writeBytes(block, 0, block.length);
     }

lucene/core/src/java/org/apache/lucene/util/fst/FST.java

@@ -435,6 +435,13 @@ public FST(FSTMetadata<T> metadata, DataInput in, Outputs<T> outputs, FSTStore f
     this.fstReader = fstReader;
   }
 
+  /**
+   * @return true if and only if this FST is readable (i.e. has a reverse BytesReader)
+   */
+  public boolean hasReverseBytesReader() {
+    return fstReader.getReverseBytesReader() != null;
+  }
+
   /**
    * Read the FST metadata from DataInput
    *
@@ -1181,7 +1188,11 @@ private void seekToNextNode(BytesReader in) throws IOException {
     }
   }
 
-  /** Returns a {@link BytesReader} for this FST, positioned at position 0. */
+  /**
+   * Returns a {@link BytesReader} for this FST, positioned at position 0.
+   *
+   * @see #hasReverseBytesReader()
+   */
   public BytesReader getBytesReader() {
     return fstReader.getReverseBytesReader();
   }