[amper][yaml] AMPER-528 Fix broken YAML due to formatting issues after adding a property

This issue seems deeper than just Amper. It seems to be caused by some formatting done after
a write action (doPostponedFormatting on writeActionFinished).

This commit is not a real fix but closer to a workaround: we pre-format the object that we
are inserting into.

GitOrigin-RevId: def33cdfad865a0af055cdd984165afd2a681c3a

plugins/yaml/src/org/jetbrains/yaml/psi/yamlPsiUtil.kt

@@ -3,8 +3,11 @@ package org.jetbrains.yaml.psi
 
 import com.intellij.codeInspection.InspectionManager
 import com.intellij.codeInspection.ProblemsHolder
+import com.intellij.openapi.editor.Document
+import com.intellij.psi.PsiElement
 import com.intellij.psi.util.parentOfType
 import com.intellij.psi.util.parents
+import com.intellij.psi.util.startOffset
 import com.intellij.util.containers.headTailOrNull
 import com.intellij.util.containers.sequenceOfNotNull
 import org.jetbrains.annotations.ApiStatus
@@ -48,3 +51,19 @@ private fun isValid(meta: YamlMetaType, value: YAMLValue): Boolean {
 
 @ApiStatus.Experimental
 fun estimatedType(scalar: YAMLScalar): YamlMetaType? = types.firstOrNull { isValid(it, scalar) }
+
+/**
+ * Returns the closest ancestor of [element] that has no indentation (is at the start of line).
+ * In short, this helps to find the containing top-level Key-Value in non-empty documents.
+ */
+internal fun Document.findClosestAncestorWithoutIndent(element: PsiElement): PsiElement {
+  var current = element
+  while (!isAtStartOfLine(current)) {
+    // It is not possible to reach the root here, because it would mean the root itself is indented - where would the indent node be then?
+    current = current.parent ?: error("the root of the PSI tree cannot be indented itself")
+  }
+  return current
+}
+
+private fun Document.isAtStartOfLine(element: PsiElement): Boolean =
+  getLineStartOffset(getLineNumber(element.startOffset)) == element.startOffset

plugins/yaml/src/org/jetbrains/yaml/schema/YamlJsonLikePsiWalkerFactory.java

@@ -8,11 +8,12 @@ import com.jetbrains.jsonSchema.impl.JsonSchemaObject;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.yaml.psi.YAMLFile;
+import org.jetbrains.yaml.psi.YAMLPsiElement;
 
 public class YamlJsonLikePsiWalkerFactory implements JsonLikePsiWalkerFactory {
   @Override
   public boolean handles(@NotNull PsiElement element) {
-    return element.getContainingFile() instanceof YAMLFile;
+    return element.getContainingFile() instanceof YAMLFile || element instanceof YAMLPsiElement;
   }
 
   @Override

plugins/yaml/src/org/jetbrains/yaml/schema/YamlJsonPsiWalker.java

@@ -6,13 +6,12 @@ import com.intellij.codeInsight.completion.CompletionUtil;
 import com.intellij.codeInsight.completion.CompletionUtilCore;
 import com.intellij.json.pointer.JsonPointerPosition;
 import com.intellij.lang.ASTNode;
+import com.intellij.openapi.editor.Document;
 import com.intellij.openapi.project.Project;
 import com.intellij.openapi.util.TextRange;
 import com.intellij.openapi.util.text.StringUtil;
-import com.intellij.psi.PsiComment;
-import com.intellij.psi.PsiElement;
-import com.intellij.psi.PsiFile;
-import com.intellij.psi.PsiWhiteSpace;
+import com.intellij.psi.*;
+import com.intellij.psi.codeStyle.CodeStyleManager;
 import com.intellij.psi.impl.source.tree.LeafPsiElement;
 import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.psi.util.PsiUtilCore;
@@ -468,5 +467,37 @@ public final class YamlJsonPsiWalker implements JsonLikePsiWalker {
       }
       return sibling;
     }
+
+    @NotNull
+    @Override
+    public PsiElement addProperty(@NotNull PsiElement contextForInsertion, @NotNull PsiElement newProperty) {
+      // Sometimes, post-write-action formatting can break the YAML structure if the area was not indented properly initially.
+      // This is why we pre-format it to avoid problems.
+      preFormatAround(contextForInsertion);
+
+      return JsonLikeSyntaxAdapter.super.addProperty(contextForInsertion, newProperty);
+    }
+
+    private static void preFormatAround(PsiElement element) {
+      PsiDocumentManager documentManager = PsiDocumentManager.getInstance(element.getProject());
+      Document document = documentManager.getDocument(element.getContainingFile());
+      if (document == null) {
+        return; // nothing to format if there is no document anyway
+      }
+      // We need to commit the pending PSI changes before triggering the formatting, otherwise it will fail.
+      // This typically happens if we have several calls to addProperty in a row, but could be with any previous PSI change too.
+      documentManager.doPostponedOperationsAndUnblockDocument(document);
+
+      // If we try to format an element that is itself indented, the formatter will not take this base indent into account.
+      // This is why we need to go up the tree to find the top-level Key-Value that contains our element.
+      PsiElement elementToFormat = YamlPsiUtilKt.findClosestAncestorWithoutIndent(document, element);
+
+      // The formatter doesn't support formatting YAMLDocument or YAMLMapping elements, but if we reach one of those, they represent the
+      // whole file anyway (because they must have an indent of 0), so we can trigger the formatting on the containing file.
+      if (elementToFormat instanceof YAMLDocument || elementToFormat instanceof YAMLMapping) {
+        elementToFormat = elementToFormat.getContainingFile();
+      }
+      CodeStyleManager.getInstance(element.getProject()).reformat(elementToFormat, true);
+    }
   }
 }

