Allow FST builder to use different writer (#12543)

Author: Anh Dung Bui

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

@@ -21,13 +21,12 @@
 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
 
-class BytesStore extends DataOutput implements Accountable {
+class BytesStore extends DataOutput implements FSTWriter {
 
   private static final long BASE_RAM_BYTES_USED =
       RamUsageEstimator.shallowSizeOfInstance(BytesStore.class)
@@ -329,6 +328,7 @@ public void skipBytes(int len) {
     }
   }
 
+  @Override
   public long getPosition() {
     return ((long) blocks.size() - 1) * blockSize + nextWrite;
   }
@@ -355,6 +355,7 @@ public void truncate(long newLen) {
     assert newLen == getPosition();
   }
 
+  @Override
   public void finish() {
     if (current != null) {
       byte[] lastBuffer = new byte[nextWrite];
@@ -364,7 +365,15 @@ public void finish() {
     }
   }
 
+  /** Writes all of our bytes to the target {@link FSTWriter}. */
+  public void writeTo(FSTWriter out) throws IOException {
+    for (byte[] block : blocks) {
+      out.writeBytes(block, 0, block.length);
+    }
+  }
+
   /** Writes all of our bytes to the target {@link DataOutput}. */
+  @Override
   public void writeTo(DataOutput out) throws IOException {
     for (byte[] block : blocks) {
       out.writeBytes(block, 0, block.length);
@@ -437,10 +446,16 @@ public boolean reversed() {
     };
   }
 
+  @Override
   public FST.BytesReader getReverseReader() {
     return getReverseReader(true);
   }
 
+  @Override
+  public FST.BytesReader getReverseReaderForSuffixSharing() {
+    return getReverseReader(false);
+  }
+
   FST.BytesReader getReverseReader(boolean allowSingle) {
     if (allowSingle && blocks.size() == 1) {
       return new ReverseBytesReader(blocks.get(0));

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

@@ -120,7 +120,7 @@ public enum INPUT_TYPE {
    * A {@link BytesStore}, used during building, or during reading when the FST is very large (more
    * than 1 GB). If the FST is less than 1 GB then bytesArray is set instead.
    */
-  final BytesStore bytes;
+  final FSTWriter fstWriter;
 
   private final FSTStore fstStore;
 
@@ -395,14 +395,18 @@ private static boolean flag(int flags, int bit) {
   }
 
   // make a new empty FST, for building; Builder invokes this
-  FST(INPUT_TYPE inputType, Outputs<T> outputs, int bytesPageBits) {
+  FST(INPUT_TYPE inputType, Outputs<T> outputs, FSTWriter fstWriter) {
     this.inputType = inputType;
     this.outputs = outputs;
     fstStore = null;
-    bytes = new BytesStore(bytesPageBits);
+    this.fstWriter = fstWriter;
     // pad: ensure no node gets address 0 which is reserved to mean
     // the stop state w/ no arcs
-    bytes.writeByte((byte) 0);
+    try {
+      this.fstWriter.writeByte((byte) 0);
+    } catch (IOException e) {
+      throw new RuntimeException(e);
+    }
     emptyOutput = null;
     this.version = VERSION_CURRENT;
   }
@@ -420,7 +424,7 @@ public FST(DataInput metaIn, DataInput in, Outputs<T> outputs) throws IOExceptio
    */
   public FST(DataInput metaIn, DataInput in, Outputs<T> outputs, FSTStore fstStore)
       throws IOException {
-    bytes = null;
+    fstWriter = null;
     this.fstStore = fstStore;
     this.outputs = outputs;
 
@@ -472,7 +476,7 @@ public long ramBytesUsed() {
     if (this.fstStore != null) {
       size += this.fstStore.ramBytesUsed();
     } else {
-      size += bytes.ramBytesUsed();
+      size += this.fstWriter.ramBytesUsed();
     }
 
     return size;
@@ -484,19 +488,19 @@ public String toString() {
   }
 
   void finish(long newStartNode) throws IOException {
-    assert newStartNode <= bytes.getPosition();
+    assert newStartNode <= fstWriter.getPosition();
     if (startNode != -1) {
       throw new IllegalStateException("already finished");
     }
     if (newStartNode == FINAL_END_NODE && emptyOutput != null) {
       newStartNode = 0;
     }
     startNode = newStartNode;
-    bytes.finish();
+    fstWriter.finish();
   }
 
   public long numBytes() {
-    return bytes.getPosition();
+    return fstWriter.getPosition();
   }
 
   public T getEmptyOutput() {
@@ -512,6 +516,21 @@ void setEmptyOutput(T v) {
   }
 
   public void save(DataOutput metaOut, DataOutput out) throws IOException {
+    saveMetadata(metaOut);
+    if (fstWriter != null) {
+      fstWriter.writeTo(out);
+    } else {
+      assert fstStore != null;
+      fstStore.writeTo(out);
+    }
+  }
+
+  /**
+   * Save the metadata to a DataOutput
+   *
+   * @param metaOut the DataOutput to save
+   */
+  public void saveMetadata(DataOutput metaOut) throws IOException {
     if (startNode == -1) {
       throw new IllegalStateException("call finish first");
     }
@@ -552,13 +571,9 @@ public void save(DataOutput metaOut, DataOutput out) throws IOException {
     }
     metaOut.writeByte(t);
     metaOut.writeVLong(startNode);
-    if (bytes != null) {
-      long numBytes = bytes.getPosition();
+    if (fstWriter != null) {
+      long numBytes = fstWriter.getPosition();
       metaOut.writeVLong(numBytes);
-      bytes.writeTo(out);
-    } else {
-      assert fstStore != null;
-      fstStore.writeTo(out);
     }
   }
 
@@ -1131,7 +1146,7 @@ public BytesReader getBytesReader() {
     if (this.fstStore != null) {
       return this.fstStore.getReverseBytesReader();
     } else {
-      return bytes.getReverseReader();
+      return fstWriter.getReverseReader();
     }
   }