IntArrayDocIdSet.java

package com.kamikaze.docidset.impl;

import java.io.IOException;
import java.io.Serializable;

import org.apache.lucene.search.DocIdSetIterator;

import com.kamikaze.docidset.api.DocSet;
import com.kamikaze.docidset.api.StatefulDSIterator;
import com.kamikaze.docidset.utils.IntArray;

public class IntArrayDocIdSet extends DocSet implements Serializable {

  private static final long serialVersionUID = 1L;

  private IntArray array = null;

  private int pos = -1;

  public IntArrayDocIdSet(int length) {
    array = new IntArray(length);
  }

  public IntArrayDocIdSet() {
    array = new IntArray();
  }

  
  public void addDoc(int docid) {
    ++pos;
    array.add(docid);
  }
  
  @Override
  public final boolean isCacheable() {
    return true;
  }

  protected int binarySearchForNearest(int val, int begin, int end) {

    int mid = (begin + end) / 2;
    int midval = array.get(mid);
    
    if(mid == end)
      return midval>=val? mid : -1;
        
    if (midval < val) {
      // Find number equal or greater than the target.
      if (array.get(mid + 1) >= val) return mid + 1;
    
      return binarySearchForNearest(val, mid + 1, end);
    }
    else {
      // Find number equal or greater than the target.
      if (midval == val) return mid;

      return binarySearchForNearest(val, begin, mid);
    }
  }

  class IntArrayDocIdSetIterator extends StatefulDSIterator {
    int lastReturn = -1;

    int cursor = -1;
    
    public IntArrayDocIdSetIterator()
    {
      if(pos == -1) lastReturn = DocIdSetIterator.NO_MORE_DOCS;
    }
    
    @Override
    public int docID() {
      return lastReturn;
    }
    
    @Override
    public int nextDoc() throws IOException {
      if (cursor < pos) {
        return (lastReturn = array.get(++cursor));
      }
      return (lastReturn = DocIdSetIterator.NO_MORE_DOCS);
    }

    @Override
    public int advance(int target) throws IOException {
      if (lastReturn == DocIdSetIterator.NO_MORE_DOCS) return DocIdSetIterator.NO_MORE_DOCS;
      
      if (target <= lastReturn) target = lastReturn + 1;
      
      int end = Math.min(cursor + (target - lastReturn), pos);
      int index = binarySearchForNearest(target, cursor + 1, end);
      
      if (index == -1) {
        cursor = pos;
        return (lastReturn = DocIdSetIterator.NO_MORE_DOCS);
      } else {
        cursor = index;
        return (lastReturn = array.get(cursor));
      }
    }

    @Override
    public int getCursor() {
      return cursor;
    }
  }

  @Override
  public IntArrayDocIdSetIterator iterator() {
     return new IntArrayDocIdSetIterator();
  }

  public int size() {
    return pos + 1;
  }
  @Override
  public int findWithIndex(int val) {
    IntArrayDocIdSetIterator dcit = new IntArrayDocIdSetIterator();
    try {
      int docid = dcit.advance(val);
      if (docid == val)
        return dcit.getCursor();
    } catch (IOException e) {
      e.printStackTrace();
    }
    return -1;
  }

  @Override
  public long sizeInBytes()
  {
    //Object Overhead
    return array.length()*4 + 64;
  }
  
  @Override
  public void optimize()
  {
    this.array.seal();
  }
   
  
}