package org.apache.lucene.index;

/**
 * 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.
 */

/** Remaps docIDs after a merge has completed, where the
 *  merged segments had at least one deletion.  This is used
 *  to renumber the buffered deletes in IndexWriter when a
 *  merge of segments with deletions commits. */

final class MergeDocIDRemapper {
  int[] starts;                                 // used for binary search of mapped docID
  int[] newStarts;                              // starts, minus the deletes
  int[][] docMaps;                              // maps docIDs in the merged set
  int minDocID;                                 // minimum docID that needs renumbering
  int maxDocID;                                 // 1+ the max docID that needs renumbering
  int docShift;                                 // total # deleted docs that were compacted by this merge

  public MergeDocIDRemapper(SegmentInfos infos, int[][] docMaps, int[] delCounts, MergePolicy.OneMerge merge, int mergedDocCount) {
    this.docMaps = docMaps;
    SegmentInfo firstSegment = merge.segments.get(0);
    int i = 0;
    while(true) {
      SegmentInfo info = infos.info(i);
      if (info.equals(firstSegment))
        break;
      minDocID += info.docCount;
      i++;
    }

    int numDocs = 0;
    for(int j=0;j<docMaps.length;i++,j++) {
      numDocs += infos.info(i).docCount;
      assert infos.info(i).equals(merge.segments.get(j));
    }
    maxDocID = minDocID + numDocs;

    starts = new int[docMaps.length];
    newStarts = new int[docMaps.length];

    starts[0] = minDocID;
    newStarts[0] = minDocID;
    for(i=1;i<docMaps.length;i++) {
      final int lastDocCount = merge.segments.get(i-1).docCount;
      starts[i] = starts[i-1] + lastDocCount;
      newStarts[i] = newStarts[i-1] + lastDocCount - delCounts[i-1];
    }
    docShift = numDocs - mergedDocCount;

    // There are rare cases when docShift is 0.  It happens
    // if you try to delete a docID that's out of bounds,
    // because the SegmentReader still allocates deletedDocs
    // and pretends it has deletions ... so we can't make
    // this assert here
    // assert docShift > 0;

    // Make sure it all adds up:
    assert docShift == maxDocID - (newStarts[docMaps.length-1] + merge.segments.get(docMaps.length-1).docCount - delCounts[docMaps.length-1]);
  }

  public int remap(int oldDocID) {
    if (oldDocID < minDocID)
      // Unaffected by merge
      return oldDocID;
    else if (oldDocID >= maxDocID)
      // This doc was "after" the merge, so simple shift
      return oldDocID - docShift;
    else {
      // Binary search to locate this document & find its new docID
      int lo = 0;                                      // search starts array
      int hi = docMaps.length - 1;                  // for first element less

      while (hi >= lo) {
        int mid = (lo + hi) >>> 1;
        int midValue = starts[mid];
        if (oldDocID < midValue)
          hi = mid - 1;
        else if (oldDocID > midValue)
          lo = mid + 1;
        else {                                      // found a match
          while (mid+1 < docMaps.length && starts[mid+1] == midValue) {
            mid++;                                  // scan to last match
          }
          if (docMaps[mid] != null)
            return newStarts[mid] + docMaps[mid][oldDocID-starts[mid]];
          else
            return newStarts[mid] + oldDocID-starts[mid];
        }
      }
      if (docMaps[hi] != null)
        return newStarts[hi] + docMaps[hi][oldDocID-starts[hi]];
      else
        return newStarts[hi] + oldDocID-starts[hi];
    }
  }
}
