content: /*
   * 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.cassandra.io.tries;
  
  import java.io.IOException;
  import java.io.PrintStream;
  import java.nio.ByteBuffer;
  import java.util.Arrays;
  import javax.annotation.concurrent.NotThreadSafe;
  
  import org.apache.cassandra.io.sstable.format.Version;
  import org.apache.cassandra.io.util.PageAware;
  import org.apache.cassandra.io.util.Rebufferer;
  import org.apache.cassandra.io.util.Rebufferer.BufferHolder;
  import org.apache.cassandra.utils.bytecomparable.ByteComparable;
  import org.apache.cassandra.utils.bytecomparable.ByteSource;
  import org.apache.lucene.util.ArrayUtil;
  
  /**
   * Thread-unsafe trie walking helper. This is analogous to {@link org.apache.cassandra.io.util.RandomAccessReader} for
   * tries -- takes an on-disk trie accessible via a supplied Rebufferer and lets user seek to nodes and work with them.
   * <p>
   * Assumes data was written using page-aware builder and thus no node crosses a page and thus a buffer boundary.
   * <p>
   * See {@code org/apache/cassandra/io/sstable/format/bti/BtiFormat.md} for a description of the mechanisms of writing
   * and reading an on-disk trie.
   */
  @NotThreadSafe
  public class Walker<CONCRETE extends Walker<CONCRETE>> implements AutoCloseable
  {
      /** Value used to indicate a branch (e.g. lesser/greaterBranch) does not exist. */
      public static int NONE = TrieNode.NONE;
  
      private final Rebufferer source;
      protected final long root;
  
      // State relating to current node.
      private BufferHolder bh;    // from Rebufferer
      private int offset;         // offset of current node within buf
      protected TrieNode nodeType;  // type of current node
      protected ByteBuffer buf;   // buffer containing the data
      protected long position;    // file position of current node
  
      // State relating to searches.
      protected long greaterBranch;
      protected long lesserBranch;
  
      // Version of the byte comparable conversion to use
      public static final ByteComparable.Version BYTE_COMPARABLE_VERSION = ByteComparable.Version.OSS50;
  
      /**
       * Creates a walker. Rebufferer must be aligned and with a buffer size that is at least 4k.
       */
      public Walker(Rebufferer source, long root)
      {
          this.source = source;
          this.root = root;
          try
          {
              bh = source.rebuffer(root);
              buf = bh.buffer();
          }
          catch (RuntimeException ex)
          {
              if (bh != null) bh.release();
              source.closeReader();
              throw ex;
          }
      }
  
      public void close()
      {
          bh.release();
          source.closeReader();
      }
  
      protected final void go(long position)
      {
          long curOffset = position - bh.offset();
          if (curOffset < 0 || curOffset >= buf.limit())
          {
              bh.release();
              bh = Rebufferer.EMPTY; // prevents double release if the call below fails
              bh = source.rebuffer(position);
              buf = bh.buffer();
              curOffset = position - bh.offset();
              assert curOffset >= 0 && curOffset < buf.limit() : String.format("Invalid offset: %d, buf: %s, bh: %s", curOffset, buf, bh);
          }
          this.offset = (int) curOffset;
          this.position = position;
          nodeType = TrieNode.at(buf, (int) curOffset);
      }
  
      protected final int payloadFlags()
      {
          return nodeType.payloadFlags(buf, offset);
      }
  
      protected final boolean hasPayload()
      {
          return payloadFlags() != 0;
      }
  
      protected final int payloadPosition()
      {
          return nodeType.payloadPosition(buf, offset);
      }
  
      protected final int search(int transitionByte)
      {
          return nodeType.search(buf, offset, transitionByte);
      }
  
      protected final long transition(int childIndex)
      {
          return nodeType.transition(buf, offset, position, childIndex);
      }
  
      protected final long lastTransition()
      {
          return nodeType.lastTransition(buf, offset, position);
      }
  
      protected final long greaterTransition(int searchIndex, long defaultValue)
      {
          return nodeType.greaterTransition(buf, offset, position, searchIndex, defaultValue);
      }
  
      protected final long lesserTransition(int searchIndex, long defaultValue)
      {
          return nodeType.lesserTransition(buf, offset, position, searchIndex, defaultValue);
      }
  
      protected final int transitionByte(int childIndex)
      {
          return nodeType.transitionByte(buf, offset, childIndex);
      }
  
      protected final int transitionRange()
      {
          return nodeType.transitionRange(buf, offset);
      }
  
      protected final boolean hasChildren()
      {
          return transitionRange() > 0;
      }
  
      protected final void goMax(long pos)
      {
          go(pos);
          while (true)
          {
              long lastChild = lastTransition();
              if (lastChild == NONE)
                  return;
              go(lastChild);
          }
      }
  
      protected final void goMin(long pos)
      {
          go(pos);
          while (true)
          {
              int payloadBits = payloadFlags();
              if (payloadBits > 0)
                  return;
  
              long firstChild = transition(0);
              if (firstChild == NONE)
                  return;
              go(firstChild);
          }
      }
  
      public interface Extractor<RESULT, VALUE>
      {
          RESULT extract(VALUE walker, int payloadPosition, int payloadFlags) throws IOException;
      }
  
      /**
       * Follows the given key while there are transitions in the trie for it.
       *
       * @return the first unmatched byte of the key, may be {@link ByteSource#END_OF_STREAM}
       */
      public int follow(ByteComparable key)
      {
          ByteSource stream = key.asComparableBytes(BYTE_COMPARABLE_VERSION);
          go(root);
          while (true)
          {
              int b = stream.next();
              int childIndex = search(b);
  
              if (childIndex < 0)
                  return b;
  
              go(transition(childIndex));
          }
      }
  
      /**
       * Follows the trie for a given key, remembering the closest greater branch.
       * On return the walker is positioned at the longest prefix that matches the input (with or without payload), and
       * min(greaterBranch) is the immediate greater neighbour.
       *
       * @return the first unmatched byte of the key, may be {@link ByteSource#END_OF_STREAM}
       */
      public int followWithGreater(ByteComparable key)
      {
          greaterBranch = NONE;
  
          ByteSource stream = key.asComparableBytes(BYTE_COMPARABLE_VERSION);
          go(root);
          while (true)
          {
              int b = stream.next();
              int searchIndex = search(b);
  
              greaterBranch = greaterTransition(searchIndex, greaterBranch);
              if (searchIndex < 0)
                  return b;
  
              go(transition(searchIndex));
          }
      }
  
      /**
       * Follows the trie for a given key, remembering the closest lesser branch.
       * On return the walker is positioned at the longest prefix that matches the input (with or without payload), and
       * max(lesserBranch) is the immediate lesser neighbour.
       *
       * @return the first unmatched byte of the key, may be {@link ByteSource#END_OF_STREAM}
       */
      public int followWithLesser(ByteComparable key)
      {
          lesserBranch = NONE;
  
          ByteSource stream = key.asComparableBytes(BYTE_COMPARABLE_VERSION);
          go(root);
          while (true)
          {
              int b = stream.next();
              int searchIndex = search(b);
  
              lesserBranch = lesserTransition(searchIndex, lesserBranch);
  
              if (searchIndex < 0)
                  return b;
  
              go(transition(searchIndex));
          }
      }
  
  
      /**
       * Takes a prefix of the given key. The prefix is in the sense of a separator key match, i.e. it is only
       * understood as valid if there are no greater entries in the trie (e.g. data at 'a' is ignored if 'ab' or 'abba'
       * is in the trie when looking for 'abc' or 'ac', but accepted when looking for 'aa').
       * In order to not have to go back to data that may have exited cache, payloads are extracted when the node is
       * visited (instead of saving the node's position), which requires an extractor to be passed as parameter.
       * @throws IOException 
       */
      @SuppressWarnings("unchecked")
      public <RESULT> RESULT prefix(ByteComparable key, Extractor<RESULT, CONCRETE> extractor) throws IOException
      {
          RESULT payload = null;
  
          ByteSource stream = key.asComparableBytes(BYTE_COMPARABLE_VERSION);
          go(root);
          while (true)
          {
              int b = stream.next();
              int childIndex = search(b);
  
              if (childIndex > 0)
                  payload = null;
              else
              {
                  int payloadBits = payloadFlags();
                  if (payloadBits > 0)
                      payload = extractor.extract((CONCRETE) this, payloadPosition(), payloadBits);
                  if (childIndex < 0)
                      return payload;
              }
  
              go(transition(childIndex));
          }
      }
  
      /**
       * Follows the trie for a given key, taking a prefix (in the sense above) and searching for neighboring values.
       * On return min(greaterBranch) and max(lesserBranch) are the immediate non-prefix neighbours for the sought value.
       * <p>
       * Note: in a separator trie the closest smaller neighbour can be another prefix of the given key. This method
       * does not take that into account. E.g. if trie contains "abba", "as" and "ask", looking for "asking" will find
       * "ask" as the match, but max(lesserBranch) will point to "abba" instead of the correct "as". This problem can
       * only occur if there is a valid prefix match.
       * @throws IOException 
       */
      @SuppressWarnings("unchecked")
      public <RESULT> RESULT prefixAndNeighbours(ByteComparable key, Extractor<RESULT, CONCRETE> extractor) throws IOException
      {
          RESULT payload = null;
          greaterBranch = NONE;
          lesserBranch = NONE;
  
          ByteSource stream = key.asComparableBytes(BYTE_COMPARABLE_VERSION);
          go(root);
          while (true)
          {
              int b = stream.next();
              int searchIndex = search(b);
  
              greaterBranch = greaterTransition(searchIndex, greaterBranch);
  
              if (searchIndex == -1 || searchIndex == 0)
              {
                  int payloadBits = payloadFlags();
                  if (payloadBits > 0)
                      payload = extractor.extract((CONCRETE) this, payloadPosition(), payloadBits);
              }
              else
              {
                  lesserBranch = lesserTransition(searchIndex, lesserBranch);
                  payload = null;
              }
  
              if (searchIndex < 0)
                  return payload;
  
              go(transition(searchIndex));
          }
      }
  
      public ByteComparable getMaxTerm()
      {
          TransitionBytesCollector collector = new TransitionBytesCollector();
          go(root);
          while (true)
          {
              int lastIdx = transitionRange() - 1;
              long lastChild = transition(lastIdx);
              if (lastIdx < 0)
              {
                  return collector.toByteComparable();
              }
              collector.add(transitionByte(lastIdx));
              go(lastChild);
          }
      }
  
      public ByteComparable getMinTerm()
      {
          TransitionBytesCollector collector = new TransitionBytesCollector();
          go(root);
          while (true)
          {
              if (hasPayload())
              {
                  return collector.toByteComparable();
              }
              collector.add(transitionByte(0));
              go(transition(0));
          }
      }
  
      /**
       * To be used only in analysis.
       */
      protected int nodeTypeOrdinal()
      {
          return nodeType.ordinal;
      }
  
      /**
       * To be used only in analysis.
       */
      protected int nodeSize()
      {
          return payloadPosition() - offset;
      }
  
      public interface PayloadToString
      {
          String payloadAsString(ByteBuffer buf, int payloadPos, int payloadFlags, Version version) throws IOException;
      }
  
      public void dumpTrie(PrintStream out, PayloadToString payloadReader, Version version) throws IOException
      {
          out.print("ROOT");
          dumpTrie(out, payloadReader, root, "", version);
      }
  
      private void dumpTrie(PrintStream out, PayloadToString payloadReader, long node, String indent, Version version) throws IOException
      {
          go(node);
          int bits = payloadFlags();
          out.format(" %s@%x %s%n", nodeType.toString(), node, bits == 0 ? "" : payloadReader.payloadAsString(buf, payloadPosition(), bits, version));
          int range = transitionRange();
          for (int i = 0; i < range; ++i)
          {
              long child = transition(i);
              if (child == NONE)
                  continue;
              out.format("%s%02x %s>", indent, transitionByte(i), PageAware.pageStart(position) == PageAware.pageStart(child) ? "--" : "==");
              dumpTrie(out, payloadReader, child, indent + "  ", version);
              go(node);
          }
      }
  
      @Override
      public String toString()
      {
          return String.format("[Trie Walker - NodeType: %s, source: %s, buffer: %s, buffer file offset: %d, Node buffer offset: %d, Node file position: %d]",
                               nodeType, source, buf, bh.offset(), offset, position);
      }
  
      public static class TransitionBytesCollector
      {
          protected byte[] bytes = new byte[32];
          protected int pos = 0;
  
          public void add(int b)
          {
              if (pos == bytes.length)
              {
                  bytes = ArrayUtil.grow(bytes, pos + 1);
              }
              bytes[pos++] = (byte) b;
          }
  
          public void pop()
          {
              assert pos >= 0;
              pos--;
          }
  
          public ByteComparable toByteComparable()
          {
              if (pos <= 0)
                  return null;
              byte[] value = new byte[pos];
              System.arraycopy(bytes, 0, value, 0, pos);
              return v -> ByteSource.fixedLength(value, 0, value.length);
          }
  
          @Override
          public String toString()
          {
              return String.format("[Bytes %s, pos %d]", Arrays.toString(bytes), pos);
          }
      }
  }
  