Merge pull request #4637 from unisonweb/cp/doc-formatter-fix

Fix for doc formatter

parser-typechecker/src/Unison/Syntax/DeclParser.hs

@@ -150,7 +150,7 @@ dataDeclaration maybeUnresolvedModifier = do
   eq <- reserved "="
   let -- go gives the type of the constructor, given the types of
       -- the constructor arguments, e.g. Cons becomes forall a . a -> List a -> List a
-      go :: L.Token v -> [Type v Ann] -> (Ann, v, Type v Ann)
+      go :: L.Token v -> [Type v Ann] -> (Ann {- Ann spanning the constructor and its args -}, (Ann, v, Type v Ann))
       go ctorName ctorArgs =
         let arrow i o = Type.arrow (ann i <> ann o) i o
             app f arg = Type.app (ann f <> ann arg) f arg
@@ -160,14 +160,16 @@ dataDeclaration maybeUnresolvedModifier = do
             --    or just `Optional a` in the case of `None`
             ctorType = foldr arrow ctorReturnType ctorArgs
             ctorAnn = ann ctorName <> maybe (ann ctorName) ann (lastMay ctorArgs)
-         in ( ann ctorName,
-              Var.namespaced [L.payload name, L.payload ctorName],
-              Type.foralls ctorAnn typeArgVs ctorType
+         in ( ctorAnn,
+              ( ann ctorName,
+                Var.namespaced [L.payload name, L.payload ctorName],
+                Type.foralls ctorAnn typeArgVs ctorType
+              )
             )
       prefixVar = TermParser.verifyRelativeVarName prefixDefinitionName
-      dataConstructor :: P v m (Ann, v, Type v Ann)
+      dataConstructor :: P v m (Ann, (Ann, v, Type v Ann))
       dataConstructor = go <$> prefixVar <*> many TypeParser.valueTypeLeaf
-      record :: P v m ([(Ann, v, Type v Ann)], [(L.Token v, [(L.Token v, Type v Ann)])], Ann)
+      record :: P v m ([(Ann, (Ann, v, Type v Ann))], [(L.Token v, [(L.Token v, Type v Ann)])], Ann)
       record = do
         _ <- openBlockWith "{"
         let field :: P v m [(L.Token v, Type v Ann)]
@@ -185,9 +187,9 @@ dataDeclaration maybeUnresolvedModifier = do
   (constructors, accessors, closingAnn) <-
     msum [Left <$> record, Right <$> sepBy (reserved "|") dataConstructor] <&> \case
       Left (constructors, accessors, closingAnn) -> (constructors, accessors, closingAnn)
-      Right constructors ->
+      Right constructors -> do
         let closingAnn :: Ann
-            closingAnn = NonEmpty.last (ann eq NonEmpty.:| ((\(_, _, t) -> ann t) <$> constructors))
+            closingAnn = NonEmpty.last (ann eq NonEmpty.:| ((\(constrSpanAnn, _) -> constrSpanAnn) <$> constructors))
          in (constructors, [], closingAnn)
   _ <- closeBlock
   case maybeUnresolvedModifier of
@@ -197,7 +199,7 @@ dataDeclaration maybeUnresolvedModifier = do
       let declSpanAnn = ann typeToken <> closingAnn
       pure
         ( L.payload name,
-          DD.mkDataDecl' modifier declSpanAnn typeArgVs constructors,
+          DD.mkDataDecl' modifier declSpanAnn typeArgVs (snd <$> constructors),
           accessors
         )
     Just unresolvedModifier -> do
@@ -207,7 +209,7 @@ dataDeclaration maybeUnresolvedModifier = do
       let declSpanAnn = ann typeToken <> ann modifier <> closingAnn
       pure
         ( L.payload name,
-          DD.mkDataDecl' (L.payload modifier) declSpanAnn typeArgVs constructors,
+          DD.mkDataDecl' (L.payload modifier) declSpanAnn typeArgVs (snd <$> constructors),
           accessors
         )
 

parser-typechecker/src/Unison/Syntax/FileParser.hs

@@ -215,7 +215,7 @@ stanza = watchExpression <|> unexpectedAction <|> binding
       binding@((_, v), _) <- TermParser.binding
       pure $ case doc of
         Nothing -> Binding binding
-        Just doc -> Bindings [((ann doc, Var.joinDot v (Var.named "doc")), doc), binding]
+        Just (spanAnn, doc) -> Bindings [((spanAnn, Var.joinDot v (Var.named "doc")), doc), binding]
 
 watched :: (Monad m, Var v) => P v m (UF.WatchKind, Text, Ann)
 watched = P.try do

parser-typechecker/src/Unison/Syntax/TermParser.hs

@@ -25,6 +25,7 @@ import Data.Set qualified as Set
 import Data.Text qualified as Text
 import Data.Tuple.Extra qualified as TupleE
 import Text.Megaparsec qualified as P
+import U.Core.ABT qualified as ABT
 import Unison.ABT qualified as ABT
 import Unison.Builtin.Decls qualified as DD
 import Unison.ConstructorReference (ConstructorReference, GConstructorReference (..))
@@ -99,7 +100,7 @@ rewriteBlock = do
     rewriteTermlike kw mk = do
       kw <- quasikeyword kw
       lhs <- term
-      rhs <- block "==>"
+      (_spanAnn, rhs) <- block "==>"
       pure (mk (ann kw <> ann rhs) lhs rhs)
     rewriteTerm = rewriteTermlike "term" DD.rewriteTerm
     rewriteCase = rewriteTermlike "case" DD.rewriteCase
@@ -206,10 +207,10 @@ matchCase = do
             [ Nothing <$ P.try (quasikeyword "otherwise"),
               Just <$> infixAppOrBooleanOp
             ]
-        t <- block "->"
+        (_spanAnn, t) <- block "->"
         pure (guard, t)
   let unguardedBlock = label "case match" do
-        t <- block "->"
+        (_spanAnn, t) <- block "->"
         pure (Nothing, t)
   -- a pattern's RHS is either one or more guards, or a single unguarded block.
   guardsAndBlocks <- guardedBlocks <|> (pure @[] <$> unguardedBlock)
@@ -349,10 +350,10 @@ lam p = label "lambda" $ mkLam <$> P.try (some prefixDefinitionName <* reserved
     mkLam vs b = Term.lam' (ann (head vs) <> ann b) (map L.payload vs) b
 
 letBlock, handle, ifthen :: (Monad m, Var v) => TermP v m
-letBlock = label "let" $ block "let"
+letBlock = label "let" $ (snd <$> block "let")
 handle = label "handle" do
-  b <- block "handle"
-  handler <- block "with"
+  (_spanAnn, b) <- block "handle"
+  (_spanAnn, handler) <- block "with"
   pure $ Term.handle (ann b) handler b
 
 checkCasesArities :: (Ord v, Annotated a) => NonEmpty (Int, a) -> P v m (Int, NonEmpty a)
@@ -382,9 +383,9 @@ lamCase = do
 
 ifthen = label "if" do
   start <- peekAny
-  c <- block "if"
-  t <- block "then"
-  f <- block "else"
+  (_spanAnn, c) <- block "if"
+  (_spanAnn, t) <- block "then"
+  (_spanAnn, f) <- block "else"
   pure $ Term.iff (ann start <> ann f) c t f
 
 text :: (Var v) => TermP v m
@@ -451,10 +452,10 @@ termLeaf =
       keywordBlock,
       list term,
       delayQuote,
-      delayBlock,
+      (snd <$> delayBlock),
       bang,
       docBlock,
-      doc2Block
+      doc2Block <&> \(spanAnn, trm) -> trm {ABT.annotation = ABT.annotation trm <> spanAnn}
     ]
 
 -- Syntax for documentation v2 blocks, which are surrounded by {{ }}.
@@ -490,41 +491,53 @@ termLeaf =
 -- variables that will be looked up in the environment like anything else. This
 -- means that the documentation syntax can have its meaning changed by
 -- overriding what functions the names `syntax.doc*` correspond to.
-doc2Block :: forall m v. (Monad m, Var v) => TermP v m
-doc2Block =
+doc2Block :: forall m v. (Monad m, Var v) => P v m (Ann {- Annotation for the whole spanning block -}, Term v Ann)
+doc2Block = do
   P.lookAhead (openBlockWith "syntax.docUntitledSection") *> elem
   where
-    elem :: TermP v m
+    -- For terms which aren't blocks the spanning annotation is the same as the
+    -- term annotation.
+    selfAnnotated :: Term v Ann -> (Ann, Term v Ann)
+    selfAnnotated t = (ann t, t)
+    elem :: P v m (Ann {- Annotation for the whole spanning block -}, Term v Ann)
     elem =
-      text <|> do
-        t <- openBlock
+      (selfAnnotated <$> text) <|> do
+        startTok <- openBlock
         let -- here, `t` will be something like `Open "syntax.docWord"`
             -- so `f` will be a term var with the name "syntax.docWord".
-            f = f' t
+            f = f' startTok
             f' t = Term.var (ann t) (Var.nameds (L.payload t))
 
             -- follows are some common syntactic forms used for parsing child elements
 
             -- regular is parsed into `f child1 child2 child3` for however many children
             regular = do
-              cs <- P.many elem <* closeBlock
-              pure $ Term.apps' f cs
+              cs <- P.many (snd <$> elem)
+              endTok <- closeBlock
+              let trm = Term.apps' f cs
+              pure (ann startTok <> ann endTok, trm)
 
             -- variadic is parsed into: `f [child1, child2, ...]`
             variadic = variadic' f
             variadic' f = do
-              cs <- P.many elem <* closeBlock
-              pure $ Term.apps' f [Term.list (ann cs) cs]
+              cs <- P.many (snd <$> elem)
+              endTok <- closeBlock
+              let trm = Term.apps' f [Term.list (ann cs) cs]
+              pure (ann startTok <> ann endTok, trm)
 
             -- sectionLike is parsed into: `f tm [child1, child2, ...]`
             sectionLike = do
-              arg1 <- elem
-              cs <- P.many elem <* closeBlock
-              pure $ Term.apps' f [arg1, Term.list (ann cs) cs]
+              arg1 <- (snd <$> elem)
+              cs <- P.many (snd <$> elem)
+              endTok <- closeBlock
+              let trm = Term.apps' f [arg1, Term.list (ann cs) cs]
+              pure (ann startTok <> ann endTok, trm)
 
             evalLike wrap = do
-              tm <- term <* closeBlock
-              pure $ Term.apps' f [wrap tm]
+              tm <- term
+              endTok <- closeBlock
+              let trm = Term.apps' f [wrap tm]
+              pure (ann startTok <> ann endTok, trm)
 
             -- converts `tm` to `'tm`
             --
@@ -533,8 +546,7 @@ doc2Block =
             -- code which renders documents. (We want the doc display to get
             -- the unevaluated expression `1 + 1` and not `2`)
             addDelay tm = Term.delay (ann tm) tm
-
-        case L.payload t of
+        case L.payload startTok of
           "syntax.docJoin" -> variadic
           "syntax.docUntitledSection" -> variadic
           "syntax.docColumn" -> variadic
@@ -545,33 +557,45 @@ doc2Block =
           "syntax.docBulletedList" -> variadic
           "syntax.docSourceAnnotations" -> variadic
           "syntax.docSourceElement" -> do
-            link <- elem
-            anns <- P.optional $ reserved "@" *> elem
-            closeBlock $> Term.apps' f [link, fromMaybe (Term.list (ann link) mempty) anns]
+            link <- (snd <$> elem)
+            anns <- P.optional $ reserved "@" *> (snd <$> elem)
+            endTok <- closeBlock
+            let trm = Term.apps' f [link, fromMaybe (Term.list (ann link) mempty) anns]
+            pure (ann startTok <> ann endTok, trm)
           "syntax.docNumberedList" -> do
-            nitems@((n, _) : _) <- P.some nitem <* closeBlock
+            nitems@((n, _) : _) <- P.some nitem
+            endTok <- closeBlock
             let items = snd <$> nitems
-            pure $ Term.apps' f [n, Term.list (ann items) items]
+            let trm = Term.apps' f [n, Term.list (ann items) items]
+            pure (ann startTok <> ann endTok, trm)
             where
               nitem = do
                 n <- number
                 t <- openBlockWith "syntax.docColumn"
                 let f = f' ("syntax.docColumn" <$ t)
-                child <- variadic' f
+                (_spanAnn, child) <- variadic' f
                 pure (n, child)
           "syntax.docSection" -> sectionLike
           -- @source{ type Blah, foo, type Bar }
           "syntax.docEmbedTermLink" -> do
             tm <- addDelay <$> (hashQualifiedPrefixTerm <|> hashQualifiedInfixTerm)
-            closeBlock $> Term.apps' f [tm]
+            endTok <- closeBlock
+            let trm = Term.apps' f [tm]
+            pure (ann startTok <> ann endTok, trm)
           "syntax.docEmbedSignatureLink" -> do
             tm <- addDelay <$> (hashQualifiedPrefixTerm <|> hashQualifiedInfixTerm)
-            closeBlock $> Term.apps' f [tm]
+            endTok <- closeBlock
+            let trm = Term.apps' f [tm]
+            pure (ann startTok <> ann endTok, trm)
           "syntax.docEmbedTypeLink" -> do
             r <- typeLink'
-            closeBlock $> Term.apps' f [Term.typeLink (ann r) (L.payload r)]
-          "syntax.docExample" ->
-            (term <* closeBlock) <&> \case
+            endTok <- closeBlock
+            let trm = Term.apps' f [Term.typeLink (ann r) (L.payload r)]
+            pure (ann startTok <> ann endTok, trm)
+          "syntax.docExample" -> do
+            trm <- term
+            endTok <- closeBlock
+            pure . (ann startTok <> ann endTok,) $ case trm of
               tm@(Term.Apps' _ xs) ->
                 let fvs = List.Extra.nubOrd $ concatMap (toList . Term.freeVars) xs
                     n = Term.nat (ann tm) (fromIntegral (length fvs))
@@ -581,11 +605,11 @@ doc2Block =
           "syntax.docTransclude" -> evalLike id
           "syntax.docEvalInline" -> evalLike addDelay
           "syntax.docExampleBlock" -> do
-            tm <- block'' False True "syntax.docExampleBlock" (pure (void t)) closeBlock
-            pure $ Term.apps' f [Term.nat (ann tm) 0, addDelay tm]
+            (spanAnn, tm) <- block'' False True "syntax.docExampleBlock" (pure (void startTok)) closeBlock
+            pure $ (spanAnn, Term.apps' f [Term.nat (ann tm) 0, addDelay tm])
           "syntax.docEval" -> do
-            tm <- block' False "syntax.docEval" (pure (void t)) closeBlock
-            pure $ Term.apps' f [addDelay tm]
+            (spanAnn, tm) <- block' False "syntax.docEval" (pure (void startTok)) closeBlock
+            pure $ (spanAnn, Term.apps' f [addDelay tm])
           _ -> regular
 
 docBlock :: (Monad m, Var v) => TermP v m
@@ -958,10 +982,10 @@ delayQuote = P.label "quote" do
   e <- termLeaf
   pure $ DD.delayTerm (ann start <> ann e) e
 
-delayBlock :: (Monad m, Var v) => TermP v m
+delayBlock :: (Monad m, Var v) => P v m (Ann {- Ann spanning the whole block -}, Term v Ann)
 delayBlock = P.label "do" do
-  b <- block "do"
-  pure $ DD.delayTerm (ann b) b
+  (spanAnn, b) <- block "do"
+  pure $ (spanAnn, DD.delayTerm (ann b) b)
 
 bang :: (Monad m, Var v) => TermP v m
 bang = P.label "bang" do
@@ -1035,7 +1059,7 @@ destructuringBind = do
     let boundVars' = snd <$> boundVars
     P.lookAhead (openBlockWith "=")
     pure (p, boundVars')
-  scrute <- block "=" -- Dwight K. Scrute ("The People's Scrutinee")
+  (_spanAnn, scrute) <- block "=" -- Dwight K. Scrute ("The People's Scrutinee")
   let guard = Nothing
   let absChain vs t = foldr (\v t -> ABT.abs' (ann t) v t) t vs
       thecase t = Term.MatchCase p (fmap (absChain boundVars) guard) $ absChain boundVars t
@@ -1073,9 +1097,11 @@ binding = label "binding" do
     Nothing -> do
       -- we haven't seen a type annotation, so lookahead to '=' before commit
       (lhsLoc, name, args) <- P.try (lhs <* P.lookAhead (openBlockWith "="))
-      body <- block "="
+      (_bodySpanAnn, body) <- block "="
       verifyRelativeName' (fmap Name.unsafeFromVar name)
       let binding = mkBinding lhsLoc args body
+      -- We don't actually use the span annotation from the block (yet) because it
+      -- may contain a bunch of white-space and comments following a top-level-definition.
       let spanAnn = ann lhsLoc <> ann binding
       pure $ ((spanAnn, (L.payload name)), binding)
     Just (nameT, typ) -> do
@@ -1084,8 +1110,10 @@ binding = label "binding" do
       when (L.payload name /= L.payload nameT) $
         customFailure $
           SignatureNeedsAccompanyingBody nameT
-      body <- block "="
+      (_bodySpanAnn, body) <- block "="
       let binding = mkBinding lhsLoc args body
+      -- We don't actually use the span annotation from the block (yet) because it
+      -- may contain a bunch of white-space and comments following a top-level-definition.
       let spanAnn = ann nameT <> ann binding
       pure $ ((spanAnn, L.payload name), Term.ann (ann nameT <> ann binding) binding typ)
   where
@@ -1097,7 +1125,7 @@ binding = label "binding" do
 customFailure :: (P.MonadParsec e s m) => e -> m a
 customFailure = P.customFailure
 
-block :: forall m v. (Monad m, Var v) => String -> TermP v m
+block :: forall m v. (Monad m, Var v) => String -> P v m (Ann, Term v Ann)
 block s = block' False s (openBlockWith s) closeBlock
 
 -- example: use Foo.bar.Baz + ++ x
@@ -1167,25 +1195,32 @@ substImports ns imports =
           Names.hasTypeNamed Names.IncludeSuffixes (Name.unsafeFromVar full) ns
       ]
 
-block' :: (Monad m, Var v) => IsTop -> String -> P v m (L.Token ()) -> P v m (L.Token ()) -> TermP v m
+block' ::
+  (Monad m, Var v) =>
+  IsTop ->
+  String ->
+  P v m (L.Token ()) ->
+  P v m (L.Token ()) ->
+  P v m (Ann {- ann which spans the whole block -}, Term v Ann)
 block' isTop = block'' isTop False
 
 block'' ::
-  forall m v b.
-  (Monad m, Var v) =>
+  forall m v end.
+  (Monad m, Var v, Annotated end) =>
   IsTop ->
   Bool -> -- `True` means insert `()` at end of block if it ends with a statement
   String ->
   P v m (L.Token ()) ->
-  P v m b ->
-  TermP v m
+  P v m end ->
+  P v m (Ann {- ann which spans the whole block -}, Term v Ann)
 block'' isTop implicitUnitAtEnd s openBlock closeBlock = do
   open <- openBlock
   (names, imports) <- imports
   _ <- optional semi
   statements <- local (\e -> e {names = names}) $ sepBy semi statement
-  _ <- closeBlock
-  substImports names imports <$> go open statements
+  end <- closeBlock
+  body <- substImports names imports <$> go open statements
+  pure (ann open <> ann end, body)
   where
     statement = asum [Binding <$> binding, DestructuringBind <$> destructuringBind, Action <$> blockTerm]
     go :: L.Token () -> [BlockElement v] -> P v m (Term v Ann)

unison-cli/src/Unison/Codebase/Editor/HandleInput.hs

@@ -1086,7 +1086,7 @@ loop e = do
                     Cli.InvalidSourceNameError -> lift $ Cli.returnEarly $ Output.InvalidSourceName filePath
                     Cli.LoadError -> lift $ Cli.returnEarly $ Output.SourceLoadFailed filePath
                     Cli.LoadSuccess contents -> pure contents
-                let updatedSource = Format.applyFormatUpdates updates source
+                let updatedSource = Format.applyTextReplacements updates source
                 liftIO $ writeSource (Text.pack filePath) updatedSource
             DebugDumpNamespacesI -> do
               let seen h = State.gets (Set.member h)