/*
 * Copyright 2012 Greplin, Inc. All Rights Reserved.
 */

package com.greplin.lucene.query;

import com.google.common.base.Objects;
import com.greplin.lucene.predicate.BitsProvider;
import org.apache.lucene.index.IndexReader;
import org.apache.lucene.index.Term;
import org.apache.lucene.search.DocIdSetIterator;
import org.apache.lucene.search.Explanation;
import org.apache.lucene.search.Query;
import org.apache.lucene.search.Scorer;
import org.apache.lucene.search.Searcher;
import org.apache.lucene.search.Similarity;
import org.apache.lucene.search.Weight;
import org.apache.lucene.util.Bits;
import org.apache.lucene.util.ToStringUtils;

import java.io.IOException;
import java.util.Set;

/**
 * Query that is filtered through a simple predicate function.
 *
 * Based on {@link org.apache.lucene.search.CachingWrapperFilter}.
 */
public class PredicateQuery extends Query {

  /**
   * The underlying query.
   */
  private final Query query;


  /**
   * The predicate to match against.
   */
  private final BitsProvider predicate;


  /**
   * Constructs a new query which applies a predicate to filter the results
   * of the original query.
   * @param query Query to be filtered.
   * @param predicate Provider of predicates to apply to the query.
   */
  public PredicateQuery(final Query query, final BitsProvider predicate) {
    this.query = query;
    this.predicate = predicate;
  }

  /**
   * Returns a Weight that applies the predicate to the enclosed query's Weight.
   * This is accomplished by overriding the Scorer returned by the Weight.
   * @param searcher the searcher to create a weight for.
   * @return a Weight that applies the predicate to the query.
   * @throws IOException if IO issues occur.
   */
  @Override
  public Weight createWeight(final Searcher searcher) throws IOException {
    final Weight weight = this.query.createWeight(searcher);
    final Similarity similarity = this.query.getSimilarity(searcher);

    return new Weight() {
      private float value;

      // pass these methods through to enclosed query's weight
      @Override
      public float getValue() {
        return this.value;
      }

      @Override
      public boolean scoresDocsOutOfOrder() {
        return false;
      }

      public float sumOfSquaredWeights() throws IOException {
        return weight.sumOfSquaredWeights() * getBoost() * getBoost();
      }

      @Override
      public void normalize(final float v) {
        weight.normalize(v * getBoost()); // incorporate boost
        this.value = weight.getValue();
      }

      @Override
      public Explanation explain(final IndexReader reader, final int i)
          throws IOException {
        Explanation inner = weight.explain(reader, i);
        Bits predicate = PredicateQuery.this.predicate.get(reader);
        if (predicate.get(i)) {
          return inner;
        } else {
          Explanation result = new Explanation(0.0f,
              "failure to match predicate: " + predicate.toString());
          result.addDetail(inner);
          return result;
        }
      }

      @Override
      public Query getQuery() {
        return PredicateQuery.this;
      }

      @Override
      public Scorer scorer(final IndexReader reader,
                           final boolean scoreDocsInOrder,
                           final boolean topScorer)
          throws IOException {
        Bits predicate = PredicateQuery.this.predicate.get(reader);
        return PredicateQuery.getScorer(
            reader, similarity, weight, this, predicate);
      }
    };
  }


  /**
   * Creates a scorer that matches documents matched by both the query and
   * the predicate.
   * @param indexReader the atomic reader
   * @param similarity the Similarity to use
   * @param weight the weight object of the underlying query
   * @param wrapperWeight the weight object.
   * @param predicate the Bits to use.
   * @return a scorer that matches the documents matched by the filter.
   * @throws IOException on IO error
   */
  private static Scorer getScorer(
      final IndexReader indexReader, final Similarity similarity,
      final Weight weight, final Weight wrapperWeight,
      final Bits predicate) throws IOException {
    // We will advance() this scorer, so we set inorder=true/toplevel=false.
    final Scorer scorer = weight.scorer(indexReader, true, false);
    return (scorer == null) ? null : new Scorer(similarity, wrapperWeight) {

      @Override
      public int nextDoc() throws IOException {
        for (;;) {
          int docId = scorer.nextDoc();
          if (docId == DocIdSetIterator.NO_MORE_DOCS
              || predicate.get(scorer.docID())) {
            return docId;
          }
        }
      }

      @Override
      public int advance(final int target) throws IOException {
        scorer.advance(target);
        if (scorer.docID() == DocIdSetIterator.NO_MORE_DOCS
            || predicate.get(scorer.docID())) {
          return scorer.docID();
        }
        return nextDoc();
      }

      @Override
      public int docID() {
        return scorer.docID();
      }

      @Override
      public float score() throws IOException {
        return scorer.score();
      }
    };
  }


  /**
   * Rewrites the wrapped query.
   * @param reader the reader to rewrite for.
   * @return the rewritten query.
   * @throws IOException if IO issues occur.
   */
  @Override
  public Query rewrite(final IndexReader reader) throws IOException {
    Query rewritten = this.query.rewrite(reader);
    if (rewritten != this.query) {
      return new PredicateQuery(rewritten, this.predicate);
    } else {
      return this;
    }
  }


  /**
   * @return the underlying query.
   */
  public Query getQuery() {
    return this.query;
  }


  /**
   * @return the predicate.
   */
  public BitsProvider getPredicate() {
    return this.predicate;
  }


  @Override
  public void extractTerms(final Set<Term> terms) {
    getQuery().extractTerms(terms);
  }


  @Override
  public int hashCode() {
    return Objects.hashCode(this.query, this.predicate);
  }


  @Override
  public boolean equals(final Object o) {
    if (this == o) {
      return true;
    }
    if (o == null || getClass() != o.getClass()) {
      return false;
    }

    PredicateQuery that = (PredicateQuery) o;
    return Objects.equal(this.query, that.query)
        && Objects.equal(this.predicate, that.predicate);
  }


  /**
   * Prints a user-readable version of this query.
   * @param s the name of the field.
   * @return a user-readable version of this query.
   */
  @Override
  public String toString(final String s) {
    StringBuilder buffer = new StringBuilder();
    buffer.append("predicate(");
    buffer.append(this.query.toString(s));
    buffer.append(")->");
    buffer.append(this.predicate);
    buffer.append(ToStringUtils.boost(getBoost()));
    return buffer.toString();
  }

}

@@ -4906,32 +4906,56 @@
 S%0A            %7C%7C
+ (scorer.docID() %3E= 0 &&
  predicate.get(s
@@ -4961,32 +4961,33 @@
 scorer.docID()))
+)
  %7B%0A          ret

Fix ArrayOutOfBoundsException