AOT and Native Compilation

/**
 * AOT processor that contributes to bean registration code generation.
 */
interface BeanRegistrationAotProcessor {
    /**
     * Process the given RegisteredBean instance ahead of time and return a contribution or null.
     * @param registeredBean the registered bean to process
     * @return a BeanRegistrationAotContribution instance or null
     */
    BeanRegistrationAotContribution processAheadOfTime(RegisteredBean registeredBean);
}

/**
 * AOT contribution from a BeanRegistrationAotProcessor, used to register a single bean definition.
 */
interface BeanRegistrationAotContribution {
    /**
     * Apply this contribution to the given BeanRegistrationCode.
     * @param generationContext the generation context
     * @param beanRegistrationCode the bean registration code
     */
    void applyTo(GenerationContext generationContext, BeanRegistrationCode beanRegistrationCode);
}

/**
 * Interface that defines the contract for code that registers a single bean.
 */
interface BeanRegistrationCode {
    /**
     * Return the class name of the bean registration.
     * @return the class name
     */
    ClassName getClassName();
    
    /**
     * Return the methods that contain the bean registration code.
     * @return the registration methods
     */
    MethodReference[] getMethods();
    
    /**
     * Return the instance supplier code generator.
     * @return the instance supplier code generator
     */
    CodeGenerator getInstanceSupplierCodeGenerator();
}

/**
 * Interface that defines the contract for code that registers all beans.
 */
interface BeanRegistrationsCode {
    /**
     * Return the class name that contains the bean registrations.
     * @return the class name
     */
    ClassName getClassName();
    
    /**
     * Return the method that contains the bean registration code.
     * @return the registration method
     */
    MethodReference getInitializerMethod();
}

/**
 * AOT processor that contributes to bean factory initialization code generation.
 */
interface BeanFactoryInitializationAotProcessor {
    /**
     * Process the given ConfigurableListableBeanFactory instance ahead of time and return a contribution or null.
     * @param beanFactory the bean factory to process
     * @return a BeanFactoryInitializationAotContribution instance or null
     */
    BeanFactoryInitializationAotContribution processAheadOfTime(ConfigurableListableBeanFactory beanFactory);
}

/**
 * AOT contribution from a BeanFactoryInitializationAotProcessor, used to initialize a bean factory.
 */
interface BeanFactoryInitializationAotContribution {
    /**
     * Apply this contribution to the given BeanFactoryInitializationCode.
     * @param generationContext the generation context
     * @param beanFactoryInitializationCode the bean factory initialization code
     */
    void applyTo(GenerationContext generationContext, BeanFactoryInitializationCode beanFactoryInitializationCode);
}

/**
 * Interface that defines the contract for code that initializes a bean factory.
 */
interface BeanFactoryInitializationCode {
    /**
     * Return the type of the bean factory.
     * @return the bean factory type
     */
    TypeSpec.Builder getTypeBuilder();
    
    /**
     * Return the methods that contain the bean factory initialization code.
     * @return the initialization methods
     */
    MethodSpec.Builder[] getMethods();
}

/**
 * Code generator for instance supplier lambda expressions.
 */
class InstanceSupplierCodeGenerator {
    /**
     * Create a new InstanceSupplierCodeGenerator instance.
     * @param generationContext the generation context to use
     * @param className the class name being generated
     * @param methods the methods that are available
     * @param allowDirectSupplierShortcut whether to allow direct supplier shortcuts
     */
    public InstanceSupplierCodeGenerator(GenerationContext generationContext, ClassName className, 
                                       MultiStatement methods, boolean allowDirectSupplierShortcut);
    
    /**
     * Generate the instance supplier code for the given registered bean and executable.
     * @param registeredBean the registered bean
     * @param constructorOrFactoryMethod the constructor or factory method
     * @return the generated code block
     */
    public CodeBlock generateCode(RegisteredBean registeredBean, Executable constructorOrFactoryMethod);
    
    /**
     * Generate the instance supplier code for the given registered bean and instantiation descriptor.
     * @param registeredBean the registered bean
     * @param instantiationDescriptor the instantiation descriptor
     * @return the generated code block
     */
    public CodeBlock generateCode(RegisteredBean registeredBean, InstantiationDescriptor instantiationDescriptor);
    
