Decouple the reader functionality from BytesStore

Author: dungba88

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

@@ -0,0 +1,45 @@
+package org.apache.lucene.util.fst;
+
+import java.io.IOException;
+import org.apache.lucene.store.ByteBuffersDataOutput;
+import org.apache.lucene.store.DataOutput;
+
+/** A {@link FSTReader} which is backed by {@link ByteBuffersDataOutput} */
+final class ByteBuffersFSTReader extends DataOutput implements FSTReader {
+
+  private final ByteBuffersDataOutput dataOutput;
+
+  public ByteBuffersFSTReader(ByteBuffersDataOutput dataOutput) {
+    this.dataOutput = dataOutput;
+  }
+
+  @Override
+  public void writeByte(byte b) {
+    dataOutput.writeByte(b);
+  }
+
+  @Override
+  public void writeBytes(byte[] b, int offset, int length) {
+    dataOutput.writeBytes(b, offset, length);
+  }
+
+  @Override
+  public long ramBytesUsed() {
+    return dataOutput.ramBytesUsed();
+  }
+
+  @Override
+  public long size() {
+    return dataOutput.size();
+  }
+
+  @Override
+  public FST.BytesReader getReverseBytesReader() {
+    return new ReverseRandomAccessReader(dataOutput.toDataInput());
+  }
+
+  @Override
+  public void writeTo(DataOutput out) throws IOException {
+    dataOutput.copyTo(out);
+  }
+}

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

@@ -21,13 +21,13 @@
 import java.util.List;
 import org.apache.lucene.store.DataInput;
 import org.apache.lucene.store.DataOutput;
+import org.apache.lucene.util.Accountable;
 import org.apache.lucene.util.RamUsageEstimator;
 
 // TODO: merge with PagedBytes, except PagedBytes doesn't
 // let you read while writing which FST needs
 
-// TODO: Separate the scratch writer and reader functionality
-class BytesStore extends DataOutput implements FSTReader {
+class BytesStore extends DataOutput implements Accountable {
 
   private static final long BASE_RAM_BYTES_USED =
       RamUsageEstimator.shallowSizeOfInstance(BytesStore.class)
@@ -333,11 +333,6 @@ public long getPosition() {
     return ((long) blocks.size() - 1) * blockSize + nextWrite;
   }
 
-  @Override
-  public long size() {
-    return getPosition();
-  }
-
   /** Similar to {@link #truncate(long)} with newLen=0 but keep the first block to reduce GC. */
   public void reset() {
     if (blocks.isEmpty()) {
@@ -371,7 +366,6 @@ public void truncate(long newLen) {
   }
 
   /** Writes all of our bytes to the target {@link DataOutput}. */
-  @Override
   public void writeTo(DataOutput out) throws IOException {
     for (byte[] block : blocks) {
       if (block == current) { // last block
@@ -443,7 +437,6 @@ public void setPosition(long pos) {
     };
   }
 
-  @Override
   public FST.BytesReader getReverseBytesReader() {
     if (blocks.size() == 1) {
       return new ReverseBytesReader(blocks.get(0));

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

@@ -16,6 +16,8 @@
  */
 package org.apache.lucene.util.fst;
 
+import static org.apache.lucene.store.ByteBuffersDataOutput.ALLOCATE_BB_ON_HEAP;
+import static org.apache.lucene.store.ByteBuffersDataOutput.NO_REUSE;
 import static org.apache.lucene.util.fst.FST.ARCS_FOR_BINARY_SEARCH;
 import static org.apache.lucene.util.fst.FST.ARCS_FOR_CONTINUOUS;
 import static org.apache.lucene.util.fst.FST.ARCS_FOR_DIRECT_ADDRESSING;
@@ -32,6 +34,7 @@
 
 import java.io.IOException;
 import org.apache.lucene.store.ByteArrayDataOutput;
+import org.apache.lucene.store.ByteBuffersDataOutput;
 import org.apache.lucene.store.DataOutput;
 import org.apache.lucene.util.Accountable;
 import org.apache.lucene.util.ArrayUtil;
@@ -140,7 +143,8 @@ public FSTCompiler(FST.INPUT_TYPE inputType, Outputs<T> outputs) {
   }
 
   static DataOutput getLegacyDataOutput(int blockBits) {
-    return new BytesStore(blockBits);
+    return new ByteBuffersFSTReader(
+        new ByteBuffersDataOutput(blockBits, blockBits, ALLOCATE_BB_ON_HEAP, NO_REUSE));
   }
 
   private FSTCompiler(
@@ -539,9 +543,9 @@ long addNode(FSTCompiler.UnCompiledNode<T> nodeIn) throws IOException {
       }
     }
 
-    scratchBytes.reverse(0, scratchBytes.size() - 1);
+    scratchBytes.reverse(0, scratchBytes.getPosition() - 1);
     scratchBytes.writeTo(dataOutput);
-    numBytesWritten += scratchBytes.size();
+    numBytesWritten += scratchBytes.getPosition();
 
     nodeCount++;
     return numBytesWritten - 1;

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

@@ -138,7 +138,7 @@ public long add(FSTCompiler.UnCompiledNode<T> nodeIn) throws IOException {
           // we use 0 as empty marker in hash table, so it better be impossible to get a frozen node
           // at 0:
           assert nodeAddress != FST.FINAL_END_NODE && nodeAddress != FST.NON_FINAL_END_NODE;
-          byte[] buf = new byte[Math.toIntExact(fstCompiler.scratchBytes.size())];
+          byte[] buf = new byte[Math.toIntExact(fstCompiler.scratchBytes.getPosition())];
           fstCompiler.scratchBytes.copyBytes(0, buf, 0, buf.length);
 
           primaryTable.setNode(hashSlot, nodeAddress, buf);

lucene/core/src/test/org/apache/lucene/util/fst/TestFSTDirectAddressing.java

@@ -65,7 +65,7 @@ public void testDeDupTails() throws Exception {
       }
       entries.add(new BytesRef(b));
     }
-    long size = buildFST(entries).ramBytesUsed();
+    long size = buildFST(entries).numBytes();
     // Size is 1648 when we use only list-encoding. We were previously failing to ever de-dup
     // direct addressing, which led this case to blow up.
     // This test will fail if there is more than 1% size increase with direct addressing.