plugins/yaml/testSrc/org/jetbrains/yaml/schema/YamlByJsonSchemaQuickFixTest.java

@@ -96,6 +96,84 @@ public class YamlByJsonSchemaQuickFixTest extends JsonSchemaQuickFixTestBase {
                  -\s""");
   }
 
+  public void testAddPropAfterObjectProp() {
+    @Language("JSON") String schema = """
+      {
+        "type": "object",
+        "properties": {
+          "obj": {
+            "type": "object",
+            "properties": {
+              "foo": {
+                "type": "number"
+              },
+              "bar": {
+                "type": "object"
+              },
+              "baz": {
+                "type": "number"
+              }
+            },
+            "required": ["foo", "bar", "baz"]
+          }
+        },
+        "required": ["obj"]
+      }""";
+    String text = """
+      obj:
+        <warning>foo: 42
+        bar:
+          aa: 42
+          ab: 42<caret></warning>""";
+    String afterFix = """
+      obj:
+        foo: 42
+        bar:
+          aa: 42
+          ab: 42
+        baz: 0""";
+    doTest(schema, text, "Add missing property 'baz'", afterFix);
+  }
+
+  public void testAddPropAfterObjectProp_wrongFormatting() {
+    @Language("JSON") String schema = """
+      {
+        "type": "object",
+        "properties": {
+          "obj": {
+            "type": "object",
+            "properties": {
+              "foo": {
+                "type": "number"
+              },
+              "bar": {
+                "type": "object"
+              },
+              "baz": {
+                "type": "number"
+              }
+            },
+            "required": ["foo", "bar", "baz"]
+          }
+        },
+        "required": ["obj"]
+      }""";
+    String text = """
+      obj:
+           <warning>foo: 42
+           bar:
+              aa: 42
+              ab: 42<caret></warning>""";
+    String afterFix = """
+      obj:
+        foo: 42
+        bar:
+          aa: 42
+          ab: 42
+        baz: 0""";
+    doTest(schema, text, "Add missing property 'baz'", afterFix);
+  }
+
   public void testEmptyObjectMultipleProps() {
     String text = """
       xyzObject: