/*
 * 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.lucene.util;

import java.io.IOException;
import org.apache.lucene.search.AbstractDocIdSetIterator;
import org.apache.lucene.search.DocIdSet;
import org.apache.lucene.search.DocIdSetIterator;

/**
 * {@link DocIdSet} implementation inspired from http://roaringbitmap.org/
 *
 * <p>The space is divided into blocks of 2^16 bits and each block is encoded independently. In each
 * block, if less than 2^12 bits are set, then documents are simply stored in a short[]. If more
 * than 2^16-2^12 bits are set, then the inverse of the set is encoded in a simple short[].
 * Otherwise a {@link FixedBitSet} is used.
 *
 * @lucene.internal
 */
public class RoaringDocIdSet extends DocIdSet {

  // Number of documents in a block
  private static final int BLOCK_SIZE = 1 << 16;
  // The maximum length for an array, beyond that point we switch to a bitset
  private static final int MAX_ARRAY_LENGTH = 1 << 12;
  private static final long BASE_RAM_BYTES_USED =
      RamUsageEstimator.shallowSizeOfInstance(RoaringDocIdSet.class);

  /** A builder of {@link RoaringDocIdSet}s. */
  public static class Builder {

    private final int maxDoc;
    private final DocIdSet[] sets;

    private int cardinality;
    private int firstDocId;
    private int lastDocId;
    private int currentBlock;
    private int currentBlockCardinality;

    // We start by filling the buffer and when it's full we copy the content of
    // the buffer to the FixedBitSet and put further documents in that bitset
    private final short[] buffer;
    private FixedBitSet denseBuffer;

    /** Sole constructor. */
    public Builder(int maxDoc) {
      this.maxDoc = maxDoc;
      sets = new DocIdSet[(maxDoc + (1 << 16) - 1) >>> 16];
      lastDocId = -1;
      currentBlock = -1;
      buffer = new short[MAX_ARRAY_LENGTH];
    }

    private void flush() {
      assert currentBlockCardinality <= BLOCK_SIZE;
      if (currentBlockCardinality == BLOCK_SIZE) {
        // all docs in the block
        sets[currentBlock] = AllDocIdSet.INSTANCE;
      } else if (currentBlockCardinality > 0
          && currentBlockCardinality == lastDocId - firstDocId + 1) {
        // doc ids are continuous, use range encoding. Even for a unique value it uses less heap.
        sets[currentBlock] = new RangeDocIdSet((short) firstDocId, (short) lastDocId);
      } else if (currentBlockCardinality <= MAX_ARRAY_LENGTH) {
        // Use sparse encoding
        assert denseBuffer == null;
        if (currentBlockCardinality > 0) {
          sets[currentBlock] =
              new ShortArrayDocIdSet(ArrayUtil.copyOfSubArray(buffer, 0, currentBlockCardinality));
        }
      } else {
        assert denseBuffer != null;
        assert denseBuffer.cardinality() == currentBlockCardinality;
        if (denseBuffer.length() == BLOCK_SIZE
            && BLOCK_SIZE - currentBlockCardinality < MAX_ARRAY_LENGTH) {
          // Doc ids are very dense, inverse the encoding
          final short[] excludedDocs = new short[BLOCK_SIZE - currentBlockCardinality];
          denseBuffer.flip(0, denseBuffer.length());
          int excludedDoc = -1;
          for (int i = 0; i < excludedDocs.length; ++i) {
            excludedDoc = denseBuffer.nextSetBit(excludedDoc + 1);
            assert excludedDoc != DocIdSetIterator.NO_MORE_DOCS;
            excludedDocs[i] = (short) excludedDoc;
          }
          assert excludedDoc + 1 == denseBuffer.length()
              || denseBuffer.nextSetBit(excludedDoc + 1) == DocIdSetIterator.NO_MORE_DOCS;
          sets[currentBlock] = new NotDocIdSet(BLOCK_SIZE, new ShortArrayDocIdSet(excludedDocs));
        } else {
          // Neither sparse nor super dense, use a fixed bit set
          sets[currentBlock] = new BitDocIdSet(denseBuffer, currentBlockCardinality);
        }
        denseBuffer = null;
      }

      cardinality += currentBlockCardinality;
      denseBuffer = null;
      currentBlockCardinality = 0;
    }

