package org.apache.lucene.search;
import java.io.IOException;
import java.util.Arrays;
import org.apache.lucene.index.BaseTermsEnum;
import org.apache.lucene.index.ImpactsEnum;
import org.apache.lucene.index.PostingsEnum;
import org.apache.lucene.index.Term;
import org.apache.lucene.index.TermState;
import org.apache.lucene.index.Terms;
import org.apache.lucene.index.TermsEnum;
import org.apache.lucene.util.Attribute;
import org.apache.lucene.util.AttributeImpl;
import org.apache.lucene.util.AttributeReflector;
import org.apache.lucene.util.AttributeSource;
import org.apache.lucene.util.BytesRef;
import org.apache.lucene.util.BytesRefBuilder;
import org.apache.lucene.util.UnicodeUtil;
import org.apache.lucene.util.automaton.Automaton;
import org.apache.lucene.util.automaton.CompiledAutomaton;
import org.apache.lucene.util.automaton.LevenshteinAutomata;
public final class FuzzyTermsEnum extends BaseTermsEnum {
private TermsEnum actualEnum;
private final BoostAttribute boostAtt;
private final MaxNonCompetitiveBoostAttribute maxBoostAtt;
private final LevenshteinAutomataAttribute dfaAtt;
private float bottom;
private BytesRef bottomTerm;
private final CompiledAutomaton automata[];
private BytesRef queuedBottom;
final int termLength;
private int maxEdits;
final Terms terms;
final Term term;
final int termText[];
final int realPrefixLength;
final boolean transpositions;
public FuzzyTermsEnum(Terms terms, AttributeSource atts, Term term,
final int maxEdits, final int prefixLength, boolean transpositions) throws IOException {
if (maxEdits < 0 || maxEdits > LevenshteinAutomata.MAXIMUM_SUPPORTED_DISTANCE) {
throw new IllegalArgumentException("max edits must be 0.." + LevenshteinAutomata.MAXIMUM_SUPPORTED_DISTANCE + ", inclusive; got: " + maxEdits);
}
if (prefixLength < 0) {
throw new IllegalArgumentException("prefixLength cannot be less than 0");
}
this.maxEdits = maxEdits;
this.terms = terms;
this.term = term;
this.termText = stringToUTF32(term.text());
this.termLength = termText.length;
this.dfaAtt = atts.addAttribute(LevenshteinAutomataAttribute.class);
this.maxBoostAtt = atts.addAttribute(MaxNonCompetitiveBoostAttribute.class);
this.boostAtt = attributes().addAttribute(BoostAttribute.class);
this.realPrefixLength = prefixLength > termLength ? termLength : prefixLength;
this.transpositions = transpositions;
CompiledAutomaton[] prevAutomata = dfaAtt.automata();
if (prevAutomata == null) {
prevAutomata = new CompiledAutomaton[maxEdits+1];
Automaton[] automata = buildAutomata(termText, prefixLength, transpositions, maxEdits);
for (int i = 0; i <= maxEdits; i++) {
prevAutomata[i] = new CompiledAutomaton(automata[i], true, false);
}
dfaAtt.setAutomata(prevAutomata);
}
this.automata = prevAutomata;
bottom = maxBoostAtt.getMaxNonCompetitiveBoost();
bottomTerm = maxBoostAtt.getCompetitiveTerm();
bottomChanged(null);
}
public static Automaton buildAutomaton(String text, int prefixLength, boolean transpositions, int maxEdits) {
int[] termText = stringToUTF32(text);
Automaton[] automata = buildAutomata(termText, prefixLength, transpositions, maxEdits);
return automata[automata.length - 1];
}
private static int[] stringToUTF32(String text) {
int[] termText = new int[text.codePointCount(0, text.length())];
for (int cp, i = 0, j = 0; i < text.length(); i += Character.charCount(cp)) {
termText[j++] = cp = text.codePointAt(i);
}
return termText;
}
private static Automaton[] buildAutomata(int[] termText, int prefixLength, boolean transpositions, int maxEdits) {
if (maxEdits < 0 || maxEdits > LevenshteinAutomata.MAXIMUM_SUPPORTED_DISTANCE) {
throw new IllegalArgumentException("max edits must be 0.." + LevenshteinAutomata.MAXIMUM_SUPPORTED_DISTANCE + ", inclusive; got: " + maxEdits);
}
if (prefixLength < 0) {
throw new IllegalArgumentException("prefixLength cannot be less than 0");
}
Automaton[] automata = new Automaton[maxEdits + 1];
int termLength = termText.length;
prefixLength = Math.min(prefixLength, termLength);
String suffix = UnicodeUtil.newString(termText, prefixLength, termText.length - prefixLength);
LevenshteinAutomata builder = new LevenshteinAutomata(suffix, transpositions);
String prefix = UnicodeUtil.newString(termText, 0, prefixLength);
for (int i = 0; i <= maxEdits; i++) {
automata[i] = builder.toAutomaton(i, prefix);
}
return automata;
}
private TermsEnum getAutomatonEnum(int editDistance, BytesRef lastTerm) throws IOException {
assert editDistance < automata.length;
final CompiledAutomaton compiled = automata[editDistance];
BytesRef initialSeekTerm;
if (lastTerm == null) {
initialSeekTerm = null;
} else {
initialSeekTerm = compiled.floor(lastTerm, new BytesRefBuilder());
}
return terms.intersect(compiled, initialSeekTerm);
}
private void bottomChanged(BytesRef lastTerm) throws IOException {
int oldMaxEdits = maxEdits;
boolean termAfter = bottomTerm == null || (lastTerm != null && lastTerm.compareTo(bottomTerm) >= 0);
while (maxEdits > 0) {
float maxBoost = 1.0f - ((float) maxEdits / (float) termLength);
if (bottom < maxBoost || (bottom == maxBoost && termAfter == false)) {
break;
}
maxEdits--;
}
if (oldMaxEdits != maxEdits || lastTerm == null) {
actualEnum = getAutomatonEnum(maxEdits, lastTerm);
}
}
@Override
public BytesRef next() throws IOException {
if (queuedBottom != null) {
bottomChanged(queuedBottom);
queuedBottom = null;
}
BytesRef term;
term = actualEnum.next();
if (term == null) {
return null;
}
int ed = maxEdits;
while (ed > 0) {
if (matches(term, ed - 1)) {
ed--;
} else {
break;
}
}
if (ed == 0) {
boostAtt.setBoost(1.0F);
} else {
final int codePointCount = UnicodeUtil.codePointCount(term);
int minTermLength = Math.min(codePointCount, termLength);
float similarity = 1.0f - (float) ed / (float) minTermLength;
boostAtt.setBoost(similarity);
}
final float bottom = maxBoostAtt.getMaxNonCompetitiveBoost();
final BytesRef bottomTerm = maxBoostAtt.getCompetitiveTerm();
if (term != null && (bottom != this.bottom || bottomTerm != this.bottomTerm)) {
this.bottom = bottom;
this.bottomTerm = bottomTerm;
queuedBottom = BytesRef.deepCopyOf(term);
}
return term;
}
private boolean matches(BytesRef termIn, int k) {
return k == 0 ? termIn.equals(term.bytes()) : automata[k].runAutomaton.run(termIn.bytes, termIn.offset, termIn.length);
}
@Override
public int docFreq() throws IOException {
return actualEnum.docFreq();
}
@Override
public long totalTermFreq() throws IOException {
return actualEnum.totalTermFreq();
}
@Override
public PostingsEnum postings(PostingsEnum reuse, int flags) throws IOException {
return actualEnum.postings(reuse, flags);
}
@Override
public ImpactsEnum impacts(int flags) throws IOException {
return actualEnum.impacts(flags);
}
@Override
public void seekExact(BytesRef term, TermState state) throws IOException {
actualEnum.seekExact(term, state);
}
@Override
public TermState termState() throws IOException {
return actualEnum.termState();
}
@Override
public long ord() throws IOException {
return actualEnum.ord();
}
@Override
public boolean seekExact(BytesRef text) throws IOException {
return actualEnum.seekExact(text);
}
@Override
public SeekStatus seekCeil(BytesRef text) throws IOException {
return actualEnum.seekCeil(text);
}
@Override
public void seekExact(long ord) throws IOException {
actualEnum.seekExact(ord);
}
@Override
public BytesRef term() throws IOException {
return actualEnum.term();
}
public static interface LevenshteinAutomataAttribute extends Attribute {
public CompiledAutomaton[] automata();
public void setAutomata(CompiledAutomaton[] automata);
}
public static final class LevenshteinAutomataAttributeImpl extends AttributeImpl implements LevenshteinAutomataAttribute {
private CompiledAutomaton[] automata;
@Override
public CompiledAutomaton[] automata() {
return automata;
}
@Override
public void setAutomata(CompiledAutomaton[] automata) {
this.automata = automata;
}
@Override
public void clear() {
automata = null;
}
@Override
public int hashCode() {
if (automata == null) {
return 0;
} else {
return automata.hashCode();
}
}
@Override
public boolean equals(Object other) {
if (this == other)
return true;
if (!(other instanceof LevenshteinAutomataAttributeImpl))
return false;
return Arrays.equals(automata, ((LevenshteinAutomataAttributeImpl) other).automata);
}
@Override
public void copyTo(AttributeImpl _target) {
LevenshteinAutomataAttribute target = (LevenshteinAutomataAttribute) _target;
if (automata == null) {
target.setAutomata(null);
} else {
target.setAutomata(automata);
}
}
@Override
public void reflectWith(AttributeReflector reflector) {
reflector.reflect(LevenshteinAutomataAttribute.class, "automata", automata);
}
}
}