PsiAugmentProvider - Augment method's parameters


I have developed an annotation called @Reified inspired by Kotlin's reified keyword. To simulate this, I create a variable of type Class<T> if the type parameter is owned by a class or a parameter of type Class<T> if it's owned by a method. The annotation processor then looks for class initializations/method calls that need to be reified and passes the correct class(for example String.class) to the constructor/method. I'm now developing a plugin to make IntelliJ work with these modifications to the source code. After brief research, I have found that the class that I need to use is PsiAugmentProvider as it allows to add "invisible" code to a class. PsiAugmentProvider works perfectly with classes as adding fields is supported, but not with methods as adding a parameter or a local variable inside of the method doesn't seem to be supported. For now, I'm simply adding a field to the class also for methods even though this is not correct. Here is my code:

import com.intellij.psi.*;
import com.intellij.psi.augment.PsiAugmentProvider;
import com.intellij.psi.impl.light.LightFieldBuilder;
import com.intellij.psi.impl.source.PsiExtensibleClass;
import it.auties.reified.annotation.Reified;
import org.apache.commons.collections.ListUtils;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.Objects;

public class ReifiedPlugin extends PsiAugmentProvider {
private static final String REIFIED_PATH = Reified.class.getName();
private static final String JAVADOC = "/**\n* Generated parameter from type parameter\n*/";

protected <Psi extends PsiElement> List<Psi> getAugments(@NotNull PsiElement element, @NotNull Class<Psi> type) {
return getAugments(element, type, null);

protected <Psi extends PsiElement> List<Psi> getAugments(@NotNull PsiElement element, @NotNull Class<Psi> type, @Nullable String nameHint) {
if (!(element instanceof PsiExtensibleClass)) {
return Collections.emptyList();

var clazz = (PsiExtensibleClass) element;
var elementFactory = PsiElementFactory.getInstance(clazz.getProject());
// The only types supported seem to be: [PsiClass, PsiMethod, PsiAnnotation, PsiField]
// Considering that the return type of this method = List<Psi> and type extends Psi
// For me to return PsiParameter one of said classes must be passed to this method
if (type != PsiField.class) {
return Collections.emptyList();

var classFields = createFieldsForConstructors(clazz, elementFactory);
// These fields should be PsiParameter, not PsiField but I cannot return them without throwing a class cast exception
var methodFields = createFieldsForMethods(clazz, elementFactory);
return (List<Psi>) ListUtils.union(classFields, methodFields);

private List<LightFieldBuilder> createFieldsForConstructors(PsiExtensibleClass clazz, PsiElementFactory elementFactory) {
.map(parameter -> createField(clazz, elementFactory, parameter))

private List<LightFieldBuilder> createFieldsForMethods(PsiExtensibleClass clazz, PsiElementFactory elementFactory) {
return clazz.getOwnMethods()
.map(parameter -> createField(clazz, elementFactory, parameter))

private boolean hasAnnotation(PsiTypeParameter parameter) {
.anyMatch(candidate -> Objects.equals(candidate.getQualifiedName(), REIFIED_PATH));

private LightFieldBuilder createField(PsiExtensibleClass clazz, PsiElementFactory elementFactory, PsiTypeParameter parameter) {
var classType = PsiType.getJavaLangClass(clazz.getManager(), clazz.getResolveScope()).resolve();
var genericType = PsiType.getTypeByName(Objects.requireNonNull(parameter.getName()), clazz.getManager().getProject(), clazz.getResolveScope());
var field = createField(clazz, elementFactory, parameter, classType, genericType);
field.setModifiers(PsiModifier.PRIVATE, PsiModifier.FINAL);
return field;

private LightFieldBuilder createField(PsiExtensibleClass clazz, PsiElementFactory elementFactory, PsiTypeParameter parameter, PsiClass classType, PsiClassType genericType) {
return new LightFieldBuilder(
constructGenericType(elementFactory, classType, genericType)

private PsiClassType constructGenericType(PsiElementFactory elementFactory, PsiClass classType, PsiClassType genericType) {
return elementFactory.createType(
1 comment
Comment actions Permalink


Adding method parameters or local variables is not supported indeed.



Please sign in to leave a comment.