rhelp/org.eclipse.statet.rhelp.core/src/org/eclipse/statet/internal/rhelp/core/index/RequestStreamCollector.java

/*=============================================================================#
 # Copyright (c) 2010, 2019 Stephan Wahlbrink and others.
 # 
 # This program and the accompanying materials are made available under the
 # terms of the Eclipse Public License 2.0 which is available at
 # https://www.eclipse.org/legal/epl-2.0, or the Apache License, Version 2.0
 # which is available at https://www.apache.org/licenses/LICENSE-2.0.
 # 
 # SPDX-License-Identifier: EPL-2.0 OR Apache-2.0
 # 
 # Contributors:
 #     Stephan Wahlbrink <sw@wahlbrink.eu> - initial API and implementation
 #=============================================================================*/

package org.eclipse.statet.internal.rhelp.core.index;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicInteger;

import org.apache.lucene.index.FieldInfo;
import org.apache.lucene.index.LeafReader;
import org.apache.lucene.search.vectorhighlight.FieldQuery;
import org.apache.lucene.search.vectorhighlight.ScoreOrderFragmentsBuilder;
import org.apache.lucene.search.vectorhighlight.SimpleFragListBuilder;

import org.eclipse.statet.jcommons.collections.ImCollections;
import org.eclipse.statet.jcommons.collections.ImSet;
import org.eclipse.statet.jcommons.lang.NonNull;
import org.eclipse.statet.jcommons.lang.NonNullByDefault;
import org.eclipse.statet.jcommons.lang.Nullable;

import org.eclipse.statet.internal.rhelp.core.RHelpSearchMatchImpl;
import org.eclipse.statet.rhelp.core.RHelpPage;
import org.eclipse.statet.rhelp.core.RHelpSearchMatch;
import org.eclipse.statet.rhelp.core.RHelpSearchMatch.MatchFragment;
import org.eclipse.statet.rhelp.core.RHelpSearchRequestor;
import org.eclipse.statet.rhelp.core.RPkgHelp;


@NonNullByDefault
public class RequestStreamCollector extends DocFieldVisitorCollector.Visitor implements REnvIndexSchema {
	
	
	private static final Highlighter HIGHLIGHTER;
	
	static {
		final SimpleFragListBuilder fragListBuilder= new SimpleFragListBuilder(10);
		final ScoreOrderFragmentsBuilder fragmentsBuilder= new ScoreOrderFragmentsBuilder(
				new HighlightBoundaryScanner() );
		fragmentsBuilder.setDiscreteMultiValueHighlighting(true);
		HIGHLIGHTER= new Highlighter(true, true, fragListBuilder, fragmentsBuilder);
	}
	
	
	private static final ImSet<String> LOAD_ID_SELECTOR= ImCollections.newSet(
			PAGE_FIELD_NAME,
			PACKAGE_FIELD_NAME );
	
	
	private final REnvIndexSearchQuery indexQuery;
	
	private final Map<String, RPkgHelp> packageMap;
	
	private final RHelpSearchRequestor requestor;
	
	private final List<MatchFragment> fragmentCollection= new ArrayList<>();
	
	private @Nullable FieldQuery fieldQuery;
	private final int maxNumFragments;
	
	private LeafReader reader;
	
	private int doc;
	private float score;
	
	private @Nullable String pkgName;
	private @Nullable String pageName;
	
	
	public RequestStreamCollector(final REnvIndexSearchQuery indexQuery, final Map<String, RPkgHelp> packageMap,
			final RHelpSearchRequestor requestor) throws IOException {
		super(LOAD_ID_SELECTOR);
		
		this.indexQuery= indexQuery;
		this.packageMap= packageMap;
		this.requestor= requestor;
		
		this.maxNumFragments= requestor.getMaxFragments();
	}
	
	
	@Override
	public void setReader(final LeafReader reader) throws IOException {
		this.reader= reader;
		if (this.indexQuery.getQuery() != null
				&& this.indexQuery.fieldNames != null && this.indexQuery.fieldNames.size() > 0 && this.maxNumFragments > 0) {
			this.fieldQuery= HIGHLIGHTER.getFieldQuery(this.indexQuery.getQuery(), this.reader);
		}
		else {
			this.fieldQuery= null;
		}
	}
	
	@Override
	public void newDocMatch(final int doc, final float score) {
		this.doc= doc;
		this.score= score;
		this.pkgName= null;
		this.pageName= null;
	}
	
	@Override
	public void stringField(final FieldInfo fieldInfo, final byte[] value) throws IOException {
		switch (fieldInfo.name) {
		case PACKAGE_FIELD_NAME:
			this.pkgName= REnvIndexUtils.toString(value);
			return;
		case PAGE_FIELD_NAME:
			this.pageName= REnvIndexUtils.toString(value);
			return;
		default:
			return;
		}
	}
	
	@Override
	public void finalizeDocMatch() throws IOException {
		if (this.pkgName != null && this.pageName != null) {
			final RPkgHelp pkgHelp= this.packageMap.get(this.pkgName);
			if (pkgHelp != null) {
				final RHelpPage page= pkgHelp.getPage(this.pageName);
				if (page != null) {
//					System.out.println(pkgName + "?" + page.getName() + ": " + score);
					final int matchCount= checkMatches();
					final RHelpSearchMatchImpl match= (matchCount >= 0) ?
							new RHelpSearchMatchImpl(page, this.score,
									matchCount, this.fragmentCollection.toArray(new @NonNull MatchFragment[this.fragmentCollection.size()]) ) :
							new RHelpSearchMatchImpl(page, this.score);
					
					this.requestor.matchFound(match);
					return;
				}
			}
		}
	}
	
	private int checkMatches() throws IOException {
		final FieldQuery fieldQuery= this.fieldQuery;
		if (fieldQuery == null) {
			return -1;
		}
		final AtomicInteger counter= new AtomicInteger();
		this.fragmentCollection.clear();
		for (final String fieldName : this.indexQuery.fieldNames) {
			final String[] fragments= HIGHLIGHTER.getBestFragments(fieldQuery, this.reader,
					this.doc, fieldName, 80, this.maxNumFragments,
					RHelpSearchMatch.PRE_TAGS, RHelpSearchMatch.POST_TAGS, Highlighter.DEFAULT_ENCODER,
					counter );
			if (fragments != null) {
				for (int j= 0; j < fragments.length; j++) {
					this.fragmentCollection.add(
							new RHelpSearchMatchImpl.Fragment(fieldName, fragments[j]) );
				}
			}
		}
		return counter.get();
	}
	
}