Make @Enhances, TraitInjectors and AST transforms work in same project using gradle custom sourceset

Groovy AST allows you to hook into compilation process and provides ability to modify the existing behaviour or add new behaviour to classes at compile time.

Grails makes it even easier by providing @Enhances annotation and TraitInjector which can be utilised to apply traits to artefacts at compile time.

However if you have previously created Groovy AST transforms, you will be well aware that the AST Transformation needs to be precompiled and be present in classpath while compiling the code. This requires either you do some hacks to precompile the AST classes within same project, or put the AST related code in a separate jar and put that jar in class path of your project.

However if you are using the gradle to build your project, the issue mentioned above can be easily solved by creating a custom ast sourceset which gets compiled before main sources, and is in classpath of main sourceset

Example of how to define additional ast sourceset.

    import org.gradle.api.tasks.JavaExec
    import org.gradle.api.tasks.Copy
    
    ......
    
    sourceSets {
                ast {
                    groovy {
                        compileClasspath += project.configurations.compile
                    }
                }
    
                main {
                    compileClasspath += sourceSets.ast.output
                }
    
                test {
                    compileClasspath += sourceSets.ast.output
                }
    }
    
    
    task copyAstClasses(type:Copy) {
        from sourceSets.ast.output
        into sourceSets.main.output.classesDir
    }
    
    tasks.withType(JavaExec) {
        classpath += sourceSets.ast.output
    }
    
    classes.dependsOn << copyAstClasses

Now you can put your ast transforms or grails TraitInjectors or traits annotated with @Enhances annotation inside src/ast/groovy folder and it will get applied to classes in the same project or any other project which has dependency on your module.

Explanation

The above code does multiple things.

1. It defines a sourceset with name ast with all compile time dependencies in its classpath
2. It puts the output of ast sourceset into classpath of main and test sourceset. This is the trick which gets the AST transforms into classpath of main sourceset.
3. It defines a new task copyAstClasses to copy the compiled ast classes in the classes directory of main sourceset so that ast classes can get bundled into jar.
4. Adds the copyAstClasses as dependency of classes task so that it gets called every time the project is compiled.
5. Adds ast classes to JavaExec task classpath 

Now when you do gradle assemble you will see that ast classes gets compiled before any thing else.

    :compileAstJava
    :compileAstGroovy 
    :processAstResources
    :astClasses
    :compileJava
    :compileGroovy
    :copyAstClasses
    :classes
``

**Note:** GrailsPluginGradlePlugin already configures ast sourceset for plugin projects. So above solution is not required for grails plugins. However the solution explained above can be used for grails applications which want to use @Enhances annotation or TraitInjectors.