package org.apache.lucene.search.uhighlight;
import java.io.Closeable;
import java.io.IOException;
import java.util.HashMap;
import java.util.List;
import java.util.Objects;
import java.util.PriorityQueue;
import java.util.function.Supplier;
import org.apache.lucene.index.PostingsEnum;
import org.apache.lucene.index.Term;
import org.apache.lucene.search.MatchesIterator;
import org.apache.lucene.search.Query;
import org.apache.lucene.search.QueryVisitor;
import org.apache.lucene.util.BytesRef;
import org.apache.lucene.util.BytesRefBuilder;
import org.apache.lucene.util.IOUtils;
public abstract class OffsetsEnum implements Comparable<OffsetsEnum>, Closeable {
@Override
public int compareTo(OffsetsEnum other) {
try {
int cmp = Integer.compare(startOffset(), other.startOffset());
if (cmp != 0) {
return cmp;
}
cmp = Integer.compare(endOffset(), other.endOffset());
if (cmp != 0) {
return cmp;
}
final BytesRef thisTerm = this.getTerm();
final BytesRef otherTerm = other.getTerm();
if (thisTerm == null || otherTerm == null) {
if (thisTerm == null && otherTerm == null) {
return 0;
} else if (thisTerm == null) {
return 1;
} else {
return -1;
}
}
return thisTerm.compareTo(otherTerm);
} catch (IOException e) {
throw new RuntimeException(e);
}
}
public abstract boolean nextPosition() throws IOException;
public abstract int freq() throws IOException;
public abstract BytesRef getTerm() throws IOException;
public abstract int startOffset() throws IOException;
public abstract int endOffset() throws IOException;
@Override
public void close() throws IOException {
}
@Override
public String toString() {
final String name = getClass().getSimpleName();
String offset = "";
try {
offset = ",[" + startOffset() + "-" + endOffset() + "]";
} catch (Exception e) {
}
try {
return name + "(term:" + getTerm().utf8ToString() + offset + ")";
} catch (Exception e) {
return name;
}
}
public static class OfPostings extends OffsetsEnum {
private final BytesRef term;
private final PostingsEnum postingsEnum;
private final int freq;
private int posCounter = -1;
public OfPostings(BytesRef term, int freq, PostingsEnum postingsEnum) throws IOException {
this.term = Objects.requireNonNull(term);
this.postingsEnum = Objects.requireNonNull(postingsEnum);
this.freq = freq;
this.posCounter = this.postingsEnum.freq();
}
public OfPostings(BytesRef term, PostingsEnum postingsEnum) throws IOException {
this(term, postingsEnum.freq(), postingsEnum);
}
public PostingsEnum getPostingsEnum() {
return postingsEnum;
}
@Override
public boolean nextPosition() throws IOException {
if (posCounter > 0) {
posCounter--;
postingsEnum.nextPosition();
return true;
} else {
return false;
}
}
@Override
public BytesRef getTerm() throws IOException {
return term;
}
@Override
public int startOffset() throws IOException {
return postingsEnum.startOffset();
}
@Override
public int endOffset() throws IOException {
return postingsEnum.endOffset();
}
@Override
public int freq() throws IOException {
return freq;
}
}
public static class OfMatchesIteratorWithSubs extends OffsetsEnum {
private final PriorityQueue<OffsetsEnum> pendingQueue = new PriorityQueue<>();
private final HashMap<Query,BytesRef> queryToTermMap = new HashMap<>();
public OfMatchesIteratorWithSubs(MatchesIterator matchesIterator) {
pendingQueue.add(new OfMatchesIterator(matchesIterator, () -> queryToTerm(matchesIterator.getQuery())));
}
@Override
public boolean nextPosition() throws IOException {
OffsetsEnum formerHeadOE = pendingQueue.poll();
if (formerHeadOE instanceof CachedOE) {
OffsetsEnum newHeadOE = pendingQueue.peek();
if (newHeadOE instanceof OfMatchesIterator) {
nextWhenMatchesIterator((OfMatchesIterator) newHeadOE);
}
} else {
OfMatchesIterator miOE = (OfMatchesIterator) formerHeadOE;
if (miOE.nextPosition()) {
nextWhenMatchesIterator(miOE);
}
}
return pendingQueue.isEmpty() == false;
}
private void nextWhenMatchesIterator(OfMatchesIterator miOE) throws IOException {
boolean isHead = miOE == pendingQueue.peek();
MatchesIterator subMatches = miOE.matchesIterator.getSubMatches();
if (subMatches != null) {
if (isHead) {
pendingQueue.poll();
}
enqueueCachedMatches(subMatches);
if (miOE.nextPosition()) {
pendingQueue.add(miOE);
assert pendingQueue.peek() != miOE;
}
} else {
if (!isHead) {
pendingQueue.add(miOE);
}
}
}
private boolean enqueueCachedMatches(MatchesIterator thisMI) throws IOException {
if (thisMI == null) {
return false;
} else {
while (thisMI.next()) {
if (false == enqueueCachedMatches(thisMI.getSubMatches())) {
pendingQueue.add(new CachedOE(queryToTerm(thisMI.getQuery()), thisMI.startOffset(), thisMI.endOffset()));
}
}
return true;
}
}
private BytesRef queryToTerm(Query query) {
return queryToTermMap.computeIfAbsent(query, (Query q) -> {
BytesRefBuilder bytesRefBuilder = new BytesRefBuilder();
q.visit(new QueryVisitor() {
@Override
public void consumeTerms(Query query, Term... terms) {
for (Term term : terms) {
if (bytesRefBuilder.length() > 0) {
bytesRefBuilder.append((byte) ' ');
}
bytesRefBuilder.append(term.bytes());
}
}
});
if (bytesRefBuilder.length() > 0) {
return bytesRefBuilder.get();
}
return new BytesRef(q.toString());
});
}
@Override
public int freq() throws IOException {
return pendingQueue.peek().freq();
}
@Override
public BytesRef getTerm() throws IOException {
return pendingQueue.peek().getTerm();
}
@Override
public int startOffset() throws IOException {
return pendingQueue.peek().startOffset();
}
@Override
public int endOffset() throws IOException {
return pendingQueue.peek().endOffset();
}
private static class CachedOE extends OffsetsEnum {
final BytesRef term;
final int startOffset;
final int endOffset;
private CachedOE(BytesRef term, int startOffset, int endOffset) {
this.term = term;
this.startOffset = startOffset;
this.endOffset = endOffset;
}
@Override
public boolean nextPosition() throws IOException {
return false;
}
@Override
public int freq() throws IOException {
return 1;
}
@Override
public BytesRef getTerm() throws IOException {
return term;
}
@Override
public int startOffset() throws IOException {
return startOffset;
}
@Override
public int endOffset() throws IOException {
return endOffset;
}
}
}
public static class OfMatchesIterator extends OffsetsEnum {
private final MatchesIterator matchesIterator;
private final Supplier<BytesRef> termSupplier;
public OfMatchesIterator(MatchesIterator matchesIterator, Supplier<BytesRef> termSupplier) {
this.matchesIterator = matchesIterator;
this.termSupplier = termSupplier;
}
@Override
public boolean nextPosition() throws IOException {
return matchesIterator.next();
}
@Override
public int freq() throws IOException {
return 1;
}
@Override
public BytesRef getTerm() throws IOException {
return termSupplier.get();
}
@Override
public int startOffset() throws IOException {
return matchesIterator.startOffset();
}
@Override
public int endOffset() throws IOException {
return matchesIterator.endOffset();
}
}
public static final OffsetsEnum EMPTY = new OffsetsEnum() {
@Override
public boolean nextPosition() throws IOException {
return false;
}
@Override
public BytesRef getTerm() throws IOException {
throw new UnsupportedOperationException();
}
@Override
public int startOffset() throws IOException {
throw new UnsupportedOperationException();
}
@Override
public int endOffset() throws IOException {
throw new UnsupportedOperationException();
}
@Override
public int freq() throws IOException {
return 0;
}
};
public static class MultiOffsetsEnum extends OffsetsEnum {
private final PriorityQueue<OffsetsEnum> queue;
private boolean started = false;
public MultiOffsetsEnum(List<OffsetsEnum> inner) throws IOException {
this.queue = new PriorityQueue<>();
for (OffsetsEnum oe : inner) {
if (oe.nextPosition())
this.queue.add(oe);
}
}
@Override
public boolean nextPosition() throws IOException {
if (started == false) {
started = true;
return this.queue.size() > 0;
}
if (this.queue.size() > 0) {
OffsetsEnum top = this.queue.poll();
if (top.nextPosition()) {
this.queue.add(top);
return true;
}
else {
top.close();
}
return this.queue.size() > 0;
}
return false;
}
@Override
public BytesRef getTerm() throws IOException {
return this.queue.peek().getTerm();
}
@Override
public int startOffset() throws IOException {
return this.queue.peek().startOffset();
}
@Override
public int endOffset() throws IOException {
return this.queue.peek().endOffset();
}
@Override
public int freq() throws IOException {
return this.queue.peek().freq();
}
@Override
public void close() throws IOException {
IOUtils.close(queue);
}
}
}