    /** Add a new doc-id to this builder. NOTE: doc ids must be added in order. */
    public Builder add(int docId) {
      if (docId <= lastDocId) {
        throw new IllegalArgumentException(
            "Doc ids must be added in-order, got " + docId + " which is <= lastDocID=" + lastDocId);
      }
      final int block = docId >>> 16;
      if (block != currentBlock) {
        // we went to a different block, let's flush what we buffered and start from fresh
        flush();
        currentBlock = block;
        firstDocId = docId;
      }

      appendDocInCurrentBlock(docId, block);
      return this;
    }

    /**
     * Add a contiguous half-open range {@code [min, max)} of doc ids. Doc ids must remain in
     * non-decreasing order relative to prior {@link #add(int)} calls; {@code min} must be strictly
     * greater than the last document added (if any). Equivalent to calling {@link #add(int)} for
     * each {@code min <= doc < max}, but uses a more efficient implementation for dense ranges.
     *
     * @param min inclusive lower bound (must be {@code >= 0})
     * @param max exclusive upper bound
     * @throws IllegalArgumentException if {@code min > max}, if {@code min <= lastDocId}, or if
     *     {@code max} exceeds {@link #Builder(int) maxDoc}
     */
    public Builder add(int min, int max) {
      if (min > max) {
        throw new IllegalArgumentException("min must be <= max, got min=" + min + " max=" + max);
      }
      if (min == max) {
        return this;
      }
      if (min <= lastDocId) {
        throw new IllegalArgumentException(
            "Doc ids must be added in-order, got range starting at "
                + min
                + " which is <= lastDocID="
                + lastDocId);
      }
      if (max > maxDoc) {
        throw new IllegalArgumentException(
            "max must be <= maxDoc (" + maxDoc + "), got max=" + max);
      }
      int doc = min;
      while (doc < max) {
        final int block = doc >>> 16;
        final int blockEnd = Math.min(max, (block + 1) << 16);
        if (block != currentBlock) {
          flush();
          currentBlock = block;
          firstDocId = doc;
        }
        appendRangeInCurrentBlock(doc, blockEnd, block);
        doc = blockEnd;
      }
      return this;
    }

    private void appendDocInCurrentBlock(int docId, int block) {
      assert docId >= firstDocId;
      if (currentBlockCardinality < MAX_ARRAY_LENGTH) {
        buffer[currentBlockCardinality] = (short) docId;
      } else {
        if (denseBuffer == null) {
          // the buffer is full, let's move to a fixed bit set
          final int numBits = Math.min(1 << 16, maxDoc - (block << 16));
          denseBuffer = new FixedBitSet(numBits);
          for (short doc : buffer) {
            denseBuffer.set(doc & 0xFFFF);
          }
        }
        denseBuffer.set(docId - (block << 16));
      }

      lastDocId = docId;
      currentBlockCardinality += 1;
    }

    private void appendRangeInCurrentBlock(int fromDoc, int toDocExclusive, int block) {
      assert fromDoc >= firstDocId;
      assert (fromDoc >>> 16) == block;
      assert (toDocExclusive >>> 16) == block || toDocExclusive == (block + 1) << 16;
      assert fromDoc < toDocExclusive;
      final int span = toDocExclusive - fromDoc;
      if (currentBlockCardinality + span <= MAX_ARRAY_LENGTH) {
        for (int d = fromDoc; d < toDocExclusive; d++) {
          buffer[currentBlockCardinality++] = (short) d;
        }
        lastDocId = toDocExclusive - 1;
      } else {
        int offset = block << 16;
        if (denseBuffer == null) {
          final int numBits = Math.min(1 << 16, maxDoc - offset);
          denseBuffer = new FixedBitSet(numBits);
          for (int i = 0; i < currentBlockCardinality; i++) {
            denseBuffer.set(buffer[i] & 0xFFFF);
          }
        }
        denseBuffer.set(fromDoc - offset, toDocExclusive - offset);
        lastDocId = toDocExclusive - 1;
        currentBlockCardinality += span;
      }
    }

