diff --git a/src/main/java/org/jastadd/tooling/aspect/AspectFoldingBuilder.java b/src/main/java/org/jastadd/tooling/aspect/AspectFoldingBuilder.java index 62b269d91cc942588cf5bb24ed6ff1e7696d38f4..9fc57f8ecf188095bada03893af729b7c79fc495 100644 --- a/src/main/java/org/jastadd/tooling/aspect/AspectFoldingBuilder.java +++ b/src/main/java/org/jastadd/tooling/aspect/AspectFoldingBuilder.java @@ -9,10 +9,7 @@ import com.intellij.openapi.editor.FoldingGroup; import com.intellij.openapi.util.TextRange; import com.intellij.psi.PsiElement; import com.intellij.psi.util.PsiTreeUtil; -import org.jastadd.tooling.aspect.psi.JastAddAspectAspectDeclaration; -import org.jastadd.tooling.aspect.psi.JastAddAspectBlock; -import org.jastadd.tooling.aspect.psi.JastAddAspectClassDeclaration; -import org.jastadd.tooling.aspect.psi.JastAddAspectConstructorDeclaration; +import org.jastadd.tooling.aspect.psi.*; import org.jetbrains.annotations.NotNull; import org.jetbrains.annotations.Nullable; @@ -37,6 +34,12 @@ public class AspectFoldingBuilder extends FoldingBuilderEx { new TextRange(classDeclaration.getClassBody().getTextRange().getStartOffset() + 1, classDeclaration.getClassBody().getTextRange().getEndOffset() - 1))); } + Collection<JastAddAspectEnumBody> enumBodies = PsiTreeUtil.findChildrenOfType(root, JastAddAspectEnumBody.class); + for (final JastAddAspectEnumBody enumBody : enumBodies) { + descriptors.add(new FoldingDescriptor(enumBody.getNode(), + new TextRange(enumBody.getTextRange().getStartOffset() + 1, + enumBody.getTextRange().getEndOffset() - 1))); + } Collection<JastAddAspectBlock> aspectBlocks = PsiTreeUtil.findChildrenOfType(root, JastAddAspectBlock.class); for (final JastAddAspectBlock block : aspectBlocks) { descriptors.add(new FoldingDescriptor(block.getNode(),