HSEARCH-3881 Expose the extended search predicate factory in extended…

… search aggregation factories

I.e. make sure that this works:

    .search( ... ).extension( LuceneExtension.get() )
        .where( ... )
        .aggregation( f -> f.terms()...
                .filter( pf -> pf.fromLuceneQuery( ... ) ) )

... and that we don't have to call .extension() again on "pf", because
the provided factory is already extended.

backend/elasticsearch/src/main/java/org/hibernate/search/backend/elasticsearch/ElasticsearchExtension.java

@@ -203,11 +203,12 @@ public Optional<ElasticsearchSearchProjectionFactory<R, E>> extendOptional(
 	@Override
 	@SuppressWarnings("unchecked") // If the factory is an instance of ElasticsearchSearchAggregationBuilderFactory, the cast is safe
 	public Optional<ElasticsearchSearchAggregationFactory> extendOptional(
-			SearchAggregationFactory original, SearchAggregationDslContext<?> dslContext) {
+			SearchAggregationFactory original, SearchAggregationDslContext<?, ?> dslContext) {
 		if ( dslContext.getBuilderFactory() instanceof ElasticsearchSearchAggregationBuilderFactory ) {
 			return Optional.of( new ElasticsearchSearchAggregationFactoryImpl(
 					original,
-					(SearchAggregationDslContext<ElasticsearchSearchAggregationBuilderFactory>) dslContext
+					((SearchAggregationDslContext<ElasticsearchSearchAggregationBuilderFactory, ?>) dslContext)
+							.withExtendedPredicateFactory( this )
 			) );
 		}
 		else {

backend/elasticsearch/src/main/java/org/hibernate/search/backend/elasticsearch/search/aggregation/dsl/ElasticsearchSearchAggregationFactory.java

@@ -6,12 +6,14 @@
  */
 package org.hibernate.search.backend.elasticsearch.search.aggregation.dsl;
 
+import org.hibernate.search.backend.elasticsearch.search.predicate.dsl.ElasticsearchSearchPredicateFactory;
 import org.hibernate.search.engine.search.aggregation.dsl.AggregationFinalStep;
-import org.hibernate.search.engine.search.aggregation.dsl.SearchAggregationFactory;
+import org.hibernate.search.engine.search.aggregation.dsl.ExtendedSearchAggregationFactory;
 
 import com.google.gson.JsonObject;
 
-public interface ElasticsearchSearchAggregationFactory extends SearchAggregationFactory {
+public interface ElasticsearchSearchAggregationFactory
+		extends ExtendedSearchAggregationFactory<ElasticsearchSearchPredicateFactory> {
 
 	/**
 	 * Create an aggregation from JSON.

backend/elasticsearch/src/main/java/org/hibernate/search/backend/elasticsearch/search/aggregation/dsl/impl/ElasticsearchSearchAggregationFactoryImpl.java

@@ -8,6 +8,7 @@
 
 import org.hibernate.search.backend.elasticsearch.search.aggregation.impl.ElasticsearchSearchAggregationBuilderFactory;
 import org.hibernate.search.backend.elasticsearch.search.aggregation.dsl.ElasticsearchSearchAggregationFactory;
+import org.hibernate.search.backend.elasticsearch.search.predicate.dsl.ElasticsearchSearchPredicateFactory;
 import org.hibernate.search.engine.search.aggregation.dsl.AggregationFinalStep;
 import org.hibernate.search.engine.search.aggregation.dsl.SearchAggregationFactory;
 import org.hibernate.search.engine.search.aggregation.dsl.spi.DelegatingSearchAggregationFactory;
@@ -16,14 +17,14 @@
 import com.google.gson.JsonObject;
 
 public class ElasticsearchSearchAggregationFactoryImpl
-		extends DelegatingSearchAggregationFactory
+		extends DelegatingSearchAggregationFactory<ElasticsearchSearchPredicateFactory>
 		implements ElasticsearchSearchAggregationFactory {
 
-	private final SearchAggregationDslContext<ElasticsearchSearchAggregationBuilderFactory> dslContext;
+	private final SearchAggregationDslContext<ElasticsearchSearchAggregationBuilderFactory, ?> dslContext;
 
 	public ElasticsearchSearchAggregationFactoryImpl(SearchAggregationFactory delegate,
-			SearchAggregationDslContext<ElasticsearchSearchAggregationBuilderFactory> dslContext) {
-		super( delegate );
+			SearchAggregationDslContext<ElasticsearchSearchAggregationBuilderFactory, ElasticsearchSearchPredicateFactory> dslContext) {
+		super( delegate, dslContext );
 		this.dslContext = dslContext;
 	}
 

backend/lucene/src/main/java/org/hibernate/search/backend/lucene/LuceneExtension.java

@@ -203,11 +203,12 @@ public Optional<LuceneSearchProjectionFactory<R, E>> extendOptional(
 	@Override
 	@SuppressWarnings("unchecked") // If the factory is an instance of LuceneSearchAggregationBuilderFactory, the cast is safe
 	public Optional<LuceneSearchAggregationFactory> extendOptional(
-			SearchAggregationFactory original, SearchAggregationDslContext<?> dslContext) {
+			SearchAggregationFactory original, SearchAggregationDslContext<?, ?> dslContext) {
 		if ( dslContext.getBuilderFactory() instanceof LuceneSearchAggregationBuilderFactory ) {
 			return Optional.of( new LuceneSearchAggregationFactoryImpl(
 					original,
-					(SearchAggregationDslContext<LuceneSearchAggregationBuilderFactory>) dslContext
+					((SearchAggregationDslContext<LuceneSearchAggregationBuilderFactory, ?>) dslContext)
+							.withExtendedPredicateFactory( this )
 			) );
 		}
 		else {

backend/lucene/src/main/java/org/hibernate/search/backend/lucene/search/aggregation/dsl/LuceneSearchAggregationFactory.java

@@ -6,7 +6,8 @@
  */
 package org.hibernate.search.backend.lucene.search.aggregation.dsl;
 
-import org.hibernate.search.engine.search.aggregation.dsl.SearchAggregationFactory;
+import org.hibernate.search.backend.lucene.search.predicate.dsl.LuceneSearchPredicateFactory;
+import org.hibernate.search.engine.search.aggregation.dsl.ExtendedSearchAggregationFactory;
 
-public interface LuceneSearchAggregationFactory extends SearchAggregationFactory {
+public interface LuceneSearchAggregationFactory extends ExtendedSearchAggregationFactory<LuceneSearchPredicateFactory> {
 }

backend/lucene/src/main/java/org/hibernate/search/backend/lucene/search/aggregation/dsl/impl/LuceneSearchAggregationFactoryImpl.java

@@ -8,20 +8,18 @@
 
 import org.hibernate.search.backend.lucene.search.aggregation.impl.LuceneSearchAggregationBuilderFactory;
 import org.hibernate.search.backend.lucene.search.aggregation.dsl.LuceneSearchAggregationFactory;
+import org.hibernate.search.backend.lucene.search.predicate.dsl.LuceneSearchPredicateFactory;
 import org.hibernate.search.engine.search.aggregation.dsl.SearchAggregationFactory;
 import org.hibernate.search.engine.search.aggregation.dsl.spi.DelegatingSearchAggregationFactory;
 import org.hibernate.search.engine.search.aggregation.dsl.spi.SearchAggregationDslContext;
 
 public class LuceneSearchAggregationFactoryImpl
-		extends DelegatingSearchAggregationFactory
+		extends DelegatingSearchAggregationFactory<LuceneSearchPredicateFactory>
 		implements LuceneSearchAggregationFactory {
 
-	private final SearchAggregationDslContext<LuceneSearchAggregationBuilderFactory> dslContext;
-
 	public LuceneSearchAggregationFactoryImpl(SearchAggregationFactory delegate,
-			SearchAggregationDslContext<LuceneSearchAggregationBuilderFactory> dslContext) {
-		super( delegate );
-		this.dslContext = dslContext;
+			SearchAggregationDslContext<LuceneSearchAggregationBuilderFactory, LuceneSearchPredicateFactory> dslContext) {
+		super( delegate, dslContext );
 	}
 
 	// Empty: no extension at the moment.

engine/src/main/java/org/hibernate/search/engine/mapper/scope/impl/MappedIndexScopeImpl.java

@@ -72,7 +72,11 @@ public SearchProjectionFactory<R, E> projection() {
 	@Override
 	public SearchAggregationFactory aggregation() {
 		return new DefaultSearchAggregationFactory(
-				SearchAggregationDslContextImpl.root( delegate.getSearchAggregationFactory(), delegate.getSearchPredicateBuilderFactory() )
+				SearchAggregationDslContextImpl.root(
+						delegate.getSearchAggregationFactory(),
+						predicate(),
+						delegate.getSearchPredicateBuilderFactory()
+				)
 		);
 	}
 }

engine/src/main/java/org/hibernate/search/engine/search/aggregation/dsl/AggregationFilterStep.java

@@ -0,0 +1,60 @@
+/*
+ * Hibernate Search, full-text search for your domain model
+ *
+ * License: GNU Lesser General Public License (LGPL), version 2.1 or later
+ * See the lgpl.txt file in the root directory or <http://www.gnu.org/licenses/lgpl-2.1.html>.
+ */
+package org.hibernate.search.engine.search.aggregation.dsl;
+
+import java.util.function.Function;
+
+import org.hibernate.search.engine.search.predicate.SearchPredicate;
+import org.hibernate.search.engine.search.predicate.dsl.PredicateFinalStep;
+import org.hibernate.search.engine.search.predicate.dsl.SearchPredicateFactory;
+
+/**
+ * The step in an aggregation definition where a filter can be set
+ * to select nested objects from which values will be extracted for this aggregation.
+ *
+ * @param <S> The "self" type (the actual exposed type of this step)
+ * @param <PDF> The type of factory used to create predicates in {@link #filter(Function)}.
+ */
+public interface AggregationFilterStep<S, PDF extends SearchPredicateFactory> {
+
+	/**
+	 * Filter nested objects from which values will be extracted for this aggregation.
+	 * <p>
+	 * The filter is based on a previously-built {@link SearchPredicate}.
+	 *
+	 * @param searchPredicate The predicate that must match.
+	 * @return {@code this}, for method chaining.
+	 */
+	S filter(SearchPredicate searchPredicate);
+
+	/**
+	 * Filter nested objects from which values will be extracted for this aggregation.
+	 * <p>
+	 * The filter is defined by the given function.
+	 * <p>
+	 * Best used with lambda expressions.
+	 *
+	 * @param clauseContributor A function that will use the factory passed in parameter to create a predicate,
+	 * returning the final step in the predicate DSL.
+	 * Should generally be a lambda expression.
+	 * @return {@code this}, for method chaining.
+	 */
+	S filter(Function<? super PDF, ? extends PredicateFinalStep> clauseContributor);
+
+	/**
+	 * Filter nested objects from which values will be extracted for this aggregation.
+	 * <p>
+	 * The filter is based on an almost-built {@link SearchPredicate}.
+	 *
+	 * @param dslFinalStep A final step in the predicate DSL allowing the retrieval of a {@link SearchPredicate}.
+	 * @return {@code this}, for method chaining.
+	 */
+	default S filter(PredicateFinalStep dslFinalStep) {
+		return filter( dslFinalStep.toPredicate() );
+	}
+
+}

engine/src/main/java/org/hibernate/search/engine/search/aggregation/dsl/ExtendedSearchAggregationFactory.java

@@ -0,0 +1,30 @@
+/*
+ * Hibernate Search, full-text search for your domain model
+ *
+ * License: GNU Lesser General Public License (LGPL), version 2.1 or later
+ * See the lgpl.txt file in the root directory or <http://www.gnu.org/licenses/lgpl-2.1.html>.
+ */
+package org.hibernate.search.engine.search.aggregation.dsl;
+
+import java.util.function.Function;
+
+import org.hibernate.search.engine.search.predicate.dsl.SearchPredicateFactory;
+
+/**
+ * A base interface for subtypes of {@link SearchAggregationFactory} allowing to
+ * easily override the predicate factory type for all relevant methods.
+ * <p>
+ * <strong>Warning:</strong> Generic parameters of this type are subject to change,
+ * so this type should not be referenced directtly in user code.
+ *
+ * @param <PDF> The type of factory used to create predicates in {@link AggregationFilterStep#filter(Function)}.
+ */
+public interface ExtendedSearchAggregationFactory<PDF extends SearchPredicateFactory>
+		extends SearchAggregationFactory {
+
+	@Override
+	RangeAggregationFieldStep<PDF> range();
+
+	@Override
+	TermsAggregationFieldStep<PDF> terms();
+}

engine/src/main/java/org/hibernate/search/engine/search/aggregation/dsl/RangeAggregationFieldStep.java

@@ -6,12 +6,17 @@
  */
 package org.hibernate.search.engine.search.aggregation.dsl;
 
+import java.util.function.Function;
+
 import org.hibernate.search.engine.search.common.ValueConvert;
+import org.hibernate.search.engine.search.predicate.dsl.SearchPredicateFactory;
 
 /**
  * The initial step in a "range" aggregation definition, where the target field can be set.
+ *
+ * @param <PDF> The type of factory used to create predicates in {@link AggregationFilterStep#filter(Function)}.
  */
-public interface RangeAggregationFieldStep {
+public interface RangeAggregationFieldStep<PDF extends SearchPredicateFactory> {
 
 	/**
 	 * Target the given field in the range aggregation.
@@ -21,7 +26,7 @@ public interface RangeAggregationFieldStep {
 	 * @param <F> The type of field values.
 	 * @return The next step.
 	 */
-	default <F> RangeAggregationRangeStep<?, F> field(String absoluteFieldPath, Class<F> type) {
+	default <F> RangeAggregationRangeStep<?, PDF, F> field(String absoluteFieldPath, Class<F> type) {
 		return field( absoluteFieldPath, type, ValueConvert.YES );
 	}
 
@@ -35,6 +40,6 @@ default <F> RangeAggregationRangeStep<?, F> field(String absoluteFieldPath, Clas
 	 * See {@link ValueConvert}.
 	 * @return The next step.
 	 */
-	<F> RangeAggregationRangeStep<?, F> field(String absoluteFieldPath, Class<F> type, ValueConvert convert);
+	<F> RangeAggregationRangeStep<?, PDF, F> field(String absoluteFieldPath, Class<F> type, ValueConvert convert);
 
 }

engine/src/main/java/org/hibernate/search/engine/search/aggregation/dsl/RangeAggregationOptionsStep.java

@@ -7,47 +7,22 @@
 package org.hibernate.search.engine.search.aggregation.dsl;
 
 import java.util.function.Function;
-import org.hibernate.search.engine.search.predicate.SearchPredicate;
-import org.hibernate.search.engine.search.predicate.dsl.PredicateFinalStep;
 import org.hibernate.search.engine.search.predicate.dsl.SearchPredicateFactory;
 
 /**
  * The final step in a "range" aggregation definition, where optional parameters can be set.
  *
  * @param <S> The "self" type (the actual exposed type of this step).
+ * @param <PDF> The type of factory used to create predicates in {@link #filter(Function)}.
  * @param <F> The type of the targeted field.
  * @param <A> The type of result for this aggregation.
  */
-public interface RangeAggregationOptionsStep<S extends RangeAggregationOptionsStep<?, F, A>, F, A>
-		extends AggregationFinalStep<A> {
+public interface RangeAggregationOptionsStep<
+				S extends RangeAggregationOptionsStep<?, PDF, F, A>,
+				PDF extends SearchPredicateFactory,
+				F,
+				A
+		>
+		extends AggregationFinalStep<A>, AggregationFilterStep<S, PDF> {
 
-	/**
-	 * Add a <a href="#filter">"filter" clause</a> based on a previously-built {@link SearchPredicate}.
-	 *
-	 * @param searchPredicate The predicate that must match.
-	 * @return {@code this}, for method chaining.
-	 */
-	S filter(SearchPredicate searchPredicate);
-
-	/**
-	 * Add a <a href="#filter">"filter" clause</a> to be defined by the given function.
-	 * <p>
-	 * Best used with lambda expressions.
-	 *
-	 * @param clauseContributor A function that will use the factory passed in parameter to create a predicate,
-	 * returning the final step in the predicate DSL.
-	 * Should generally be a lambda expression.
-	 * @return {@code this}, for method chaining.
-	 */
-	S filter(Function<? super SearchPredicateFactory, ? extends PredicateFinalStep> clauseContributor);
-
-	/**
-	 * Add a <a href="#filter">"filter" clause</a> based on an almost-built {@link SearchPredicate}.
-	 *
-	 * @param dslFinalStep A final step in the predicate DSL allowing the retrieval of a {@link SearchPredicate}.
-	 * @return {@code this}, for method chaining.
-	 */
-	default S filter(PredicateFinalStep dslFinalStep) {
-		return filter( dslFinalStep.toPredicate() );
-	}
 }

engine/src/main/java/org/hibernate/search/engine/search/aggregation/dsl/RangeAggregationRangeMoreStep.java

@@ -7,7 +7,9 @@
 package org.hibernate.search.engine.search.aggregation.dsl;
 
 import java.util.Map;
+import java.util.function.Function;
 
+import org.hibernate.search.engine.search.predicate.dsl.SearchPredicateFactory;
 import org.hibernate.search.util.common.data.Range;
 
 /**
@@ -16,15 +18,17 @@
  * or more ranges can be added.
  *
  * @param <S> The "self" type (the actual exposed type of this step).
+ * @param <PDF> The type of factory used to create predicates in {@link #filter(Function)}.
  * @param <N> The type of the next step.
  * @param <F> The type of the targeted field.
  */
 public interface RangeAggregationRangeMoreStep<
-				S extends RangeAggregationRangeMoreStep<?, ?, F>,
-				N extends RangeAggregationOptionsStep<?, F, Map<Range<F>, Long>>,
+				S extends RangeAggregationRangeMoreStep<?, ?, PDF, F>,
+				N extends RangeAggregationOptionsStep<?, PDF, F, Map<Range<F>, Long>>,
+				PDF extends SearchPredicateFactory,
 				F
 		>
-		extends RangeAggregationOptionsStep<N, F, Map<Range<F>, Long>>,
-				RangeAggregationRangeStep<S, F> {
+		extends RangeAggregationOptionsStep<N, PDF, F, Map<Range<F>, Long>>,
+				RangeAggregationRangeStep<S, PDF, F> {
 
 }

engine/src/main/java/org/hibernate/search/engine/search/aggregation/dsl/RangeAggregationRangeStep.java

@@ -7,16 +7,23 @@
 package org.hibernate.search.engine.search.aggregation.dsl;
 
 import java.util.Collection;
+import java.util.function.Function;
 
+import org.hibernate.search.engine.search.predicate.dsl.SearchPredicateFactory;
 import org.hibernate.search.util.common.data.Range;
 
 /**
  * The step in a "range" aggregation definition where the ranges can be set.
  *
  * @param <N> The type of the next step.
+ * @param <PDF> The type of factory used to create predicates in {@link AggregationFilterStep#filter(Function)}.
  * @param <F> The type of the targeted field.
  */
-public interface RangeAggregationRangeStep<N extends RangeAggregationRangeMoreStep<?, ?, F>, F> {
+public interface RangeAggregationRangeStep<
+				N extends RangeAggregationRangeMoreStep<?, ?, PDF, F>,
+				PDF extends SearchPredicateFactory,
+				F
+		> {
 
 	/**
 	 * Add a bucket for the range {@code [lowerBound, upperBound)} (lower bound included, upper bound excluded),

engine/src/main/java/org/hibernate/search/engine/search/aggregation/dsl/SearchAggregationFactory.java

@@ -28,7 +28,7 @@ public interface SearchAggregationFactory {
 	 *
 	 * @return The next step.
 	 */
-	RangeAggregationFieldStep range();
+	RangeAggregationFieldStep<?> range();
 
 	/**
 	 * Perform aggregation in term buckets.
@@ -43,7 +43,7 @@ public interface SearchAggregationFactory {
 	 *
 	 * @return The next step.
 	 */
-	TermsAggregationFieldStep terms();
+	TermsAggregationFieldStep<?> terms();
 
 	/**
 	 * Extend the current factory with the given extension,

engine/src/main/java/org/hibernate/search/engine/search/aggregation/dsl/SearchAggregationFactoryExtension.java

@@ -38,6 +38,6 @@ public interface SearchAggregationFactoryExtension<T> {
 	 * @return An optional containing the extended aggregation factory ({@link T}) in case
 	 * of success, or an empty optional otherwise.
 	 */
-	Optional<T> extendOptional(SearchAggregationFactory original, SearchAggregationDslContext<?> dslContext);
+	Optional<T> extendOptional(SearchAggregationFactory original, SearchAggregationDslContext<?, ?> dslContext);
 
 }