format: retain comments in case and if expressions

compiler/src/AST/Source.hs

@@ -8,6 +8,8 @@ module AST.Source
     Expr,
     Expr_ (..),
     VarType (..),
+    IfBranch,
+    CaseBranch,
     Def (..),
     Pattern,
     Pattern_ (..),
@@ -63,9 +65,9 @@ data Expr_
   | Binops [(Expr, [Comment], A.Located Name)] Expr
   | Lambda [([Comment], Pattern)] Expr SC.LambdaComments
   | Call Expr [([Comment], Expr)]
-  | If [(Expr, Expr)] Expr
+  | If [IfBranch] Expr SC.IfComments
   | Let [([Comment], A.Located Def)] Expr SC.LetComments
-  | Case Expr [([Comment], Pattern, Expr)]
+  | Case Expr [CaseBranch] SC.CaseComments
   | Accessor Name
   | Access Expr (A.Located Name)
   | Update Expr [(A.Located Name, Expr)]
@@ -76,6 +78,12 @@ data Expr_
 data VarType = LowVar | CapVar
   deriving (Show)
 
+type IfBranch =
+  (Expr, Expr, SC.IfBranchComments)
+
+type CaseBranch =
+  (Pattern, Expr, SC.CaseBranchComments)
+
 -- DEFINITIONS
 
 data Def
@@ -92,8 +100,8 @@ data Pattern_
   | PVar Name
   | PRecord [RecordFieldPattern]
   | PAlias Pattern (A.Located Name)
-  | PCtor A.Region Name [Pattern]
-  | PCtorQual A.Region Name Name [Pattern]
+  | PCtor A.Region Name [([Comment], Pattern)]
+  | PCtorQual A.Region Name Name [([Comment], Pattern)]
   | PArray [Pattern]
   | PChr ES.String
   | PStr ES.String

compiler/src/AST/SourceComments.hs

@@ -95,8 +95,36 @@ data LambdaComments = LambdaComments
   }
   deriving (Show)
 
+data IfComments = IfComments
+  { _beforeElseBody :: [Comment],
+    _afterElseBody :: [Comment]
+  }
+  deriving (Show)
+
+data IfBranchComments = IfBranchComments
+  { _afterIfKeyword :: [Comment],
+    _beforeThenKeyword :: [Comment],
+    _beforeThenBody :: [Comment],
+    _afterThenBody :: [Comment]
+  }
+  deriving (Show)
+
 data LetComments = LetComments
   { _afterLetDecls :: [Comment],
     _afterIn :: [Comment]
   }
   deriving (Show)
+
+data CaseComments = CaseComments
+  { _afterCaseKeyword :: [Comment],
+    _beforeOfKeyword :: [Comment]
+  }
+  deriving (Show)
+
+data CaseBranchComments = CaseBranchComments
+  { _beforeBranch :: [Comment],
+    _beforeBranchArrow :: [Comment],
+    _beforeBranchBody :: [Comment],
+    _afterBranchBody :: [Comment]
+  }
+  deriving (Show)

compiler/src/Canonicalize/Expression.hs

@@ -94,13 +94,13 @@ canonicalize env (A.At region expression) =
         Can.Call
           <$> canonicalize env func
           <*> traverse (canonicalize env) (fmap snd args)
-      Src.If branches finally ->
+      Src.If branches finally _ ->
         Can.If
           <$> traverse (canonicalizeIfBranch env) branches
           <*> canonicalize env finally
       Src.Let defs expr _ ->
         A.toValue <$> canonicalizeLet region env (fmap snd defs) expr
-      Src.Case expr branches ->
+      Src.Case expr branches _ ->
         Can.Case
           <$> canonicalize env expr
           <*> traverse (canonicalizeCaseBranch env) branches
@@ -125,16 +125,16 @@ canonicalize env (A.At region expression) =
 
 -- CANONICALIZE IF BRANCH
 
-canonicalizeIfBranch :: Env.Env -> (Src.Expr, Src.Expr) -> Result FreeLocals [W.Warning] (Can.Expr, Can.Expr)
-canonicalizeIfBranch env (condition, branch) =
+canonicalizeIfBranch :: Env.Env -> Src.IfBranch -> Result FreeLocals [W.Warning] (Can.Expr, Can.Expr)
+canonicalizeIfBranch env (condition, branch, _) =
   (,)
     <$> canonicalize env condition
     <*> canonicalize env branch
 
 -- CANONICALIZE CASE BRANCH
 
