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.utils;
  
  import java.util.ArrayList;
  import java.util.Comparator;
  import java.util.Iterator;
  import java.util.List;
  import java.util.Map;
  import java.util.function.Consumer;
  import java.util.function.Function;
  import java.util.function.Predicate;
  import java.util.stream.Stream;
  import java.util.stream.StreamSupport;
  import javax.annotation.CheckForNull;
  
  import com.google.common.collect.AbstractIterator;
  
  public class RTree<Token, Range, Value> implements RangeTree<Token, Range, Value>
  {
      /**
       * Tuning size target can be tricky as it is based on expected access patterns and expected matche sizes.  There is also
       * a memory cost to account for as large tree sizes will have far more nodes with a small target than a large target.
       *
       * If matching most of the data then larger sizes leads to fewer hops
       * If matching few elements then tree depth maters the most, if walking a long tree is more costly than walking the
       * element list, then shrinking depth (by having larger size target) can improve performance.
       */
      private static final int DEFAULT_SIZE_TARGET = 1 << 7;
      private static final int DEFAULT_NUMBER_OF_CHILDREN = 6;
  
      private final Comparator<Token> comparator;
      private final Accessor<Token, Range> accessor;
      private final int sizeTarget;
      private final int numChildren;
      private Node node = new Node();
  
      public RTree(Comparator<Token> comparator, Accessor<Token, Range> accessor)
      {
          this(comparator, accessor, DEFAULT_SIZE_TARGET, DEFAULT_NUMBER_OF_CHILDREN);
      }
  
      public RTree(Comparator<Token> comparator, Accessor<Token, Range> accessor, int sizeTarget, int numChildren)
      {
          if (sizeTarget <= 1)
              throw new IllegalArgumentException("size target must be 2 or more");
          if (numChildren <= 1)
              throw new IllegalArgumentException("Number of children must be 2 or more");
          if (sizeTarget < numChildren)
              throw new IllegalArgumentException("Size target (" + sizeTarget + ") was less than number of children (" + numChildren + ")");
          this.comparator = comparator;
          this.accessor = accessor;
          this.sizeTarget = sizeTarget;
          this.numChildren = numChildren;
      }
  
      public static <Token extends Comparable<? super Token>, Range, Value> RTree<Token, Range, Value> create(Accessor<Token, Range> accessor)
      {
          return new RTree<>(Comparator.naturalOrder(), accessor);
      }
  
      @Override
      public List<Value> get(Range range)
      {
          List<Value> matches = new ArrayList<>();
          get(range, e -> matches.add(e.getValue()));
          return matches;
      }
  
      @Override
      public void get(Range range, Consumer<Map.Entry<Range, Value>> onMatch)
      {
          node.search(range, onMatch, e -> e.getKey().equals(range), Function.identity());
      }
  
      @Override
      public List<Map.Entry<Range, Value>> search(Range range)
      {
          List<Map.Entry<Range, Value>> matches = new ArrayList<>();
          search(range, matches::add);
          return matches;
      }
  
      @Override
      public void search(Range range, Consumer<Map.Entry<Range, Value>> onMatch)
      {
          node.search(range, onMatch, ignore -> true, Function.identity());
      }
  
      public List<Value> find(Range range)
      {
          List<Value> matches = new ArrayList<>();
          find(range, matches::add);
          return matches;
      }
  
      public void find(Range range, Consumer<Value> onMatch)
      {
          node.search(range, onMatch, ignore -> true, Map.Entry::getValue);
      }
  
      @Override
      public List<Map.Entry<Range, Value>> searchToken(Token token)
      {
          List<Map.Entry<Range, Value>> matches = new ArrayList<>();
          searchToken(token, matches::add);
          return matches;
      }
  
      @Override
      public void searchToken(Token token, Consumer<Map.Entry<Range, Value>> onMatch)
      {
          node.searchToken(token, onMatch, ignore -> true, Function.identity());
      }
  
      public List<Value> findToken(Token token)
      {
          List<Value> matches = new ArrayList<>();
          findToken(token, matches::add);
          return matches;
      }
  
      public void findToken(Token token, Consumer<Value> onMatch)
      {
          node.searchToken(token, onMatch, ignore -> true, Map.Entry::getValue);
      }
  
      @Override
      public boolean add(Range key, Value value)
      {
          node.add(key, value);
          return true;
      }
  
      @Override
      public int remove(Range key)
      {
          return node.removeIf(e -> e.getKey().equals(key));
      }
  
      public int remove(Range key, Value value)
      {
          Map.Entry<Range, Value> match = Map.entry(key, value);
          return node.removeIf(match::equals);
      }
  
      @Override
      public void clear()
      {
          node = new Node();
      }
  
      @Override
      public int size()
      {
          return node.size;
      }
  
      @Override
      public boolean isEmpty()
      {
          return node.size == 0;
      }
  
      public String displayTree()
      {
          StringBuilder sb = new StringBuilder();
          node.displayTree(0, sb);
          return sb.toString();
      }
  
      @Override
      public Iterator<Map.Entry<Range, Value>> iterator()
      {
          return node.iterator();
      }
  
      @Override
      public Stream<Map.Entry<Range, Value>> stream()
      {
          return StreamSupport.stream(spliterator(), false);
      }
  
      private class Node implements Iterable<Map.Entry<Range, Value>>
      {
          private List<Map.Entry<Range, Value>> values = new ArrayList<>();
          private List<Node> children = null;
          private int size = 0;
          private Token minStart, maxStart, minEnd, maxEnd;
  
          int removeIf(Predicate<Map.Entry<Range, Value>> condition)
          {
              if (minStart == null)
                  return 0;
              if (children != null)
              {
                  int sum = 0;
                  for (Node node : children)
                      sum += node.removeIf(condition);
                  size -= sum;
                  return sum;
              }
              class Counter {int value;}
              Counter counter = new Counter();
              values.removeIf(e -> {
                  if (condition.test(e))
                  {
                      counter.value++;
                      return true;
                  }
                  return false;
              });
              size -= counter.value;
              if (values.isEmpty())
                  minStart = maxStart = minEnd = maxEnd = null;
              return counter.value;
          }
  
          void add(Range range, Value value)
          {
              size++;
              if (minStart == null)
              {
                  minStart = maxStart = accessor.start(range);
                  minEnd = maxEnd = accessor.end(range);
              }
              else
              {
                  Token start = accessor.start(range);
                  minStart = min(minStart, start);
                  maxStart = max(maxStart, start);
                  Token end = accessor.end(range);
                  minEnd = min(minEnd, end);
                  maxEnd = max(maxEnd, end);
              }
              if (children != null)
              {
                  findBestMatch(range).add(range, value);
                  return;
              }
              values.add(new MutableEntry(range, value));
              if (shouldSplit())
                  split();
          }
  
          private Node findBestMatch(Range range)
          {
              int topIdx = 0;
              Node node = children.get(0);
              int topScore = node.score(range);
              int size = node.size;
              for (int i = 1; i < children.size(); i++)
              {
                  node = children.get(i);
                  int score = node.score(range);
                  if (score > topScore || (score == topScore && size > node.size))
                  {
                      topIdx = i;
                      size = node.size;
                  }
              }
              return children.get(topIdx);
          }
  
          private int score(Range range)
          {
              if (minStart == null)
                  return 0;
              if (!intersects(range))
                  return -10;
              int score = 5; // overlapps
              if (values != null) // is leaf
                  score += 5;
  
              int startScore = 0;
              if (comparator.compare(maxStart, accessor.start(range)) <= 0)
                  startScore += 10;
              else if (comparator.compare(minStart, accessor.start(range)) <= 0)
                  startScore += 5;
  
              int endScore = 0;
              if (comparator.compare(minEnd, accessor.end(range)) >= 0)
                  endScore += 10;
              else if (comparator.compare(maxEnd, accessor.end(range)) >= 0)
                  endScore += 5;
              // if fully contained, then add the scores: 10 for largest bounds, 20 for smallest bounds
              if (!(startScore == 0 || endScore == 0))
                  score += startScore + endScore;
              return score;
          }
  
          boolean shouldSplit()
          {
              return values.size() > sizeTarget
                     // if the same range is used over and over again, splitting doesn't do much
                     && !(comparator.compare(minStart, maxStart) == 0
                          && comparator.compare(minEnd, maxEnd) == 0);
          }
  
          List<List<Map.Entry<Range, Value>>> partitionByEnd()
          {
              List<Token> allEndpoints = new ArrayList<>(values.size() * 2);
              for (Map.Entry<Range, Value> a : values)
              {
                  allEndpoints.add(accessor.start(a.getKey()));
                  allEndpoints.add(accessor.end(a.getKey()));
              }
              allEndpoints.sort(comparator);
              List<Token> maxToken = new ArrayList<>(numChildren);
              int tick = allEndpoints.size() / numChildren;
              int offset = tick;
              for (int i = 0; i < numChildren; i++)
              {
                  maxToken.add(allEndpoints.get(offset));
                  offset += tick;
                  if (offset >= allEndpoints.size())
                  {
                      maxToken.add(allEndpoints.get(allEndpoints.size() - 1));
                      break;
                  }
              }
  
              List<List<Map.Entry<Range, Value>>> partitions = new ArrayList<>(numChildren);
              for (int i = 0; i < numChildren; i++)
                  partitions.add(new ArrayList<>());
  
              for (Map.Entry<Range, Value> a : values)
              {
                  Token end = accessor.end(a.getKey());
                  List<Map.Entry<Range, Value>> selected = null;
                  for (int i = 0; i < numChildren; i++)
                  {
                      if (comparator.compare(end, maxToken.get(i)) < 0)
                      {
                          selected = partitions.get(i);
                          break;
                      }
                  }
                  if (selected == null)
                      selected = partitions.get(partitions.size() - 1);
                  selected.add(a);
              }
              int[] sizes = partitions.stream().mapToInt(List::size).toArray();
              return goodEnough(sizes) ? partitions : null;
          }
  
          private boolean goodEnough(int[] sizes)
          {
              double sum = 0.0;
              for (int i : sizes)
                  sum += i;
              double mean = sum / sizes.length;
              double stddev = 0.0;
              for (int i : sizes)
                  stddev += Math.pow(i - mean, 2);
              stddev = Math.sqrt(stddev / sizes.length);
              return stddev < 1.5;
          }
  
          void split()
          {
              children = new ArrayList<>(numChildren);
              for (int i = 0; i < numChildren; i++)
                  children.add(new Node());
  
              List<List<Map.Entry<Range, Value>>> partitions = partitionByEnd();
              if (partitions == null)
                  partitions = partitionEven();
              for (int i = 0; i < children.size(); i++)
              {
                  Node c = children.get(i);
                  List<Map.Entry<Range, Value>> entries = partitions.get(i);
                  entries.forEach(e -> c.add(e.getKey(), e.getValue()));
              }
  
              values.clear();
              values = null;
          }
  
          private List<List<Map.Entry<Range, Value>>> partitionEven()
          {
              values.sort((a, b) -> {
                  Range left = a.getKey();
                  Range right = b.getKey();
                  int rc = comparator.compare(accessor.start(left), accessor.start(right));
                  if (rc == 0)
                      rc = comparator.compare(accessor.end(left), accessor.end(right));
                  return rc;
              });
              List<List<Map.Entry<Range, Value>>> partition = new ArrayList<>(numChildren);
              int size = Math.max(1, values.size() / numChildren);
              int offset = 0;
              for (int i = 0; i < numChildren - 1; i++)
              {
                  int total = size;
                  partition.add(new ArrayList<>(values.subList(offset, offset + total)));
                  offset += total;
              }
              partition.add(new ArrayList<>(values.subList(offset, values.size())));
              return partition;
          }
  
          <T> void search(Range range, Consumer<T> matches, Predicate<Map.Entry<Range, Value>> predicate, Function<Map.Entry<Range, Value>, T> transformer)
          {
              if (minStart == null)
                  return;
              if (!intersects(range))
                  return;
              if (children != null)
              {
                  children.forEach(n -> n.search(range, matches, predicate, transformer));
                  return;
              }
              values.forEach(e -> {
                  if (accessor.intersects(e.getKey(), range) && predicate.test(e))
                      matches.accept(transformer.apply(e));
              });
          }
  
          <T> void searchToken(Token token, Consumer<T> matches, Predicate<Map.Entry<Range, Value>> predicate, Function<Map.Entry<Range, Value>, T> transformer)
          {
              if (minStart == null)
                  return;
              if (!contains(minStart, maxEnd, token))
                  return;
              if (children != null)
              {
                  for (int i = 0, size = children.size(); i < size; i++)
                  {
                      Node node = children.get(i);
                      node.searchToken(token, matches, predicate, transformer);
                  }
                  return;
              }
              values.forEach(e -> {
                  if (accessor.contains(e.getKey(), token) && predicate.test(e))
                      matches.accept(transformer.apply(e));
              });
          }
  
          boolean intersects(Range range)
          {
              return accessor.intersects(range, minStart, maxEnd);
          }
  
          boolean contains(Token start, Token end, Token value)
          {
              return accessor.contains(start, end, value);
          }
  
          private void displayTree(int level, StringBuilder sb)
          {
              for (int i = 0; i < level; i++)
                  sb.append('\t');
              sb.append("start:(").append(minStart).append(", ").append(maxStart).append("), end:(").append(minEnd).append(", ").append(maxEnd).append("):");
              if (children != null)
              {
                  sb.append('\n');
                  children.forEach(n -> n.displayTree(level + 1, sb));
              }
              else
              {
                  sb.append(' ').append(size).append('\n');
              }
          }
  
          @Override
          public String toString()
          {
              return "Node{" +
                     "minStart=" + minStart +
                     ", maxStart=" + maxStart +
                     ", minEnd=" + minEnd +
                     ", maxEnd=" + maxEnd +
                     ", values=" + values +
                     ", children=" + children +
                     '}';
          }
  
          private Token min(Token a, Token b)
          {
              return comparator.compare(a, b) < 0 ? a : b;
          }
  
          private Token max(Token a, Token b)
          {
              return comparator.compare(a, b) < 0 ? b : a;
          }
  
          @Override
          public Iterator<Map.Entry<Range, Value>> iterator()
          {
              if (values != null)
                  return values.iterator();
              return new AbstractIterator<>()
              {
                  private int index = 0;
                  private Iterator<Map.Entry<Range, Value>> it = null;
                  @CheckForNull
                  @Override
                  protected Map.Entry<Range, Value> computeNext()
                  {
                      while (true)
                      {
                          if (it == null)
                          {
                              if (index == children.size())
                                  return endOfData();
                              it = children.get(index++).iterator();
                          }
                          if (it.hasNext())
                              return it.next();
                          it = null;
                      }
                  }
              };
          }
      }
  }
  