    /**
     * Generate the instance supplier code for the given registered bean and instance supplier.
     * @param registeredBean the registered bean
     * @param instanceSupplier the instance supplier
     * @return the generated code block
     */
    public CodeBlock generateCode(RegisteredBean registeredBean, InstanceSupplier<?> instanceSupplier);
    
    /**
     * Check whether the given bean class has a constructor with an optional parameter.
     * @param beanClass the bean class to check
     * @return true if the bean class has a constructor with an optional parameter
     */
    public static boolean hasConstructorWithOptionalParameter(Class<?> beanClass);
    
    /**
     * Register reflection hints for the given executable.
     * @param runtimeHints the runtime hints
     * @param executable the executable to register hints for
     */
    public void registerRuntimeHints(RuntimeHints runtimeHints, Executable executable);
}

/**
 * Interface for resolving autowired arguments in generated code.
 */
interface AutowiredArguments {
    /**
     * Resolve the arguments for the given parameter types.
     * @param parameterTypes the parameter types
     * @return the resolved arguments
     */
    Object[] resolve(Class<?>[] parameterTypes);
    
    /**
     * Resolve arguments for the given parameter types starting from the specified index.
     * @param parameterTypes the parameter types
     * @param startIndex the start index
     * @return the resolved arguments
     */
    Object[] resolve(Class<?>[] parameterTypes, int startIndex);
}

/**
 * Code generator for autowired arguments resolution.
 */
class AutowiredArgumentsCodeGenerator {
    /**
     * Create a new AutowiredArgumentsCodeGenerator.
     * @param className the class name being generated
     * @param methods the methods that are available
     */
    public AutowiredArgumentsCodeGenerator(ClassName className, MultiStatement methods);
    
    /**
     * Generate code to resolve arguments for the given parameter types.
     * @param parameterTypes the parameter types
     * @return the generated code block
     */
    public CodeBlock generateCode(Class<?>[] parameterTypes);
    
    /**
     * Generate code to resolve arguments for the given parameter types starting from the specified index.
     * @param parameterTypes the parameter types
     * @param startIndex the start index
     * @return the generated code block
     */
    public CodeBlock generateCode(Class<?>[] parameterTypes, int startIndex);
    
    /**
     * Generate code to resolve arguments for the given parameter types with a custom variable name.
     * @param parameterTypes the parameter types
     * @param startIndex the start index
     * @param variableName the variable name to use
     * @return the generated code block
     */
    public CodeBlock generateCode(Class<?>[] parameterTypes, int startIndex, String variableName);
}

/**
 * Interface that defines the various code fragments that are used during bean registration code generation.
 */
interface BeanRegistrationCodeFragments {
    /**
     * Return the target class name for the given registered bean.
     * @param registeredBean the registered bean
     * @return the target class name
     */
    ClassName getTarget(RegisteredBean registeredBean);
    
    /**
     * Generate the instance supplier code.
     * @param generationContext the generation context
     * @param beanRegistrationCode the bean registration code
     * @param instanceSupplierCode the instance supplier code
     * @return the generated instance supplier code
     */
    CodeBlock generateInstanceSupplierCode(GenerationContext generationContext, 
                                         BeanRegistrationCode beanRegistrationCode, 
                                         CodeBlock instanceSupplierCode);
    
    /**
     * Generate the return statement for the bean registration method.
     * @param generationContext the generation context
     * @param beanRegistrationCode the bean registration code
     * @return the generated return statement
     */
    CodeBlock generateReturnStatement(GenerationContext generationContext, 
                                    BeanRegistrationCode beanRegistrationCode);
    
    /**
     * Generate the set bean definition properties code.
     * @param generationContext the generation context
     * @param beanRegistrationCode the bean registration code
     * @param assignedVariable the assigned variable name
     * @param beanDefinition the bean definition
     * @return the generated properties code
     */
    CodeBlock generateSetBeanDefinitionPropertiesCode(GenerationContext generationContext,
                                                    BeanRegistrationCode beanRegistrationCode,
                                                    CodeBlock assignedVariable,
                                                    BeanDefinition beanDefinition);
    
    /**
     * Generate the set bean instance supplier code.
     * @param generationContext the generation context
     * @param beanRegistrationCode the bean registration code
     * @param assignedVariable the assigned variable name
     * @param instanceSupplierCode the instance supplier code
     * @return the generated instance supplier code
     */
    CodeBlock generateSetBeanInstanceSupplierCode(GenerationContext generationContext,
                                                 BeanRegistrationCode beanRegistrationCode,
                                                 CodeBlock assignedVariable,
                                                 CodeBlock instanceSupplierCode);
}