-canonicalizeCaseBranch :: Env.Env -> ([Src.Comment], Src.Pattern, Src.Expr) -> Result FreeLocals [W.Warning] Can.CaseBranch
-canonicalizeCaseBranch env (_, pattern, expr) =
+canonicalizeCaseBranch :: Env.Env -> Src.CaseBranch -> Result FreeLocals [W.Warning] Can.CaseBranch
+canonicalizeCaseBranch env (pattern, expr, _) =
   directUsage $
     do
       (cpattern, bindings) <-
@@ -246,9 +246,9 @@ addBindingsHelp bindings (A.At region pattern) =
     Src.PRecord fields ->
       List.foldl' addBindingsHelp bindings (map extractRecordFieldPattern fields)
     Src.PCtor _ _ patterns ->
-      List.foldl' addBindingsHelp bindings patterns
+      List.foldl' addBindingsHelp bindings (fmap snd patterns)
     Src.PCtorQual _ _ _ patterns ->
-      List.foldl' addBindingsHelp bindings patterns
+      List.foldl' addBindingsHelp bindings (fmap snd patterns)
     Src.PArray patterns ->
       List.foldl' addBindingsHelp bindings patterns
     Src.PAlias aliasPattern (A.At nameRegion name) ->
@@ -356,8 +356,8 @@ getPatternNames names (A.At region pattern) =
     Src.PRecord fields ->
       List.foldl' (\n f -> getPatternNames n (extractRecordFieldPattern f)) names fields
     Src.PAlias ptrn name -> getPatternNames (name : names) ptrn
-    Src.PCtor _ _ args -> List.foldl' getPatternNames names args
-    Src.PCtorQual _ _ _ args -> List.foldl' getPatternNames names args
+    Src.PCtor _ _ args -> List.foldl' getPatternNames names (fmap snd args)
+    Src.PCtorQual _ _ _ args -> List.foldl' getPatternNames names (fmap snd args)
     Src.PArray patterns -> List.foldl' getPatternNames names patterns
     Src.PChr _ -> names
     Src.PStr _ -> names

compiler/src/Canonicalize/Pattern.hs

@@ -65,9 +65,9 @@ canonicalize env (A.At region pattern) =
       Src.PRecord fields ->
         Can.PRecord <$> canonicalizeRecordFields env fields
       Src.PCtor nameRegion name patterns ->
-        canonicalizeCtor env region name patterns =<< Env.findCtor nameRegion env name
+        canonicalizeCtor env region name (fmap snd patterns) =<< Env.findCtor nameRegion env name
       Src.PCtorQual nameRegion home name patterns ->
-        canonicalizeCtor env region name patterns =<< Env.findCtorQual nameRegion env home name
+        canonicalizeCtor env region name (fmap snd patterns) =<< Env.findCtorQual nameRegion env home name
       Src.PArray patterns ->
         Can.PArray <$> canonicalizeList env patterns
       Src.PAlias ptrn (A.At reg name) ->

compiler/src/Gren/Format.hs

@@ -579,33 +579,44 @@ formatExpr = \case
     where
       formatArg (commentsBefore, arg) =
         exprParensProtectSpaces (formatExpr $ A.toValue arg)
-  Src.If [] else_ ->
+  Src.If [] else_ _ ->
     formatExpr $ A.toValue else_
-  Src.If (if_ : elseifs) else_ ->
+  Src.If (if_ : elseifs) else_ (SC.IfComments commentsBeforeElseBody commentsAfterElseBody) ->
     ExpressionHasAmbiguousEnd $
       Block.stack $
         NonEmpty.fromList $
           mconcat
-            [ List.singleton $ formatIfClause "if" if_,
-              fmap (formatIfClause "else if") elseifs,
+            [ List.singleton $ formatIfBranch "if" if_,
+              fmap (formatIfBranch "else if") elseifs,
               List.singleton $
                 Block.stack
                   [ Block.line $ Block.string7 "else",
-                    Block.indent $ exprParensNone $ formatExpr $ A.toValue else_
+                    Block.indent $
+                      withCommentsStackAround commentsBeforeElseBody commentsAfterElseBody $
+                        exprParensNone $
+                          formatExpr $
+                            A.toValue else_
                   ]
             ]
     where
-      formatIfClause :: String -> (Src.Expr, Src.Expr) -> Block
-      formatIfClause keyword (predicate, body) =
+      formatIfBranch :: String -> Src.IfBranch -> Block
+      formatIfBranch keyword (predicate, body, SC.IfBranchComments commentsAfterIf commentsBeforeThen commentsBeforeBody commentsAfterBody) =
         Block.stack
           [ spaceOrStack
               [ spaceOrIndent
                   [ Block.line $ Block.string7 keyword,
-                    exprParensNone $ formatExpr $ A.toValue predicate
+                    withCommentsAround commentsAfterIf commentsBeforeThen $
+                      exprParensNone $
+                        formatExpr $
+                          A.toValue predicate
                   ],
                 Block.line $ Block.string7 "then"
               ],
-            Block.indent $ exprParensNone $ formatExpr $ A.toValue body
+            Block.indent $
+              withCommentsStackAround commentsBeforeBody commentsAfterBody $
+                exprParensNone $
+                  formatExpr $
+                    A.toValue body
           ]
   Src.Let [] body _ ->
     formatExpr $ A.toValue body
@@ -624,26 +635,35 @@ formatExpr = \case
                 ],
           withCommentsStackBefore commentsAfterIn $ exprParensNone $ formatExpr (A.toValue body)
         ]
