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.db.tries;
  
  import java.util.ArrayList;
  import java.util.Collection;
  import java.util.List;
  
  import com.google.common.collect.Iterables;
  
  /**
   * A merged view of multiple tries.
   *
   * This is accomplished by walking the cursors in parallel; the merged cursor takes the position and features of the
   * smallest and advances with it; when multiple cursors are equal, all of them are advanced. The ordered view of the
   * cursors is maintained using a custom binary min-heap, built for efficiently reforming the heap when the top elements
   * are advanced (see {@link CollectionMergeCursor}).
   *
   * Crucial for the efficiency of this is the fact that when they are advanced like this, we can compare cursors'
   * positions by their depth descending and then incomingTransition ascending.
   *
   * See Trie.md for further details.
   */
  class CollectionMergeTrie<T> extends Trie<T>
  {
      private final CollectionMergeResolver<T> resolver;  // only called on more than one input
      protected final Collection<? extends Trie<T>> inputs;
  
      CollectionMergeTrie(Collection<? extends Trie<T>> inputs, CollectionMergeResolver<T> resolver)
      {
          this.resolver = resolver;
          this.inputs = inputs;
      }
  
      @Override
      protected Cursor<T> cursor(Direction direction)
      {
          return new CollectionMergeCursor<>(resolver, direction, inputs);
      }
  
      /**
       * Compare the positions of two cursors. One is before the other when
       * - its depth is greater, or
       * - its depth is equal, and the incoming transition is smaller.
       */
      static <T> boolean greaterCursor(Direction direction, Cursor<T> c1, Cursor<T> c2)
      {
          int c1depth = c1.depth();
          int c2depth = c2.depth();
          if (c1depth != c2depth)
              return c1depth < c2depth;
          return direction.lt(c2.incomingTransition(), c1.incomingTransition());
      }
  
      static <T> boolean equalCursor(Cursor<T> c1, Cursor<T> c2)
      {
          return c1.depth() == c2.depth() && c1.incomingTransition() == c2.incomingTransition();
      }
  
      /*
       * The merge cursor is a variation of the idea of a merge iterator with one key observation: because we advance
       * the source iterators together, we can compare them just by depth and incoming transition.
       *
       * The most straightforward way to implement merging of iterators is to use a {@code PriorityQueue},
       * {@code poll} it to find the next item to consume, then {@code add} the iterator back after advancing.
       * This is not very efficient as {@code poll} and {@code add} in all cases require at least
       * {@code log(size)} comparisons and swaps (usually more than {@code 2*log(size)}) per consumed item, even
       * if the input is suitable for fast iteration.
       *
       * The implementation below makes use of the fact that replacing the top element in a binary heap can be
       * done much more efficiently than separately removing it and placing it back, especially in the cases where
       * the top iterator is to be used again very soon (e.g. when there are large sections of the output where
       * only a limited number of input iterators overlap, which is normally the case in many practically useful
       * situations, e.g. levelled compaction).
       *
       * The implementation builds and maintains a binary heap of sources (stored in an array), where we do not
       * add items after the initial construction. Instead we advance the smallest element (which is at the top
       * of the heap) and push it down to find its place for its new position. Should this source be exhausted,
       * we swap it with the last source in the heap and proceed by pushing that down in the heap.
       *
       * In the case where we have multiple sources with matching positions, the merging algorithm
       * must be able to merge all equal values. To achieve this {@code content} walks the heap to
       * find all equal cursors without advancing them, and separately {@code advance} advances
       * all equal sources and restores the heap structure.
       *
       * The latter is done equivalently to the process of initial construction of a min-heap using back-to-front
       * heapification as done in the classic heapsort algorithm. It only needs to heapify subheaps whose top item
       * is advanced (i.e. one whose position matches the current), and we can do that recursively from
       * bottom to top. Should a source be exhausted when advancing, it can be thrown away by swapping in the last
       * source in the heap (note: we must be careful to advance that source too if required).
       *
       * To make it easier to advance efficienty in single-sourced branches of tries, we extract the current smallest
       * cursor (the head) separately, and start any advance with comparing that to the heap's first. When the smallest
       * cursor remains the same (e.g. in branches coming from a single source) this makes it possible to advance with
       * just one comparison instead of two at the expense of increasing the number by one in the general case.
       *
       * Note: This is a simplification of the MergeIterator code from CASSANDRA-8915, without the leading ordered
       * section and equalParent flag since comparisons of cursor positions are cheap.
       */
      static class CollectionMergeCursor<T> implements Cursor<T>
      {
          private final CollectionMergeResolver<T> resolver;
          private final Direction direction;
  
          /**
           * The smallest cursor, tracked separately to improve performance in single-source sections of the trie.
           */
          private Cursor<T> head;
  
          /**
           * Binary heap of the remaining cursors. The smallest element is at position 0.
           * Every element i is smaller than or equal to its two children, i.e.
           *     heap[i] <= heap[i*2 + 1] && heap[i] <= heap[i*2 + 2]
           */
          private final Cursor<T>[] heap;
  
          /**
           * A list used to collect contents during content() calls.
           */
          private final List<T> contents;
  
          public CollectionMergeCursor(CollectionMergeResolver<T> resolver, Direction direction, Collection<? extends Trie<T>> inputs)
          {
              this.resolver = resolver;
              this.direction = direction;
              int count = inputs.size();
              // Get cursors for all inputs. Put one of them in head and the rest in the heap.
              heap = new Cursor[count - 1];
              contents = new ArrayList<>(count);
              int i = -1;
              for (Trie<T> trie : inputs)
              {
                  Cursor<T> cursor = trie.cursor(direction);
                  assert cursor.depth() == 0;
                  if (i >= 0)
                      heap[i] = cursor;
                  else
                      head = cursor;
                  ++i;
              }
              // The cursors are all currently positioned on the root and thus in valid heap order.
          }
  
          /**
           * Interface for internal operations that can be applied to the equal top elements of the heap.
           */
          interface HeapOp<T>
          {
              void apply(CollectionMergeCursor<T> self, Cursor<T> cursor, int index);
          }
  
          /**
           * Apply a non-interfering operation, i.e. one that does not change the cursor state, to all inputs in the heap
           * that are on equal position to the head.
           * For interfering operations like advancing the cursors, use {@link #advanceEqualAndRestoreHeap(AdvancingHeapOp)}.
           */
          private void applyToEqualOnHeap(HeapOp<T> action)
          {
              applyToEqualElementsInHeap(action, 0);
          }
  
          /**
           * Interface for internal advancing operations that can be applied to the heap cursors. This interface provides
           * the code to restore the heap structure after advancing the cursors.
           */
          interface AdvancingHeapOp<T> extends HeapOp<T>
          {
              void apply(Cursor<T> cursor);
  
              default void apply(CollectionMergeCursor<T> self, Cursor<T> cursor, int index)
              {
                  // Apply the operation, which should advance the position of the element.
                  apply(cursor);
  
                  // This method is called on the back path of the recursion. At this point the heaps at both children are
                  // advanced and well-formed.
                  // Place current node in its proper position.
                  self.heapifyDown(cursor, index);
                  // The heap rooted at index is now advanced and well-formed.
              }
          }
  
  
          /**
           * Advance the state of all inputs in the heap that are on equal position as the head and restore the heap
           * invariant.
           */
          private void advanceEqualAndRestoreHeap(AdvancingHeapOp<T> action)
          {
              applyToEqualElementsInHeap(action, 0);
          }
  
          /**
           * Apply an operation to all elements on the heap that are equal to the head. Descends recursively in the heap
           * structure to all equal children and applies the operation on the way back.
           *
           * This operation can be something that does not change the cursor state (see {@link #content}) or an operation
           * that advances the cursor to a new state, wrapped in a {@link AdvancingHeapOp} ({@link #advance} or
           * {@link #skipChildren}). The latter interface takes care of pushing elements down in the heap after advancing
           * and restores the subheap state on return from each level of the recursion.
           */
          private void applyToEqualElementsInHeap(HeapOp<T> action, int index)
          {
              if (index >= heap.length)
                  return;
              Cursor<T> item = heap[index];
              if (!equalCursor(item, head))
                  return;
  
              // If the children are at the same position, they also need advancing and their subheap
              // invariant to be restored.
              applyToEqualElementsInHeap(action, index * 2 + 1);
              applyToEqualElementsInHeap(action, index * 2 + 2);
  
              // Apply the action. This is done on the reverse direction to give the action a chance to form proper
              // subheaps and combine them on processing the parent.
              action.apply(this, item, index);
          }
  
          /**
           * Push the given state down in the heap from the given index until it finds its proper place among
           * the subheap rooted at that position.
           */
          private void heapifyDown(Cursor<T> item, int index)
          {
              while (true)
              {
                  int next = index * 2 + 1;
                  if (next >= heap.length)
                      break;
                  // Select the smaller of the two children to push down to.
                  if (next + 1 < heap.length && greaterCursor(direction, heap[next], heap[next + 1]))
                      ++next;
                  // If the child is greater or equal, the invariant has been restored.
                  if (!greaterCursor(direction, item, heap[next]))
                      break;
                  heap[index] = heap[next];
                  index = next;
              }
              heap[index] = item;
          }
  
          /**
           * Check if the head is greater than the top element in the heap, and if so, swap them and push down the new
           * top until its proper place.
           * @param headDepth the depth of the head cursor (as returned by e.g. advance).
           * @return the new head element's depth
           */
          private int maybeSwapHead(int headDepth)
          {
              int heap0Depth = heap[0].depth();
              if (headDepth > heap0Depth ||
                  (headDepth == heap0Depth && direction.le(head.incomingTransition(), heap[0].incomingTransition())))
                  return headDepth;   // head is still smallest
  
              // otherwise we need to swap heap and heap[0]
              Cursor<T> newHeap0 = head;
              head = heap[0];
              heapifyDown(newHeap0, 0);
              return heap0Depth;
          }
  
          @Override
          public int advance()
          {
              advanceEqualAndRestoreHeap(Cursor::advance);
              return maybeSwapHead(head.advance());
          }
  
          @Override
          public int advanceMultiple(TransitionsReceiver receiver)
          {
              // If the current position is present in just one cursor, we can safely descend multiple levels within
              // its branch as no one of the other tries has content for it.
              if (equalCursor(heap[0], head))
                  return advance();   // More than one source at current position, do single-step advance.
  
              // If there are no children, i.e. the cursor ascends, we have to check if it's become larger than some
              // other candidate.
              return maybeSwapHead(head.advanceMultiple(receiver));
          }
  
          @Override
          public int skipChildren()
          {
              advanceEqualAndRestoreHeap(Cursor::skipChildren);
              return maybeSwapHead(head.skipChildren());
          }
  
          @Override
          public int depth()
          {
              return head.depth();
          }
  
          @Override
          public int incomingTransition()
          {
              return head.incomingTransition();
          }
  
          @Override
          public T content()
          {
              applyToEqualOnHeap(CollectionMergeCursor::collectContent);
              collectContent(head, -1);
  
              T toReturn;
              switch (contents.size())
              {
                  case 0:
                      toReturn = null;
                      break;
                  case 1:
                      toReturn = contents.get(0);
                      break;
                  default:
                      toReturn = resolver.resolve(contents);
                      break;
              }
              contents.clear();
              return toReturn;
          }
  
          private void collectContent(Cursor<T> item, int index)
          {
              T itemContent = item.content();
              if (itemContent != null)
                  contents.add(itemContent);
          }
      }
  
      /**
       * Special instance for sources that are guaranteed distinct. The main difference is that we can form unordered
       * value list by concatenating sources.
       */
      static class Distinct<T> extends CollectionMergeTrie<T>
      {
          Distinct(Collection<? extends Trie<T>> inputs)
          {
              super(inputs, throwingResolver());
          }
  
          @Override
          public Iterable<T> valuesUnordered()
          {
              return Iterables.concat(Iterables.transform(inputs, Trie::valuesUnordered));
          }
      }
  }
  