/**
 * A BeanRegistrationCodeFragments decorator implementation.
 */
class BeanRegistrationCodeFragmentsDecorator implements BeanRegistrationCodeFragments {
    /**
     * Create a new BeanRegistrationCodeFragmentsDecorator.
     * @param delegate the delegate to use
     */
    public BeanRegistrationCodeFragmentsDecorator(BeanRegistrationCodeFragments delegate);
    
    /**
     * Return the target class name for the given registered bean.
     * @param registeredBean the registered bean
     * @return the target class name
     */
    public ClassName getTarget(RegisteredBean registeredBean);
    
    /**
     * Generate the instance supplier code.
     * @param generationContext the generation context
     * @param beanRegistrationCode the bean registration code
     * @param instanceSupplierCode the instance supplier code
     * @return the generated instance supplier code
     */
    public CodeBlock generateInstanceSupplierCode(GenerationContext generationContext,
                                                 BeanRegistrationCode beanRegistrationCode,
                                                 CodeBlock instanceSupplierCode);
    
    /**
     * Generate the return statement for the bean registration method.
     * @param generationContext the generation context
     * @param beanRegistrationCode the bean registration code
     * @return the generated return statement
     */
    public CodeBlock generateReturnStatement(GenerationContext generationContext,
                                           BeanRegistrationCode beanRegistrationCode);
    
    /**
     * Generate the set bean definition properties code.
     * @param generationContext the generation context
     * @param beanRegistrationCode the bean registration code
     * @param assignedVariable the assigned variable name
     * @param beanDefinition the bean definition
     * @return the generated properties code
     */
    public CodeBlock generateSetBeanDefinitionPropertiesCode(GenerationContext generationContext,
                                                           BeanRegistrationCode beanRegistrationCode,
                                                           CodeBlock assignedVariable,
                                                           BeanDefinition beanDefinition);
    
    /**
     * Generate the set bean instance supplier code.
     * @param generationContext the generation context
     * @param beanRegistrationCode the bean registration code
     * @param assignedVariable the assigned variable name
     * @param instanceSupplierCode the instance supplier code
     * @return the generated instance supplier code
     */
    public CodeBlock generateSetBeanInstanceSupplierCode(GenerationContext generationContext,
                                                        BeanRegistrationCode beanRegistrationCode,
                                                        CodeBlock assignedVariable,
                                                        CodeBlock instanceSupplierCode);
}

/**
 * Utility class for tracking and managing code warnings during AOT processing.
 */
class CodeWarnings {
    /**
     * Create a new CodeWarnings instance.
     */
    public CodeWarnings();
    
    /**
     * Register a warning message.
     * @param warning the warning message
     */
    public void register(String warning);
    
    /**
     * Detect deprecation warnings for the given annotated elements.
     * @param elements the annotated elements to check
     * @return this CodeWarnings instance for method chaining
     */
    public CodeWarnings detectDeprecation(AnnotatedElement... elements);
    
    /**
     * Detect deprecation warnings for the given stream of annotated elements.
     * @param elements the stream of annotated elements to check
     * @return this CodeWarnings instance for method chaining
     */
    public CodeWarnings detectDeprecation(Stream<AnnotatedElement> elements);
    
    /**
     * Detect deprecation warnings for the given resolvable type.
     * @param resolvableType the resolvable type to check
     * @return this CodeWarnings instance for method chaining
     */
    public CodeWarnings detectDeprecation(ResolvableType resolvableType);
    
    /**
     * Suppress warnings on the given method specification builder.
     * @param method the method specification builder
     */
    public void suppress(MethodSpec.Builder method);
    
    /**
     * Suppress warnings on the given type specification builder.
     * @param type the type specification builder
     */
    public void suppress(TypeSpec.Builder type);
    
    /**
     * Check if there are any warnings registered.
     * @return true if there are warnings
     */
    public boolean hasWarnings();
    
    /**
     * Return a string representation of all warnings.
     * @return the warnings as a string
     */
    public String toString();
}

/**
 * Filter that can be used to exclude bean definitions from AOT processing.
 */
interface BeanRegistrationExcludeFilter {
    /**
     * Determine if the given registered bean should be excluded from AOT processing.
     * @param registeredBean the registered bean to check
     * @return true if the bean should be excluded from AOT processing
     */
    boolean isExcluded(RegisteredBean registeredBean);
}

