Improve formatting of comments (#2022)

- Closes #2016 
This pr is a work in progress towards respecting empty lines with
respect to comments.

While the algorithm is supposed to work in general, I've focused on
testing it on comments between statements.

See
https://github.com/anoma/juvix/blob/13442eee34e230463fb4b91455a0f552531626d8/tests/positive/Format.juvix
to see some examples

src/Juvix/Compiler/Abstract/Translation/FromConcrete.hs

@@ -342,7 +342,7 @@ goInductive ty@InductiveDef {..} = do
             _inductiveType = fromMaybe (Abstract.ExpressionUniverse (smallUniverse loc)) _inductiveType',
             _inductiveConstructors = toList _inductiveConstructors',
             _inductiveExamples = _inductiveExamples',
-            _inductivePositive = ty ^. inductivePositive
+            _inductivePositive = isJust (ty ^. inductivePositive)
           }
   whenJust ((^. withLocParam) <$> _inductiveBuiltin) (registerBuiltinInductive indDef)
   inductiveInfo <- registerInductive indDef

src/Juvix/Compiler/Backend/Geb/Translation/FromSource/Lexer.hs

@@ -13,7 +13,7 @@ import Text.Megaparsec as P hiding (sepBy1, sepEndBy1, some)
 import Text.Megaparsec.Char.Lexer qualified as L
 
 space :: forall r. ParsecS r ()
-space = L.space space1 lineCmnt blockCmnt
+space = L.space whiteSpace1 lineCmnt blockCmnt
   where
     lineCmnt :: ParsecS r () = L.skipLineComment ";"
     blockCmnt :: ParsecS r () = L.skipBlockComment "#|" "|#"

src/Juvix/Compiler/Concrete/Data/Highlight.hs

@@ -69,8 +69,8 @@ goFaceParsedItem i = WithLoc (i ^. parsedLoc) (PropertyFace f)
       ParsedTagLiteralInt -> FaceNumber
       ParsedTagLiteralString -> FaceString
       ParsedTagComment -> FaceComment
-      ParsedTagDelimiter -> FaceDelimiter
       ParsedTagJudoc -> FaceJudoc
+      ParsedTagDelimiter -> FaceDelimiter
 
 goFaceName :: AName -> Maybe (WithLoc PropertyFace)
 goFaceName n = do

src/Juvix/Compiler/Concrete/Data/ParsedInfoTableBuilder.hs

@@ -4,7 +4,7 @@ module Juvix.Compiler.Concrete.Data.ParsedInfoTableBuilder
     registerDelimiter,
     registerKeyword,
     registerJudocText,
-    registerComment,
+    registerSpaceSpan,
     registerModule,
     moduleVisited,
     visitModule,
@@ -23,7 +23,7 @@ import Juvix.Prelude
 
 data InfoTableBuilder m a where
   RegisterItem :: ParsedItem -> InfoTableBuilder m ()
-  RegisterComment :: Comment -> InfoTableBuilder m ()
+  RegisterSpaceSpan :: SpaceSpan -> InfoTableBuilder m ()
   RegisterModule :: Module 'Parsed 'ModuleTop -> InfoTableBuilder m ()
   VisitModule :: TopModulePath -> InfoTableBuilder m ()
   ModuleVisited :: TopModulePath -> InfoTableBuilder m Bool
@@ -71,7 +71,7 @@ registerLiteral l =
 
 data BuilderState = BuilderState
   { _stateItems :: [ParsedItem],
-    _stateComments :: [Comment],
+    _stateComments :: [SpaceSpan],
     _stateVisited :: HashSet TopModulePath,
     _stateModules :: HashMap TopModulePath (Module 'Parsed 'ModuleTop)
   }
@@ -111,11 +111,12 @@ runParserInfoTableBuilder =
             modify' (over stateModules (HashMap.insert (m ^. modulePath) m))
           RegisterItem i ->
             modify' (over stateItems (i :))
