package org.apache.lucene.search.uhighlight;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import org.apache.lucene.analysis.Analyzer;
import org.apache.lucene.analysis.FilteringTokenFilter;
import org.apache.lucene.analysis.TokenStream;
import org.apache.lucene.analysis.tokenattributes.CharTermAttribute;
import org.apache.lucene.index.LeafReader;
import org.apache.lucene.index.memory.MemoryIndex;
import org.apache.lucene.search.spans.SpanQuery;
import org.apache.lucene.util.automaton.Automata;
import org.apache.lucene.util.automaton.CharacterRunAutomaton;
public class MemoryIndexOffsetStrategy extends AnalysisOffsetStrategy {
private final MemoryIndex memoryIndex;
private final LeafReader memIndexLeafReader;
private final CharacterRunAutomaton preMemIndexFilterAutomaton;
public MemoryIndexOffsetStrategy(UHComponents components, Analyzer analyzer) {
super(components, analyzer);
boolean storePayloads = components.getPhraseHelper().hasPositionSensitivity();
memoryIndex = new MemoryIndex(true, storePayloads);
memIndexLeafReader = (LeafReader) memoryIndex.createSearcher().getIndexReader();
preMemIndexFilterAutomaton = buildCombinedAutomaton(components);
}
private static CharacterRunAutomaton buildCombinedAutomaton(UHComponents components) {
if (components.getTerms() == null || components.getAutomata() == null) {
return null;
}
List<CharacterRunAutomaton> allAutomata = new ArrayList<>();
if (components.getTerms().length > 0) {
allAutomata.add(new CharacterRunAutomaton(Automata.makeStringUnion(Arrays.asList(components.getTerms()))));
}
Collections.addAll(allAutomata, components.getAutomata());
for (SpanQuery spanQuery : components.getPhraseHelper().getSpanQueries()) {
Collections.addAll(allAutomata,
MultiTermHighlighting.extractAutomata(spanQuery, components.getFieldMatcher(), true));
}
if (allAutomata.size() == 1) {
return allAutomata.get(0);
}
return new CharacterRunAutomaton(Automata.makeEmpty()) {
@Override
public boolean run(char[] chars, int offset, int length) {
for (int i = 0; i < allAutomata.size(); i++) {
if (allAutomata.get(i).run(chars, offset, length)) {
return true;
}
}
return false;
}
};
}
@Override
public OffsetsEnum getOffsetsEnum(LeafReader reader, int docId, String content) throws IOException {
TokenStream tokenStream = tokenStream(content);
if (preMemIndexFilterAutomaton != null) {
tokenStream = newKeepWordFilter(tokenStream, preMemIndexFilterAutomaton);
}
memoryIndex.reset();
memoryIndex.addField(getField(), tokenStream);
if (reader == null) {
return createOffsetsEnumFromReader(memIndexLeafReader, 0);
} else {
return createOffsetsEnumFromReader(
new OverlaySingleDocTermsLeafReader(
reader,
memIndexLeafReader,
getField(),
docId),
docId);
}
}
private static FilteringTokenFilter newKeepWordFilter(final TokenStream tokenStream,
final CharacterRunAutomaton charRunAutomaton) {
return new FilteringTokenFilter(tokenStream) {
final CharTermAttribute charAtt = addAttribute(CharTermAttribute.class);
@Override
protected boolean accept() throws IOException {
return charRunAutomaton.run(charAtt.buffer(), 0, charAtt.length());
}
};
}
}