EPA: order of semicolons and comments for top-level decls is wrong

A comment followed by a semicolon at the top level resulted in the
preceding comments being attached to the following declaration.

Capture the comments as belonging to the declaration preceding the
semicolon instead.

Closes #20258

9 files changed, 233 insertions, 56 deletions

diff --git a/compiler/GHC/Parser.y b/compiler/GHC/Parser.y
index 1c0c65bb96..732a03f7d5 100644
--- a/compiler/GHC/Parser.y
+++ b/compiler/GHC/Parser.y
@@ -1066,9 +1066,9 @@ qcname  :: { LocatedN RdrName }  -- Variable or type constructor
 -- top handles the fact that these may be optional.
 
 -- One or more semicolons
-semis1  :: { [TrailingAnn] }
-semis1  : semis1 ';'  { if isZeroWidthSpan (gl $2) then $1 else (AddSemiAnn (glAA $2) : $1) }
-        | ';'         { msemi $1 }
+semis1  :: { Located [TrailingAnn] }
+semis1  : semis1 ';'  { sLL $1 $> $ if isZeroWidthSpan (gl $2) then (unLoc $1) else (AddSemiAnn (glAA $2) : (unLoc $1)) }
+        | ';'         { sL1 $1 $ msemi $1 }
 
 -- Zero or more semicolons
 semis   :: { [TrailingAnn] }
@@ -1085,7 +1085,7 @@ importdecls
 importdecls_semi :: { [LImportDecl GhcPs] }
 importdecls_semi
         : importdecls_semi importdecl semis1