-          RegisterComment c -> do
-            modify' (over stateComments (c :))
-            registerItem'
-              ParsedItem
-                { _parsedLoc = c ^. commentInterval,
-                  _parsedTag = ParsedTagComment
-                }
+          RegisterSpaceSpan g -> do
+            modify' (over stateComments (g :))
+            forM_ (g ^.. spaceSpan . each . _SpaceComment) $ \c ->
+              registerItem'
+                ParsedItem
+                  { _parsedLoc = getLoc c,
+                    _parsedTag = ParsedTagComment
+                  }
       )

src/Juvix/Compiler/Concrete/Data/ParsedItem.hs

@@ -13,8 +13,8 @@ data ParsedItemTag
   | ParsedTagLiteralInt
   | ParsedTagLiteralString
   | ParsedTagComment
-  | ParsedTagDelimiter
   | ParsedTagJudoc
+  | ParsedTagDelimiter
   deriving stock (Eq, Show, Generic)
 
 makeLenses ''ParsedItem

src/Juvix/Compiler/Concrete/Language.hs

@@ -264,7 +264,7 @@ data InductiveDef (s :: Stage) = InductiveDef
     _inductiveParameters :: [InductiveParameters s],
     _inductiveType :: Maybe (ExpressionType s),
     _inductiveConstructors :: NonEmpty (InductiveConstructorDef s),
-    _inductivePositive :: Bool
+    _inductivePositive :: Maybe KeywordRef
   }
 
 deriving stock instance (Show (ExpressionType s), Show (SymbolType s)) => Show (InductiveDef s)
@@ -1054,6 +1054,33 @@ instance HasLoc ScopedIden where
     ScopedFunction a -> getLoc a
     ScopedVar a -> getLoc a
 
