From 2beb85e6950ad6b95dc6be951d621fcbeda046d0 Mon Sep 17 00:00:00 2001
From: Gavin King <gavin@hibernate.org>
Date: Fri, 1 Mar 2024 21:00:07 +0100
Subject: [PATCH] HHH-17772 support for returning the argument from lifecycle
 methods

---
 .../test/data/data/BookAuthorRepository.java  |  6 +++++
 .../annotation/AnnotationMetaEntity.java      | 13 +++++++++--
 .../annotation/LifecycleMethod.java           | 22 +++++++++++++++----
 3 files changed, 35 insertions(+), 6 deletions(-)

diff --git a/tooling/metamodel-generator/src/jakartaData/java/org/hibernate/jpamodelgen/test/data/data/BookAuthorRepository.java b/tooling/metamodel-generator/src/jakartaData/java/org/hibernate/jpamodelgen/test/data/data/BookAuthorRepository.java
index f1fbf9c606..71d0a8db77 100644
--- a/tooling/metamodel-generator/src/jakartaData/java/org/hibernate/jpamodelgen/test/data/data/BookAuthorRepository.java
+++ b/tooling/metamodel-generator/src/jakartaData/java/org/hibernate/jpamodelgen/test/data/data/BookAuthorRepository.java
@@ -140,4 +140,10 @@ public interface BookAuthorRepository {
 
 	@Find
 	List<Book> allBooksWithLotsOfSorting(Sort<? super Book> s1, Order<? super Book> order, Sort<? super Book>... s3);
+
+	@Save
+	Book write(Book book);
+
+	@Insert
+	Iterable<Book> createAll(Iterable<Book> books);
 }
diff --git a/tooling/metamodel-generator/src/main/java/org/hibernate/jpamodelgen/annotation/AnnotationMetaEntity.java b/tooling/metamodel-generator/src/main/java/org/hibernate/jpamodelgen/annotation/AnnotationMetaEntity.java
index 669c2dce22..64733dbb6a 100644
--- a/tooling/metamodel-generator/src/main/java/org/hibernate/jpamodelgen/annotation/AnnotationMetaEntity.java
+++ b/tooling/metamodel-generator/src/main/java/org/hibernate/jpamodelgen/annotation/AnnotationMetaEntity.java
@@ -824,12 +824,13 @@ else if ( method.getParameters().size() != 1 ) {
 					Diagnostic.Kind.ERROR );
 
 		}
-		else if ( returnType == null || returnType.getKind() != TypeKind.VOID ) {
+		else if ( returnType == null ) {
 			context.message( method,
 					"must be declared 'void'",
 					Diagnostic.Kind.ERROR );
 		}
 		else {
+			final boolean returnArgument = returnType.getKind() != TypeKind.VOID;
 			final String operation = lifecycleOperation( method );
 			final VariableElement parameter = method.getParameters().get(0);
 			final TypeMirror parameterType = parameter.asType();
@@ -844,6 +845,13 @@ else if ( !containsAnnotation( declaredType.asElement(), ENTITY ) ) {
 						"incorrect parameter type '" + parameterType + "' is not annotated '@Entity'",
 						Diagnostic.Kind.ERROR );
 			}
+			else if ( returnArgument
+					&& !context.getTypeUtils().isSameType( returnType, parameterType ) ) {
+				context.message( parameter,
+						"return type '" + returnType
+								+ "' disagrees with parameter type '" + parameterType + "'",
+						Diagnostic.Kind.ERROR );
+			}
 			else {
 				putMember(
 						method.getSimpleName().toString()
@@ -856,7 +864,8 @@ else if ( !containsAnnotation( declaredType.asElement(), ENTITY ) ) {
 								getSessionVariableName(),
 								operation,
 								context.addNonnullAnnotation(),
-								declaredType != parameterType
+								declaredType != parameterType,
+								returnArgument
 						)
 				);
 			}
diff --git a/tooling/metamodel-generator/src/main/java/org/hibernate/jpamodelgen/annotation/LifecycleMethod.java b/tooling/metamodel-generator/src/main/java/org/hibernate/jpamodelgen/annotation/LifecycleMethod.java
index 0f1708257c..bed61849f5 100644
--- a/tooling/metamodel-generator/src/main/java/org/hibernate/jpamodelgen/annotation/LifecycleMethod.java
+++ b/tooling/metamodel-generator/src/main/java/org/hibernate/jpamodelgen/annotation/LifecycleMethod.java
@@ -18,6 +18,7 @@ public class LifecycleMethod implements MetaAttribute {
 	private final String operationName;
 	private final boolean addNonnullAnnotation;
 	private final boolean iterateParameter;
+	private final boolean returnArgument;
 
 	public LifecycleMethod(
 			AnnotationMetaEntity annotationMetaEntity,
@@ -27,7 +28,8 @@ public LifecycleMethod(
 			String sessionName,
 			String operationName,
 			boolean addNonnullAnnotation,
-			boolean iterateParameter) {
+			boolean iterateParameter,
+			boolean returnArgument) {
 		this.annotationMetaEntity = annotationMetaEntity;
 		this.entity = entity;
 		this.methodName = methodName;
@@ -36,6 +38,7 @@ public LifecycleMethod(
 		this.operationName = operationName;
 		this.addNonnullAnnotation = addNonnullAnnotation;
 		this.iterateParameter = iterateParameter;
+		this.returnArgument = returnArgument;
 	}
 
 	@Override
@@ -55,6 +58,8 @@ public String getAttributeDeclarationString() {
 		nullCheck(declaration);
 		declaration.append("\ttry {\n");
 		delegateCall(declaration);
+		returnArgument(declaration);
+		declaration.append("\t}\n");
 		if ( operationName.equals("insert") ) {
 			convertException(declaration,
 					"jakarta.persistence.EntityExistsException",
@@ -72,6 +77,15 @@ public String getAttributeDeclarationString() {
 		return declaration.toString();
 	}
 
+	private void returnArgument(StringBuilder declaration) {
+		if ( returnArgument ) {
+			declaration
+					.append("\t\treturn ")
+					.append(parameterName)
+					.append(";\n");
+		}
+	}
+
 	private void delegateCall(StringBuilder declaration) {
 		if ( iterateParameter ) {
 			declaration
@@ -92,13 +106,13 @@ private void delegateCall(StringBuilder declaration) {
 			declaration
 					.append("\t\t}\n");
 		}
-		declaration
-				.append("\t}\n");
 	}
 
 	private void preamble(StringBuilder declaration) {
 		declaration
-				.append("\n@Override\npublic void ")
+				.append("\n@Override\npublic ")
+				.append(returnArgument ? annotationMetaEntity.importType(entity) : "void")
+				.append(' ')
 				.append(methodName)
 				.append('(');
 		notNull(declaration);