    /** Add the content of the provided {@link DocIdSetIterator}. */
    public Builder add(DocIdSetIterator disi) throws IOException {
      for (int doc = disi.nextDoc(); doc != DocIdSetIterator.NO_MORE_DOCS; doc = disi.nextDoc()) {
        add(doc);
      }
      return this;
    }

    /** Build an instance. */
    public RoaringDocIdSet build() {
      flush();
      return new RoaringDocIdSet(sets, cardinality);
    }
  }

  private static class AllDocIdSet extends DocIdSet {

    private static final AllDocIdSet INSTANCE = new AllDocIdSet();

    private AllDocIdSet() {}

    @Override
    public DocIdSetIterator iterator() {
      return DocIdSetIterator.range(0, BLOCK_SIZE);
    }

    @Override
    public long ramBytesUsed() {
      // we use a static instance
      return 0;
    }
  }

  private static class RangeDocIdSet extends DocIdSet {

    private static final long BASE_RAM_BYTES_USED =
        RamUsageEstimator.shallowSizeOfInstance(RangeDocIdSet.class);

    final short minDocId;
    final short maxDocId;

    private RangeDocIdSet(short minDocId, short maxDocId) {
      this.minDocId = minDocId;
      this.maxDocId = maxDocId;
    }

    @Override
    public DocIdSetIterator iterator() {
      return DocIdSetIterator.range(minDocId & 0xFFFF, (maxDocId & 0xFFFF) + 1);
    }

    @Override
    public long ramBytesUsed() {
      return BASE_RAM_BYTES_USED;
    }
  }

  /** {@link DocIdSet} implementation that can store documents up to 2^16-1 in a short[]. */
  private static class ShortArrayDocIdSet extends DocIdSet {

    private static final long BASE_RAM_BYTES_USED =
        RamUsageEstimator.shallowSizeOfInstance(ShortArrayDocIdSet.class);

    private final short[] docIDs;

    private ShortArrayDocIdSet(short[] docIDs) {
      this.docIDs = docIDs;
    }

    @Override
    public long ramBytesUsed() {
      return BASE_RAM_BYTES_USED + RamUsageEstimator.sizeOf(docIDs);
    }

    @Override
    public DocIdSetIterator iterator() throws IOException {
      return new AbstractDocIdSetIterator() {

        int i = -1; // this is the index of the current document in the array

        private int docId(int i) {
          return docIDs[i] & 0xFFFF;
        }

        @Override
        public int nextDoc() throws IOException {
          if (++i >= docIDs.length) {
            return doc = NO_MORE_DOCS;
          }
          return doc = docId(i);
        }

        @Override
        public long cost() {
          return docIDs.length;
        }

        @Override
        public int advance(int target) throws IOException {
          // binary search
          int lo = i + 1;
          int hi = docIDs.length - 1;
          while (lo <= hi) {
            final int mid = (lo + hi) >>> 1;
            final int midDoc = docId(mid);
            if (midDoc < target) {
              lo = mid + 1;
            } else {
              hi = mid - 1;
            }
          }
          if (lo == docIDs.length) {
            i = docIDs.length;
            return doc = NO_MORE_DOCS;
          } else {
            i = lo;
            return doc = docId(i);
          }
        }

        @Override
        public void intoBitSet(int upTo, FixedBitSet bitSet, int offset) throws IOException {
          if (doc >= upTo) {
            return;
          }

          int from = i;
          advance(upTo);
          int to = i;
          for (int i = from; i < to; ++i) {
            bitSet.set(docId(i) - offset);
          }
        }

        @Override
        public int docIDRunEnd() {
          int runEnd = doc + 1;
          int j = i + 1;
          while (j < docIDs.length && (docIDs[j] & 0xFFFF) == runEnd) {
            runEnd++;
            j++;
          }
          return runEnd;
        }
      };
    }
  }

  private final DocIdSet[] docIdSets;
  private final int cardinality;
  private final long ramBytesUsed;