-                                {% do { i <- amsA $2 $3
+                                {% do { i <- amsAl $2 (comb2 (reLoc $2) $3) (unLoc $3)
                                       ; return (i : $1)} }
         | {- empty -}           { [] }
 
@@ -1187,7 +1187,7 @@ topdecls :: { OrdList (LHsDecl GhcPs) }
 
 -- May have trailing semicolons, can be empty
 topdecls_semi :: { OrdList (LHsDecl GhcPs) }
-        : topdecls_semi topdecl semis1 {% do { t <- amsA $2 $3
+        : topdecls_semi topdecl semis1 {% do { t <- amsAl $2 (comb2 (reLoc $2) $3) (unLoc $3)
                                              ; return ($1 `snocOL` t) }}
         | {- empty -}                  { nilOL }
 
@@ -1200,7 +1200,7 @@ topdecls_cs :: { OrdList (LHsDecl GhcPs) }
 
 -- May have trailing semicolons, can be empty
 topdecls_cs_semi :: { OrdList (LHsDecl GhcPs) }
-        : topdecls_cs_semi topdecl_cs semis1 {% do { t <- amsA $2 $3
+        : topdecls_cs_semi topdecl_cs semis1 {% do { t <- amsAl $2 (comb2 (reLoc $2) $3) (unLoc $3)
                                                    ; return ($1 `snocOL` t) }}
         | {- empty -}                  { nilOL }
 
@@ -4254,6 +4254,11 @@ amsA (L l a) bs = do
   cs <- getCommentsFor (locA l)
   return (L (addAnnsA l bs cs) a)
 
+amsAl :: MonadP m => LocatedA a -> SrcSpan -> [TrailingAnn] -> m (LocatedA a)
+amsAl (L l a) loc bs = do
+  cs <- getCommentsFor loc
+  return (L (addAnnsA l bs cs) a)
+
 amsrc :: MonadP m => Located a -> AnnContext -> m (LocatedC a)
 amsrc a@(L l _) bs = do
   cs <- getCommentsFor l
diff --git a/testsuite/tests/ghc-api/exactprint/AddDecl2.expected.hs b/testsuite/tests/ghc-api/exactprint/AddDecl2.expected.hs
index 2bbbcf5b37..5e134b5234 100644
--- a/testsuite/tests/ghc-api/exactprint/AddDecl2.expected.hs
+++ b/testsuite/tests/ghc-api/exactprint/AddDecl2.expected.hs
@@ -8,6 +8,6 @@ foo a b = a + b
 -- | Do bar
 bar x y = {- baz -} foo (x+y) x
 
-nn = n2
-
 -- end of file
+
+nn = n2
diff --git a/testsuite/tests/ghc-api/exactprint/RmDecl4.expected.hs b/testsuite/tests/ghc-api/exactprint/RmDecl4.expected.hs
index e7c71dbd08..9a36673062 100644
--- a/testsuite/tests/ghc-api/exactprint/RmDecl4.expected.hs
+++ b/testsuite/tests/ghc-api/exactprint/RmDecl4.expected.hs
@@ -7,4 +7,5 @@ ff y = y + zz + xx
 
 zz = 1
 
+
 -- EOF
diff --git a/testsuite/tests/printer/Makefile b/testsuite/tests/printer/Makefile
index e2081aee9c..4d2dad7b86 100644
--- a/testsuite/tests/printer/Makefile
+++ b/testsuite/tests/printer/Makefile
@@ -728,3 +728,8 @@ PprUnicodeSyntax:
 PprCommentPlacement2:
 	$(CHECK_PPR)   $(LIBDIR) PprCommentPlacement2.hs
 	$(CHECK_EXACT) $(LIBDIR) PprCommentPlacement2.hs
+
+.PHONY: Test20258
+Test20258:
+	$(CHECK_PPR)   $(LIBDIR) Test20258.hs
+	$(CHECK_EXACT) $(LIBDIR) Test20258.hs
diff --git a/testsuite/tests/printer/Test20258.hs b/testsuite/tests/printer/Test20258.hs
new file mode 100644
index 0000000000..bfcad743ae
--- /dev/null
+++ b/testsuite/tests/printer/Test20258.hs
@@ -0,0 +1,79 @@
+{-# LANGUAGE StandaloneDeriving #-}
+{-# LANGUAGE StandaloneKindSignatures #-}
+{-# LANGUAGE TemplateHaskell #-}
+module Test20258 where
+
+x = 1
+
+-- Comment
+;
+
+data Foo = Foo
+
+-- After TyClD
+;
+
+instance Monoid CIRB where
+  mempty = CIRB mempty mempty mempty mempty
+
+-- After InstD
+;
+
+deriving instance Eq (GenTickish 'TickishPassCore)
+
+-- After DerivD
+;
+
+transferCodingStr DeflateTransferCoding  = "deflate"
+
+-- After ValD
+;
+
+getContentType :: Int
+
+-- After SigD
+;
+
+type MyMaybe :: Type -> Type
+
+-- After KindSigD
+;
+
+default (Integer)
+
+-- After DefD
+;
+
+foreign import ccall unsafe "isDoubleFinite" isDoubleFinite :: Double -> Int
+
+-- After ForD
+;
+
+{-# DEPRECATED foo2 [] #-}
+
+-- After WarningD
+;
+
+{-# ANN module FromA #-}
+
+-- After AnnD
+;
+
+{-# RULES "myrule2" id f = f #-}
+
+-- After RuleD
+;
+
+$foo
+
+-- After SpliceD
+;
+
+type role Representational representational
+
+-- After RoleAnnotD
+;
+
+getContentType = 1
+
+-- Note: skipping DocD, only generated in haddock mode
diff --git a/testsuite/tests/printer/all.T b/testsuite/tests/printer/all.T
index 2a477cf83b..4b9c04c01c 100644
--- a/testsuite/tests/printer/all.T
+++ b/testsuite/tests/printer/all.T
@@ -169,3 +169,5 @@ test('PprT13747', ignore_stderr, makefile_test, ['PprT13747'])
 test('PprBracesSemiDataDecl', ignore_stderr, makefile_test, ['PprBracesSemiDataDecl'])
 test('PprUnicodeSyntax', ignore_stderr, makefile_test, ['PprUnicodeSyntax'])
 test('PprCommentPlacement2', ignore_stderr, makefile_test, ['PprCommentPlacement2'])
+
+test('Test20258', ignore_stderr, makefile_test, ['Test20258'])
diff --git a/utils/check-exact/ExactPrint.hs b/utils/check-exact/ExactPrint.hs
index 74135cb9f6..fc04e24332 100644
--- a/utils/check-exact/ExactPrint.hs
+++ b/utils/check-exact/ExactPrint.hs
@@ -268,15 +268,31 @@ enterAnn (Entry anchor' cs) a = do
   withOffset st (advance edp >> exact a)
 
   when ((getFollowingComments cs) /= []) $ do
-    -- debugM $ "starting trailing comments:" ++ showAst (getFollowingComments cs)
+    debugM $ "starting trailing comments:" ++ showAst (getFollowingComments cs)
     mapM_ printOneComment (map tokComment $ getFollowingComments cs)
-    -- debugM $ "ending trailing comments"
+    debugM $ "ending trailing comments"
 
 -- ---------------------------------------------------------------------
 
 addCommentsA :: [LEpaComment] -> EPP ()
 addCommentsA csNew = addComments (map tokComment csNew)
 
+{-
+TODO: When we addComments, some may have an anchor that is no longer
+valid, as it has been moved and has an anchor_op.
+
+Does an Anchor even make sense for a comment, perhaps it should be an
+EpaLocation?
+
+How do we sort them? do we assign a location based on when we add them
+to the list, based on the current output pos?  Except the offset is a
+delta compared to a reference location.  Need to nail the concept of
+the reference location.
+
+By definition it is the current anchor, so work against that. And that
+also means that the first entry comment that has moved should not have
+a line offset.
+-}
 addComments :: [Comment] -> EPP ()
 addComments csNew = do
   debugM $ "addComments:" ++ show csNew
@@ -1188,12 +1204,6 @@ instance (ExactPrint tm, ExactPrint ty, Outputable tm, Outputable ty)
 
 -- ---------------------------------------------------------------------
 
--- instance ExactPrint [LHsTyVarBndr () GhcPs] where
---   getAnnotationEntry = const NoEntryVal
---   exact bs = mapM_ markAnnotated bs
-
--- ---------------------------------------------------------------------
-
 instance ExactPrint (ClsInstDecl GhcPs) where
   getAnnotationEntry cid = fromAnn (fst $ cid_ext cid)
 
@@ -3911,7 +3921,7 @@ printString layout str = do
       cr = getDeltaLine strDP
   p <- getPosP
   colOffset <- getLayoutOffsetP
-  debugM $ "printString:(p,colOffset,strDP,cr)="  ++ show (p,colOffset,strDP,cr)
+  -- debugM $ "printString:(p,colOffset,strDP,cr)="  ++ show (p,colOffset,strDP,cr)
   if cr == 0
     then setPosP (undelta p strDP colOffset)
     else setPosP (undelta p strDP 1)
diff --git a/utils/check-exact/Main.hs b/utils/check-exact/Main.hs
index 0d79249398..2b93f2553e 100644
--- a/utils/check-exact/Main.hs
+++ b/utils/check-exact/Main.hs
@@ -72,7 +72,7 @@ _tt = testOneFile changers "/home/alanz/mysrc/git.haskell.org/worktree/exactprin
  -- "../../testsuite/tests/ghc-api/exactprint/RmDecl1.hs" (Just rmDecl1)
  -- "../../testsuite/tests/ghc-api/exactprint/RmDecl2.hs" (Just rmDecl2)
  -- "../../testsuite/tests/ghc-api/exactprint/RmDecl3.hs" (Just rmDecl3)
- -- "../../testsuite/tests/ghc-api/exactprint/RmDecl4.hs" (Just rmDecl4)
+ "../../testsuite/tests/ghc-api/exactprint/RmDecl4.hs" (Just rmDecl4)
  -- "../../testsuite/tests/ghc-api/exactprint/RmDecl5.hs" (Just rmDecl5)
  -- "../../testsuite/tests/ghc-api/exactprint/RmDecl6.hs" (Just rmDecl6)
  -- "../../testsuite/tests/ghc-api/exactprint/RmDecl7.hs" (Just rmDecl7)
@@ -198,7 +198,7 @@ _tt = testOneFile changers "/home/alanz/mysrc/git.haskell.org/worktree/exactprin
  -- "../../testsuite/tests/printer/Test19850.hs" Nothing
  -- "../../testsuite/tests/printer/PprLinearArrow.hs" Nothing
  -- "../../testsuite/tests/printer/PprSemis.hs" Nothing
- "../../testsuite/tests/printer/PprEmptyMostly.hs" Nothing
+ -- "../../testsuite/tests/printer/PprEmptyMostly.hs" Nothing
 
 -- cloneT does not need a test, function can be retired
 
@@ -462,6 +462,7 @@ changeAddDecl1 libdir top = do
   return p'
 
 -- ---------------------------------------------------------------------
+
 changeAddDecl2 :: Changer
 changeAddDecl2 libdir top = do
   Right decl <- withDynFlags libdir (\df -> parseDecl df "<interactive>" "nn = n2")
@@ -483,9 +484,10 @@ changeAddDecl3 libdir top = do
 
   let (p',(_,_),_) = runTransform mempty doAddDecl
       doAddDecl = everywhereM (mkM replaceTopLevelDecls) top
-      f d (l1:l2:ls) = l1:d:l2':ls
+      f d (l1:l2:ls) = (l1:d:l2':ls)
         where
           l2' = setEntryDP' l2 (DifferentLine 2 0)
+
       replaceTopLevelDecls :: ParsedSource -> Transform ParsedSource
       replaceTopLevelDecls m = insertAt f m decl'
   return p'
@@ -557,9 +559,9 @@ changeWhereIn3a :: Changer
 changeWhereIn3a _libdir (L l p) = do
   let decls0 = hsmodDecls p
       (decls,(_,_),w) = runTransform mempty (balanceCommentsList decls0)
-      (_de0:_:de1:_d2:_) = decls
+  --     (_de0:_:de1:_d2:_) = decls
   debugM $ unlines w
-  debugM $ "changeWhereIn3a:de1:" ++ showAst de1
+  -- debugM $ "changeWhereIn3a:de1:" ++ showAst de1
   let p2 = p { hsmodDecls = decls}
   return (L l p2)
 
@@ -710,9 +712,10 @@ rmDecl1 _libdir lp = do
   let doRmDecl = do
          tlDecs0 <- hsDecls lp
          tlDecs <- balanceCommentsList $ captureLineSpacing tlDecs0
-         let (de1:_s1:_d2:ds) = tlDecs
+         let (de1:_s1:_d2:d3:ds) = tlDecs
+         let d3' = setEntryDP' d3 (DifferentLine 2 0)
 
-         replaceDecls lp (de1:ds)
+         replaceDecls lp (de1:d3':ds)
 
   (lp',(_,_),_w) <- runTransformT mempty doRmDecl
   debugM $ "log:[\n" ++ intercalate "\n" _w ++ "]log end\n"
diff --git a/utils/check-exact/Transform.hs b/utils/check-exact/Transform.hs
index b9e400613f..867133f195 100644
--- a/utils/check-exact/Transform.hs
+++ b/utils/check-exact/Transform.hs
@@ -271,7 +271,7 @@ captureMatchLineSpacing (L l (ValD x (FunBind a b (MG c (L d ms ) e) f)))
 captureMatchLineSpacing d = d
 
 captureLineSpacing :: Monoid t
-                   => [LocatedAn t e] -> [GenLocated (SrcSpanAnn' (EpAnn t)) e]
+                   => [LocatedAn t e] -> [LocatedAn t e]
 captureLineSpacing [] = []
 captureLineSpacing [d] = [d]
 captureLineSpacing (de1:d2:ds) = de1:captureLineSpacing (d2':ds)
@@ -600,7 +600,7 @@ balanceCommentsList (a:b:ls) = do
   r <- balanceCommentsList (b':ls)
   return (a':r)
 
--- |The relatavise phase puts all comments appearing between the end of one AST
+-- |The GHC parser puts all comments appearing between the end of one AST
 -- item and the beginning of the next as 'annPriorComments' for the second one.
 -- This function takes two adjacent AST items and moves any 'annPriorComments'
 -- from the second one to the 'annFollowingComments' of the first if they belong
@@ -610,7 +610,6 @@ balanceComments :: (Monad m)
   => LHsDecl GhcPs -> LHsDecl GhcPs
   -> TransformT m (LHsDecl GhcPs, LHsDecl GhcPs)
 balanceComments first second = do
-  -- ++AZ++ : replace the nested casts with appropriate gmapM
   -- logTr $ "balanceComments entered"
   -- logDataWithAnnsTr "first" first
   case first of
@@ -626,14 +625,37 @@ balanceCommentsFB :: (Monad m)
   => LHsBind GhcPs -> LocatedA b -> TransformT m (LHsBind GhcPs, LocatedA b)
 balanceCommentsFB (L lf (FunBind x n (MG mx (L lm matches) o) t)) second = do
   logTr $ "balanceCommentsFB entered: " ++ showGhc (ss2range $ locA lf)
-  matches' <- balanceCommentsList' matches
-  let (m,ms) = case reverse matches' of
-                 (m':ms') -> (m',ms')
+  -- There are comments on lf.  We need to
+  -- + Keep the prior ones here
+  -- + move the interior ones to the first match,
+  -- + move the trailing ones to the last match.
+  let
+    split = splitCommentsEnd (realSrcSpan $ locA lf) (epAnnComments $ ann lf)
+    split2 = splitCommentsStart (realSrcSpan $ locA lf)  (EpaComments (sort $ priorComments split))
+
+    before = sort $ priorComments split2
+    middle = sort $ getFollowingComments split2
+    after  = sort $ getFollowingComments split
+
+    lf' = setCommentsSrcAnn lf (EpaComments before)
+  logTr $ "balanceCommentsFB (before, after): " ++ showAst (before, after)
+  let matches' = case matches of
+                    (L lm' m':ms') ->
+                      (L (addCommentsToSrcAnn lm' (EpaComments middle )) m':ms')
+                    _ -> error "balanceCommentsFB"
+  matches'' <- balanceCommentsList' matches'
+  let (m,ms) = case reverse matches'' of
+                 (L lm' m':ms') ->
+                   (L (addCommentsToSrcAnn lm' (EpaCommentsBalanced [] after)) m',ms')
                  _ -> error "balanceCommentsFB"
   (m',second') <- balanceComments' m second
   m'' <- balanceCommentsMatch m'
+  let (m''',lf'') = case ms of
+        [] -> moveLeadingComments m'' lf'
+        _  -> (m'',lf')
   logTr $ "balanceCommentsMatch done"
-  return (L lf (FunBind x n (MG mx (L lm (reverse (m'':ms))) o) t), second')
+  -- return (L lf'' (FunBind x n (MG mx (L lm (reverse (m''':ms))) o) t), second')
+  balanceComments' (L lf'' (FunBind x n (MG mx (L lm (reverse (m''':ms))) o) t)) second'
 balanceCommentsFB f s = balanceComments' f s
 
 -- | Move comments on the same line as the end of the match into the
@@ -642,10 +664,11 @@ balanceCommentsMatch :: (Monad m)
   => LMatch GhcPs (LHsExpr GhcPs) -> TransformT m (LMatch GhcPs (LHsExpr GhcPs))
 balanceCommentsMatch (L l (Match am mctxt pats (GRHSs xg grhss binds))) = do
   logTr $ "balanceCommentsMatch: (loc1)=" ++ showGhc (ss2range (locA l))
-  logTr $ "balanceCommentsMatch: (move',stay')=" ++ showAst (move',stay')
+  -- logTr $ "balanceCommentsMatch: (move',stay')=" ++ showAst (move',stay')
   logTr $ "balanceCommentsMatch: (logInfo)=" ++ showAst (logInfo)
-  logTr $ "balanceCommentsMatch: (loc1)=" ++ showGhc (ss2range (locA l))
+  -- logTr $ "balanceCommentsMatch: (loc1)=" ++ showGhc (ss2range (locA l))
   logTr $ "balanceCommentsMatch: (anc1,cs1f)=" ++ showAst (anc1,cs1f)
+  logTr $ "balanceCommentsMatch: (move,stay)=" ++ showAst (move,stay)
   logTr $ "balanceCommentsMatch: (l'', grhss')=" ++ showAst (l'', grhss')
   return (L l'' (Match am mctxt pats (GRHSs xg grhss' binds')))
   where
@@ -670,7 +693,7 @@ balanceCommentsMatch (L l (Match am mctxt pats (GRHSs xg grhss binds))) = do
               -- ---------------------------------
 
               (EpAnn anc an lgc) = ag
-              lgc' = splitComments (realSrcSpan lg) $ addCommentOrigDeltas lgc
+              lgc' = splitCommentsEnd (realSrcSpan lg) $ addCommentOrigDeltas lgc
               ag' = if moved
                       then EpAnn anc an lgc'
                       else EpAnn anc an (lgc' <> (EpaCommentsBalanced [] move))
@@ -688,7 +711,7 @@ pushTrailingComments w cs lb@(HsValBinds an _)
     (an', decls') = case reverse decls of
       [] -> (addCommentsToEpAnn (spanHsLocaLBinds lb) an cs, decls)
       (L la d:ds) -> (an, L (addCommentsToSrcAnn la cs) d:ds)
-    (vb,_ws2) = case runTransform mempty (replaceDeclsValbinds w lb decls') of
+    (vb,_ws2) = case runTransform mempty (replaceDeclsValbinds w lb (reverse decls')) of
       ((HsValBinds _ vb'), _, ws2') -> (vb', ws2')
       _ -> (ValBinds NoAnnSortKey emptyBag [], [])
 
@@ -711,11 +734,11 @@ balanceCommentsList' (a:b:ls) = do
 balanceComments' :: (Monad m) => LocatedA a -> LocatedA b -> TransformT m (LocatedA a, LocatedA b)
 balanceComments' la1 la2 = do
   logTr $ "balanceComments': (loc1,loc2)=" ++ showGhc (ss2range loc1,ss2range loc2)
-  logTr $ "balanceComments': (anchorFromLocatedA la1)=" ++ showGhc (anchorFromLocatedA la1)
-  logTr $ "balanceComments': (sort cs2b)=" ++ showAst (sort cs2b)
-  logTr $ "balanceComments': (move',stay')=" ++ showAst (move',stay')
-  logTr $ "balanceComments': (move'',stay'')=" ++ showAst (move'',stay'')
-  logTr $ "balanceComments': (move,stay)=" ++ showAst (move,stay)
+  logTr $ "balanceComments': (anc1)=" ++ showAst (anc1)
+  logTr $ "balanceComments': (cs1s)=" ++ showAst (cs1s)
+  logTr $ "balanceComments': (sort cs1f)=" ++ showAst (sort cs1f)
+  logTr $ "balanceComments': (cs1stay,cs1move)=" ++ showAst (cs1stay,cs1move)
+  logTr $ "balanceComments': (an1',an2')=" ++ showAst (an1',an2')
   return (la1', la2')
   where
     simpleBreak n (r,_) = r > n
@@ -723,19 +746,27 @@ balanceComments' la1 la2 = do
     L (SrcSpanAnn an2 loc2) s = la2
     anc1 = addCommentOrigDeltas $ epAnnComments an1
     anc2 = addCommentOrigDeltas $ epAnnComments an2
-    cs1f = getFollowingComments anc1
-    cs2b = priorComments anc2
-    (stay'',move') = break (simpleBreak 1) (priorCommentsDeltas (anchorFromLocatedA la2) cs2b)
+
+    cs1s = splitCommentsEnd (anchorFromLocatedA la1) anc1
+    cs1p = priorCommentsDeltas    (anchorFromLocatedA la1) (priorComments        cs1s)
+    cs1f = trailingCommentsDeltas (anchorFromLocatedA la1) (getFollowingComments cs1s)
+
+    cs2s = splitCommentsEnd (anchorFromLocatedA la2) anc2
+    cs2p = priorCommentsDeltas    (anchorFromLocatedA la2) (priorComments        cs2s)
+    cs2f = trailingCommentsDeltas (anchorFromLocatedA la2) (getFollowingComments cs2s)
+
+    -- Split cs1f into those that belong on an1 and ones that must move to an2
+    (cs1move,cs1stay) = break (simpleBreak 1) cs1f
+
+    (stay'',move') = break (simpleBreak 1) cs2p
     -- Need to also check for comments more closely attached to la1,
     -- ie trailing on the same line
     (move'',stay') = break (simpleBreak 0) (trailingCommentsDeltas (anchorFromLocatedA la1) (map snd stay''))
-    move = map snd (move'' ++ move')
-    stay = map snd stay'
-    cs1 = setFollowingComments anc1 (sort $ cs1f ++ move)
-    cs2 = setPriorComments anc2 stay
+    move = sort $ map snd (cs1move ++ move'' ++ move')
+    stay = sort $ map snd (cs1stay ++ stay')
 
-    an1' = setCommentsSrcAnn (getLoc la1) cs1
-    an2' = setCommentsSrcAnn (getLoc la2) cs2
+    an1' = setCommentsSrcAnn (getLoc la1) (EpaCommentsBalanced (map snd cs1p) move)
+    an2' = setCommentsSrcAnn (getLoc la2) (EpaCommentsBalanced stay (map snd cs2f))
     la1' = L an1' f
     la2' = L an2' s
 
@@ -769,21 +800,54 @@ priorCommentsDeltas anc cs = go anc (reverse $ sort cs)
 
 -- | Split comments into ones occuring before the end of the reference
 -- span, and those after it.
-splitComments :: RealSrcSpan -> EpAnnComments -> EpAnnComments
-splitComments p (EpaComments cs) = cs'
+splitCommentsEnd :: RealSrcSpan -> EpAnnComments -> EpAnnComments
+splitCommentsEnd p (EpaComments cs) = cs'
+  where
+    cmp (L (Anchor l _) _) = ss2pos l > ss2posEnd p
+    (before, after) = break cmp cs
+    cs' = case after of
+      [] -> EpaComments cs
+      _ -> EpaCommentsBalanced before after
+splitCommentsEnd p (EpaCommentsBalanced cs ts) = EpaCommentsBalanced cs' ts'
+  where
+    cmp (L (Anchor l _) _) = ss2pos l > ss2posEnd p
+    (before, after) = break cmp cs
+    cs' = before
+    ts' = after <> ts
+
+-- | Split comments into ones occuring before the start of the reference
+-- span, and those after it.
+splitCommentsStart :: RealSrcSpan -> EpAnnComments -> EpAnnComments
+splitCommentsStart p (EpaComments cs) = cs'
   where
-    cmp (L (Anchor l _) _) = ss2pos l < ss2posEnd p
+    cmp (L (Anchor l _) _) = ss2pos l > ss2pos p
     (before, after) = break cmp cs
     cs' = case after of
       [] -> EpaComments cs
       _ -> EpaCommentsBalanced before after
-splitComments p (EpaCommentsBalanced cs ts) = EpaCommentsBalanced cs' ts'
+splitCommentsStart p (EpaCommentsBalanced cs ts) = EpaCommentsBalanced cs' ts'
   where
-    cmp (L (Anchor l _) _) = ss2pos l < ss2posEnd p
+    cmp (L (Anchor l _) _) = ss2pos l > ss2pos p
     (before, after) = break cmp cs
     cs' = before
     ts' = after <> ts
 
+moveLeadingComments :: (Data t, Data u, Monoid t, Monoid u)
+  => LocatedAn t a -> SrcAnn u -> (LocatedAn t a, SrcAnn u)
+moveLeadingComments from@(L (SrcSpanAnn EpAnnNotUsed _) _) to = (from, to)
+moveLeadingComments (L la a) lb = (L la' a, lb')
+  `debug` ("moveLeadingComments: (before, after, la', lb'):" ++ showAst (before, after, la', lb'))
+  where
+    split = splitCommentsEnd (realSrcSpan $ locA la) (epAnnComments $ ann la)
+    before = sort $ priorComments split
+    after = sort $ getFollowingComments split
+
+    -- TODO: need to set an entry delta on lb' to zero, and move the
+    -- original spacing to the first comment.
+
+    la' = setCommentsSrcAnn la (EpaComments after)
+    lb' = addCommentsToSrcAnn lb (EpaCommentsBalanced before [])
+
 -- | A GHC comment includes the span of the preceding (non-comment)
 -- token.  Takes an original list of comments, and converts the
 -- 'Anchor's to have a have a `MovedAnchor` operation based on the
@@ -814,7 +878,6 @@ anchorFromLocatedA (L (SrcSpanAnn an loc) _)
 -- `MovedAnchor` operation based on the original location, only if it
 -- does not already have one.
 commentOrigDelta :: LEpaComment -> LEpaComment
--- commentOrigDelta c@(L (GHC.Anchor _ (GHC.MovedAnchor _)) _) = c
 commentOrigDelta (L (GHC.Anchor la _) (GHC.EpaComment t pp))
   = (L (GHC.Anchor la op) (GHC.EpaComment t pp))
   where
@@ -922,6 +985,13 @@ anchorEof (L l m@(HsModule an _lo _mn _exps _imps _decls _ _)) = L l (m { hsmodA
 
 -- ---------------------------------------------------------------------
 
+commentsOrigDeltasDecl :: LHsDecl GhcPs -> LHsDecl GhcPs
+commentsOrigDeltasDecl (L (SrcSpanAnn an l) d) = L (SrcSpanAnn an' l) d
+  where
+    an' = addCommentOrigDeltasAnn an
+
+-- ---------------------------------------------------------------------
+
 -- | Take an anchor and a preceding location, and generate an
 -- equivalent one with a 'MovedAnchor' delta.
 deltaAnchor :: Anchor -> RealSrcSpan -> Anchor
@@ -983,7 +1053,9 @@ insertAt :: (HasDecls ast)
               -> Transform ast
 insertAt f t decl = do
   oldDecls <- hsDecls t
-  replaceDecls t (f decl oldDecls)
+  oldDeclsb <- balanceCommentsList oldDecls
+  let oldDecls' = map commentsOrigDeltasDecl oldDeclsb
+  replaceDecls t (f decl oldDecls')
 
 -- |Insert a declaration at the beginning or end of the subdecls of the given
 -- AST item