indentation of classes, collection contributions, expressions after '='

src/main/java/org/jastadd/tooling/aspect/AspectBlock.java

@@ -4,11 +4,13 @@ package org.jastadd.tooling.aspect;
 import com.intellij.formatting.*;
 import com.intellij.lang.ASTNode;
 import com.intellij.openapi.util.TextRange;
+import com.intellij.psi.PsiElement;
 import com.intellij.psi.TokenType;
 import com.intellij.psi.codeStyle.CodeStyleSettings;
 import com.intellij.psi.formatter.common.AbstractBlock;
 import com.intellij.psi.formatter.common.InjectedLanguageBlockBuilder;
 import com.intellij.psi.formatter.java.LeafBlock;
+import com.intellij.psi.tree.IElementType;
 import org.jastadd.tooling.aspect.psi.*;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
@@ -32,7 +34,9 @@ public class AspectBlock extends AbstractBlock {
   @Override
   protected List<Block> buildChildren() {
     List<Block> blocks = new ArrayList<>();
-    if (myNode.getPsi() instanceof JastAddAspectBlock) {
+    if (myNode.getPsi() instanceof JastAddAspectAspectClassDeclaration) {
+      myInjectedBlockBuilder.addInjectedBlocks(blocks, myNode, Wrap.createWrap(WrapType.NONE, false), null, Indent.getIndent(Indent.Type.NONE, false, true));
+    } else if (myNode.getPsi() instanceof JastAddAspectBlock) {
       myInjectedBlockBuilder.addInjectedBlocks(blocks, myNode, Wrap.createWrap(WrapType.NONE, false), null, Indent.getIndent(Indent.Type.NONE, false, true));
     } else {
       ASTNode child = myNode.getFirstChildNode();
@@ -49,12 +53,25 @@ public class AspectBlock extends AbstractBlock {
 
   @Override
   public Indent getIndent() {
-
-    if (myNode.getTreeParent() != null && myNode.getTreeParent().getPsi() instanceof JastAddAspectAspectBody && myNode.getElementType() != AspectTypes.RBRACE) {
+    if (myNode.getTreeParent() != null) {
+      PsiElement parentPsi = myNode.getTreeParent().getPsi();
+      IElementType childType =  myNode.getElementType();
+      if (parentPsi instanceof JastAddAspectAspectBody && childType != AspectTypes.RBRACE) {
+        return Indent.getNormalIndent();
+      } else if (parentPsi instanceof JastAddAspectCollectionContribution && (childType == AspectTypes.WHEN || childType == AspectTypes.TO || childType == AspectTypes.FOR)) {
         return Indent.getNormalIndent();
-    } else {
-      return Indent.getNoneIndent();
       }
+      ASTNode nonBlankSuccecssor = myNode.getTreePrev();
+      while (nonBlankSuccecssor != null && nonBlankSuccecssor.getElementType() == TokenType.WHITE_SPACE) {
+        nonBlankSuccecssor = nonBlankSuccecssor.getTreePrev();
+      }
+      if (nonBlankSuccecssor != null && nonBlankSuccecssor.getElementType() == AspectTypes.ASSIGN) {
+        return Indent.getContinuationIndent();
+      } else if (nonBlankSuccecssor != null) {
+        //System.out.println(nonBlankSuccecssor.getElementType().getDebugName());
+      }
+    }
+    return Indent.getNoneIndent();
   }
 
   @Nullable