From cb02ae515966439e29a677a1c3d95e96faab5d6e Mon Sep 17 00:00:00 2001
From: Maksim Kurnikov <maxim.kurnikov@gmail.com>
Date: Thu, 27 Jun 2024 15:14:31 +0300
Subject: [PATCH] fix indentation in spec statements

---
 src/main/kotlin/org/move/ide/formatter/impl/indent.kt     | 4 ++++
 src/main/kotlin/org/move/ide/formatter/impl/utils.kt      | 8 ++++++++
 .../resources/org/move/ide/formatter.fixtures/specs.move  | 5 +++++
 .../org/move/ide/formatter.fixtures/specs_after.move      | 5 +++++
 4 files changed, 22 insertions(+)

diff --git a/src/main/kotlin/org/move/ide/formatter/impl/indent.kt b/src/main/kotlin/org/move/ide/formatter/impl/indent.kt
index 40736d8cf..8b9fd2747 100644
--- a/src/main/kotlin/org/move/ide/formatter/impl/indent.kt
+++ b/src/main/kotlin/org/move/ide/formatter/impl/indent.kt
@@ -5,6 +5,7 @@ import com.intellij.lang.ASTNode
 import org.move.ide.formatter.MoveFmtBlock
 import org.move.lang.MvElementTypes.*
 import org.move.lang.core.psi.MvExpr
+import org.move.lang.core.psi.MvPragmaSpecStmt
 
 fun MoveFmtBlock.computeChildIndent(childNode: ASTNode): Indent? {
     val parentNode = node
@@ -52,6 +53,9 @@ fun MoveFmtBlock.computeChildIndent(childNode: ASTNode): Indent? {
         //     .myotherfield
         parentPsi is MvExpr -> Indent.getContinuationWithoutFirstIndent()
 
+        // same thing as previous one, but for spec statements
+        parentPsi.isSpecStmt -> Indent.getContinuationWithoutFirstIndent()
+
         else -> Indent.getNoneIndent()
     }
 }
diff --git a/src/main/kotlin/org/move/ide/formatter/impl/utils.kt b/src/main/kotlin/org/move/ide/formatter/impl/utils.kt
index a35507fdb..6f62e2283 100644
--- a/src/main/kotlin/org/move/ide/formatter/impl/utils.kt
+++ b/src/main/kotlin/org/move/ide/formatter/impl/utils.kt
@@ -43,6 +43,14 @@ val DELIMITED_BLOCKS = orSet(
 
 fun ASTNode?.isWhitespaceOrEmpty() = this == null || textLength == 0 || elementType == TokenType.WHITE_SPACE
 
+val PsiElement.isSpecStmt: Boolean
+    get() = this is MvSchemaFieldStmt
+            || this is MvGlobalVariableStmt
+            || this is MvPragmaSpecStmt
+            || this is MvUpdateSpecStmt
+            || this is MvIncludeStmt
+            || this is MvApplySchemaStmt
+
 val PsiElement.isTopLevelItem: Boolean
     get() = (this is MvModule || this is MvAddressDef || this is MvScript || this is MvModuleSpec)
             && parent is MoveFile
diff --git a/src/test/resources/org/move/ide/formatter.fixtures/specs.move b/src/test/resources/org/move/ide/formatter.fixtures/specs.move
index 485aa833e..1fadd081e 100644
--- a/src/test/resources/org/move/ide/formatter.fixtures/specs.move
+++ b/src/test/resources/org/move/ide/formatter.fixtures/specs.move
@@ -2,6 +2,11 @@ module M {
     spec initialize {
 assert true;
 
+        pragma hello = world,
+        hello2 = world2;
+        include
+        MySchema;
+
     /// After genesis, time progresses monotonically.
     invariant update
     old(is_operating()) ==> old(spec_now_microseconds()) <= spec_now_microseconds();
diff --git a/src/test/resources/org/move/ide/formatter.fixtures/specs_after.move b/src/test/resources/org/move/ide/formatter.fixtures/specs_after.move
index 150b99992..e9cd1494b 100644
--- a/src/test/resources/org/move/ide/formatter.fixtures/specs_after.move
+++ b/src/test/resources/org/move/ide/formatter.fixtures/specs_after.move
@@ -2,6 +2,11 @@ module M {
     spec initialize {
         assert true;
 
+        pragma hello = world,
+            hello2 = world2;
+        include
+            MySchema;
+
         /// After genesis, time progresses monotonically.
         invariant update
             old(is_operating()) ==> old(spec_now_microseconds()) <= spec_now_microseconds();