/**
 * Abstract base class for AOT processing related exceptions.
 */
abstract class AotException extends RuntimeException {
    /**
     * Create a new AotException with the specified message.
     * @param msg the detail message
     */
    public AotException(String msg);
    
    /**
     * Create a new AotException with the specified message and root cause.
     * @param msg the detail message
     * @param cause the root cause
     */
    public AotException(String msg, Throwable cause);
}

/**
 * Exception thrown when AOT processing fails.
 */
class AotProcessingException extends AotException {
    /**
     * Create a new AotProcessingException with the specified message.
     * @param msg the detail message
     */
    public AotProcessingException(String msg);
    
    /**
     * Create a new AotProcessingException with the specified message and root cause.
     * @param msg the detail message
     * @param cause the root cause
     */
    public AotProcessingException(String msg, Throwable cause);
}

/**
 * Exception thrown when bean processing fails during AOT.
 */
class AotBeanProcessingException extends AotException {
    private final RootBeanDefinition beanDefinition;
    
    /**
     * Create a new AotBeanProcessingException for the given bean definition.
     * @param beanDefinition the bean definition that failed processing
     * @param msg the detail message
     */
    public AotBeanProcessingException(RootBeanDefinition beanDefinition, String msg);
    
    /**
     * Create a new AotBeanProcessingException for the given bean definition.
     * @param beanDefinition the bean definition that failed processing
     * @param msg the detail message
     * @param cause the root cause
     */
    public AotBeanProcessingException(RootBeanDefinition beanDefinition, String msg, Throwable cause);
    
    /**
     * Return the bean definition that failed processing.
     * @return the bean definition
     */
    public RootBeanDefinition getBeanDefinition();
}

import org.springframework.aot.generate.GenerationContext;
import org.springframework.beans.factory.aot.BeanRegistrationAotProcessor;
import org.springframework.beans.factory.aot.BeanRegistrationAotContribution;
import org.springframework.beans.factory.support.RegisteredBean;

// Custom AOT processor example
public class CustomBeanRegistrationAotProcessor implements BeanRegistrationAotProcessor {
    
    @Override
    public BeanRegistrationAotContribution processAheadOfTime(RegisteredBean registeredBean) {
        // Only process beans of specific types
        if (!isProcessable(registeredBean)) {
            return null;
        }
        
        return new CustomBeanRegistrationAotContribution(registeredBean);
    }
    
    private boolean isProcessable(RegisteredBean registeredBean) {
        Class<?> beanClass = registeredBean.getBeanClass();
        return MyProcessableInterface.class.isAssignableFrom(beanClass);
    }
    
    private static class CustomBeanRegistrationAotContribution implements BeanRegistrationAotContribution {
        private final RegisteredBean registeredBean;
        
        public CustomBeanRegistrationAotContribution(RegisteredBean registeredBean) {
            this.registeredBean = registeredBean;
        }
        
        @Override
        public void applyTo(GenerationContext generationContext, BeanRegistrationCode beanRegistrationCode) {
            // Generate custom initialization code
            CodeBlock customCode = CodeBlock.builder()
                .addStatement("$T bean = $L", registeredBean.getBeanClass(), "instanceSupplier.get()")
                .addStatement("bean.initializeForAot()")
                .addStatement("return bean")
                .build();
            
            // Apply the custom code to the registration
            beanRegistrationCode.getInstanceSupplierCodeGenerator().addCode(customCode);
        }
    }
}

// Bean factory initialization AOT processor example
public class CustomBeanFactoryInitializationAotProcessor implements BeanFactoryInitializationAotProcessor {
    
    @Override
    public BeanFactoryInitializationAotContribution processAheadOfTime(ConfigurableListableBeanFactory beanFactory) {
        // Check if we need to contribute initialization code
        if (!needsCustomInitialization(beanFactory)) {
            return null;
        }
        
        return new CustomBeanFactoryInitializationAotContribution();
    }
    
    private boolean needsCustomInitialization(ConfigurableListableBeanFactory beanFactory) {
        // Check for specific beans or configuration that requires custom initialization
        return beanFactory.containsBean("customConfigurationBean");
    }
    
    private static class CustomBeanFactoryInitializationAotContribution implements BeanFactoryInitializationAotContribution {
        
