HSEARCH-3395 Make the .reference() and .object() projections type-safe in the Projection DSL

engine/src/main/java/org/hibernate/search/engine/common/impl/MappedIndexManagerImpl.java

@@ -6,6 +6,8 @@
  */
 package org.hibernate.search.engine.common.impl;
 
+import java.util.function.Function;
+
 import org.hibernate.search.engine.backend.document.DocumentElement;
 import org.hibernate.search.engine.backend.index.IndexManager;
 import org.hibernate.search.engine.backend.index.spi.IndexManagerImplementor;
@@ -14,6 +16,7 @@
 import org.hibernate.search.engine.mapper.mapping.context.spi.MappingContextImplementor;
 import org.hibernate.search.engine.mapper.session.context.spi.SessionContextImplementor;
 import org.hibernate.search.engine.mapper.mapping.spi.MappedIndexManager;
+import org.hibernate.search.engine.search.DocumentReference;
 
 class MappedIndexManagerImpl<D extends DocumentElement> implements MappedIndexManager<D> {
 
@@ -34,12 +37,16 @@ public IndexWorkPlan<D> createWorkPlan(SessionContextImplementor sessionContext)
 	}
 
 	@Override
-	public MappedIndexSearchTargetBuilder createSearchTargetBuilder(MappingContextImplementor mappingContext) {
-		return new MappedIndexSearchTargetBuilderImpl( implementor, mappingContext );
+	public <R, O> MappedIndexSearchTargetBuilder<R, O> createSearchTargetBuilder(MappingContextImplementor mappingContext,
+			Function<DocumentReference, R> documentReferenceTransformer) {
+		return new MappedIndexSearchTargetBuilderImpl<>(
+				implementor, mappingContext,
+				documentReferenceTransformer
+		);
 	}
 
 	@Override
-	public void addToSearchTarget(MappedIndexSearchTargetBuilder builder) {
-		((MappedIndexSearchTargetBuilderImpl) builder).add( implementor );
+	public void addToSearchTarget(MappedIndexSearchTargetBuilder<?, ?> builder) {
+		((MappedIndexSearchTargetBuilderImpl<?, ?>) builder).add( implementor );
 	}
 }

engine/src/main/java/org/hibernate/search/engine/common/impl/MappedIndexSearchTargetBuilderImpl.java

@@ -6,26 +6,32 @@
  */
 package org.hibernate.search.engine.common.impl;
 
+import java.util.function.Function;
+
 import org.hibernate.search.engine.backend.index.spi.IndexManagerImplementor;
 import org.hibernate.search.engine.backend.index.spi.IndexSearchTargetContextBuilder;
 import org.hibernate.search.engine.mapper.mapping.context.spi.MappingContextImplementor;
 import org.hibernate.search.engine.mapper.mapping.spi.MappedIndexSearchTarget;
 import org.hibernate.search.engine.mapper.mapping.spi.MappedIndexSearchTargetBuilder;
+import org.hibernate.search.engine.search.DocumentReference;
 
-class MappedIndexSearchTargetBuilderImpl implements MappedIndexSearchTargetBuilder {
+class MappedIndexSearchTargetBuilderImpl<R, O> implements MappedIndexSearchTargetBuilder<R, O> {
 	private final IndexSearchTargetContextBuilder delegate;
+	private final Function<DocumentReference, R> documentReferenceTransformer;
 
 	MappedIndexSearchTargetBuilderImpl(IndexManagerImplementor<?> firstIndexManager,
-			MappingContextImplementor mappingContext) {
+			MappingContextImplementor mappingContext,
+			Function<DocumentReference, R> documentReferenceTransformer) {
 		this.delegate = firstIndexManager.createSearchTargetContextBuilder( mappingContext );
+		this.documentReferenceTransformer = documentReferenceTransformer;
 	}
 
 	void add(IndexManagerImplementor<?> indexManager) {
 		indexManager.addToSearchTarget( delegate );
 	}
 
 	@Override
-	public MappedIndexSearchTarget build() {
-		return new MappedIndexSearchTargetImpl<>( delegate.build() );
+	public MappedIndexSearchTarget<R, O> build() {
+		return new MappedIndexSearchTargetImpl<>( delegate.build(), documentReferenceTransformer );
 	}
 }

engine/src/main/java/org/hibernate/search/engine/common/impl/MappedIndexSearchTargetImpl.java

@@ -33,12 +33,15 @@
 import org.hibernate.search.engine.search.query.spi.ReferenceHitCollector;
 import org.hibernate.search.engine.search.query.spi.SearchQueryBuilder;
 
-class MappedIndexSearchTargetImpl<C> implements MappedIndexSearchTarget {
+class MappedIndexSearchTargetImpl<C, R, O> implements MappedIndexSearchTarget<R, O> {
 
 	private final SearchTargetContext<C> searchTargetContext;
+	private final Function<DocumentReference, R> documentReferenceTransformer;
 
-	MappedIndexSearchTargetImpl(SearchTargetContext<C> searchTargetContext) {
+	MappedIndexSearchTargetImpl(SearchTargetContext<C> searchTargetContext,
+			Function<DocumentReference, R> documentReferenceTransformer) {
 		this.searchTargetContext = searchTargetContext;
+		this.documentReferenceTransformer = documentReferenceTransformer;
 	}
 
 	@Override
@@ -51,14 +54,13 @@ public String toString() {
 	}
 
 	@Override
-	public <R, O, Q> SearchQueryResultContext<Q> queryAsLoadedObjects(SessionContextImplementor sessionContext,
-			Function<DocumentReference, R> documentReferenceTransformer,
-			ObjectLoader<R, O> objectLoader,
-			Function<SearchQuery<O>, Q> searchQueryWrapperFactory) {
-		HitAggregator<LoadingHitCollector, List<O>> hitAggregator =
+	public <T, Q> SearchQueryResultContext<Q> queryAsLoadedObjects(SessionContextImplementor sessionContext,
+			ObjectLoader<R, T> objectLoader,
+			Function<SearchQuery<T>, Q> searchQueryWrapperFactory) {
+		HitAggregator<LoadingHitCollector, List<T>> hitAggregator =
 				new ObjectHitAggregator<>( documentReferenceTransformer, objectLoader );
 
-		SearchQueryBuilder<O, C> builder = searchTargetContext.getSearchQueryBuilderFactory()
+		SearchQueryBuilder<T, C> builder = searchTargetContext.getSearchQueryBuilderFactory()
 				.asObjects( sessionContext, hitAggregator );
 
 		return new SearchQueryResultContextImpl<>(
@@ -67,8 +69,7 @@ public <R, O, Q> SearchQueryResultContext<Q> queryAsLoadedObjects(SessionContext
 	}
 
 	@Override
-	public <R, T, Q> SearchQueryResultContext<Q> queryAsReferences(SessionContextImplementor sessionContext,
-			Function<DocumentReference, R> documentReferenceTransformer,
+	public <T, Q> SearchQueryResultContext<Q> queryAsReferences(SessionContextImplementor sessionContext,
 			Function<R, T> hitTransformer,
 			Function<SearchQuery<T>, Q> searchQueryWrapperFactory) {
 		HitAggregator<ReferenceHitCollector, List<T>> hitAggregator =
@@ -83,9 +84,8 @@ public <R, T, Q> SearchQueryResultContext<Q> queryAsReferences(SessionContextImp
 	}
 
 	@Override
-	public <R, O, T, Q> SearchQueryResultContext<Q> queryAsProjections(
+	public <T, Q> SearchQueryResultContext<Q> queryAsProjections(
 			SessionContextImplementor sessionContext,
-			Function<DocumentReference, R> documentReferenceTransformer,
 			ObjectLoader<R, O> objectLoader,
 			Function<List<?>, T> hitTransformer,
 			Function<SearchQuery<T>, Q> searchQueryWrapperFactory,
@@ -113,7 +113,7 @@ public SearchSortContainerContext sort() {
 	}
 
 	@Override
-	public SearchProjectionFactoryContext projection() {
-		return new SearchProjectionFactoryContextImpl( searchTargetContext.getSearchProjectionFactory() );
+	public SearchProjectionFactoryContext<R, O> projection() {
+		return new SearchProjectionFactoryContextImpl<>( searchTargetContext.getSearchProjectionFactory() );
 	}
 }

engine/src/main/java/org/hibernate/search/engine/mapper/mapping/spi/MappedIndexManager.java

@@ -6,11 +6,14 @@
  */
 package org.hibernate.search.engine.mapper.mapping.spi;
 
+import java.util.function.Function;
+
 import org.hibernate.search.engine.backend.document.DocumentElement;
 import org.hibernate.search.engine.backend.index.IndexManager;
 import org.hibernate.search.engine.backend.index.spi.IndexWorkPlan;
 import org.hibernate.search.engine.mapper.mapping.context.spi.MappingContextImplementor;
 import org.hibernate.search.engine.mapper.session.context.spi.SessionContextImplementor;
+import org.hibernate.search.engine.search.DocumentReference;
 
 /**
  * The object responsible for applying works and searches to a full-text index.
@@ -23,7 +26,8 @@ public interface MappedIndexManager<D extends DocumentElement> {
 
 	IndexWorkPlan<D> createWorkPlan(SessionContextImplementor sessionContext);
 
-	MappedIndexSearchTargetBuilder createSearchTargetBuilder(MappingContextImplementor mappingContext);
+	<R, O> MappedIndexSearchTargetBuilder<R, O> createSearchTargetBuilder(MappingContextImplementor mappingContext,
+			Function<DocumentReference, R> documentReferenceTransformer);
 
-	void addToSearchTarget(MappedIndexSearchTargetBuilder builder);
+	void addToSearchTarget(MappedIndexSearchTargetBuilder<?, ?> builder);
 }

engine/src/main/java/org/hibernate/search/engine/mapper/mapping/spi/MappedIndexSearchTarget.java

@@ -10,7 +10,6 @@
 import java.util.function.Function;
 
 import org.hibernate.search.engine.mapper.session.context.spi.SessionContextImplementor;
-import org.hibernate.search.engine.search.DocumentReference;
 import org.hibernate.search.engine.search.SearchProjection;
 import org.hibernate.search.engine.search.SearchQuery;
 import org.hibernate.search.engine.search.dsl.query.SearchQueryResultContext;
@@ -19,23 +18,36 @@
 import org.hibernate.search.engine.search.dsl.projection.SearchProjectionFactoryContext;
 import org.hibernate.search.engine.search.dsl.sort.SearchSortContainerContext;
 
-public interface MappedIndexSearchTarget {
+/**
+ * @param <R> The type of references, i.e. the type of hits returned by
+ * {@link #queryAsReferences(SessionContextImplementor, Function, Function) reference queries}
+ * when not using any hit transformer,
+ * or the type of objects returned for {@link SearchProjectionFactoryContext#reference() reference projections}.
+ * @param <O> The type of loaded objects, i.e. the type of hits returned by
+ * {@link #queryAsLoadedObjects(SessionContextImplementor, ObjectLoader, Function) loaded object queries}
+ * when not using any hit transformer,
+ * or the type of objects returned for {@link SearchProjectionFactoryContext#object() loaded object projections}.
+ */
+public interface MappedIndexSearchTarget<R, O> {
 
-	<R, O, Q> SearchQueryResultContext<Q> queryAsLoadedObjects(
+	<T, Q> SearchQueryResultContext<Q> queryAsLoadedObjects(
 			SessionContextImplementor sessionContext,
-			Function<DocumentReference, R> documentReferenceTransformer,
-			ObjectLoader<R, O> objectLoader,
-			Function<SearchQuery<O>, Q> searchQueryWrapperFactory);
+			ObjectLoader<R, T> objectLoader,
+			Function<SearchQuery<T>, Q> searchQueryWrapperFactory);
 
-	<R, T, Q> SearchQueryResultContext<Q> queryAsReferences(
+	<T, Q> SearchQueryResultContext<Q> queryAsReferences(
 			SessionContextImplementor sessionContext,
-			Function<DocumentReference, R> documentReferenceTransformer,
 			Function<R, T> hitTransformer,
 			Function<SearchQuery<T>, Q> searchQueryWrapperFactory);
 
-	<R, O, T, Q> SearchQueryResultContext<Q> queryAsProjections(
+	/*
+	 * IMPLEMENTATION NOTE: we *must* only accept an object loader with the same R/O type parameters as this class,
+	 * otherwise some casts in ObjectProjectionContextImpl and ReferenceProjectionContextImpl
+	 * will be wrong.
+	 * In particular, we cannot accept an ObjectLoader<R, T> like we do in queryAsLoadedObjects(...).
+	 */
+	<T, Q> SearchQueryResultContext<Q> queryAsProjections(
 			SessionContextImplementor sessionContext,
-			Function<DocumentReference, R> documentReferenceTransformer,
 			ObjectLoader<R, O> objectLoader,
 			Function<List<?>, T> hitTransformer,
 			Function<SearchQuery<T>, Q> searchQueryWrapperFactory,
@@ -45,6 +57,11 @@ <R, O, T, Q> SearchQueryResultContext<Q> queryAsProjections(
 
 	SearchSortContainerContext sort();
 
-	SearchProjectionFactoryContext projection();
+	/*
+	 * IMPLEMENTATION NOTE: we *must* return a factory with the same R/O type arguments as this class,
+	 * otherwise some casts in ObjectProjectionContextImpl and ReferenceProjectionContextImpl
+	 * will be wrong.
+	 */
+	SearchProjectionFactoryContext<R, O> projection();
 
 }

engine/src/main/java/org/hibernate/search/engine/mapper/mapping/spi/MappedIndexSearchTargetBuilder.java

@@ -9,8 +9,8 @@
 /**
  * @author Yoann Rodiere
  */
-public interface MappedIndexSearchTargetBuilder {
+public interface MappedIndexSearchTargetBuilder<R, O> {
 
-	MappedIndexSearchTarget build();
+	MappedIndexSearchTarget<R, O> build();
 
 }

engine/src/main/java/org/hibernate/search/engine/search/dsl/projection/ObjectProjectionContext.java

@@ -9,6 +9,6 @@
 /**
  * The context used when starting to define an object projection.
  */
-public interface ObjectProjectionContext extends SearchProjectionTerminalContext<Object> {
+public interface ObjectProjectionContext<O> extends SearchProjectionTerminalContext<O> {
 
 }

engine/src/main/java/org/hibernate/search/engine/search/dsl/projection/ReferenceProjectionContext.java

@@ -9,6 +9,6 @@
 /**
  * The context used when starting to define a reference projection.
  */
-public interface ReferenceProjectionContext extends SearchProjectionTerminalContext<Object> {
+public interface ReferenceProjectionContext<R> extends SearchProjectionTerminalContext<R> {
 
 }

engine/src/main/java/org/hibernate/search/engine/search/dsl/projection/SearchProjectionFactoryContext.java

@@ -11,8 +11,12 @@
 
 /**
  * A context allowing to create a projection.
+ *
+ * @param <R> The type of references, i.e. the type of objects returned for {@link #reference() reference projections}.
+ * @param <O> The type of loaded objects, i.e. the type of objects returned for
+ * {@link #object() object projections}.
  */
-public interface SearchProjectionFactoryContext {
+public interface SearchProjectionFactoryContext<R, O> {
 
 	/**
 	 * Project the match to a {@link DocumentReference}.
@@ -29,7 +33,7 @@ public interface SearchProjectionFactoryContext {
 	 *
 	 * @return A context allowing to define the projection more precisely.
 	 */
-	ReferenceProjectionContext reference();
+	ReferenceProjectionContext<R> reference();
 
 	/**
 	 * Project to an object representing the match.
@@ -41,7 +45,7 @@ public interface SearchProjectionFactoryContext {
 	 *
 	 * @return A context allowing to define the projection more precisely.
 	 */
-	ObjectProjectionContext object();
+	ObjectProjectionContext<O> object();
 
 	/**
 	 * Project to a field of the indexed document.

engine/src/main/java/org/hibernate/search/engine/search/dsl/projection/impl/ObjectProjectionContextImpl.java

@@ -12,7 +12,7 @@
 import org.hibernate.search.engine.search.projection.spi.SearchProjectionBuilderFactory;
 
 
-public class ObjectProjectionContextImpl implements ObjectProjectionContext {
+public class ObjectProjectionContextImpl<O> implements ObjectProjectionContext<O> {
 
 	private ObjectSearchProjectionBuilder objectProjectionBuilder;
 
@@ -21,8 +21,15 @@ public class ObjectProjectionContextImpl implements ObjectProjectionContext {
 	}
 
 	@Override
-	public SearchProjection<Object> toProjection() {
-		return objectProjectionBuilder.build();
+	/*
+	 * The backend has no control over the type of loaded objects.
+	 * This cast is only safe because we make sure to only use SearchProjectionFactoryContext
+	 * with generic type arguments that are consistent with the type of object loaders.
+	 * See comments in MappedIndexSearchTarget.
+	 */
+	@SuppressWarnings("unchecked")
+	public SearchProjection<O> toProjection() {
+		return (SearchProjection<O>) objectProjectionBuilder.build();
 	}
 
 }

engine/src/main/java/org/hibernate/search/engine/search/dsl/projection/impl/ReferenceProjectionContextImpl.java

@@ -12,7 +12,7 @@
 import org.hibernate.search.engine.search.projection.spi.SearchProjectionBuilderFactory;
 
 
-public class ReferenceProjectionContextImpl implements ReferenceProjectionContext {
+public class ReferenceProjectionContextImpl<R> implements ReferenceProjectionContext<R> {
 
 	private ReferenceSearchProjectionBuilder referenceProjectionBuilder;
 
@@ -21,8 +21,15 @@ public class ReferenceProjectionContextImpl implements ReferenceProjectionContex
 	}
 
 	@Override
-	public SearchProjection<Object> toProjection() {
-		return referenceProjectionBuilder.build();
+	/*
+	 * The backend has no control over the type of loaded objects.
+	 * This cast is only safe because we make sure to only use SearchProjectionFactoryContext
+	 * with generic type arguments that are consistent with the type of object loaders.
+	 * See comments in MappedIndexSearchTarget.
+	 */
+	@SuppressWarnings("unchecked")
+	public SearchProjection<R> toProjection() {
+		return (SearchProjection<R>) referenceProjectionBuilder.build();
 	}
 
 }

engine/src/main/java/org/hibernate/search/engine/search/dsl/projection/impl/SearchProjectionFactoryContextImpl.java

@@ -18,7 +18,7 @@
 import org.hibernate.search.util.impl.common.Contracts;
 
 
-public class SearchProjectionFactoryContextImpl implements SearchProjectionFactoryContext {
+public class SearchProjectionFactoryContextImpl<R, O> implements SearchProjectionFactoryContext<R, O> {
 
 	private final SearchProjectionBuilderFactory factory;
 
@@ -44,13 +44,13 @@ public FieldProjectionContext<Object> field(String absoluteFieldPath) {
 	}
 
 	@Override
-	public ReferenceProjectionContext reference() {
-		return new ReferenceProjectionContextImpl( factory );
+	public ReferenceProjectionContext<R> reference() {
+		return new ReferenceProjectionContextImpl<>( factory );
 	}
 
 	@Override
-	public ObjectProjectionContext object() {
-		return new ObjectProjectionContextImpl( factory );
+	public ObjectProjectionContext<O> object() {
+		return new ObjectProjectionContextImpl<>( factory );
 	}
 
 	@Override