simple language: SimpleFoldingBuilder use visitor, cleanup

code_samples/simple_language_plugin/src/main/java/org/intellij/sdk/language/SimpleFoldingBuilder.java

@@ -11,16 +11,16 @@
 import com.intellij.openapi.project.Project;
 import com.intellij.openapi.util.TextRange;
 import com.intellij.openapi.util.text.StringUtil;
+import com.intellij.psi.JavaRecursiveElementWalkingVisitor;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiLiteralExpression;
-import com.intellij.psi.util.PsiTreeUtil;
+import com.intellij.psi.util.PsiLiteralUtil;
 import com.intellij.util.containers.ContainerUtil;
 import org.intellij.sdk.language.psi.SimpleProperty;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
 import java.util.ArrayList;
-import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
 
@@ -32,28 +32,32 @@ public class SimpleFoldingBuilder extends FoldingBuilderEx implements DumbAware
     FoldingGroup group = FoldingGroup.newGroup(SimpleAnnotator.SIMPLE_PREFIX_STR);
     // Initialize the list of folding regions
     List<FoldingDescriptor> descriptors = new ArrayList<>();
-    // Get a collection of the literal expressions in the document below root
-    Collection<PsiLiteralExpression> literalExpressions =
-            PsiTreeUtil.findChildrenOfType(root, PsiLiteralExpression.class);
-    // Evaluate the collection
-    for (final PsiLiteralExpression literalExpression : literalExpressions) {
-      String value = literalExpression.getValue() instanceof String ? (String) literalExpression.getValue() : null;
-      if (value != null && value.startsWith(SimpleAnnotator.SIMPLE_PREFIX_STR + SimpleAnnotator.SIMPLE_SEPARATOR_STR)) {
-        Project project = literalExpression.getProject();
-        String key = value.substring(
-                SimpleAnnotator.SIMPLE_PREFIX_STR.length() + SimpleAnnotator.SIMPLE_SEPARATOR_STR.length()
-        );
-        // find SimpleProperty for the given key in the project
-        SimpleProperty simpleProperty = ContainerUtil.getOnlyItem(SimpleUtil.findProperties(project, key));
-        if (simpleProperty != null) {
-          // Add a folding descriptor for the literal expression at this node.
-          descriptors.add(new FoldingDescriptor(literalExpression.getNode(),
-                  new TextRange(literalExpression.getTextRange().getStartOffset() + 1,
-                          literalExpression.getTextRange().getEndOffset() - 1),
-                  group, Collections.singleton(simpleProperty)));
+
+    root.accept(new JavaRecursiveElementWalkingVisitor() {
+
+      @Override
+      public void visitLiteralExpression(@NotNull PsiLiteralExpression literalExpression) {
+        super.visitLiteralExpression(literalExpression);
+
+        String value = PsiLiteralUtil.getStringLiteralContent(literalExpression);
+        if (value != null && value.startsWith(SimpleAnnotator.SIMPLE_PREFIX_STR + SimpleAnnotator.SIMPLE_SEPARATOR_STR)) {
+          Project project = literalExpression.getProject();
+          String key = value.substring(
+                  SimpleAnnotator.SIMPLE_PREFIX_STR.length() + SimpleAnnotator.SIMPLE_SEPARATOR_STR.length()
+          );
+          // find SimpleProperty for the given key in the project
+          SimpleProperty simpleProperty = ContainerUtil.getOnlyItem(SimpleUtil.findProperties(project, key));
+          if (simpleProperty != null) {
+            // Add a folding descriptor for the literal expression at this node.
+            descriptors.add(new FoldingDescriptor(literalExpression.getNode(),
+                    new TextRange(literalExpression.getTextRange().getStartOffset() + 1,
+                            literalExpression.getTextRange().getEndOffset() - 1),
+                    group, Collections.singleton(simpleProperty)));
+          }
         }
       }
-    }
+    });
+
     return descriptors.toArray(FoldingDescriptor.EMPTY);
   }
 
@@ -67,22 +71,34 @@ public class SimpleFoldingBuilder extends FoldingBuilderEx implements DumbAware
   @Nullable
   @Override
   public String getPlaceholderText(@NotNull ASTNode node) {
-    if (!(node.getPsi() instanceof PsiLiteralExpression nodeElement)) {
-      return null ;
-    }
+    if (node.getPsi() instanceof PsiLiteralExpression psiLiteralExpression) {
+      String text = PsiLiteralUtil.getStringLiteralContent(psiLiteralExpression);
+      if (text == null) {
+        return null;
+      }
+
+      String key = text.substring(SimpleAnnotator.SIMPLE_PREFIX_STR.length() +
+              SimpleAnnotator.SIMPLE_SEPARATOR_STR.length());
+
+      SimpleProperty simpleProperty = ContainerUtil.getOnlyItem(
+              SimpleUtil.findProperties(psiLiteralExpression.getProject(), key)
+      );
+      if (simpleProperty == null) {
+        return StringUtil.THREE_DOTS;
+      }
 
-    String key = ((String) nodeElement.getValue()).substring(
-            SimpleAnnotator.SIMPLE_PREFIX_STR.length() + SimpleAnnotator.SIMPLE_SEPARATOR_STR.length()
-    );
+      String propertyValue = simpleProperty.getValue();
+      // IMPORTANT: keys can come with no values, so a test for null is needed
+      // IMPORTANT: Convert embedded \n to backslash n, so that the string will look
+      // like it has LF embedded in it and embedded " to escaped "
+      if (propertyValue == null) {
+        return StringUtil.THREE_DOTS;
+      }
 
-    SimpleProperty simpleProperty = ContainerUtil.getOnlyItem(SimpleUtil.findProperties(nodeElement.getProject(), key));
-    if (simpleProperty == null) return StringUtil.THREE_DOTS;
+      return propertyValue.replaceAll("\n", "\\n").replaceAll("\"", "\\\\\"");
+    }
 
-    String propertyValue = simpleProperty.getValue();
-    // IMPORTANT: keys can come with no values, so a test for null is needed
-    // IMPORTANT: Convert embedded \n to backslash n, so that the string will look
-    // like it has LF embedded in it and embedded " to escaped "
-    return propertyValue == null ? StringUtil.THREE_DOTS : propertyValue.replaceAll("\n", "\\n").replaceAll("\"", "\\\\\"");
+    return null;
   }
 
   @Override