        @Override
        public void applyTo(GenerationContext generationContext, BeanFactoryInitializationCode beanFactoryInitializationCode) {
            // Generate custom bean factory initialization code
            MethodSpec initMethod = MethodSpec.methodBuilder("initializeCustomConfiguration")
                .addModifiers(Modifier.PRIVATE, Modifier.STATIC)
                .addParameter(ConfigurableListableBeanFactory.class, "beanFactory")
                .addStatement("// Custom initialization logic here")
                .addStatement("$T customBean = beanFactory.getBean($S, $T.class)", 
                             CustomConfigBean.class, "customConfigurationBean", CustomConfigBean.class)
                .addStatement("customBean.performAotInitialization()")
                .build();
            
            beanFactoryInitializationCode.getTypeBuilder().addMethod(initMethod);
        }
    }
}

// Bean exclusion filter example
public class CustomBeanRegistrationExcludeFilter implements BeanRegistrationExcludeFilter {
    
    @Override
    public boolean isExcluded(RegisteredBean registeredBean) {
        // Exclude prototype beans from AOT processing
        if (registeredBean.getBeanDefinition().isPrototype()) {
            return true;
        }
        
        // Exclude beans with specific annotations
        Class<?> beanClass = registeredBean.getBeanClass();
        if (beanClass.isAnnotationPresent(ExcludeFromAot.class)) {
            return true;
        }
        
        // Exclude test-related beans
        String packageName = beanClass.getPackage().getName();
        if (packageName.contains(".test.") || packageName.endsWith(".test")) {
            return true;
        }
        
        return false;
    }
}

// Custom annotation for excluding beans from AOT
@Target(ElementType.TYPE)
@Retention(RetentionPolicy.RUNTIME)
@Documented
public @interface ExcludeFromAot {
    String reason() default "";
}

// Usage in Spring configuration
@Configuration
public class AotProcessingConfiguration {
    
    @Bean
    public BeanRegistrationAotProcessor customBeanRegistrationAotProcessor() {
        return new CustomBeanRegistrationAotProcessor();
    }
    
    @Bean
    public BeanFactoryInitializationAotProcessor customBeanFactoryInitializationAotProcessor() {
        return new CustomBeanFactoryInitializationAotProcessor();
    }
    
    @Bean
    public BeanRegistrationExcludeFilter customBeanRegistrationExcludeFilter() {
        return new CustomBeanRegistrationExcludeFilter();
    }
}

/**
 * Enumeration of different sources for AOT processing.
 */
enum Source {
    /** Source from Spring factories. */
    SPRING_FACTORIES,
    /** Source from bean factory. */
    BEAN_FACTORY;
    
    /**
     * Create a loader for Spring factories.
     * @return the loader
     */
    public static Loader factories();
    
    /**
     * Create a loader for Spring factories with a specific ClassLoader.
     * @param classLoader the ClassLoader to use
     * @return the loader
     */
    public static Loader factories(ClassLoader classLoader);
    
    /**
     * Create a loader for Spring factories with a specific SpringFactoriesLoader.
     * @param springFactoriesLoader the SpringFactoriesLoader to use
     * @return the loader
     */
    public static Loader factories(SpringFactoriesLoader springFactoriesLoader);
    
    /**
     * Create a loader for both Spring factories and beans.
     * @param beanFactory the bean factory
     * @return the loader
     */
    public static Loader factoriesAndBeans(ListableBeanFactory beanFactory);
    
    /**
     * Create a loader for both Spring factories and beans.
     * @param springFactoriesLoader the SpringFactoriesLoader to use
     * @param beanFactory the bean factory
     * @return the loader
     */
    public static Loader factoriesAndBeans(SpringFactoriesLoader springFactoriesLoader, ListableBeanFactory beanFactory);
    
    /**
     * Loader interface for loading instances from different sources.
     */
    interface Loader<T> extends Iterable<T> {
        /**
         * Return an iterator over the loaded instances.
         * @return the iterator
         */
        Iterator<T> iterator();
        
        /**
         * Return a stream of the loaded instances.
         * @return the stream
         */
        Stream<T> stream();
        
        /**
         * Return the loaded instances as a list.
         * @return the list of instances
         */
        List<T> asList();
    }
    
    /** The location of the AOT factories resource. */
    public static final String FACTORIES_RESOURCE_LOCATION = "META-INF/spring/aot.factories";
}