Consolidate FSTStore and BytesStore in FST (#12709)

* Remove direct dependency of NodeHash to FST

* Fix index out of bounds when writing FST to different metaOut (#12697)

* Tidify code

* Update CHANGES.txt

* Re-add assertion

* Remove direct dependency of NodeHash to FST

* Hold off the FSTTraversal changes

* Rename variable

* Add Javadoc

* Add @OverRide

* tidy

* tidy

* Change to FSTReader

* Update CHANGES.txt

Author: dungba88

lucene/CHANGES.txt

@@ -18,9 +18,6 @@ API Changes
 * GITHUB#12592: Add RandomAccessInput#length method to the RandomAccessInput interface. In addition deprecate
   ByteBuffersDataInput#size in favour of this new method. (Ignacio Vera)
 
-* GITHUB#12646: Move FST#addNode to FSTCompiler to avoid a circular dependency between FST and FSTCompiler
-  (Anh Dung Bui via Alan Woodward)
-
 * GITHUB#12718: Make IndexSearcher#getSlices final as it is not expected to be overridden (Luca Cavanna)
 
 * GITHUB#12427: Automata#makeStringUnion #makeBinaryStringUnion now accept Iterable<BytesRef> instead of
@@ -32,6 +29,12 @@ API Changes
 
 * GITHUB#12816: Add HumanReadableQuery which takes a description parameter for debugging purposes. (Jakub Slowinski)
 
+* GITHUB#12646, GITHUB#12690: Move FST#addNode to FSTCompiler to avoid a circular dependency
+  between FST and FSTCompiler (Anh Dung Bui)
+
+* GITHUB#12709 Consolidate FSTStore and BytesStore in FST. Created FSTReader which contains the common methods
+  of the two (Anh Dung Bui)
+
 New Features
 ---------------------
 
@@ -108,7 +111,7 @@ Optimizations
 
 * GITHUB#12382: Faster top-level conjunctions on term queries when sorting by
   descending score. (Adrien Grand)
-  
+
 * GITHUB#12591: Use stable radix sort to speed up the sorting of update terms. (Guo Feng)
 
 * GITHUB#12587: Use radix sort to speed up the sorting of terms in TermInSetQuery. (Guo Feng)
@@ -123,7 +126,7 @@ Optimizations
 
 * GITHUB#12668: ImpactsEnums now decode frequencies lazily like PostingsEnums.
   (Adrien Grand)
-  
+
 * GITHUB#12651: Use 2d array for OnHeapHnswGraph representation. (Patrick Zhai)
 
 * GITHUB#12653: Optimize computing number of levels in MultiLevelSkipListWriter#bufferSkip. (Shubham Chaudhary)

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 FSTReader {
 
   private static final long BASE_RAM_BYTES_USED =
       RamUsageEstimator.shallowSizeOfInstance(BytesStore.class)
@@ -333,6 +332,11 @@ public long getPosition() {
     return ((long) blocks.size() - 1) * blockSize + nextWrite;
   }
 
+  @Override
+  public long size() {
+    return getPosition();
+  }
+
   /**
    * Pos must be less than the max position written so far! Ie, you cannot "grow" the file with
    * this!
@@ -365,6 +369,7 @@ public void finish() {
   }
 
   /** 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,7 +442,8 @@ public boolean reversed() {
     };
   }
 
-  public FST.BytesReader getReverseReader() {
+  @Override
+  public FST.BytesReader getReverseBytesReader() {
     return getReverseReader(true);
   }
 

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

@@ -123,9 +123,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;
-
-  private final FSTStore fstStore;
+  private final FSTReader fstReader;
 
   private long startNode = -1;
 
@@ -398,15 +396,11 @@ 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, FSTReader fstReader) {
     this.inputType = inputType;
     this.outputs = outputs;
-    fstStore = null;
-    bytes = new BytesStore(bytesPageBits);
-    // pad: ensure no node gets address 0 which is reserved to mean
-    // the stop state w/ no arcs
-    bytes.writeByte((byte) 0);
     emptyOutput = null;
+    this.fstReader = fstReader;
     this.version = VERSION_CURRENT;
   }
 
@@ -423,8 +417,6 @@ 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;
-    this.fstStore = fstStore;
     this.outputs = outputs;
 
     // NOTE: only reads formats VERSION_START up to VERSION_CURRENT; we don't have
@@ -438,7 +430,7 @@ public FST(DataInput metaIn, DataInput in, Outputs<T> outputs, FSTStore fstStore
       emptyBytes.copyBytes(metaIn, numBytes);
 
       // De-serialize empty-string output:
-      BytesReader reader = emptyBytes.getReverseReader();
+      BytesReader reader = emptyBytes.getReverseBytesReader();
       // NoOutputs uses 0 bytes when writing its output,
       // so we have to check here else BytesStore gets
       // angry:
@@ -466,19 +458,13 @@ public FST(DataInput metaIn, DataInput in, Outputs<T> outputs, FSTStore fstStore
     startNode = metaIn.readVLong();
 
     long numBytes = metaIn.readVLong();
-    this.fstStore.init(in, numBytes);
+    fstStore.init(in, numBytes);
+    this.fstReader = fstStore;
   }
 
   @Override
   public long ramBytesUsed() {
-    long size = BASE_RAM_BYTES_USED;
-    if (this.fstStore != null) {
-      size += this.fstStore.ramBytesUsed();
-    } else {
-      size += bytes.ramBytesUsed();
-    }
-
-    return size;
+    return BASE_RAM_BYTES_USED + fstReader.ramBytesUsed();
   }
 
   @Override
@@ -487,19 +473,18 @@ public String toString() {
   }
 
   void finish(long newStartNode) throws IOException {
-    assert newStartNode <= bytes.getPosition();
+    assert newStartNode <= fstReader.size();
     if (startNode != -1) {
       throw new IllegalStateException("already finished");
     }
     if (newStartNode == FINAL_END_NODE && emptyOutput != null) {
       newStartNode = 0;
     }
     startNode = newStartNode;
-    bytes.finish();
   }
 
   public long numBytes() {
-    return bytes.getPosition();
+    return fstReader.size();
   }
 
   public T getEmptyOutput() {
@@ -555,16 +540,8 @@ public void save(DataOutput metaOut, DataOutput out) throws IOException {
     }
     metaOut.writeByte(t);
     metaOut.writeVLong(startNode);
-    if (bytes != null) {
-      long numBytes = bytes.getPosition();
-      metaOut.writeVLong(numBytes);
-      bytes.writeTo(out);
-    } else {
-      assert fstStore != null;
-      long numBytes = fstStore.size();
-      metaOut.writeVLong(numBytes);
-      fstStore.writeTo(out);
-    }
+    metaOut.writeVLong(numBytes());
+    fstReader.writeTo(out);
   }
 
   /** Writes an automaton to a file. */
@@ -1141,11 +1118,7 @@ private void seekToNextNode(BytesReader in) throws IOException {
 
   /** Returns a {@link BytesReader} for this FST, positioned at position 0. */
   public BytesReader getBytesReader() {
-    if (this.fstStore != null) {
-      return this.fstStore.getReverseBytesReader();
-    } else {
-      return bytes.getReverseReader();
-    }
+    return fstReader.getReverseBytesReader();
   }
 
   /** Reads bytes stored in an FST. */

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

@@ -137,9 +137,11 @@ private FSTCompiler(
       float directAddressingMaxOversizingFactor) {
     this.allowFixedLengthArcs = allowFixedLengthArcs;
     this.directAddressingMaxOversizingFactor = directAddressingMaxOversizingFactor;
-    fst = new FST<>(inputType, outputs, bytesPageBits);
-    bytes = fst.bytes;
-    assert bytes != null;
+    bytes = new BytesStore(bytesPageBits);
+    // pad: ensure no node gets address 0 which is reserved to mean
+    // the stop state w/ no arcs
+    bytes.writeByte((byte) 0);
+    fst = new FST<>(inputType, outputs, bytes);
     if (suffixRAMLimitMB < 0) {
       throw new IllegalArgumentException("ramLimitMB must be >= 0; got: " + suffixRAMLimitMB);
     } else if (suffixRAMLimitMB > 0) {
@@ -317,8 +319,6 @@ private CompiledNode compileNode(UnCompiledNode<T> nodeIn, int tailLength) throw
   // serializes new node by appending its bytes to the end
   // of the current byte[]
   long addNode(FSTCompiler.UnCompiledNode<T> nodeIn) throws IOException {
-    T NO_OUTPUT = fst.outputs.getNoOutput();
-
     // System.out.println("FST.addNode pos=" + bytes.getPosition() + " numArcs=" + nodeIn.numArcs);
     if (nodeIn.numArcs == 0) {
       if (nodeIn.isFinal) {
@@ -859,6 +859,7 @@ public FST<T> compile() throws IOException {
     // if (DEBUG) System.out.println("  builder.finish root.isFinal=" + root.isFinal + "
     // root.output=" + root.output);
     fst.finish(compileNode(root, lastInput.length()).node);
+    bytes.finish();
 
     return fst;
   }

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

@@ -0,0 +1,47 @@
+/*
+ * 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.fst;
+
+import java.io.IOException;
+import org.apache.lucene.store.DataOutput;
+import org.apache.lucene.util.Accountable;
+
+/** Abstraction for reading bytes necessary for FST. */
+public interface FSTReader extends Accountable {
+
+  /**
+   * The raw size in bytes of the FST
+   *
+   * @return the FST size
+   */
+  long size();
+
+  /**
+   * Get the reverse BytesReader for this FST
+   *
+   * @return the reverse BytesReader
+   */
+  FST.BytesReader getReverseBytesReader();
+
+  /**
+   * Write this FST to another DataOutput
+   *
+   * @param out the DataOutput
+   * @throws IOException if exception occurred during writing
+   */
+  void writeTo(DataOutput out) throws IOException;
+}

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

@@ -18,16 +18,8 @@
 
 import java.io.IOException;
 import org.apache.lucene.store.DataInput;
-import org.apache.lucene.store.DataOutput;
-import org.apache.lucene.util.Accountable;
 
-/** Abstraction for reading/writing bytes necessary for FST. */
-public interface FSTStore extends Accountable {
+/** A type of {@link FSTReader} which needs data to be initialized before use */
+public interface FSTStore extends FSTReader {
   void init(DataInput in, long numBytes) throws IOException;
-
-  long size();
-
-  FST.BytesReader getReverseBytesReader();
-
-  void writeTo(DataOutput out) throws IOException;
 }

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

@@ -88,7 +88,7 @@ public FST.BytesReader getReverseBytesReader() {
     if (bytesArray != null) {
       return new ReverseBytesReader(bytesArray);
     } else {
-      return bytes.getReverseReader();
+      return bytes.getReverseBytesReader();
     }
   }
 

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

@@ -271,7 +271,7 @@ private void verify(BytesStore bytes, byte[] expected, int totalLength) throws E
         System.out.println("    bulk: reversed");
       }
       // reversed
-      FST.BytesReader r = bytes.getReverseReader();
+      FST.BytesReader r = bytes.getReverseBytesReader();
       assertTrue(r.reversed());
       r.setPosition(totalLength - 1);
       r.readBytes(actual, 0, actual.length);
@@ -306,7 +306,7 @@ private void verify(BytesStore bytes, byte[] expected, int totalLength) throws E
       if (VERBOSE) {
         System.out.println("    ops: reversed");
       }
-      r = bytes.getReverseReader();
+      r = bytes.getReverseBytesReader();
     } else {
       if (VERBOSE) {
         System.out.println("    ops: forward");