-  Src.Case subject branches ->
+  Src.Case subject branches (SC.CaseComments commentsAfterCaseKeyword commentsBeforeOfKeyword) ->
     ExpressionHasAmbiguousEnd $
       Block.stack $
         spaceOrStack
           [ spaceOrIndent
               [ Block.line (Block.string7 "case"),
-                exprParensNone $ formatExpr (A.toValue subject)
+                withCommentsAround commentsAfterCaseKeyword commentsBeforeOfKeyword $
+                  exprParensNone $
+                    formatExpr (A.toValue subject)
               ],
             Block.line (Block.string7 "of")
           ]
           :| List.intersperse Block.blankLine (fmap (Block.indent . formatCaseBranch) branches)
     where
-      formatCaseBranch (commentsBefore, pat, expr) =
-        Block.stack
-          [ spaceOrStack
-              [ patternParensNone $ formatPattern (A.toValue pat),
-                Block.line $ Block.string7 "->"
-              ],
-            Block.indent $ exprParensNone $ formatExpr $ A.toValue expr
-          ]
+      formatCaseBranch (pat, expr, SC.CaseBranchComments commentsBefore commentsAfterPattern commentsBeforeBody commentsAfterBody) =
+        withCommentsStackBefore commentsBefore $
+          Block.stack
+            [ spaceOrStack
+                [ withCommentsAround [] commentsAfterPattern $
+                    patternParensNone $
+                      formatPattern (A.toValue pat),
+                  Block.line $ Block.string7 "->"
+                ],
+              Block.indent $
+                withCommentsStackAround commentsBeforeBody commentsAfterBody $
+                  exprParensNone $
+                    formatExpr $
+                      A.toValue expr
+            ]
   Src.Accessor field ->
     NoExpressionParens $
       Block.line $
@@ -857,15 +877,15 @@ formatPattern = \case
     PatternContainsSpaces $
       spaceOrIndent $
         Block.line (utf8 name)
-          :| fmap (patternParensProtectSpaces . formatPattern . A.toValue) args
+          :| fmap (patternParensProtectSpaces . formatPatternConstructorArg) args
   Src.PCtorQual _ ns name [] ->
     NoPatternParens $
       Block.line (utf8 ns <> Block.char7 '.' <> utf8 name)
   Src.PCtorQual _ ns name args ->
     PatternContainsSpaces $
       spaceOrIndent $
         Block.line (utf8 ns <> Block.char7 '.' <> utf8 name)
-          :| fmap (patternParensProtectSpaces . formatPattern . A.toValue) args
+          :| fmap (patternParensProtectSpaces . formatPatternConstructorArg) args
   Src.PArray items ->
     NoPatternParens $
       group '[' ',' ']' False $
@@ -881,6 +901,10 @@ formatPattern = \case
       Block.line $
         Block.string7 (show int)
 
+formatPatternConstructorArg :: ([Src.Comment], Src.Pattern) -> PatternBlock
+formatPatternConstructorArg (commentsBefore, pat) =
+  formatPattern (A.toValue pat)
+
 data StringStyle
   = StringStyleChar
   | StringStyleSingleQuoted

compiler/src/Parse/Declaration.hs

@@ -100,7 +100,7 @@ chompDefArgsAndBody maybeDocs start name tipe revArgs commentsBefore =
         word1 0x3D {-=-} E.DeclDefEquals
         commentsAfterEquals <- Space.chompAndCheckIndent E.DeclDefSpace E.DeclDefIndentBody
         ((body, commentsAfter), end) <- specialize E.DeclDefBody Expr.expression
-        let (commentsAfterBody, commentsAfterDef) = List.span (A.isIndentedAtLeast 2) commentsAfter
+        let (commentsAfterBody, commentsAfterDef) = List.span (A.isIndentedMoreThan 1) commentsAfter
         let comments = SC.ValueComments commentsBefore commentsAfterEquals commentsAfterBody
         let value = Src.Value name (reverse revArgs) body tipe comments
         let avalue = A.at start end value