Fix: Format for comprehension expressions (#3916)

Author: MikaelMayer

Source/DafnyCore/AST/Expressions/ComprehensionExpr.cs

@@ -473,6 +473,7 @@ public virtual bool SetIndent(int indentBefore, TokenNewIndentCollector formatte
             indentBefore = formatter.ReduceBlockiness ? indentBefore : formatter.GetNewTokenVisualIndent(token, indentBefore);
             assignOpIndent = formatter.ReduceBlockiness ? indentBefore + formatter.SpaceTab : indentBefore;
             formatter.SetOpeningIndentedRegion(token, indentBefore);
+
             break;
           }
         case ":=":

Source/DafnyCore/AST/Expressions/Expressions.cs

@@ -5,6 +5,7 @@
 using System.Numerics;
 using System.Linq;
 using System.Diagnostics;
+using System.Reflection.Emit;
 using System.Security.AccessControl;
 using Microsoft.Boogie;
 
@@ -2274,7 +2275,11 @@ public bool SetIndent(int indentBefore, TokenNewIndentCollector formatter) {
           formatter.SetIndentations(ownedTokens[0], formatter.binOpIndent, formatter.binOpIndent, formatter.binOpArgIndent);
         } else {
           var startToken = this.StartToken;
-          var newIndent = formatter.GetNewTokenVisualIndent(startToken, formatter.GetIndentInlineOrAbove(startToken));
+          //"," in a comprehension is an "&&", except that it should not try to keep a visual indentation between components.
+          var newIndent =
+            ownedTokens[0].val == "," ?
+              formatter.GetIndentInlineOrAbove(startToken)
+              : formatter.GetNewTokenVisualIndent(startToken, formatter.GetIndentInlineOrAbove(startToken));
           formatter.SetIndentations(ownedTokens[0], newIndent, newIndent, newIndent);
         }
       }
@@ -2339,6 +2344,33 @@ public bool SetIndent(int indentBefore, TokenNewIndentCollector formatter) {
       formatter.Visit(E1, item2Indent);
       formatter.SetIndentations(EndToken, below: indent);
       return false;
+    } else if (Op is Opcode.In or Opcode.NotIn) {
+      var itemIndent = formatter.GetNewTokenVisualIndent(
+        E0.StartToken, indent);
+      var item2Indent = itemIndent + formatter.SpaceTab;
+      formatter.Visit(E0, itemIndent);
+      foreach (var token in this.OwnedTokens) {
+        switch (token.val) {
+          case "<": {
+              if (token.Prev.line != token.line) {
+                itemIndent = formatter.GetNewTokenVisualIndent(token, indent);
+              }
+
+              break;
+            }
+          case "in":
+          case "-": {
+              if (TokenNewIndentCollector.IsFollowedByNewline(token)) {
+                formatter.SetOpeningIndentedRegion(token, itemIndent);
+              } else {
+                formatter.SetAlign(itemIndent, token, out item2Indent, out _);
+              }
+              break;
+            }
+        }
+      }
+      formatter.Visit(E1, item2Indent);
+      return false;
     } else {
       foreach (var token in OwnedTokens) {
         formatter.SetIndentations(token, indent, indent, indent);

Source/DafnyPipeline.Test/FormatterIssues.cs

@@ -4,6 +4,83 @@ namespace DafnyPipeline.Test;
 
 [Collection("Singleton Test Collection - FormatterForTopLevelDeclarations")]
 public class FormatterIssues : FormatterBaseTest {
+  [Fact]
+  public void GitIssue3912FormatterCollectionArrow() {
+    FormatterWorksFor(@"
+const i :=
+  a in
+    B + // Previously it was not indented
+    C +
+    D &&
+  b in
+    B +
+    C
+");
+  }
+
+  [Fact]
+  public void GitIssue3912FormatterCollectionArrowA() {
+    FormatterWorksFor(@"
+const newline :=
+  set
+    i <-
+      PartA +
+      PartB +
+      PartC,
+    j <-
+      PartD
+    ::
+      f(i,j)
+
+const sameline :=
+  set i <-
+        PartA +
+        PartB +
+        PartC,
+    j <-
+      PartD
+    ::
+      f(i,j)
+
+");
+  }
+
+  [Fact]
+  public void GitIssue3912FormatterCollectionArrowB() {
+    FormatterWorksFor(@"
+const newlineOp :=
+  set
+    i
+    <-
+      PartA +
+      PartB +
+      PartC,
+    j
+    <-
+      PartD
+    ::
+      f(i,j)
+
+");
+  }
+
+  [Fact]
+  public void GitIssue3912FormatterCollectionArrowC() {
+    FormatterWorksFor(@"
+const sameLineOp :=
+  set i
+    <-
+      PartA +
+      PartB +
+      PartC,
+    j
+    <- PartD +
+       PartE
+    ::
+      f(i,j)
+");
+  }
+
   [Fact]
   public void GitIssue3960FormattingIssueForallStatements() {
     FormatterWorksFor(@"
@@ -50,4 +127,4 @@ public void FormatterWorksForEmptyDocument() {
     FormatterWorksFor(@"
 ", null, true);
   }
-}
+}

docs/dev/news/3912.fix

@@ -0,0 +1 @@
+Format for comprehension expressions