+instance HasLoc (InductiveDef 'Scoped) where
+  getLoc i = (getLoc <$> i ^. inductivePositive) ?<> getLoc (i ^. inductiveKw)
+
+instance HasLoc (FunctionClause 'Scoped) where
+  getLoc c = getLoc (c ^. clauseOwnerFunction) <> getLoc (c ^. clauseBody)
+
+instance HasLoc ModuleRef where
+  getLoc (ModuleRef' (_ :&: r)) = getLoc r
+
+instance HasLoc (AxiomDef 'Scoped) where
+  getLoc m = getLoc (m ^. axiomKw) <> getLoc (m ^. axiomType)
+
+instance HasLoc (OpenModule 'Scoped) where
+  getLoc m = getLoc (m ^. openModuleKw) <> getLoc (m ^. openModuleName)
+
+instance HasLoc (Statement 'Scoped) where
+  getLoc :: Statement 'Scoped -> Interval
+  getLoc = \case
+    StatementOperator t -> getLoc t
+    StatementTypeSignature t -> getLoc t
+    StatementImport t -> getLoc t
+    StatementInductive t -> getLoc t
+    StatementModule t -> getLoc t
+    StatementOpenModule t -> getLoc t
+    StatementFunctionClause t -> getLoc t
+    StatementAxiom t -> getLoc t
+
 instance HasLoc Application where
   getLoc (Application l r) = getLoc l <> getLoc r
 

src/Juvix/Compiler/Concrete/Pretty/Base.hs

@@ -50,17 +50,18 @@ ppSymbol = case sing :: SStage s of
   SParsed -> ppCode
   SScoped -> ppCode
 
-groupStatements :: forall s. (SingI s) => [Statement s] -> [[Statement s]]
-groupStatements = reverse . map reverse . uncurry cons . foldl' aux ([], [])
+groupStatements :: forall s. (SingI s) => [Statement s] -> [NonEmpty (Statement s)]
+groupStatements = \case
+  [] -> []
+  s : ss -> reverse . map NonEmpty.reverse . uncurry cons . foldl' aux (pure s, []) $ ss
   where
     aux ::
-      ([Statement s], [[Statement s]]) ->
+      (NonEmpty (Statement s), [NonEmpty (Statement s)]) ->
       Statement s ->
-      ([Statement s], [[Statement s]])
-    aux ([], acc) s = ([s], acc)
-    aux (gr@(a : _), acc) b
-      | g a b = (b : gr, acc)
-      | otherwise = ([b], gr : acc)
+      (NonEmpty (Statement s), [NonEmpty (Statement s)])
+    aux (gr@(a :| _), acc) b
+      | g a b = (NonEmpty.cons b gr, acc)
+      | otherwise = (pure b, gr : acc)
     -- Decides if statements a and b should be next to each other without a
     -- blank line
     g :: Statement s -> Statement s -> Bool

src/Juvix/Compiler/Concrete/Print/Base.hs

@@ -49,7 +49,7 @@ ppModulePathType ::
   Sem r ()
 ppModulePathType x = case sing :: SStage s of
   SParsed -> case sing :: SModuleIsTop t of
-    SModuleLocal -> noLoc (pretty x)
+    SModuleLocal -> ppCode x
     SModuleTop -> ppCode x
   SScoped -> case sing :: SModuleIsTop t of
     SModuleLocal -> P.ppCode x >>= morpheme (getLoc x) . P.annSDef x
@@ -61,20 +61,24 @@ instance SingI t => PrettyPrint (Module 'Scoped t) where
     let moduleBody' = localIndent (ppCode _moduleBody)
         modulePath' = ppModulePathType _modulePath
         moduleDoc' :: Sem r () = maybe (return ()) ppCode _moduleDoc
+        body'
+          | null _moduleBody = ensureEmptyLine
+          | otherwise =
+              topSpace
+                <> moduleBody'
+                <> line
     moduleDoc'
       <> ppCode _moduleKw
       <+> modulePath'
         <> ppCode kwSemicolon
         <> line
-        <> topSpace
-        <> moduleBody'
-        <> line
+        <> body'
         <> ending
     where
       topSpace :: Sem r ()
       topSpace = case sing :: SModuleIsTop t of
         SModuleLocal -> mempty
-        SModuleTop -> line
+        SModuleTop -> ensureEmptyLine
 
       localIndent :: Sem r () -> Sem r ()
       localIndent = case sing :: SModuleIsTop t of
@@ -88,10 +92,10 @@ instance SingI t => PrettyPrint (Module 'Scoped t) where
 
 instance PrettyPrint [Statement 'Scoped] where
   ppCode :: forall r. Members '[ExactPrint, Reader Options] r => [Statement 'Scoped] -> Sem r ()
-  ppCode ss = vsep2 (map ppGroup (P.groupStatements ss))
+  ppCode ss = paragraphs (ppGroup <$> P.groupStatements ss)
     where
-      ppGroup :: [Statement 'Scoped] -> Sem r ()
-      ppGroup = vsep . endSemicolon . map ppCode
+      ppGroup :: NonEmpty (Statement 'Scoped) -> Sem r ()
+      ppGroup = vsep . endSemicolon . fmap ppCode
 
 instance PrettyPrint TopModulePath where
   ppCode t@TopModulePath {..} =
@@ -255,7 +259,7 @@ instance PrettyPrint UsingHiding where
             (noLoc P.kwBraceR)
             (noLoc P.kwSemicolon)
             (ppUnkindedSymbol <$> syms)
-    noLoc (pretty word) <+> bracedList
+    ppCode word <+> bracedList
     where
       (word, syms) = case uh of
         Using s -> (kwUsing, s)
@@ -333,7 +337,7 @@ ppInductiveSignature InductiveDef {..} = do
         Nothing -> Nothing
         Just e -> Just (noLoc P.kwColon <+> ppCode e)
       positive'
-        | _inductivePositive = (<> line) <$> Just (noLoc P.kwPositive)
+        | Just k <- _inductivePositive = (<> line) <$> Just (ppCode k)
         | otherwise = Nothing
   builtin'
     ?<> positive'

src/Juvix/Compiler/Concrete/Translation/FromSource.hs

@@ -517,7 +517,7 @@ lambda = do
 
 inductiveDef :: Members '[InfoTableBuilder, JudocStash, NameIdGen] r => Maybe (WithLoc BuiltinInductive) -> ParsecS r (InductiveDef 'Parsed)
 inductiveDef _inductiveBuiltin = do
-  _inductivePositive <- isJust <$> optional (kw kwPositive)
+  _inductivePositive <- optional (kw kwPositive)
   _inductiveKw <- kw kwInductive
   _inductiveDoc <- getJudoc
   _inductiveName <- symbol P.<?> "<type name>"

src/Juvix/Compiler/Concrete/Translation/FromSource/Lexer.hs

@@ -9,7 +9,7 @@ where
 import Data.Text qualified as Text
 import GHC.Unicode
 import Juvix.Compiler.Concrete.Data.ParsedInfoTableBuilder
-import Juvix.Compiler.Concrete.Extra hiding (Pos, space, string')
+import Juvix.Compiler.Concrete.Extra hiding (Pos, hspace, space, string')
 import Juvix.Compiler.Concrete.Extra qualified as P
 import Juvix.Compiler.Concrete.Keywords
 import Juvix.Data.Keyword
@@ -30,7 +30,7 @@ judocText_ :: Members '[InfoTableBuilder] r => ParsecS r a -> ParsecS r ()
 judocText_ = void . judocText
 
 space :: forall r. Members '[InfoTableBuilder] r => ParsecS r ()
-space = space' True >>= mapM_ (P.lift . registerComment)
+space = space' True >>= mapM_ (P.lift . registerSpaceSpan)
 
 lexeme :: (Members '[InfoTableBuilder] r) => ParsecS r a -> ParsecS r a
 lexeme = L.lexeme space
@@ -72,10 +72,10 @@ string :: Members '[InfoTableBuilder] r => ParsecS r (Text, Interval)
 string = lexemeInterval string'
 
 judocExampleStart :: ParsecS r ()
-judocExampleStart = P.chunk Str.judocExample >> hspace
+judocExampleStart = P.chunk Str.judocExample >> hspace_
 
-judocStart :: ParsecS r ()
-judocStart = P.chunk Str.judocStart >> hspace
+judocStart :: Members '[InfoTableBuilder] r => ParsecS r ()
+judocStart = judocText_ (P.chunk Str.judocStart) >> hspace_
 
 judocEmptyLine :: (Members '[InfoTableBuilder] r) => ParsecS r ()
 judocEmptyLine = lexeme (void (P.try (judocStart >> P.newline)))

src/Juvix/Data/Comment.hs

@@ -13,7 +13,7 @@ newtype Comments = Comments
 
 data FileComments = FileComments
   { -- | sorted by position
-    _fileCommentsSorted :: [Comment],
+    _fileCommentsSorted :: [SpaceSpan],
     _fileCommentsFile :: Path Abs File
   }
   deriving stock (Eq, Show, Generic, Data)
@@ -30,9 +30,73 @@ data Comment = Comment
   }
   deriving stock (Show, Eq, Ord, Generic, Data)
 
+data SpaceSection
+  = SpaceComment Comment
+  | SpaceLines EmptyLines
+  deriving stock (Show, Eq, Ord, Generic, Data)
+
+-- | One or more empty lines
+data EmptyLines = EmptyLines
+  { _emptyLinesLoc :: Interval,
+    -- | The number of empty lines. Always positive
+    _emptyLinesNum :: Int
+  }
+  deriving stock (Show, Eq, Ord, Generic, Data)
+
+newtype SpaceSpan = SpaceSpan
+  { _spaceSpan :: NonEmpty SpaceSection
+  }
+  deriving stock (Show, Eq, Ord, Generic, Data)
+
 makeLenses ''Comment
+makeLenses ''SpaceSpan
 makeLenses ''FileComments
 makeLenses ''Comments
+makeLenses ''EmptyLines
+
+instance Semigroup EmptyLines where
+  a <> b =
+    EmptyLines
+      { _emptyLinesLoc = a ^. emptyLinesLoc <> b ^. emptyLinesLoc,
+        _emptyLinesNum = a ^. emptyLinesNum + b ^. emptyLinesNum
+      }
+
+instance Semigroup SpaceSpan where
+  SpaceSpan a <> SpaceSpan b@(headb :| tailb)
+    | (inia, SpaceLines emptya) <- nonEmptyUnsnoc a,
+      SpaceLines emptyb <- headb =
+        SpaceSpan (nonEmpty' $ maybe [] toList inia <> (pure (SpaceLines (emptya <> emptyb))) <> tailb)
+    | otherwise = SpaceSpan (a <> b)
+
+_SpaceComment :: Traversal' SpaceSection Comment
+_SpaceComment f s = case s of
+  SpaceComment l -> SpaceComment <$> f l
+  SpaceLines {} -> pure s
+
+_SpaceLines :: Traversal' SpaceSection EmptyLines
+_SpaceLines f s = case s of
+  SpaceComment {} -> pure s
+  SpaceLines l -> SpaceLines <$> f l
+
+hasComments :: SpaceSpan -> Bool
+hasComments = any (has _SpaceComment) . (^. spaceSpan)
+
+hasEmptyLines :: SpaceSpan -> Bool
+hasEmptyLines = any (has _SpaceLines) . (^. spaceSpan)
+
+instance HasLoc SpaceSpan where
+  getLoc = getLocSpan . (^. spaceSpan)
+
+instance HasLoc EmptyLines where
+  getLoc = (^. emptyLinesLoc)
+
+instance HasLoc SpaceSection where
+  getLoc = \case
+    SpaceComment g -> getLoc g
+    SpaceLines w -> getLoc w
+
+instance HasLoc Comment where
+  getLoc = (^. commentInterval)
 
 instance Pretty Comment where
   pretty :: Comment -> Doc ann
@@ -44,20 +108,22 @@ instance Pretty Comment where
         CommentBlock -> enclose "{-" "-}"
 
 allComments :: Comments -> [Comment]
-allComments c = concat [f ^. fileCommentsSorted | f <- toList (c ^. commentsByFile)]
+allComments c =
+  [ m | f <- toList (c ^. commentsByFile), s <- f ^. fileCommentsSorted, SpaceComment m <- toList (s ^. spaceSpan)
+  ]
 
-mkComments :: [Comment] -> Comments
+mkComments :: [SpaceSpan] -> Comments
 mkComments cs = Comments {..}
   where
-    commentFile :: Comment -> Path Abs File
-    commentFile = (^. commentInterval . intervalFile)
+    spSpanFile :: SpaceSpan -> Path Abs File
+    spSpanFile = (^. intervalFile) . getLoc
     _commentsByFile :: HashMap (Path Abs File) FileComments
     _commentsByFile =
       HashMap.fromList
         [ (_fileCommentsFile, FileComments {..})
-          | filecomments :: NonEmpty Comment <- groupSortOn commentFile cs,
-            let _fileCommentsFile = commentFile (head filecomments),
-            let _fileCommentsSorted = sortOn (^. commentInterval) (toList filecomments)
+          | filecomments :: NonEmpty SpaceSpan <- groupSortOn spSpanFile cs,
+            let _fileCommentsFile = spSpanFile (head filecomments),
+            let _fileCommentsSorted = sortOn getLoc (toList filecomments)
         ]
 
 emptyComments :: Comments
@@ -73,11 +139,14 @@ emptyFileComments _fileCommentsFile =
 fileComments :: Path Abs File -> Comments -> FileComments
 fileComments f cs = HashMap.lookupDefault (emptyFileComments f) f (cs ^. commentsByFile)
 
+flattenComments :: [SpaceSpan] -> [Comment]
+flattenComments m = [c | s <- m, SpaceComment c <- toList (s ^. spaceSpan)]
+
 instance Pretty FileComments where
   pretty fc =
     pretty (fc ^. fileCommentsFile)
       <> line
-      <> vsep [pretty c | c <- toList (fc ^. fileCommentsSorted)]
+      <> vsep [pretty c | c <- flattenComments (fc ^. fileCommentsSorted)]
 
 instance Pretty Comments where
   pretty c