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.btree;
  
  import java.util.ArrayList;
  import java.util.Arrays;
  import java.util.Collections;
  import java.util.Comparator;
  import java.util.List;
  import java.util.Random;
  import java.util.Set;
  import java.util.TreeSet;
  import java.util.stream.Collectors;
  
  import org.junit.Assert;
  import org.junit.Test;
  
  import accord.utils.DefaultRandom;
  import accord.utils.Invariants;
  import accord.utils.QuadFunction;
  import accord.utils.RandomSource;
  import accord.utils.SymmetricComparator;
  import org.apache.cassandra.config.CassandraRelevantProperties;
  
  import static org.apache.cassandra.utils.btree.BTree.getChildCount;
  import static org.apache.cassandra.utils.btree.BTree.getChildStart;
  import static org.apache.cassandra.utils.btree.BTree.getKeyEnd;
  import static org.apache.cassandra.utils.btree.BTree.isLeaf;
  import static org.apache.cassandra.utils.btree.IntervalBTree.InclusiveEndHelper.endWithStart;
  import static org.apache.cassandra.utils.btree.IntervalBTree.InclusiveEndHelper.startWithEnd;
  import static org.apache.cassandra.utils.btree.IntervalBTree.InclusiveEndHelper.startWithStart;
  
  public class IntervalBTreeTest
  {
      static
      {
          // build deeper trees to explore more behaviours
          CassandraRelevantProperties.BTREE_BRANCH_SHIFT.setInt(2);
      }
  
      static class TestInterval implements Comparable<TestInterval>
      {
          final int start, end;
          final int value;
  
          TestInterval(int start, int end, int value)
          {
              this.start = start;
              this.end = end;
              this.value = value;
          }
  
          @Override
          public String toString()
          {
              return start + "," + end + ':' + value;
          }
  
          @Override
          public int compareTo(TestInterval that)
          {
              int c = Integer.compare(this.start, that.start);
              if (c == 0) c = Integer.compare(this.end, that.end);
              if (c == 0) c = Integer.compare(this.value, that.value);
              return c;
          }
      }
  
      static class TestComparators implements IntervalBTree.IntervalComparators<TestInterval>
      {
          static final TestComparators INSTANCE = new TestComparators();
  
          @Override public Comparator<TestInterval> totalOrder() { return TestInterval::compareTo; }
          @Override public SymmetricComparator<TestInterval> endWithEndSorter() { return (a, b) -> Integer.compare(a.end, b.end); }
          @Override public SymmetricComparator<TestInterval> startWithStartSeeker() { return (a, b) -> startWithStart(Integer.compare(a.start, b.start)); }
          @Override public SymmetricComparator<TestInterval> startWithEndSeeker() { return (a, b) -> startWithEnd(Integer.compare(a.start, b.end)); }
          @Override public SymmetricComparator<TestInterval> endWithStartSeeker() { return (a, b) -> endWithStart(Integer.compare(a.end, b.start)); }
      }
  
      @Test
      public void testN()
      {
          testOne(-415085593308080411L);
          Random seeds = new Random();
          for (int i = 0 ; i < 200 ; ++i)
              testOne(seeds.nextLong());
      }
  
      // TODO (expected): test extreme patterns
      public static void testOne(long seed)
      {
          try
          {
              List<TestInterval> tmp = new ArrayList<>();
              List<TestInterval> list = new ArrayList<>();
              RandomSource random = new DefaultRandom();
              random.setSeed(seed);
              System.out.println(seed);
  
              int count = random.nextFloat() < 0.001 ? random.nextInt(1, 4) : 4 << random.nextInt(10);
              int keyDomain = count == 1 ? 2 : 1 + log2uniform(random, count);
              int valueDomain = count == 1 ? 1 : random.nextInt((int)Math.sqrt(count), count);
              int maxRemoveSize = log2uniform(random, count);
              int maxRunLength = log2uniform(random, count);
              float runChance = random.nextFloat();
              count = count + random.nextInt(count);
  
              TreeSet<TestInterval> unique = new TreeSet<>();
              for (int i = 0 ; i < count ; ++i)
              {
                  TestInterval interval = newInterval(random, keyDomain, valueDomain);
                  if (unique.add(interval))
                      list.add(interval);
  
                  if (maxRunLength > 1 && random.decide(runChance) && count - i > 1)
                  {
                      boolean startRun = random.nextBoolean();
                      int runLength = random.nextBoolean() ? log2uniform(random, maxRunLength - 1)
                                                           : random.nextInt(0, maxRunLength - 1);
                      runLength = Math.min(runLength, count - (i + 1));
                      runLength = Math.min(runLength, startRun ? keyDomain - (1 + interval.start) : interval.end);
                      while (--runLength >= 0)
                      {
                          int start = startRun ? interval.start : random.nextInt(interval.end);
                          int end = startRun ? random.nextInt(interval.start + 1, keyDomain) : interval.end;
                          int value = random.nextInt(valueDomain);
                          TestInterval interval2 = new TestInterval(start, end, value);
                          if (unique.add(interval2))
                          {
                              list.add(interval2);
                              ++i;
                          }
                      }
                  }
              }
  
              Object[] tree = BTree.empty();
              if (random.nextBoolean())
              {
                  for (TestInterval v : list)
                  {
                      tree = IntervalBTree.update(tree, BTree.singleton(v), TestComparators.INSTANCE);
                  }
              }
              else
              {
                  tmp.addAll(list);
                  tmp.sort(TestComparators.INSTANCE.totalOrder());
                  try (BTree.FastBuilder<TestInterval> builder = IntervalBTree.fastBuilder(TestComparators.INSTANCE))
                  {
                      for (TestInterval add : tmp)
                          builder.add(add);
                      tree = builder.build();
                  }
              }
              validate(tree, TestComparators.INSTANCE.endWithEndSorter());
  
              for (int i = 0 ; i < list.size() ; ++i)
              {
                  TestInterval iv = list.get(i);
                  TreeSet<TestInterval> collect = collect(list, 0, iv);
                  remove(tree, collect, iv);
                  Invariants.require(collect.isEmpty());
                  iv = newInterval(random, keyDomain, valueDomain);
                  collect = collect(list, 0, iv);
                  remove(tree, collect, iv);
                  Invariants.require(collect.isEmpty());
              }
  
              Collections.shuffle(list, random.asJdkRandom());
              for (int i = 0 ; i < list.size() ;)
              {
                  int remaining = list.size() - i;
                  int removeCount = remaining == 1 ? 1 : 1 + random.nextInt(Math.min(remaining, maxRemoveSize));
                  remaining -= removeCount;
  
                  TestInterval iv = list.get(i++);
                  Object[] remove;
                  {
                      tmp.clear();
                      tmp.add(iv);
                      int c = removeCount;
                      while (--c > 0)
                          tmp.add(list.get(i++));
  
                      int notPresentCount;
                      switch (random.nextInt(4))
                      {
                          default: throw new IllegalStateException();
                          case 0: notPresentCount = 0; break;
                          case 1: notPresentCount = random.nextInt(5); break;
                          case 2: notPresentCount = random.nextInt(Math.max(2, count/2)); break;
                          case 3: notPresentCount = random.nextInt(count*2); break;
                      }
  
                      while (--notPresentCount > 0)
                      {
                          TestInterval add = newInterval(random, keyDomain, valueDomain);
                          if (!unique.contains(add))
                              tmp.add(add);
                      }
  
                      tmp.sort(TestComparators.INSTANCE.totalOrder());
                      try (BTree.FastBuilder<TestInterval> builder = IntervalBTree.fastBuilder(TestComparators.INSTANCE))
                      {
                          for (TestInterval add : tmp)
                              builder.add(add);
                          remove = builder.build();
                      }
                  }
                  tree = IntervalBTree.subtract(tree, remove, TestComparators.INSTANCE);
                  validate(tree, TestComparators.INSTANCE.endWithEndSorter());
                  TreeSet<TestInterval> collect = collect(list, i, iv);
                  remove(tree, collect, iv);
                  Invariants.require(collect.isEmpty());
                  iv = newInterval(random, keyDomain, valueDomain);
                  collect = collect(list, i, iv);
                  remove(tree, collect, iv);
                  Invariants.require(collect.isEmpty());
              }
          }
          catch (Throwable t)
          {
              throw new AssertionError("Failed with seed " + seed, t);
          }
      }
  
      @Test
      public void printTest()
      {
          Random rng = new Random(0);
          List<TestInterval> intervals = new ArrayList<>();
          for (int i = 0; i < 20; i++)
          {
              intervals.add(interval(i, i + 1, i));
          }
          Collections.shuffle(intervals, rng);
          Object[] tree = IntervalBTree.empty();
          for (TestInterval v : intervals)
              tree = IntervalBTree.update(tree, IntervalBTree.singleton(v), TestComparators.INSTANCE);
  
          String printed = BTreePrinter.print(tree);
          for (TestInterval interval : intervals)
              Invariants.require(printed.contains(interval.toString()));
      }
  
      @Test
      public void simpleTest()
      {
          Random rng = new Random(0);
          List<TestInterval> intervals = Arrays.asList(interval(10, 40, 100), interval(20, 50, 200), interval(30, 60, 300));
          Collections.shuffle(intervals, rng);
          Object[] tree = IntervalBTree.empty();
          for (TestInterval v : intervals)
          {
              tree = IntervalBTree.update(tree, IntervalBTree.singleton(v), TestComparators.INSTANCE);
          }
  
          {
              Set<TestInterval> actual = new TreeSet<>();
              IntervalBTree.accumulate(tree, TestComparators.INSTANCE, interval(20, 55, 400), new QuadFunction<Object, Object, TestInterval, Object, Object>()
              {
                  @Override
                  public Object apply(Object o, Object o2, TestInterval match, Object o3)
                  {
                      actual.add(match);
                      return null;
                  }
              }, null, null, null);
              Assert.assertEquals(actual, new TreeSet<>(intervals));
          }
  
          {
              Set<TestInterval> actual = new TreeSet<>();
              IntervalBTree.accumulate(tree, TestComparators.INSTANCE, interval(10, 15, 400), new QuadFunction<Object, Object, TestInterval, Object, Object>()
              {
                  @Override
                  public Object apply(Object o, Object o2, TestInterval match, Object o3)
                  {
                      actual.add(match);
                      return null;
                  }
              }, null, null, null);
              Assert.assertEquals(actual, new TreeSet<>(Arrays.asList(interval(10, 40, 100))));
          }
      }
  
      @Test
      public void midSizeTest()
      {
          for (int i = 0; i < 10_000; i++)
          {
              midSizeTestIteration(i);
          }
      }
  
      public void midSizeTestIteration(int iteration)
      {
          RandomSource rng = new DefaultRandom(iteration);
          List<TestInterval> intervals = new ArrayList<>();
          TreeSet<TestInterval> unique = new TreeSet<>();
          for (int i = 0; i < rng.nextInt(10_000) + 1; i++)
          {
              TestInterval interval = newInterval(rng, 2, 10000);
              if (unique.add(interval))
                  intervals.add(interval);
          }
  
          Collections.shuffle(intervals, rng.asJdkRandom());
  
          Object[] tree = IntervalBTree.empty();
          for (TestInterval v : intervals)
              tree = IntervalBTree.update(tree, IntervalBTree.singleton(v), TestComparators.INSTANCE);
  
          for (int i = 0; i < 100; i++)
          {
              TestInterval searched = newInterval(rng, 2, 10000);
              Set<TestInterval> actual = new TreeSet<>();
              IntervalBTree.accumulate(tree, TestComparators.INSTANCE, searched, (o, o2, match, o3) -> {
                  actual.add(match);
                  return null;
              }, null, null, null);
              Assert.assertEquals(String.format("Searching for %s", searched),
                                  actual, new TreeSet<>(intervals.stream().filter(match -> {
                  if ((match.start <= searched.start && searched.start < match.end) ||
                      (searched.start <= match.start && match.start < searched.end))
                      return true;
                  return false;
              }).collect(Collectors.toList())));
          }
      }
  
      @Test
      public void subtractTest()
      {
          for (int j = 0; j < 10_000; j++)
              subtractTestIteration(j);
      }
  
      public void subtractTestIteration(int iteration)
      {
          RandomSource rng = new DefaultRandom(iteration);
          List<TestInterval> intervals = new ArrayList<>();
          TreeSet<TestInterval> unique = new TreeSet<>();
          for (int i = 0; i < rng.nextInt(10_000) + 1; i++)
          {
              TestInterval interval = newInterval(rng, 2, 10000);
              if (unique.add(interval))
                  intervals.add(interval);
          }
  
          Collections.shuffle(intervals, rng.asJdkRandom());
  
          Object[] tree = IntervalBTree.empty();
          for (TestInterval v : intervals)
              tree = IntervalBTree.update(tree, IntervalBTree.singleton(v), TestComparators.INSTANCE);
  
          for (int i = 0; i < 100 && BTree.size(tree) > 0; i++)
          {
              int subtractCount = Math.max(1, rng.nextInt(intervals.size() / 10 + 1));
              List<TestInterval> toSubtract = new ArrayList<>();
              TreeSet<TestInterval> subtractSet = new TreeSet<>();
  
              for (int j = 0; j < subtractCount; j++)
              {
                  TestInterval interval;
                  // pick an existing or a new interval
                  if (rng.nextBoolean() && !intervals.isEmpty())
                      interval = intervals.get(rng.nextInt(intervals.size()));
                  else
                      interval = newInterval(rng, 2, 10000);
  
                  if (subtractSet.add(interval))
                      toSubtract.add(interval);
              }
  
              Object[] resultTree;
              {
                  Object[] subtractTree = IntervalBTree.empty();
                  for (TestInterval v : toSubtract)
                      subtractTree = IntervalBTree.update(subtractTree, IntervalBTree.singleton(v), TestComparators.INSTANCE);
  
                  resultTree = IntervalBTree.subtract(tree, subtractTree, TestComparators.INSTANCE);
              }
              // Collect all intervals remaining in the result tree
              Set<TestInterval> remaining = new TreeSet<>();
              IntervalBTree.accumulate(resultTree, (o1, o2, interval, o3) -> {
                  remaining.add((TestInterval)interval);
                  return null;
              }, null, null, null);
  
              Set<TestInterval> expectedRemaining = new TreeSet<>();
              for (TestInterval original : intervals)
              {
                  boolean shouldBeSubtracted = subtractSet.contains(original);
                  if (!shouldBeSubtracted)
                      expectedRemaining.add(original);
              }
  
              Assert.assertEquals(String.format("Subtraction iteration %d: expected %d remaining, got %d. Subtracted: %s",
                                                i, expectedRemaining.size(), remaining.size(), subtractSet),
                                  expectedRemaining, remaining);
  
              intervals.retainAll(expectedRemaining);
              tree = resultTree;
          }
      }
  
      public static TestInterval interval(int start, int end, int value)
      {
          return new TestInterval(start, end, value);
      }
  
      private static int log2uniform(RandomSource random, int max)
      {
          int logn = 31 - Integer.numberOfLeadingZeros(max);
          int loglogn = 31 - Integer.numberOfLeadingZeros(logn);
          int scale = 1 << random.nextInt(loglogn, logn);
          if (scale == max)
              return max;
          return random.nextInt(scale, Math.min(scale * 2, max));
      }
  
      private static TreeSet<TestInterval> collect(List<TestInterval> list, int from, TestInterval intersects)
      {
          TreeSet<TestInterval> collect = new TreeSet<>();
          for (int i = from; i < list.size() ; ++i)
          {
              TestInterval iv2 = list.get(i);
              if (intersects.start < iv2.end && iv2.start < intersects.end)
                  collect.add(iv2);
          }
          return collect;
      }
  
      private static void remove(Object[] tree, TreeSet<TestInterval> removeFrom, TestInterval intersects)
      {
          IntervalBTree.accumulate(tree, TestComparators.INSTANCE, intersects, (c, v, i, s) -> {
              Invariants.require(c.remove(i));
              return null;
          }, removeFrom, null, null);
      }
  
      private static TestInterval newInterval(RandomSource random, int keyDomain, int valueDomain)
      {
          int end = keyDomain == 1 ? 1 : 1 + random.nextInt(Math.max(1, keyDomain - 1));
          int start = random.nextInt(end);
          return new TestInterval(start, end, random.nextInt(valueDomain));
      }
  
      static Object validate(Object[] tree, Comparator endSorter)
      {
          if (isLeaf(tree))
              return max(tree, endSorter);
  
          IntervalBTree.IntervalMaxIndex index = (IntervalBTree.IntervalMaxIndex) tree[tree.length - 1];
          Object[] tmp = new Object[getChildCount(tree)];
          for (int i = 0 ; i < index.sortedByEndIndex.length ; ++i)
              tmp[index.sortedByEndIndex[i]] = index.sortedByEnd[i];
  
          Object max = max(tree, endSorter);
          for (int i = 0 ; i < getChildCount(tree) ; ++i)
          {
              Object childMax = validate((Object[])tree[getChildStart(tree) + i], endSorter);
              if (i < getChildStart(tree) && endSorter.compare(childMax, tree[i]) < 0)
                  childMax = tree[i];
              if (endSorter.compare(childMax, max) > 0)
                  max = childMax;
              Invariants.require(endSorter.compare(childMax, tmp[i]) == 0);
          }
          return max;
      }
  
      private static Object max(Object[] tree, Comparator endSorter)
      {
          Object max = tree[0];
          for (int i = 1; i < getKeyEnd(tree) ; ++i)
          {
              if (endSorter.compare(tree[i], max) > 0)
                  max = tree[i];
          }
          return max;
      }
  }
  