change to package-private

Author: dungba88

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

@@ -31,9 +31,11 @@
 import java.io.IOException;
 import org.apache.lucene.store.ByteArrayDataOutput;
 import org.apache.lucene.store.DataOutput;
+import org.apache.lucene.util.Accountable;
 import org.apache.lucene.util.ArrayUtil;
 import org.apache.lucene.util.IntsRef;
 import org.apache.lucene.util.IntsRefBuilder;
+import org.apache.lucene.util.RamUsageEstimator;
 import org.apache.lucene.util.fst.FST.INPUT_TYPE; // javadoc
 
 // TODO: could we somehow stream an FST to disk while we
@@ -83,7 +85,7 @@ public class FSTCompiler<T> {
    */
   private static final float DIRECT_ADDRESSING_MAX_OVERSIZE_WITH_CREDIT_FACTOR = 1.66f;
 
-  private final NodeHash<T> dedupHash;
+  final NodeHash<T> dedupHash;
   final FST<T> fst;
   private final T NO_OUTPUT;
 
@@ -145,7 +147,7 @@ private FSTCompiler(
     if (suffixRAMLimitMB < 0) {
       throw new IllegalArgumentException("ramLimitMB must be >= 0; got: " + suffixRAMLimitMB);
     } else if (suffixRAMLimitMB > 0) {
-      dedupHash = new NodeHash<>(this, suffixRAMLimitMB, bytes.getReverseReader(false));
+      dedupHash = new NodeHash<>(this, suffixRAMLimitMB);
     } else {
       dedupHash = null;
     }
@@ -270,10 +272,6 @@ public float getDirectAddressingMaxOversizingFactor() {
     return directAddressingMaxOversizingFactor;
   }
 
-  public long getTermCount() {
-    return frontier[0].inputCount;
-  }
-
   public long getNodeCount() {
     // 1+ in order to count the -1 implicit final node
     return 1 + nodeCount;
@@ -292,13 +290,13 @@ private CompiledNode compileNode(UnCompiledNode<T> nodeIn, int tailLength) throw
     long bytesPosStart = bytes.getPosition();
     if (dedupHash != null) {
       if (nodeIn.numArcs == 0) {
-        node = addNode(nodeIn);
+        node = addNode(nodeIn).nodeAddress;
         lastFrozenNode = node;
       } else {
         node = dedupHash.add(nodeIn);
       }
     } else {
-      node = addNode(nodeIn);
+      node = addNode(nodeIn).nodeAddress;
     }
     assert node != -2;
 
@@ -318,13 +316,13 @@ 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 {
+  NodeAndBuffer addNode(FSTCompiler.UnCompiledNode<T> nodeIn) throws IOException {
     // System.out.println("FST.addNode pos=" + bytes.getPosition() + " numArcs=" + nodeIn.numArcs);
     if (nodeIn.numArcs == 0) {
       if (nodeIn.isFinal) {
-        return FINAL_END_NODE;
+        return new NodeAndBuffer(FINAL_END_NODE, null);
       } else {
-        return NON_FINAL_END_NODE;
+        return new NodeAndBuffer(NON_FINAL_END_NODE, null);
       }
     }
     final long startAddress = bytes.getPosition();
@@ -461,7 +459,13 @@ long addNode(FSTCompiler.UnCompiledNode<T> nodeIn) throws IOException {
     final long thisNodeAddress = bytes.getPosition() - 1;
     bytes.reverse(startAddress, thisNodeAddress);
     nodeCount++;
-    return thisNodeAddress;
+    byte[] buf = new byte[Math.toIntExact(thisNodeAddress - startAddress + 1)];
+    bytes.copyBytes(startAddress, buf, 0, buf.length);
+    return new NodeAndBuffer(thisNodeAddress, buf);
+  }
+
+  record NodeAndBuffer(long nodeAddress, byte[] bytes) {
+
   }
 
   private void writeLabel(DataOutput out, int v) throws IOException {
@@ -749,7 +753,6 @@ public void add(IntsRef input, T output) throws IOException {
       // format cannot represent the empty input since
       // 'finalness' is stored on the incoming arc, not on
       // the node
-      frontier[0].inputCount++;
       frontier[0].isFinal = true;
       fst.setEmptyOutput(output);
       return;
@@ -760,9 +763,6 @@ public void add(IntsRef input, T output) throws IOException {
     int pos2 = input.offset;
     final int pos1Stop = Math.min(lastInput.length(), input.length);
     while (true) {
-      frontier[pos1].inputCount++;
-      // System.out.println("  incr " + pos1 + " ct=" + frontier[pos1].inputCount + " n=" +
-      // frontier[pos1]);
       if (pos1 >= pos1Stop || lastInput.intAt(pos1) != input.ints[pos2]) {
         break;
       }
@@ -786,7 +786,6 @@ public void add(IntsRef input, T output) throws IOException {
     // init tail states for current input
     for (int idx = prefixLenPlus1; idx <= input.length; idx++) {
       frontier[idx - 1].addArc(input.ints[input.offset + idx - 1], frontier[idx]);
-      frontier[idx].inputCount++;
     }
 
     final UnCompiledNode<T> lastNode = frontier[input.length];
@@ -835,8 +834,6 @@ public void add(IntsRef input, T output) throws IOException {
 
     // save last input
     lastInput.copyInts(input);
-
-    // System.out.println("  count[0]=" + frontier[0].inputCount);
   }
 
   private boolean validOutput(T output) {
@@ -866,6 +863,7 @@ public FST<T> compile() throws IOException {
 
   /** Expert: holds a pending (seen but not yet serialized) arc. */
   static class Arc<T> {
+
     int label; // really an "unsigned" byte
     Node target;
     boolean isFinal;
@@ -895,7 +893,13 @@ public boolean isCompiled() {
   }
 
   /** Expert: holds a pending (seen but not yet serialized) Node. */
-  static final class UnCompiledNode<T> implements Node {
+  static final class UnCompiledNode<T> implements Node, Accountable {
+
+    private static final long BASE_RAM_BYTES_USED =
+        RamUsageEstimator.shallowSizeOfInstance(UnCompiledNode.class);
+    private static final long BASE_ARC_RAM_BYTES_USED =
+        RamUsageEstimator.shallowSizeOfInstance(Arc.class);
+
     final FSTCompiler<T> owner;
     int numArcs;
     Arc<T>[] arcs;
@@ -906,10 +910,6 @@ static final class UnCompiledNode<T> implements Node {
     T output;
     boolean isFinal;
 
-    // TODO: remove this tracking?  we used to use it for confusingly pruning NodeHash, but
-    // we switched to LRU by RAM usage instead:
-    long inputCount;
-
     /** This node's depth, starting from the automaton root. */
     final int depth;
 
@@ -931,11 +931,79 @@ public boolean isCompiled() {
       return false;
     }
 
+    @Override
+    public int hashCode() {
+      final int PRIME = 31;
+      long h = 0;
+      // TODO: maybe if number of arcs is high we can safely subsample?
+      for (int arcIdx = 0; arcIdx < numArcs; arcIdx++) {
+        final FSTCompiler.Arc<T> arc = arcs[arcIdx];
+        h = PRIME * h + arc.label;
+        long n = ((FSTCompiler.CompiledNode) arc.target).node;
+        h = PRIME * h + (int) (n ^ (n >> 32));
+        h = PRIME * h + arc.output.hashCode();
+        h = PRIME * h + arc.nextFinalOutput.hashCode();
+        if (arc.isFinal) {
+          h += 17;
+        }
+      }
+
+      return Long.hashCode(h);
+    }
+
+    @Override
+    public boolean equals(Object obj) {
+      if ((obj instanceof FSTCompiler.UnCompiledNode) == false) {
+        return false;
+      }
+      @SuppressWarnings("unchecked")
+      UnCompiledNode<T> other = (UnCompiledNode<T>) obj;
+
+      if (numArcs != other.numArcs || isFinal != other.isFinal) {
+        return false;
+      }
+
+      for (int arcUpto = 0; arcUpto < numArcs; arcUpto++) {
+        final FSTCompiler.Arc<T> arc = arcs[arcUpto];
+        final FSTCompiler.Arc<T> otherArc = other.arcs[arcUpto];
+        if (arc.label != otherArc.label
+            || arc.output.equals(otherArc.output) == false
+            || ((FSTCompiler.CompiledNode) arc.target).node
+                != ((FSTCompiler.CompiledNode) otherArc.target).node
+            || arc.nextFinalOutput.equals(otherArc.nextFinalOutput) == false
+            || arc.isFinal != otherArc.isFinal) {
+          return false;
+        }
+      }
+
+      return true;
+    }
+
+    @Override
+    public String toString() {
+      String arcString = "";
+      for (int arcIndex = 0; arcIndex < numArcs; arcIndex++) {
+        Arc<T> arc = arcs[arcIndex];
+        arcString +=
+            "{label="
+                + arc.label
+                + ", output="
+                + arc.output
+                + ", isFinal="
+                + arc.isFinal
+                + ", nextFinalOutput="
+                + arc.nextFinalOutput
+                + ", node="
+                + (arc.target != null ? ((CompiledNode) arc.target).node : "")
+                + "},";
+      }
+      return "{numArcs=" + numArcs + ", isFinal=" + isFinal + ", arcs=[" + arcString + "]}";
+    }
+
     void clear() {
       numArcs = 0;
       isFinal = false;
       output = owner.NO_OUTPUT;
-      inputCount = 0;
 
       // We don't clear the depth here because it never changes
       // for nodes on the frontier (even when reused).
@@ -1009,6 +1077,31 @@ void prependOutput(T outputPrefix) {
         assert owner.validOutput(output);
       }
     }
+
+    @Override
+    public long ramBytesUsed() {
+      return BASE_RAM_BYTES_USED + BASE_ARC_RAM_BYTES_USED * numArcs;
+    }
+
+    @SuppressWarnings({"rawtypes", "unchecked"})
+    UnCompiledNode<T> copyToNewNode() {
+      FSTCompiler.UnCompiledNode<T> cloned = new FSTCompiler.UnCompiledNode<>(owner, depth);
+      cloned.numArcs = numArcs;
+      cloned.output = output;
+      cloned.isFinal = isFinal;
+
+      cloned.arcs = new FSTCompiler.Arc[numArcs];
+      for (int arcIndex = 0; arcIndex < numArcs; arcIndex++) {
+        FSTCompiler.Arc<T> nodeArc = arcs[arcIndex];
+        cloned.arcs[arcIndex] = new FSTCompiler.Arc<>();
+        cloned.arcs[arcIndex].target = nodeArc.target;
+        cloned.arcs[arcIndex].label = nodeArc.label;
+        cloned.arcs[arcIndex].output = nodeArc.output;
+        cloned.arcs[arcIndex].isFinal = nodeArc.isFinal;
+        cloned.arcs[arcIndex].nextFinalOutput = nodeArc.nextFinalOutput;
+      }
+      return cloned;
+    }
   }
 
   /**