package mr.intellij.plugin.autofactory.augment;

import com.intellij.psi.PsiElement;
import com.intellij.psi.augment.PsiAugmentProvider;
import lombok.AccessLevel;
import lombok.RequiredArgsConstructor;
import lombok.SneakyThrows;
import org.jetbrains.annotations.NotNull;

import java.util.List;

@RequiredArgsConstructor(access = AccessLevel.PROTECTED)
public abstract class GuardedPsiAugmentProvider<T extends PsiElement> extends PsiAugmentProvider {

    private boolean suppress = false;

    private final Class<T> supportedType;

    @NotNull
    @Override
    @SneakyThrows
    @SuppressWarnings("unchecked")
    protected final <Psi extends PsiElement> List<Psi> getAugments(@NotNull PsiElement element,
                                                                   @NotNull Class<Psi> type) {

        if (suppress || supportedType != type) {
            return super.getAugments(element, type);
        }

        try (Suppressor ignored = new Suppressor()) {
            return (List<Psi>) doGetAugments(element);
        }
    }

    protected abstract List<T> doGetAugments(@NotNull PsiElement element);

    private class Suppressor implements AutoCloseable {

        public Suppressor() {
            suppress = true;
        }

        @Override
        public void close() throws Exception {
            suppress = false;
        }
    }
}

@@ -288,14 +288,640 @@
 til.
-List;%0A
+Collections;%0Aimport java.util.List;%0A%0A/**%0A * PsiAugment provider that is stack overflow resistant.%0A *%0A * The rational for such class is that %7B@link InjectAugmentProvider%7D is supposed to (sometimes) add %7B@code @Inject%7D%0A * annotations on constructors, but in its %7B@code getAugments()%7D logic is check for other annotations presence, which%0A * triggers another %7B@code getAugments()%7D and so on, causing stack overflow.%0A *%0A * This class fixes this issue by introducing a %7B@link Suppressor%7D that disables further %7B@code getAugment()%7D calls.%0A *%0A * @param %3CT%3E The supported %7B@link PsiElement%7D (i.e. %7B@link com.intellij.psi.PsiAnnotation%7D).%0A */
 %0A@Re
@@ -1090,35 +1090,38 @@
 ate 
-boolean suppress = fals
+final Class%3CT%3E supportedTyp
 e;%0A
-%0A
     
@@ -1132,35 +1132,31 @@
 ate 
-final Class%3CT%3E supportedTyp
+boolean suppress = fals
 e;%0A%0A
@@ -1499,39 +1499,30 @@
 urn 
-super.getAugments(element, type
+Collections.emptyList(
 );%0A 

Add documentation