  private RoaringDocIdSet(DocIdSet[] docIdSets, int cardinality) {
    this.docIdSets = docIdSets;
    long ramBytesUsed = BASE_RAM_BYTES_USED + RamUsageEstimator.shallowSizeOf(docIdSets);
    for (DocIdSet set : this.docIdSets) {
      if (set != null) {
        ramBytesUsed += set.ramBytesUsed();
      }
    }
    this.ramBytesUsed = ramBytesUsed;
    this.cardinality = cardinality;
  }

  @Override
  public long ramBytesUsed() {
    return ramBytesUsed;
  }

  @Override
  public DocIdSetIterator iterator() throws IOException {
    if (cardinality == 0) {
      return null;
    }
    return new Iterator();
  }

  private class Iterator extends AbstractDocIdSetIterator {

    int block;
    DocIdSetIterator sub;

    Iterator() throws IOException {
      block = -1;
      sub = DocIdSetIterator.empty();
    }

    @Override
    public int nextDoc() throws IOException {
      final int subNext = sub.nextDoc();
      if (subNext == NO_MORE_DOCS) {
        return firstDocFromNextBlock();
      }
      return doc = (block << 16) | subNext;
    }

    @Override
    public int advance(int target) throws IOException {
      final int targetBlock = target >>> 16;
      if (targetBlock != block) {
        block = targetBlock;
        if (block >= docIdSets.length) {
          sub = null;
          return doc = NO_MORE_DOCS;
        }
        if (docIdSets[block] == null) {
          return firstDocFromNextBlock();
        }
        sub = docIdSets[block].iterator();
      }
      final int subNext = sub.advance(target & 0xFFFF);
      if (subNext == NO_MORE_DOCS) {
        return firstDocFromNextBlock();
      }
      return doc = (block << 16) | subNext;
    }

    private int firstDocFromNextBlock() throws IOException {
      while (true) {
        block += 1;
        if (block >= docIdSets.length) {
          sub = null;
          return doc = NO_MORE_DOCS;
        } else if (docIdSets[block] != null) {
          sub = docIdSets[block].iterator();
          final int subNext = sub.nextDoc();
          assert subNext != NO_MORE_DOCS;
          return doc = (block << 16) | subNext;
        }
      }
    }

    @Override
    public void intoBitSet(int upTo, FixedBitSet bitSet, int offset) throws IOException {
      for (; ; ) {
        if (doc == NO_MORE_DOCS) {
          break;
        }
        int subUpto = upTo - (block << 16);
        if (subUpto < 0) {
          break;
        }
        int subOffset = offset - (block << 16);
        sub.intoBitSet(subUpto, bitSet, subOffset);
        if (sub.docID() == NO_MORE_DOCS) {
          if (firstDocFromNextBlock() == NO_MORE_DOCS) {
            break;
          }
        } else {
          doc = (block << 16) | sub.docID();
          break;
        }
      }
    }

    @Override
    public int docIDRunEnd() throws IOException {
      assert doc != NO_MORE_DOCS && sub != null;
      int b = block;
      long globalEnd = ((long) b << 16) + (long) sub.docIDRunEnd();
      // Merge with the next block when the sub-run ends exactly on a 64K boundary and the next
      // block's first doc is the immediate successor (local doc 0).
      while (globalEnd == ((long) (b + 1) << 16) && globalEnd < (long) Integer.MAX_VALUE) {
        int nb = b + 1;
        if (nb >= docIdSets.length || docIdSets[nb] == null) {
          break;
        }
        DocIdSetIterator nextIt = docIdSets[nb].iterator();
        int first = nextIt.nextDoc();
        if (first != 0) {
          break;
        }
        b = nb;
        globalEnd = ((long) b << 16) + (long) nextIt.docIDRunEnd();
      }
      return globalEnd >= Integer.MAX_VALUE ? Integer.MAX_VALUE : (int) globalEnd;
    }

    @Override
    public long cost() {
      return cardinality;
    }
  }

  /** Return the exact number of documents that are contained in this set. */
  public int cardinality() {
    return cardinality;
  }

  @Override
  public String toString() {
    return "RoaringDocIdSet(cardinality=" + cardinality + ")";
  }
}
