From b8f9913aa1dcfd52caa08e12063ae5066a2e2f3b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?St=C3=A9phane=20Nicoll?= <stephane.nicoll@broadcom.com>
Date: Wed, 24 Jul 2024 13:58:36 +0200
Subject: [PATCH] Add AOT/Native support

This commit adds AOT/Native support for beans that are contributed by
the Kotlin DSL.

Since they use an instance supplier, such beans are now configured to
be ignored by AOT generation. They are part of the bean factory still
so any hint generation works.

This commit removes a previous attempt at fixing this issue when we
were not checking for instance suppliers. Rather than skipping the
initializr at runtime, it runs again as intended since their state
can't be stored in AOT-generated code.

Closes gh-29555
---
 .../context/support/BeanDefinitionDsl.kt      |  9 ++++-----
 .../context/support/BeanDefinitionDslTests.kt | 20 +++++++++++++++++++
 2 files changed, 24 insertions(+), 5 deletions(-)

diff --git a/spring-context/src/main/kotlin/org/springframework/context/support/BeanDefinitionDsl.kt b/spring-context/src/main/kotlin/org/springframework/context/support/BeanDefinitionDsl.kt
index c0cc8b573ac..8077fe8a182 100644
--- a/spring-context/src/main/kotlin/org/springframework/context/support/BeanDefinitionDsl.kt
+++ b/spring-context/src/main/kotlin/org/springframework/context/support/BeanDefinitionDsl.kt
@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2023 the original author or authors.
+ * Copyright 2002-2024 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -16,8 +16,8 @@
 
 package org.springframework.context.support
 
-import org.springframework.aot.AotDetector
 import org.springframework.beans.factory.ObjectProvider
+import org.springframework.beans.factory.aot.BeanRegistrationAotProcessor
 import org.springframework.beans.factory.config.BeanDefinition
 import org.springframework.beans.factory.config.BeanDefinitionCustomizer
 import org.springframework.beans.factory.getBeanProvider
@@ -191,6 +191,7 @@ open class BeanDefinitionDsl internal constructor (private val init: BeanDefinit
 			description?.let { bd.description = description }
 			role?.let { bd.role = role.ordinal }
 			order?.let { bd.setAttribute(AbstractBeanDefinition.ORDER_ATTRIBUTE, order) }
+			bd.setAttribute(BeanRegistrationAotProcessor.IGNORE_REGISTRATION_ATTRIBUTE, true)
 		}
 
 		val beanName = name ?: BeanDefinitionReaderUtils.uniqueBeanName(T::class.java.name, context);
@@ -237,6 +238,7 @@ open class BeanDefinitionDsl internal constructor (private val init: BeanDefinit
 			description?.let { bd.description = description }
 			role?.let { bd.role = role.ordinal }
 			order?.let { bd.setAttribute(AbstractBeanDefinition.ORDER_ATTRIBUTE, order) }
+			bd.setAttribute(BeanRegistrationAotProcessor.IGNORE_REGISTRATION_ATTRIBUTE, true)
 		}
 
 
@@ -1199,9 +1201,6 @@ open class BeanDefinitionDsl internal constructor (private val init: BeanDefinit
 	 * @param context The `ApplicationContext` to use for registering the beans
 	 */
 	override fun initialize(context: GenericApplicationContext) {
-		if (AotDetector.useGeneratedArtifacts()) {
-			return
-		}
 		this.context = context
 		init()
 		for (child in children) {
diff --git a/spring-context/src/test/kotlin/org/springframework/context/support/BeanDefinitionDslTests.kt b/spring-context/src/test/kotlin/org/springframework/context/support/BeanDefinitionDslTests.kt
index 50e22b7012e..544c93ffba0 100644
--- a/spring-context/src/test/kotlin/org/springframework/context/support/BeanDefinitionDslTests.kt
+++ b/spring-context/src/test/kotlin/org/springframework/context/support/BeanDefinitionDslTests.kt
@@ -21,6 +21,7 @@ import org.assertj.core.api.Assertions.assertThatExceptionOfType
 import org.junit.jupiter.api.Test
 import org.junit.jupiter.api.fail
 import org.springframework.beans.factory.NoSuchBeanDefinitionException
+import org.springframework.beans.factory.aot.BeanRegistrationAotProcessor
 import org.springframework.beans.factory.getBean
 import org.springframework.beans.factory.getBeanProvider
 import org.springframework.context.support.BeanDefinitionDsl.*
@@ -217,6 +218,25 @@ class BeanDefinitionDslTests {
 
 		assertThat(context.getBeanProvider<FooFoo>().orderedStream().map { it.name }).containsExactly("highest", "lowest")
 	}
+
+	@Test
+	fun `Declare beans with the functional Kotlin DSL flag them as to be ignored by AOT`() {
+		val beans = beans {
+			bean<Foo>("one")
+			bean<Bar>("two")
+		}
+
+		val context = GenericApplicationContext().apply {
+			beans.initialize(this)
+		}
+
+		assertThat(context.beanDefinitionNames).contains("one", "two")
+		assertThat(context.getBeanDefinition("one").getAttribute(
+			BeanRegistrationAotProcessor.IGNORE_REGISTRATION_ATTRIBUTE)).isEqualTo(true)
+		assertThat(context.getBeanDefinition("two").getAttribute(
+			BeanRegistrationAotProcessor.IGNORE_REGISTRATION_ATTRIBUTE)).isEqualTo(true)
+	}
+
 }
 
 class Foo