HSEARCH-3395 Make the .reference() and .object() projections type-saf…

…e in the Projection DSL

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

@@ -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/MappedIndexSearchTarget.java

@@ -40,6 +40,12 @@ <T, Q> SearchQueryResultContext<Q> queryAsReferences(
 			Function<R, T> hitTransformer,
 			Function<SearchQuery<T>, Q> searchQueryWrapperFactory);
 
+	/*
+	 * 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,
 			ObjectLoader<R, O> objectLoader,
@@ -51,6 +57,11 @@ <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/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

integrationtest/backend/tck/src/main/java/org/hibernate/search/integrationtest/backend/tck/search/StubDocumentReferenceTransformer.java

@@ -13,6 +13,6 @@
 /**
  * The only purpose of this class is to avoid unchecked cast warnings when creating mocks.
  */
-interface StubDocumentReferenceTransformer
+public interface StubDocumentReferenceTransformer
 		extends Function<DocumentReference, StubTransformedReference> {
 }

integrationtest/backend/tck/src/main/java/org/hibernate/search/integrationtest/backend/tck/search/StubLoadedObject.java

@@ -8,10 +8,10 @@
 
 import org.hibernate.search.engine.search.DocumentReference;
 
-class StubLoadedObject {
+public class StubLoadedObject {
 	private final DocumentReference documentReference;
 
-	StubLoadedObject(DocumentReference documentReference) {
+	public StubLoadedObject(DocumentReference documentReference) {
 		this.documentReference = documentReference;
 	}
 

integrationtest/backend/tck/src/main/java/org/hibernate/search/integrationtest/backend/tck/search/StubObjectLoader.java

@@ -11,6 +11,6 @@
 /**
  * The only purpose of this class is to avoid unchecked cast warnings when creating mocks.
  */
-interface StubObjectLoader
+public interface StubObjectLoader
 		extends ObjectLoader<StubTransformedReference, StubLoadedObject> {
 }

integrationtest/backend/tck/src/main/java/org/hibernate/search/integrationtest/backend/tck/search/StubTransformedReference.java

@@ -8,10 +8,10 @@
 
 import org.hibernate.search.engine.search.DocumentReference;
 
-class StubTransformedReference {
+public class StubTransformedReference {
 	private final DocumentReference documentReference;
 
-	StubTransformedReference(DocumentReference documentReference) {
+	public StubTransformedReference(DocumentReference documentReference) {
 		this.documentReference = documentReference;
 	}
 

integrationtest/backend/tck/src/main/java/org/hibernate/search/integrationtest/backend/tck/search/projection/SearchProjectionIT.java

@@ -6,6 +6,7 @@
  */
 package org.hibernate.search.integrationtest.backend.tck.search.projection;
 
+import static org.hibernate.search.util.impl.integrationtest.common.EasyMockUtils.referenceMatcher;
 import static org.hibernate.search.util.impl.integrationtest.common.NormalizationUtils.reference;
 import static org.hibernate.search.util.impl.integrationtest.common.assertion.ProjectionsSearchResultAssert.assertThat;
 import static org.hibernate.search.util.impl.integrationtest.common.stub.mapper.StubMapperUtils.referenceProvider;
@@ -27,25 +28,36 @@
 import org.hibernate.search.engine.backend.document.model.dsl.Projectable;
 import org.hibernate.search.engine.backend.index.spi.IndexWorkPlan;
 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.SearchResult;
+import org.hibernate.search.engine.search.loading.spi.ObjectLoader;
 import org.hibernate.search.engine.spatial.GeoPoint;
 import org.hibernate.search.integrationtest.backend.tck.configuration.DefaultAnalysisDefinitions;
+import org.hibernate.search.integrationtest.backend.tck.search.StubDocumentReferenceTransformer;
+import org.hibernate.search.integrationtest.backend.tck.search.StubLoadedObject;
+import org.hibernate.search.integrationtest.backend.tck.search.StubObjectLoader;
+import org.hibernate.search.integrationtest.backend.tck.search.StubTransformedReference;
 import org.hibernate.search.integrationtest.backend.tck.util.StandardFieldMapper;
 import org.hibernate.search.integrationtest.backend.tck.util.ValueWrapper;
 import org.hibernate.search.integrationtest.backend.tck.util.rule.SearchSetupHelper;
 import org.hibernate.search.util.SearchException;
+import org.hibernate.search.util.impl.integrationtest.common.EasyMockUtils;
 import org.hibernate.search.util.impl.integrationtest.common.assertion.DocumentReferencesSearchResultAssert;
+import org.hibernate.search.util.impl.integrationtest.common.stub.mapper.GenericStubMappingSearchTarget;
 import org.hibernate.search.util.impl.integrationtest.common.stub.mapper.StubMappingIndexManager;
 import org.hibernate.search.util.impl.integrationtest.common.stub.mapper.StubMappingSearchTarget;
 import org.hibernate.search.util.impl.test.SubTest;
+import org.hibernate.search.util.impl.test.annotation.TestForIssue;
+
 import org.junit.Assume;
 import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.rules.ExpectedException;
 
 import org.assertj.core.api.Assertions;
+import org.easymock.EasyMock;
 
 public class SearchProjectionIT {
 	private static final String INDEX_NAME = "IndexName";
@@ -240,11 +252,22 @@ public void projectionConstants_references() {
 		DocumentReference document3Reference = reference( INDEX_NAME, DOCUMENT_3 );
 		DocumentReference emptyReference = reference( INDEX_NAME, EMPTY );
 
+		/*
+		 * Note to test writers: make sure to assign these projections to variables,
+		 * just so that tests do not compile if someone changes the APIs in an incorrect way.
+		 */
+		SearchProjection<DocumentReference> documentReferenceProjection =
+				searchTarget.projection().documentReference().toProjection();
+		SearchProjection<DocumentReference> referenceProjection =
+				searchTarget.projection().reference().toProjection();
+		SearchProjection<DocumentReference> objectProjection =
+				searchTarget.projection().object().toProjection();
+
 		query = searchTarget.query()
 				.asProjections(
-						searchTarget.projection().documentReference().toProjection(),
-						searchTarget.projection().reference().toProjection(),
-						searchTarget.projection().object().toProjection()
+						documentReferenceProjection,
+						referenceProjection,
+						objectProjection
 				)
 				.predicate( f -> f.matchAll().toPredicate() )
 				.build();
@@ -256,6 +279,88 @@ public void projectionConstants_references() {
 		} );
 	}
 
+	/**
+	 * Test documentReference/reference/object projections as they are likely to be used by mappers,
+	 * i.e. with a custom reference transformer and a custom object loader.
+	 */
+	@Test
+	@TestForIssue(jiraKey = "HSEARCH-3395")
+	public void projectionConstants_references_transformed() {
+		DocumentReference document1Reference = reference( INDEX_NAME, DOCUMENT_1 );
+		DocumentReference document2Reference = reference( INDEX_NAME, DOCUMENT_2 );
+		DocumentReference document3Reference = reference( INDEX_NAME, DOCUMENT_3 );
+		DocumentReference emptyReference = reference( INDEX_NAME, EMPTY );
+		StubTransformedReference document1TransformedReference = new StubTransformedReference( document1Reference );
+		StubTransformedReference document2TransformedReference = new StubTransformedReference( document2Reference );
+		StubTransformedReference document3TransformedReference = new StubTransformedReference( document3Reference );
+		StubTransformedReference emptyTransformedReference = new StubTransformedReference( emptyReference );
+		StubLoadedObject document1LoadedObject = new StubLoadedObject( document1Reference );
+		StubLoadedObject document2LoadedObject = new StubLoadedObject( document2Reference );
+		StubLoadedObject document3LoadedObject = new StubLoadedObject( document3Reference );
+		StubLoadedObject emptyLoadedObject = new StubLoadedObject( emptyReference );
+
+		Function<DocumentReference, StubTransformedReference> referenceTransformerMock =
+				EasyMock.createMock( StubDocumentReferenceTransformer.class );
+		ObjectLoader<StubTransformedReference, StubLoadedObject> objectLoaderMock =
+				EasyMock.createMock( StubObjectLoader.class );
+
+		EasyMock.expect( referenceTransformerMock.apply( referenceMatcher( document1Reference ) ) )
+				.andReturn( document1TransformedReference )
+				.times( 2 );
+		EasyMock.expect( referenceTransformerMock.apply( referenceMatcher( document2Reference ) ) )
+				.andReturn( document2TransformedReference )
+				.times( 2 );
+		EasyMock.expect( referenceTransformerMock.apply( referenceMatcher( document3Reference ) ) )
+				.andReturn( document3TransformedReference )
+				.times( 2 );
+		EasyMock.expect( referenceTransformerMock.apply( referenceMatcher( emptyReference ) ) )
+				.andReturn( emptyTransformedReference )
+				.times( 2 );
+		EasyMock.expect( objectLoaderMock.load(
+				EasyMockUtils.collectionAnyOrderMatcher( Arrays.asList(
+						document1TransformedReference, document2TransformedReference,
+						document3TransformedReference, emptyTransformedReference
+				) )
+		) )
+				.andReturn( Arrays.asList(
+						document1LoadedObject, document2LoadedObject,
+						document3LoadedObject, emptyLoadedObject
+				) );
+		EasyMock.replay( referenceTransformerMock, objectLoaderMock );
+
+		GenericStubMappingSearchTarget<StubTransformedReference, StubLoadedObject> searchTarget =
+				indexManager.createSearchTarget( referenceTransformerMock );
+		SearchQuery<List<?>> query;
+
+		/*
+		 * Note to test writers: make sure to assign these projections to variables,
+		 * just so that tests do not compile if someone changes the APIs in an incorrect way.
+		 */
+		SearchProjection<DocumentReference> documentReferenceProjection =
+				searchTarget.projection().documentReference().toProjection();
+		SearchProjection<StubTransformedReference> referenceProjection =
+				searchTarget.projection().reference().toProjection();
+		SearchProjection<StubLoadedObject> objectProjection =
+				searchTarget.projection().object().toProjection();
+
+		query = searchTarget.query( objectLoaderMock )
+				.asProjections(
+						documentReferenceProjection,
+						referenceProjection,
+						objectProjection
+				)
+				.predicate( f -> f.matchAll().toPredicate() )
+				.build();
+		assertThat( query ).hasProjectionsHitsAnyOrder( b -> {
+			b.projection( document1Reference, document1TransformedReference, document1LoadedObject );
+			b.projection( document2Reference, document2TransformedReference, document2LoadedObject );
+			b.projection( document3Reference, document3TransformedReference, document3LoadedObject );
+			b.projection( emptyReference, emptyTransformedReference, emptyLoadedObject );
+		} );
+
+		EasyMock.verify( referenceTransformerMock, objectLoaderMock );
+	}
+
 	@Test
 	public void score() {
 		StubMappingSearchTarget searchTarget = indexManager.createSearchTarget();

mapper/javabean/src/main/java/org/hibernate/search/mapper/javabean/search/JavaBeanSearchTarget.java

@@ -10,6 +10,7 @@
 import org.hibernate.search.engine.search.dsl.projection.SearchProjectionFactoryContext;
 import org.hibernate.search.engine.search.dsl.sort.SearchSortContainerContext;
 import org.hibernate.search.mapper.javabean.search.dsl.query.JavaBeanQueryResultDefinitionContext;
+import org.hibernate.search.mapper.pojo.search.PojoReference;
 
 public interface JavaBeanSearchTarget {
 
@@ -19,6 +20,6 @@ public interface JavaBeanSearchTarget {
 
 	SearchSortContainerContext sort();
 
-	SearchProjectionFactoryContext projection();
+	SearchProjectionFactoryContext<PojoReference, PojoReference> projection();
 
 }

mapper/javabean/src/main/java/org/hibernate/search/mapper/javabean/search/impl/JavaBeanSearchTargetImpl.java

@@ -39,7 +39,7 @@ public SearchSortContainerContext sort() {
 	}
 
 	@Override
-	public SearchProjectionFactoryContext projection() {
+	public SearchProjectionFactoryContext<PojoReference, PojoReference> projection() {
 		return delegate.projection();
 	}
 }