format: retain comments in expression bodies (value decls, let, lambda)

builder/src/Build.hs

@@ -282,7 +282,7 @@ crawlFile env@(Env _ root projectType _ _ buildID _ _) mvar docsNeed expectedNam
               else return $ SBadSyntax path time source (Syntax.ModuleNameMismatch expectedName name)
 
 isMain :: A.Located Src.Value -> Bool
-isMain (A.At _ (Src.Value (A.At _ name) _ _ _)) =
+isMain (A.At _ (Src.Value (A.At _ name) _ _ _ _)) =
   name == Name._main
 
 -- CHECK MODULE

compiler/src/AST/Source.hs

@@ -2,7 +2,8 @@
 {-# OPTIONS_GHC -Wall #-}
 
 module AST.Source
-  ( Comment (..),
+  ( Comment,
+    Comment_ (..),
     GREN_COMMENT,
     Expr,
     Expr_ (..),
@@ -34,7 +35,7 @@ module AST.Source
   )
 where
 
-import AST.SourceComments (Comment, GREN_COMMENT)
+import AST.SourceComments (Comment, Comment_, GREN_COMMENT)
 import AST.SourceComments qualified as SC
 import AST.Utils.Binop qualified as Binop
 import Data.List.NonEmpty (NonEmpty)
@@ -60,10 +61,10 @@ data Expr_
   | Op Name
   | Negate Expr
   | Binops [(Expr, [Comment], A.Located Name)] Expr
-  | Lambda [Pattern] Expr
+  | Lambda [([Comment], Pattern)] Expr SC.LambdaComments
   | Call Expr [([Comment], Expr)]
   | If [(Expr, Expr)] Expr
-  | Let [A.Located Def] Expr
+  | Let [([Comment], A.Located Def)] Expr SC.LetComments
   | Case Expr [([Comment], Pattern, Expr)]
   | Accessor Name
   | Access Expr (A.Located Name)
@@ -78,7 +79,7 @@ data VarType = LowVar | CapVar
 -- DEFINITIONS
 
 data Def
-  = Define (A.Located Name) [Pattern] Expr (Maybe Type)
+  = Define (A.Located Name) [([Comment], Pattern)] Expr (Maybe Type) SC.ValueComments
   | Destruct Pattern Expr
   deriving (Show)
 
@@ -157,7 +158,7 @@ data Import = Import
   }
   deriving (Show)
 
-data Value = Value (A.Located Name) [Pattern] Expr (Maybe Type)
+data Value = Value (A.Located Name) [([Comment], Pattern)] Expr (Maybe Type) SC.ValueComments
   deriving (Show)
 
 data Union = Union (A.Located Name) [A.Located Name] [(A.Located Name, [([Comment], Type)])]

compiler/src/AST/SourceComments.hs

@@ -1,10 +1,13 @@
 module AST.SourceComments where
 
 import Data.Utf8 qualified as Utf8
+import Reporting.Annotation qualified as A
 
 data GREN_COMMENT
 
-data Comment
+type Comment = A.Located Comment_
+
+data Comment_
   = BlockComment (Utf8.Utf8 GREN_COMMENT)
   | LineComment (Utf8.Utf8 GREN_COMMENT)
   deriving (Eq, Show)
@@ -67,10 +70,33 @@ data ImportAliasComments = ImportAliasComments
   { _afterAs :: [Comment],
     _afterAliasName :: [Comment]
   }
-  deriving (Eq, Show)
+  deriving (Show)
 
 data ImportExposingComments = ImportExposingComments
   { _afterExposing :: [Comment],
     _afterExposingListing :: [Comment]
   }
   deriving (Show)
+
+-- Declarations
+
+data ValueComments = ValueComments
+  { _beforeValueEquals :: [Comment],
+    _beforeValueBody :: [Comment],
+    _afterValueBody :: [Comment]
+  }
+  deriving (Show)
+
+-- Expressions
+
+data LambdaComments = LambdaComments
+  { _beforeArrow :: [Comment],
+    _afterArrow :: [Comment]
+  }
+  deriving (Show)
+
+data LetComments = LetComments
+  { _afterLetDecls :: [Comment],
+    _afterIn :: [Comment]
+  }
+  deriving (Show)

compiler/src/Canonicalize/Effects.hs

@@ -113,7 +113,7 @@ verifyEffectType (A.At region name) unions =
     else Result.throw (Error.EffectNotFound region name)
 
 toNameRegion :: A.Located Src.Value -> (Name.Name, A.Region)
-toNameRegion (A.At _ (Src.Value (A.At region name) _ _ _)) =
+toNameRegion (A.At _ (Src.Value (A.At region name) _ _ _ _)) =
   (name, region)
 
 verifyManager :: A.Region -> Map.Map Name.Name A.Region -> Name.Name -> Result i w A.Region

compiler/src/Canonicalize/Environment/Local.hs

@@ -47,7 +47,7 @@ addVars module_ (Env.Env home vs ts cs bs qvs qts qcs) =
 
 collectVars :: Src.Module -> Result i w (Map.Map Name.Name Env.Var)
 collectVars (Src.Module _ _ _ _ values _ _ _ _ _ effects) =
-  let addDecl dict (A.At _ (Src.Value (A.At region name) _ _ _)) =
+  let addDecl dict (A.At _ (Src.Value (A.At region name) _ _ _ _)) =
         Dups.insert name region (Env.TopLevel region) dict
    in Dups.detect Error.DuplicateDecl $
         List.foldl' addDecl (toEffectDups effects) (fmap snd values)

compiler/src/Canonicalize/Expression.hs

@@ -76,12 +76,12 @@ canonicalize env (A.At region expression) =
         Can.Negate <$> canonicalize env expr
       Src.Binops ops final ->
         A.toValue <$> canonicalizeBinops region env ops final
-      Src.Lambda srcArgs body ->
+      Src.Lambda srcArgs body _ ->
         delayedUsage $
           do
             (args, bindings) <-
               Pattern.verify Error.DPLambdaArgs $
-                traverse (Pattern.canonicalize env) srcArgs
+                traverse (Pattern.canonicalize env) (fmap snd srcArgs)
 
             newEnv <-
               Env.addLocals bindings env
@@ -98,8 +98,8 @@ canonicalize env (A.At region expression) =
         Can.If
           <$> traverse (canonicalizeIfBranch env) branches
           <*> canonicalize env finally
-      Src.Let defs expr ->
-        A.toValue <$> canonicalizeLet region env defs expr
+      Src.Let defs expr _ ->
+        A.toValue <$> canonicalizeLet region env (fmap snd defs) expr
       Src.Case expr branches ->
         Can.Case
           <$> canonicalize env expr
@@ -231,7 +231,7 @@ canonicalizeLet letRegion env defs body =
 addBindings :: Dups.Dict A.Region -> A.Located Src.Def -> Dups.Dict A.Region
 addBindings bindings (A.At _ def) =
   case def of
-    Src.Define (A.At region name) _ _ _ ->
+    Src.Define (A.At region name) _ _ _ _ ->
       Dups.insert name region region bindings
     Src.Destruct pattern _ ->
       addBindingsHelp bindings pattern
@@ -274,13 +274,13 @@ data Binding
 addDefNodes :: Env.Env -> [Node] -> A.Located Src.Def -> Result FreeLocals [W.Warning] [Node]
 addDefNodes env nodes (A.At _ def) =
   case def of
-    Src.Define aname@(A.At _ name) srcArgs body maybeType ->
+    Src.Define aname@(A.At _ name) srcArgs body maybeType _ ->
       case maybeType of
         Nothing ->
           do
             (args, argBindings) <-
               Pattern.verify (Error.DPFuncArgs name) $
-                traverse (Pattern.canonicalize env) srcArgs
+                traverse (Pattern.canonicalize env . snd) srcArgs
 
             newEnv <-
               Env.addLocals argBindings env
@@ -296,7 +296,7 @@ addDefNodes env nodes (A.At _ def) =
             (Can.Forall freeVars ctipe) <- Type.toAnnotation env tipe
             ((args, resultType), argBindings) <-
               Pattern.verify (Error.DPFuncArgs name) $
-                gatherTypedArgs env name srcArgs ctipe Index.first []
+                gatherTypedArgs env name (fmap snd srcArgs) ctipe Index.first []
 
             newEnv <-
               Env.addLocals argBindings env

compiler/src/Canonicalize/Module.hs

@@ -121,13 +121,13 @@ type NodeTwo =
   (Can.Def, Name.Name, [Name.Name])
 
 toNodeOne :: Env.Env -> A.Located Src.Value -> Result i [W.Warning] NodeOne
-toNodeOne env (A.At _ (Src.Value aname@(A.At _ name) srcArgs body maybeType)) =
+toNodeOne env (A.At _ (Src.Value aname@(A.At _ name) srcArgs body maybeType _)) =
   case maybeType of
     Nothing ->
       do
         (args, argBindings) <-
           Pattern.verify (Error.DPFuncArgs name) $
-            traverse (Pattern.canonicalize env) srcArgs
+            traverse (Pattern.canonicalize env . snd) srcArgs
 
         newEnv <-
           Env.addLocals argBindings env
@@ -147,7 +147,7 @@ toNodeOne env (A.At _ (Src.Value aname@(A.At _ name) srcArgs body maybeType)) =
 
         ((args, resultType), argBindings) <-
           Pattern.verify (Error.DPFuncArgs name) $
-            Expr.gatherTypedArgs env name srcArgs tipe Index.first []
+            Expr.gatherTypedArgs env name (fmap snd srcArgs) tipe Index.first []
 
         newEnv <-
           Env.addLocals argBindings env
@@ -197,7 +197,7 @@ canonicalizeExports values unions aliases binops effects (A.At region exposing)
         Can.Export <$> Dups.detect Error.ExportDuplicate (Dups.unions infos)
 
 valueToName :: A.Located Src.Value -> (Name.Name, ())
-valueToName (A.At _ (Src.Value (A.At _ name) _ _ _)) =
+valueToName (A.At _ (Src.Value (A.At _ name) _ _ _ _)) =
   (name, ())
 
 checkExposed ::

compiler/src/Gren/Format.hs

@@ -143,17 +143,29 @@ withCommentsAround before after block =
     (Just beforeBlock, Just afterBlock) ->
       spaceOrStack [beforeBlock, spaceOrIndent [block, afterBlock]]
 
+withCommentsStackBefore :: [Src.Comment] -> Block -> Block
+withCommentsStackBefore before = withCommentsStackAround before []
+
+withCommentsStackAround :: [Src.Comment] -> [Src.Comment] -> Block -> Block
+withCommentsStackAround [] [] block = block
+withCommentsStackAround before after block =
+  case (formatCommentBlock before, formatCommentBlock after) of
+    (Nothing, Nothing) -> block
+    (Just beforeBlock, Nothing) -> Block.stack [beforeBlock, block]
+    (Nothing, Just afterBlock) -> Block.stack [block, afterBlock]
+    (Just beforeBlock, Just afterBlock) -> Block.stack [beforeBlock, block, afterBlock]
+
 --
 -- AST -> Block
 --
 
 formatComment :: Src.Comment -> Block
 formatComment = \case
-  Src.BlockComment text ->
+  A.At _ (Src.BlockComment text) ->
     let open = if Utf8.startsWithChar (== ' ') text then "{-" else "{- "
         close = if Utf8.endsWithWord8 0x20 {- space -} text then "-}" else " -}"
      in Block.line $ Block.string7 open <> utf8 text <> Block.string7 close
-  Src.LineComment text ->
+  A.At _ (Src.LineComment text) ->
     let open = if Utf8.startsWithChar (== ' ') text then "--" else "-- "
      in Block.mustBreak $ Block.string7 open <> utf8 text
 
@@ -223,7 +235,7 @@ formatModule (Src.Module moduleName exports docs imports values unions aliases b
         Src.NoDocs _ -> Map.empty
         Src.YesDocs _ defs -> Map.fromList defs
 
-    valueName (Src.Value name _ _ _) = A.toValue name
+    valueName (Src.Value name _ _ _ _) = A.toValue name
     unionName (Src.Union name _ _) = A.toValue name
     aliasName (Src.Alias name _ _) = A.toValue name
     portName (Src.Port name _) = A.toValue name
@@ -368,23 +380,33 @@ formatAssociativity = \case
   Binop.Right -> Block.string7 "right"
 
 formatValue :: Src.Value -> Block
-formatValue (Src.Value name args body type_) =
-  formatBasicDef (A.toValue name) (fmap A.toValue args) (A.toValue body) (fmap A.toValue type_)
+formatValue (Src.Value name args body type_ comments) =
+  formatBasicDef (A.toValue name) args (A.toValue body) (fmap A.toValue type_) comments
 
-formatBasicDef :: Name -> [Src.Pattern_] -> Src.Expr_ -> Maybe Src.Type_ -> Block
-formatBasicDef name args body type_ =
+formatBasicDef :: Name -> [([Src.Comment], Src.Pattern)] -> Src.Expr_ -> Maybe Src.Type_ -> SC.ValueComments -> Block
+formatBasicDef name args body type_ (SC.ValueComments commentsBeforeEquals commentsBeforeBody commentsAfterBody) =
   Block.stack $
     NonEmpty.fromList $
       catMaybes
         [ fmap (formatTypeAnnotation Nothing name) type_,
           Just $
             spaceOrIndent $
               Block.line (utf8 name)
-                :| fmap (patternParensProtectSpaces . formatPattern) args
-                ++ [ Block.line $ Block.char7 '='
+                :| fmap formatPat args
+                ++ [ withCommentsBefore commentsBeforeEquals $
+                       Block.line (Block.char7 '=')
                    ],
-          Just $ Block.indent $ exprParensNone $ formatExpr body
+          Just $
+            Block.indent $
+              withCommentsStackAround commentsBeforeBody commentsAfterBody $
+                exprParensNone $
+                  formatExpr body
         ]
+  where
+    formatPat (comments, pat) =
+      withCommentsBefore comments $
+        patternParensProtectSpaces $
+          formatPattern (A.toValue pat)
 
 formatTypeAnnotation :: Maybe String -> Name -> Src.Type_ -> Block
 formatTypeAnnotation prefix name t =
@@ -524,19 +546,29 @@ formatExpr = \case
           4
           (utf8 (A.toValue op) <> Block.space)
           (exprParensProtectInfixOps $ formatExpr $ A.toValue expr)
-  Src.Lambda [] body ->
+  Src.Lambda [] body _ ->
     formatExpr $ A.toValue body
-  Src.Lambda (arg1 : args) body ->
+  Src.Lambda (arg1 : args) body (SC.LambdaComments commentsBeforeArrow commentsAfterArrow) ->
     ExpressionHasAmbiguousEnd $
       spaceOrIndent
         [ Block.prefix 1 (Block.char7 '\\') $
             spaceOrStack $
               join
-                [ fmap (patternParensProtectSpaces . formatPattern . A.toValue) (arg1 :| args),
-                  pure $ Block.line $ Block.string7 "->"
+                [ fmap formatArg (arg1 :| args),
+                  pure $
+                    withCommentsBefore commentsBeforeArrow $
+                      Block.line $
+                        Block.string7 "->"
                 ],
-          exprParensNone $ formatExpr $ A.toValue body
+          withCommentsBefore commentsAfterArrow $
+            exprParensNone $
+              formatExpr $
+                A.toValue body
         ]
+    where
+      formatArg (commentsBefore, arg) =
+        withCommentsBefore commentsBefore $
+          patternParensProtectSpaces (formatPattern $ A.toValue arg)
   Src.Call fn [] ->
     formatExpr $ A.toValue fn
   Src.Call fn args ->
@@ -575,15 +607,22 @@ formatExpr = \case
               ],
             Block.indent $ exprParensNone $ formatExpr $ A.toValue body
           ]
-  Src.Let [] body ->
+  Src.Let [] body _ ->
     formatExpr $ A.toValue body
-  Src.Let (def1 : defs) body ->
+  Src.Let (def1 : defs) body (SC.LetComments commentsBeforeIn commentsAfterIn) ->
     ExpressionHasAmbiguousEnd $
       Block.stack
         [ Block.line (Block.string7 "let"),
-          Block.indent $ Block.stack $ NonEmpty.intersperse Block.blankLine $ fmap (formatDef . A.toValue) (def1 :| defs),
-          Block.line (Block.string7 "in"),
-          exprParensNone $ formatExpr (A.toValue body)
+          Block.indent $ Block.stack $ NonEmpty.intersperse Block.blankLine $ fmap formatDef (def1 :| defs),
+          case formatCommentBlock commentsBeforeIn of
+            Nothing -> Block.line (Block.string7 "in")
+            Just comments ->
+              Block.stack
+                [ Block.blankLine,
+                  Block.indent comments,
+                  Block.line (Block.string7 "in")
+                ],
+          withCommentsStackBefore commentsAfterIn $ exprParensNone $ formatExpr (A.toValue body)
         ]
   Src.Case subject branches ->
     ExpressionHasAmbiguousEnd $
@@ -657,18 +696,20 @@ opForcesMultiline op =
   op == Utf8.fromChars "|>"
     || op == Utf8.fromChars "<|"
 
-formatDef :: Src.Def -> Block
-formatDef = \case
-  Src.Define name args body ann ->
-    formatBasicDef (A.toValue name) (fmap A.toValue args) (A.toValue body) (fmap A.toValue ann)
-  Src.Destruct pat body ->
-    Block.stack
-      [ spaceOrIndent
-          [ patternParensProtectSpaces $ formatPattern $ A.toValue pat,
-            Block.line $ Block.char7 '='
-          ],
-        Block.indent $ exprParensNone $ formatExpr $ A.toValue body
-      ]
+formatDef :: ([Src.Comment], A.Located Src.Def) -> Block
+formatDef (commentsBefore, def) =
+  withCommentsStackBefore commentsBefore $
+    case A.toValue def of
+      Src.Define name args body ann comments ->
+        formatBasicDef (A.toValue name) args (A.toValue body) (fmap A.toValue ann) comments
+      Src.Destruct pat body ->
+        Block.stack
+          [ spaceOrIndent
+              [ patternParensProtectSpaces $ formatPattern $ A.toValue pat,
+                Block.line $ Block.char7 '='
+              ],
+            Block.indent $ exprParensNone $ formatExpr $ A.toValue body
+          ]
 
 data TypeBlock
   = NoTypeParens Block