[Git][ghc/ghc][wip/marge_bot_batch_merge_job] 2 commits: finder: Add `IsBootInterface` to finder cache keys

Marge Bot (@marge-bot) gitlab at gitlab.haskell.org
Thu Sep 12 20:41:16 UTC 2024



Marge Bot pushed to branch wip/marge_bot_batch_merge_job at Glasgow Haskell Compiler / GHC


Commits:
bf856255 by Torsten Schmits at 2024-09-12T16:40:26-04:00
finder: Add `IsBootInterface` to finder cache keys

- - - - -
dda9c763 by Alan Zimmerman at 2024-09-12T16:40:27-04:00
EPA: Sync ghc-exactprint to GHC

- - - - -


18 changed files:

- compiler/GHC/Driver/Backpack.hs
- compiler/GHC/Driver/Make.hs
- compiler/GHC/Driver/Pipeline/Execute.hs
- compiler/GHC/Unit/Finder.hs
- compiler/GHC/Unit/Finder/Types.hs
- compiler/GHC/Unit/Module/Env.hs
- compiler/GHC/Unit/Types.hs
- + testsuite/tests/driver/boot-target/A.hs
- + testsuite/tests/driver/boot-target/A.hs-boot
- + testsuite/tests/driver/boot-target/B.hs
- + testsuite/tests/driver/boot-target/Makefile
- + testsuite/tests/driver/boot-target/all.T
- utils/check-exact/ExactPrint.hs
- utils/check-exact/Main.hs
- utils/check-exact/Parsers.hs
- utils/check-exact/Transform.hs
- utils/check-exact/Types.hs
- utils/check-exact/Utils.hs


Changes:

=====================================
compiler/GHC/Driver/Backpack.hs
=====================================
@@ -781,7 +781,7 @@ summariseRequirement pn mod_name = do
     let loc = srcLocSpan (mkSrcLoc (mkFastString (bkp_filename env)) 1 1)
 
     let fc = hsc_FC hsc_env
-    mod <- liftIO $ addHomeModuleToFinder fc home_unit mod_name location
+    mod <- liftIO $ addHomeModuleToFinder fc home_unit (notBoot mod_name) location
 
     extra_sig_imports <- liftIO $ findExtraSigImports hsc_env HsigFile mod_name
 
@@ -893,7 +893,7 @@ hsModuleToModSummary home_keys pn hsc_src modname
     this_mod <- liftIO $ do
       let home_unit = hsc_home_unit hsc_env
       let fc        = hsc_FC hsc_env
-      addHomeModuleToFinder fc home_unit modname location
+      addHomeModuleToFinder fc home_unit (GWIB modname (hscSourceToIsBoot hsc_src)) location
     let ms = ModSummary {
             ms_mod = this_mod,
             ms_hsc_src = hsc_src,


=====================================
compiler/GHC/Driver/Make.hs
=====================================
@@ -2044,25 +2044,43 @@ summariseFile hsc_env' home_unit old_summaries src_fn mb_phase maybe_buf
             <- getPreprocessedImports hsc_env src_fn mb_phase maybe_buf
 
         let fopts = initFinderOpts (hsc_dflags hsc_env)
-
-        -- Make a ModLocation for this file
-        let location = mkHomeModLocation fopts pi_mod_name (unsafeEncodeUtf src_fn)
+            src_path = unsafeEncodeUtf src_fn
+
+            is_boot = case takeExtension src_fn of
+              ".hs-boot" -> IsBoot
+              ".lhs-boot" -> IsBoot
+              _ -> NotBoot
+
+            (path_without_boot, hsc_src)
+              | isHaskellSigFilename src_fn = (src_path, HsigFile)
+              | IsBoot <- is_boot = (removeBootSuffix src_path, HsBootFile)
+              | otherwise = (src_path, HsSrcFile)
+
+            -- Make a ModLocation for the Finder, who only has one entry for
+            -- each @ModuleName@, and therefore needs to use the locations for
+            -- the non-boot files.
+            location_without_boot =
+              mkHomeModLocation fopts pi_mod_name path_without_boot
+
+            -- Make a ModLocation for this file, adding the @-boot@ suffix to
+            -- all paths if the original was a boot file.
+            location
+              | IsBoot <- is_boot
+              = addBootSuffixLocn location_without_boot
+              | otherwise
+              = location_without_boot
 
         -- Tell the Finder cache where it is, so that subsequent calls
         -- to findModule will find it, even if it's not on any search path
         mod <- liftIO $ do
           let home_unit = hsc_home_unit hsc_env
           let fc        = hsc_FC hsc_env
-          addHomeModuleToFinder fc home_unit pi_mod_name location
+          addHomeModuleToFinder fc home_unit (GWIB pi_mod_name is_boot) location
 
         liftIO $ makeNewModSummary hsc_env $ MakeNewModSummary
             { nms_src_fn = src_fn
             , nms_src_hash = src_hash
-            , nms_is_boot = NotBoot
-            , nms_hsc_src =
-                if isHaskellSigFilename src_fn
-                   then HsigFile
-                   else HsSrcFile
+            , nms_hsc_src = hsc_src
             , nms_location = location
             , nms_mod = mod
             , nms_preimps = preimps
@@ -2090,9 +2108,10 @@ checkSummaryHash
            -- Also, only add to finder cache for non-boot modules as the finder cache
            -- makes sure to add a boot suffix for boot files.
            _ <- do
-              let fc        = hsc_FC hsc_env
+              let fc = hsc_FC hsc_env
+                  gwib = GWIB (ms_mod old_summary) (isBootSummary old_summary)
               case ms_hsc_src old_summary of
-                HsSrcFile -> addModuleToFinder fc (ms_mod old_summary) location
+                HsSrcFile -> addModuleToFinder fc gwib location
                 _ -> return ()
 
            hi_timestamp <- modificationTimeIfExists (ml_hi_file location)
@@ -2230,7 +2249,6 @@ summariseModule hsc_env' home_unit old_summary_map is_boot (L _ wanted_mod) mb_p
         liftIO $ makeNewModSummary hsc_env $ MakeNewModSummary
             { nms_src_fn = src_fn
             , nms_src_hash = src_hash
-            , nms_is_boot = is_boot
             , nms_hsc_src = hsc_src
             , nms_location = location
             , nms_mod = mod
@@ -2243,7 +2261,6 @@ data MakeNewModSummary
   = MakeNewModSummary
       { nms_src_fn :: FilePath
       , nms_src_hash :: Fingerprint
-      , nms_is_boot :: IsBootInterface
       , nms_hsc_src :: HscSource
       , nms_location :: ModLocation
       , nms_mod :: Module


=====================================
compiler/GHC/Driver/Pipeline/Execute.hs
=====================================
@@ -734,7 +734,7 @@ runHscPhase pipe_env hsc_env0 input_fn src_flavour = do
   mod <- do
     let home_unit = hsc_home_unit hsc_env
     let fc        = hsc_FC hsc_env
-    addHomeModuleToFinder fc home_unit mod_name location
+    addHomeModuleToFinder fc home_unit (GWIB mod_name (hscSourceToIsBoot src_flavour)) location
 
   -- Make the ModSummary to hand to hscMain
   let


=====================================
compiler/GHC/Unit/Finder.hs
=====================================
@@ -89,23 +89,23 @@ type BaseName = OsPath  -- Basename of file
 
 initFinderCache :: IO FinderCache
 initFinderCache = do
-  mod_cache <- newIORef emptyInstalledModuleEnv
+  mod_cache <- newIORef emptyInstalledModuleWithIsBootEnv
   file_cache <- newIORef M.empty
   let flushFinderCaches :: UnitEnv -> IO ()
       flushFinderCaches ue = do
-        atomicModifyIORef' mod_cache $ \fm -> (filterInstalledModuleEnv is_ext fm, ())
+        atomicModifyIORef' mod_cache $ \fm -> (filterInstalledModuleWithIsBootEnv is_ext fm, ())
         atomicModifyIORef' file_cache $ \_ -> (M.empty, ())
        where
-        is_ext mod _ = not (isUnitEnvInstalledModule ue mod)
+        is_ext mod _ = not (isUnitEnvInstalledModule ue (gwib_mod mod))
 
-      addToFinderCache :: InstalledModule -> InstalledFindResult -> IO ()
+      addToFinderCache :: InstalledModuleWithIsBoot -> InstalledFindResult -> IO ()
       addToFinderCache key val =
-        atomicModifyIORef' mod_cache $ \c -> (extendInstalledModuleEnv c key val, ())
+        atomicModifyIORef' mod_cache $ \c -> (extendInstalledModuleWithIsBootEnv c key val, ())
 
-      lookupFinderCache :: InstalledModule -> IO (Maybe InstalledFindResult)
+      lookupFinderCache :: InstalledModuleWithIsBoot -> IO (Maybe InstalledFindResult)
       lookupFinderCache key = do
          c <- readIORef mod_cache
-         return $! lookupInstalledModuleEnv c key
+         return $! lookupInstalledModuleWithIsBootEnv c key
 
       lookupFileCache :: FilePath -> IO Fingerprint
       lookupFileCache key = do
@@ -255,7 +255,7 @@ orIfNotFound this or_this = do
 homeSearchCache :: FinderCache -> UnitId -> ModuleName -> IO InstalledFindResult -> IO InstalledFindResult
 homeSearchCache fc home_unit mod_name do_this = do
   let mod = mkModule home_unit mod_name
-  modLocationCache fc mod do_this
+  modLocationCache fc (notBoot mod) do_this
 
 findExposedPackageModule :: FinderCache -> FinderOpts -> UnitState -> ModuleName -> PkgQual -> IO FindResult
 findExposedPackageModule fc fopts units mod_name mb_pkg =
@@ -312,7 +312,7 @@ findLookupResult fc fopts r = case r of
                        , fr_unusables = []
                        , fr_suggestions = suggest' })
 
-modLocationCache :: FinderCache -> InstalledModule -> IO InstalledFindResult -> IO InstalledFindResult
+modLocationCache :: FinderCache -> InstalledModuleWithIsBoot -> IO InstalledFindResult -> IO InstalledFindResult
 modLocationCache fc mod do_this = do
   m <- lookupFinderCache fc mod
   case m of
@@ -322,17 +322,17 @@ modLocationCache fc mod do_this = do
         addToFinderCache fc mod result
         return result
 
-addModuleToFinder :: FinderCache -> Module -> ModLocation -> IO ()
+addModuleToFinder :: FinderCache -> ModuleWithIsBoot -> ModLocation -> IO ()
 addModuleToFinder fc mod loc = do
-  let imod = toUnitId <$> mod
-  addToFinderCache fc imod (InstalledFound loc imod)
+  let imod = fmap toUnitId <$> mod
+  addToFinderCache fc imod (InstalledFound loc (gwib_mod imod))
 
 -- This returns a module because it's more convenient for users
-addHomeModuleToFinder :: FinderCache -> HomeUnit -> ModuleName -> ModLocation -> IO Module
+addHomeModuleToFinder :: FinderCache -> HomeUnit -> ModuleNameWithIsBoot -> ModLocation -> IO Module
 addHomeModuleToFinder fc home_unit mod_name loc = do
-  let mod = mkHomeInstalledModule home_unit mod_name
-  addToFinderCache fc mod (InstalledFound loc mod)
-  return (mkHomeModule home_unit mod_name)
+  let mod = mkHomeInstalledModule home_unit <$> mod_name
+  addToFinderCache fc mod (InstalledFound loc (gwib_mod mod))
+  return (mkHomeModule home_unit (gwib_mod mod_name))
 
 -- -----------------------------------------------------------------------------
 --      The internal workers
@@ -466,7 +466,7 @@ findPackageModule_ :: FinderCache -> FinderOpts -> InstalledModule -> UnitInfo -
 findPackageModule_ fc fopts mod pkg_conf = do
   massertPpr (moduleUnit mod == unitId pkg_conf)
              (ppr (moduleUnit mod) <+> ppr (unitId pkg_conf))
-  modLocationCache fc mod $
+  modLocationCache fc (notBoot mod) $
 
     -- special case for GHC.Prim; we won't find it in the filesystem.
     if mod `installedModuleEq` gHC_PRIM


=====================================
compiler/GHC/Unit/Finder/Types.hs
=====================================
@@ -30,9 +30,9 @@ data FinderCache = FinderCache { flushFinderCaches :: UnitEnv -> IO ()
                                -- ^ remove all the home modules from the cache; package modules are
                                -- assumed to not move around during a session; also flush the file hash
                                -- cache.
-                               , addToFinderCache  :: InstalledModule -> InstalledFindResult -> IO ()
+                               , addToFinderCache  :: InstalledModuleWithIsBoot -> InstalledFindResult -> IO ()
                                -- ^ Add a found location to the cache for the module.
-                               , lookupFinderCache :: InstalledModule -> IO (Maybe InstalledFindResult)
+                               , lookupFinderCache :: InstalledModuleWithIsBoot -> IO (Maybe InstalledFindResult)
                                -- ^ Look for a location in the cache.
                                , lookupFileCache   :: FilePath -> IO Fingerprint
                                -- ^ Look for the hash of a file in the cache. This should add it to the


=====================================
compiler/GHC/Unit/Module/Env.hs
=====================================
@@ -33,6 +33,17 @@ module GHC.Unit.Module.Env
    , mergeInstalledModuleEnv
    , plusInstalledModuleEnv
    , installedModuleEnvElts
+
+     -- * InstalledModuleWithIsBootEnv
+   , InstalledModuleWithIsBootEnv
+   , emptyInstalledModuleWithIsBootEnv
+   , lookupInstalledModuleWithIsBootEnv
+   , extendInstalledModuleWithIsBootEnv
+   , filterInstalledModuleWithIsBootEnv
+   , delInstalledModuleWithIsBootEnv
+   , mergeInstalledModuleWithIsBootEnv
+   , plusInstalledModuleWithIsBootEnv
+   , installedModuleWithIsBootEnvElts
    )
 where
 
@@ -283,3 +294,56 @@ plusInstalledModuleEnv :: (elt -> elt -> elt)
 plusInstalledModuleEnv f (InstalledModuleEnv xm) (InstalledModuleEnv ym) =
   InstalledModuleEnv $ Map.unionWith f xm ym
 
+
+
+--------------------------------------------------------------------
+-- InstalledModuleWithIsBootEnv
+--------------------------------------------------------------------
+
+-- | A map keyed off of 'InstalledModuleWithIsBoot'
+newtype InstalledModuleWithIsBootEnv elt = InstalledModuleWithIsBootEnv (Map InstalledModuleWithIsBoot elt)
+
+instance Outputable elt => Outputable (InstalledModuleWithIsBootEnv elt) where
+  ppr (InstalledModuleWithIsBootEnv env) = ppr env
+
+
+emptyInstalledModuleWithIsBootEnv :: InstalledModuleWithIsBootEnv a
+emptyInstalledModuleWithIsBootEnv = InstalledModuleWithIsBootEnv Map.empty
+
+lookupInstalledModuleWithIsBootEnv :: InstalledModuleWithIsBootEnv a -> InstalledModuleWithIsBoot -> Maybe a
+lookupInstalledModuleWithIsBootEnv (InstalledModuleWithIsBootEnv e) m = Map.lookup m e
+
+extendInstalledModuleWithIsBootEnv :: InstalledModuleWithIsBootEnv a -> InstalledModuleWithIsBoot -> a -> InstalledModuleWithIsBootEnv a
+extendInstalledModuleWithIsBootEnv (InstalledModuleWithIsBootEnv e) m x = InstalledModuleWithIsBootEnv (Map.insert m x e)
+
+filterInstalledModuleWithIsBootEnv :: (InstalledModuleWithIsBoot -> a -> Bool) -> InstalledModuleWithIsBootEnv a -> InstalledModuleWithIsBootEnv a
+filterInstalledModuleWithIsBootEnv f (InstalledModuleWithIsBootEnv e) =
+  InstalledModuleWithIsBootEnv (Map.filterWithKey f e)
+
+delInstalledModuleWithIsBootEnv :: InstalledModuleWithIsBootEnv a -> InstalledModuleWithIsBoot -> InstalledModuleWithIsBootEnv a
+delInstalledModuleWithIsBootEnv (InstalledModuleWithIsBootEnv e) m = InstalledModuleWithIsBootEnv (Map.delete m e)
+
+installedModuleWithIsBootEnvElts :: InstalledModuleWithIsBootEnv a -> [(InstalledModuleWithIsBoot, a)]
+installedModuleWithIsBootEnvElts (InstalledModuleWithIsBootEnv e) = Map.assocs e
+
+mergeInstalledModuleWithIsBootEnv
+  :: (elta -> eltb -> Maybe eltc)
+  -> (InstalledModuleWithIsBootEnv elta -> InstalledModuleWithIsBootEnv eltc)  -- map X
+  -> (InstalledModuleWithIsBootEnv eltb -> InstalledModuleWithIsBootEnv eltc) -- map Y
+  -> InstalledModuleWithIsBootEnv elta
+  -> InstalledModuleWithIsBootEnv eltb
+  -> InstalledModuleWithIsBootEnv eltc
+mergeInstalledModuleWithIsBootEnv f g h (InstalledModuleWithIsBootEnv xm) (InstalledModuleWithIsBootEnv ym)
+  = InstalledModuleWithIsBootEnv $ Map.mergeWithKey
+      (\_ x y -> (x `f` y))
+      (coerce g)
+      (coerce h)
+      xm ym
+
+plusInstalledModuleWithIsBootEnv :: (elt -> elt -> elt)
+  -> InstalledModuleWithIsBootEnv elt
+  -> InstalledModuleWithIsBootEnv elt
+  -> InstalledModuleWithIsBootEnv elt
+plusInstalledModuleWithIsBootEnv f (InstalledModuleWithIsBootEnv xm) (InstalledModuleWithIsBootEnv ym) =
+  InstalledModuleWithIsBootEnv $ Map.unionWith f xm ym
+


=====================================
compiler/GHC/Unit/Types.hs
=====================================
@@ -86,6 +86,8 @@ module GHC.Unit.Types
    , GenWithIsBoot (..)
    , ModuleNameWithIsBoot
    , ModuleWithIsBoot
+   , InstalledModuleWithIsBoot
+   , notBoot
    )
 where
 
@@ -723,6 +725,8 @@ type ModuleNameWithIsBoot = GenWithIsBoot ModuleName
 
 type ModuleWithIsBoot = GenWithIsBoot Module
 
+type InstalledModuleWithIsBoot = GenWithIsBoot InstalledModule
+
 instance Binary a => Binary (GenWithIsBoot a) where
   put_ bh (GWIB { gwib_mod, gwib_isBoot }) = do
     put_ bh gwib_mod
@@ -736,3 +740,6 @@ instance Outputable a => Outputable (GenWithIsBoot a) where
   ppr (GWIB  { gwib_mod, gwib_isBoot }) = hsep $ ppr gwib_mod : case gwib_isBoot of
     IsBoot -> [ text "{-# SOURCE #-}" ]
     NotBoot -> []
+
+notBoot :: mod -> GenWithIsBoot mod
+notBoot gwib_mod = GWIB {gwib_mod, gwib_isBoot = NotBoot}


=====================================
testsuite/tests/driver/boot-target/A.hs
=====================================
@@ -0,0 +1,5 @@
+module A where
+
+import B
+
+data A = A B


=====================================
testsuite/tests/driver/boot-target/A.hs-boot
=====================================
@@ -0,0 +1,3 @@
+module A where
+
+data A


=====================================
testsuite/tests/driver/boot-target/B.hs
=====================================
@@ -0,0 +1,5 @@
+module B where
+
+import {-# source #-} A
+
+data B = B A


=====================================
testsuite/tests/driver/boot-target/Makefile
=====================================
@@ -0,0 +1,8 @@
+boot1:
+	$(TEST_HC) -c A.hs-boot B.hs
+
+boot2:
+	$(TEST_HC) A.hs-boot A.hs B.hs -v0
+
+boot3:
+	$(TEST_HC) A.hs-boot B.hs -v0
\ No newline at end of file


=====================================
testsuite/tests/driver/boot-target/all.T
=====================================
@@ -0,0 +1,10 @@
+def test_boot(name):
+    return test(name,
+     [extra_files(['A.hs', 'A.hs-boot', 'B.hs']),
+      ],
+     makefile_test,
+     [])
+
+test_boot('boot1')
+test_boot('boot2')
+test_boot('boot3')


=====================================
utils/check-exact/ExactPrint.hs
=====================================
@@ -1,3 +1,4 @@
+{-# LANGUAGE BangPatterns         #-}
 {-# LANGUAGE DeriveDataTypeable   #-}
 {-# LANGUAGE FlexibleContexts     #-}
 {-# LANGUAGE FlexibleInstances    #-}
@@ -25,7 +26,7 @@ module ExactPrint
   , makeDeltaAst
 
   -- * Configuration
-  , EPOptions(epRigidity, epAstPrint, epTokenPrint, epWhitespacePrint, epUpdateAnchors)
+  , EPOptions(epTokenPrint, epWhitespacePrint)
   , stringOptions
   , epOptions
   , deltaOptions
@@ -43,10 +44,11 @@ import GHC.Types.ForeignCall
 import GHC.Types.Name.Reader
 import GHC.Types.PkgQual
 import GHC.Types.SourceText
+import GHC.Types.SrcLoc
 import GHC.Types.Var
-import GHC.Utils.Outputable hiding ( (<>) )
 import GHC.Unit.Module.Warnings
 import GHC.Utils.Misc
+import GHC.Utils.Outputable hiding ( (<>) )
 import GHC.Utils.Panic
 
 import Language.Haskell.Syntax.Basic (FieldLabelString(..))
@@ -77,8 +79,7 @@ import Types
 exactPrint :: ExactPrint ast => ast -> String
 exactPrint ast = snd $ runIdentity (runEP stringOptions (markAnnotated ast))
 
--- | The additional option to specify the rigidity and printing
--- configuration.
+-- | The additional option to specify the printing configuration.
 exactPrintWithOptions :: (ExactPrint ast, Monoid b, Monad m)
                       => EPOptions m b
                       -> ast
@@ -86,9 +87,8 @@ exactPrintWithOptions :: (ExactPrint ast, Monoid b, Monad m)
 exactPrintWithOptions r ast =
     runEP r (markAnnotated ast)
 
--- | Transform concrete annotations into relative annotations which
--- are more useful when transforming an AST. This corresponds to the
--- earlier 'relativiseApiAnns'.
+-- | Transform concrete annotations into relative annotations.
+-- This should be unnecessary from GHC 9.10
 makeDeltaAst :: ExactPrint ast => ast -> ast
 makeDeltaAst ast = fst $ runIdentity (runEP deltaOptions (markAnnotated ast))
 
@@ -115,6 +115,7 @@ defaultEPState = EPState
              , dPriorEndPosition = (1,1)
              , uAnchorSpan = badRealSrcSpan
              , uExtraDP = Nothing
+             , uExtraDPReturn = Nothing
              , pAcceptSpan = False
              , epComments = []
              , epCommentsApplied = []
@@ -128,39 +129,27 @@ defaultEPState = EPState
 -- | The R part of RWS. The environment. Updated via 'local' as we
 -- enter a new AST element, having a different anchor point.
 data EPOptions m a = EPOptions
-            {
-              epAstPrint :: forall ast . Data ast => GHC.Located ast -> a -> m a
-            , epTokenPrint :: String -> m a
+            { epTokenPrint :: String -> m a
             , epWhitespacePrint :: String -> m a
-            , epRigidity :: Rigidity
-            , epUpdateAnchors :: Bool
             }
 
 -- | Helper to create a 'EPOptions'
-epOptions ::
-      (forall ast . Data ast => GHC.Located ast -> a -> m a)
-      -> (String -> m a)
-      -> (String -> m a)
-      -> Rigidity
-      -> Bool
-      -> EPOptions m a
-epOptions astPrint tokenPrint wsPrint rigidity delta = EPOptions
-             {
-               epAstPrint = astPrint
-             , epWhitespacePrint = wsPrint
+epOptions :: (String -> m a)
+          -> (String -> m a)
+          -> EPOptions m a
+epOptions tokenPrint wsPrint = EPOptions
+             { epWhitespacePrint = wsPrint
              , epTokenPrint = tokenPrint
-             , epRigidity = rigidity
-             , epUpdateAnchors = delta
              }
 
 -- | Options which can be used to print as a normal String.
 stringOptions :: EPOptions Identity String
-stringOptions = epOptions (\_ b -> return b) return return NormalLayout False
+stringOptions = epOptions return return
 
 -- | Options which can be used to simply update the AST to be in delta
 -- form, without generating output
 deltaOptions :: EPOptions Identity ()
-deltaOptions = epOptions (\_ _ -> return ()) (\_ -> return ()) (\_ -> return ()) NormalLayout True
+deltaOptions = epOptions (\_ -> return ()) (\_ -> return ())
 
 data EPWriter a = EPWriter
               { output :: !a }
@@ -177,6 +166,8 @@ data EPState = EPState
                                           -- Annotation
              , uExtraDP :: !(Maybe EpaLocation) -- ^ Used to anchor a
                                                 -- list
+             , uExtraDPReturn :: !(Maybe DeltaPos)
+                  -- ^ Used to return Delta version of uExtraDP
              , pAcceptSpan :: Bool -- ^ When we have processed an
                                    -- entry of EpaDelta, accept the
                                    -- next `EpaSpan` start as the
@@ -213,7 +204,7 @@ class HasTrailing a where
   trailing :: a -> [TrailingAnn]
   setTrailing :: a -> [TrailingAnn] -> a
 
-setAnchorEpa :: (HasTrailing an, NoAnn an)
+setAnchorEpa :: (HasTrailing an)
              => EpAnn an -> EpaLocation -> [TrailingAnn] -> EpAnnComments -> EpAnn an
 setAnchorEpa (EpAnn _ an _) anc ts cs = EpAnn anc (setTrailing an ts)          cs
 
@@ -223,7 +214,7 @@ setAnchorHsModule hsmod anc cs = hsmod { hsmodExt = (hsmodExt hsmod) {hsmodAnn =
     anc' = anc
     an' = setAnchorEpa (hsmodAnn $ hsmodExt hsmod) anc' [] cs
 
-setAnchorAn :: (HasTrailing an, NoAnn an)
+setAnchorAn :: (HasTrailing an)
              => LocatedAn an a -> EpaLocation -> [TrailingAnn] -> EpAnnComments -> LocatedAn an a
 setAnchorAn (L (EpAnn _ an _) a) anc ts cs = (L (EpAnn anc (setTrailing an ts) cs) a)
      -- `debug` ("setAnchorAn: anc=" ++ showAst anc)
@@ -248,7 +239,7 @@ data FlushComments = FlushComments
 data CanUpdateAnchor = CanUpdateAnchor
                      | CanUpdateAnchorOnly
                      | NoCanUpdateAnchor
-                   deriving (Eq, Show)
+                   deriving (Eq, Show, Data)
 
 data Entry = Entry EpaLocation [TrailingAnn] EpAnnComments FlushComments CanUpdateAnchor
            | NoEntryVal
@@ -402,7 +393,7 @@ enterAnn NoEntryVal a = do
   r <- exact a
   debugM $ "enterAnn:done:NO ANN:p =" ++ show (p, astId a)
   return r
-enterAnn (Entry anchor' trailing_anns cs flush canUpdateAnchor) a = do
+enterAnn !(Entry anchor' trailing_anns cs flush canUpdateAnchor) a = do
   acceptSpan <- getAcceptSpan
   setAcceptSpan False
   case anchor' of
@@ -421,9 +412,11 @@ enterAnn (Entry anchor' trailing_anns cs flush canUpdateAnchor) a = do
     _ -> return ()
   case anchor' of
     EpaDelta _ _ dcs -> do
-      debugM $ "enterAnn:Printing comments:" ++ showGhc (priorComments cs)
+      debugM $ "enterAnn:Delta:Flushing comments"
+      flushComments []
+      debugM $ "enterAnn:Delta:Printing prior comments:" ++ showGhc (priorComments cs)
       mapM_ printOneComment (concatMap tokComment $ priorComments cs)
-      debugM $ "enterAnn:Printing EpaDelta comments:" ++ showGhc dcs
+      debugM $ "enterAnn:Delta:Printing EpaDelta comments:" ++ showGhc dcs
       mapM_ printOneComment (concatMap tokComment dcs)
     _ -> do
       debugM $ "enterAnn:Adding comments:" ++ showGhc (priorComments cs)
@@ -465,7 +458,7 @@ enterAnn (Entry anchor' trailing_anns cs flush canUpdateAnchor) a = do
   -- The first part corresponds to the delta phase, so should only use
   -- delta phase variables -----------------------------------
   -- Calculate offset required to get to the start of the SrcSPan
-  off <- getLayoutOffsetD
+  !off <- getLayoutOffsetD
   let spanStart = ss2pos curAnchor
   priorEndAfterComments <- getPriorEndD
   let edp' = adjustDeltaForOffset
@@ -480,17 +473,18 @@ enterAnn (Entry anchor' trailing_anns cs flush canUpdateAnchor) a = do
   -- ---------------------------------------------
   med <- getExtraDP
   setExtraDP Nothing
-  let edp = case med of
-        Nothing -> edp''
-        Just (EpaDelta _ dp _) -> dp
+  let (edp, medr) = case med of
+        Nothing -> (edp'', Nothing)
+        Just (EpaDelta _ dp _) -> (dp, Nothing)
                    -- Replace original with desired one. Allows all
                    -- list entry values to be DP (1,0)
-        Just (EpaSpan (RealSrcSpan r _)) -> dp
+        Just (EpaSpan (RealSrcSpan r _)) -> (dp, Just dp)
           where
             dp = adjustDeltaForOffset
                    off (ss2delta priorEndAfterComments r)
         Just (EpaSpan (UnhelpfulSpan r)) -> panic $ "enterAnn: UnhelpfulSpan:" ++ show r
   when (isJust med) $ debugM $ "enterAnn:(med,edp)=" ++ showAst (med,edp)
+  when (isJust medr) $ setExtraDPReturn medr
   -- ---------------------------------------------
   -- Preparation complete, perform the action
   when (priorEndAfterComments < spanStart) (do
@@ -511,12 +505,15 @@ enterAnn (Entry anchor' trailing_anns cs flush canUpdateAnchor) a = do
   debugM $ "enterAnn:exact a starting:" ++ show (showAst anchor')
   a' <- exact a
   debugM $ "enterAnn:exact a done:" ++ show (showAst anchor')
+
+  -- Core recursive exactprint done, start end of Entry processing
+
   when (flush == FlushComments) $ do
-    debugM $ "flushing comments in enterAnn:" ++ showAst cs
+    debugM $ "flushing comments in enterAnn:" ++ showAst (cs, getFollowingComments cs)
     flushComments (getFollowingComments cs)
     debugM $ "flushing comments in enterAnn done"
 
-  eof <- getEofPos
+  !eof <- getEofPos
   case eof of
     Nothing -> return ()
     Just (pos, prior) -> do
@@ -544,28 +541,50 @@ enterAnn (Entry anchor' trailing_anns cs flush canUpdateAnchor) a = do
 
   -- Outside the anchor, mark any trailing
   postCs <- cua canUpdateAnchor takeAppliedCommentsPop
-  when (flush == NoFlushComments) $ do
-    when ((getFollowingComments cs) /= []) $ do
-
-      -- debugM $ "enterAnn:in:(anchor') =" ++ show (eloc2str anchor')
-      debugM $ "starting trailing comments:" ++ showAst (getFollowingComments cs)
-      mapM_ printOneComment (concatMap tokComment $ getFollowingComments cs)
-      debugM $ "ending trailing comments"
-  trailing' <- markTrailing trailing_anns
+  following <- if (flush == NoFlushComments)
+          then do
+              let (before, after) = splitAfterTrailingAnns trailing_anns
+                                                           (getFollowingComments cs)
+              addCommentsA before
+              return after
+           else return []
+  !trailing' <- markTrailing trailing_anns
+  -- mapM_ printOneComment (concatMap tokComment $ following)
+  addCommentsA following
 
   -- Update original anchor, comments based on the printing process
   -- TODO:AZ: probably need to put something appropriate in instead of noSrcSpan
-  let newAchor = EpaDelta noSrcSpan edp []
+  let newAnchor = EpaDelta noSrcSpan edp []
   let r = case canUpdateAnchor of
-            CanUpdateAnchor -> setAnnotationAnchor a' newAchor trailing' (mkEpaComments (priorCs ++ postCs) [])
-            CanUpdateAnchorOnly -> setAnnotationAnchor a' newAchor [] emptyComments
+            CanUpdateAnchor -> setAnnotationAnchor a' newAnchor trailing' (mkEpaComments priorCs postCs)
+            CanUpdateAnchorOnly -> setAnnotationAnchor a' newAnchor [] emptyComments
             NoCanUpdateAnchor -> a'
   return r
 
 -- ---------------------------------------------------------------------
 
+-- | Split the span following comments into ones that occur prior to
+-- the last trailing ann, and ones after.
+splitAfterTrailingAnns :: [TrailingAnn] -> [LEpaComment] -> ([LEpaComment], [LEpaComment])
+splitAfterTrailingAnns [] cs = ([], cs)
+splitAfterTrailingAnns tas cs = (before, after)
+  where
+    trailing_loc ta = case ta_location ta of
+        EpaSpan (RealSrcSpan s _) -> [s]
+        _ -> []
+    (before, after) = case reverse (concatMap trailing_loc tas) of
+        [] -> ([],cs)
+        (s:_) -> (b,a)
+          where
+            s_pos = ss2pos s
+            (b,a) = break (\(L ll _) -> (ss2pos $ anchor ll) > s_pos)
+                          cs
+
+
+-- ---------------------------------------------------------------------
+
 addCommentsA :: (Monad m, Monoid w) => [LEpaComment] -> EP w m ()
-addCommentsA csNew = addComments (concatMap tokComment csNew)
+addCommentsA csNew = addComments False (concatMap tokComment csNew)
 
 {-
 TODO: When we addComments, some may have an anchor that is no longer
@@ -583,24 +602,36 @@ 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 :: (Monad m, Monoid w) => [Comment] -> EP w m ()
-addComments csNew = do
-  -- debugM $ "addComments:" ++ show csNew
+addComments :: (Monad m, Monoid w) => Bool -> [Comment] -> EP w m ()
+addComments sortNeeded csNew = do
+  debugM $ "addComments:csNew" ++ show csNew
   cs <- getUnallocatedComments
+  debugM $ "addComments:cs" ++ show cs
+  -- We can only sort the comments if we are in the first phase,
+  -- where all comments have locations. If any have EpaDelta the
+  -- sort will fail, so we do not try.
+  if sortNeeded && all noDelta (csNew ++ cs)
+    then putUnallocatedComments (sort (cs ++ csNew))
+    else putUnallocatedComments (cs ++ csNew)
 
-  putUnallocatedComments (sort (cs ++ csNew))
+noDelta :: Comment -> Bool
+noDelta c = case commentLoc c of
+    EpaSpan _ -> True
+    _ -> False
 
 -- ---------------------------------------------------------------------
 
 -- | Just before we print out the EOF comments, flush the remaining
 -- ones in the state.
 flushComments :: (Monad m, Monoid w) => [LEpaComment] -> EP w m ()
-flushComments trailing_anns = do
+flushComments !trailing_anns = do
+  debugM $ "flushComments entered: " ++ showAst trailing_anns
   addCommentsA trailing_anns
+  debugM $ "flushComments after addCommentsA"
   cs <- getUnallocatedComments
-  debugM $ "flushing comments starting"
-    -- AZ:TODO: is the sort still needed?
-  mapM_ printOneComment (sortComments cs)
+  debugM $ "flushComments: got cs"
+  debugM $ "flushing comments starting: cs" ++ showAst cs
+  mapM_ printOneComment cs
   putUnallocatedComments []
   debugM $ "flushing comments done"
 
@@ -612,7 +643,7 @@ annotationsToComments :: (Monad m, Monoid w)
   => a -> Lens a [AddEpAnn] -> [AnnKeywordId] -> EP w m a
 annotationsToComments a l kws = do
   let (newComments, newAnns) = go ([],[]) (view l a)
-  addComments newComments
+  addComments True newComments
   return (set l (reverse newAnns) a)
   where
     keywords = Set.fromList kws
@@ -654,14 +685,11 @@ printSourceText (NoSourceText) txt   =  printStringAdvance txt >> return ()
 printSourceText (SourceText   txt) _ =  printStringAdvance (unpackFS txt) >> return ()
 
 printSourceTextAA :: (Monad m, Monoid w) => SourceText -> String -> EP w m ()
-printSourceTextAA (NoSourceText) txt   = printStringAtAA noAnn txt >> return ()
-printSourceTextAA (SourceText   txt) _ =  printStringAtAA noAnn (unpackFS txt) >> return ()
+printSourceTextAA (NoSourceText) txt   = printStringAdvanceA  txt >> return ()
+printSourceTextAA (SourceText   txt) _ = printStringAdvanceA  (unpackFS txt) >> return ()
 
 -- ---------------------------------------------------------------------
 
-printStringAtSs :: (Monad m, Monoid w) => SrcSpan -> String -> EP w m ()
-printStringAtSs ss str = printStringAtRs (realSrcSpan ss) str >> return ()
-
 printStringAtRs :: (Monad m, Monoid w) => RealSrcSpan -> String -> EP w m EpaLocation
 printStringAtRs pa str = printStringAtRsC CaptureComments pa str
 
@@ -676,7 +704,7 @@ printStringAtRsC capture pa str = do
   p' <- adjustDeltaForOffsetM p
   debugM $ "printStringAtRsC:(p,p')=" ++ show (p,p')
   printStringAtLsDelta p' str
-  setPriorEndASTD True pa
+  setPriorEndASTD pa
   cs' <- case capture of
     CaptureComments -> takeAppliedComments
     NoCaptureComments -> return []
@@ -709,6 +737,9 @@ printStringAtMLocL (EpAnn anc an cs) l s = do
       printStringAtLsDelta (SameLine 1) str
       return (Just (EpaDelta noSrcSpan (SameLine 1) []))
 
+printStringAdvanceA :: (Monad m, Monoid w) => String -> EP w m ()
+printStringAdvanceA str = printStringAtAA (EpaDelta noSrcSpan (SameLine 0) []) str >> return ()
+
 printStringAtAA :: (Monad m, Monoid w) => EpaLocation -> String -> EP w m EpaLocation
 printStringAtAA el str = printStringAtAAC CaptureComments el str
 
@@ -735,7 +766,7 @@ printStringAtAAC capture (EpaDelta ss d cs) s = do
   p2 <- getPosP
   pe2 <- getPriorEndD
   debugM $ "printStringAtAA:(pe1,pe2,p1,p2)=" ++ show (pe1,pe2,p1,p2)
-  setPriorEndASTPD True (pe1,pe2)
+  setPriorEndASTPD (pe1,pe2)
   cs' <- case capture of
     CaptureComments -> takeAppliedComments
     NoCaptureComments -> return []
@@ -883,8 +914,7 @@ markAnnOpenP' :: (Monad m, Monoid w) => AnnPragma -> SourceText -> String -> EP
 markAnnOpenP' an NoSourceText txt   = markEpAnnLMS0 an lapr_open AnnOpen (Just txt)
 markAnnOpenP' an (SourceText txt) _ = markEpAnnLMS0 an lapr_open AnnOpen (Just $ unpackFS txt)
 
-markAnnOpen :: (Monad m, Monoid w)
-  => [AddEpAnn] -> SourceText -> String -> EP w m [AddEpAnn]
+markAnnOpen :: (Monad m, Monoid w) => [AddEpAnn] -> SourceText -> String -> EP w m [AddEpAnn]
 markAnnOpen an NoSourceText txt   = markEpAnnLMS'' an lidl AnnOpen (Just txt)
 markAnnOpen an (SourceText txt) _ = markEpAnnLMS'' an lidl AnnOpen (Just $ unpackFS txt)
 
@@ -1589,7 +1619,7 @@ markTopLevelList ls = mapM (\a -> setLayoutTopLevelP $ markAnnotated a) ls
 instance (ExactPrint a) => ExactPrint (Located a) where
   getAnnotationEntry (L l _) = case l of
     UnhelpfulSpan _ -> NoEntryVal
-    _ -> Entry (hackSrcSpanToAnchor l) [] emptyComments NoFlushComments CanUpdateAnchorOnly
+    _ -> Entry (EpaSpan l) [] emptyComments NoFlushComments CanUpdateAnchorOnly
 
   setAnnotationAnchor (L l a) _anc _ts _cs = L l a
 
@@ -1664,16 +1694,10 @@ instance ExactPrint (HsModule GhcPs) where
         _ -> return lo
 
     am_decls' <- markTrailing (am_decls $ anns an0)
-    imports' <- markTopLevelList imports
-
-    case lo of
-        EpExplicitBraces _ _ -> return ()
-        _ -> do
-          -- Get rid of the balance of the preceding comments before starting on the decls
-          flushComments []
-          putUnallocatedComments []
 
-    decls' <- markTopLevelList (filter removeDocDecl decls)
+    mid <- markAnnotated (HsModuleImpDecls (am_cs $ anns an0) imports decls)
+    let imports' = id_imps mid
+    let decls' = id_decls mid
 
     lo1 <- case lo0 of
         EpExplicitBraces open close -> do
@@ -1688,15 +1712,32 @@ instance ExactPrint (HsModule GhcPs) where
         debugM $ "am_eof:" ++ showGhc (pos, prior)
         setEofPos (Just (pos, prior))
 
-    let anf = an0 { anns = (anns an0) { am_decls = am_decls' }}
+    let anf = an0 { anns = (anns an0) { am_decls = am_decls', am_cs = [] }}
     debugM $ "HsModule, anf=" ++ showAst anf
 
     return (HsModule (XModulePs anf lo1 mdeprec' mbDoc') mmn' mexports' imports' decls')
 
+-- ---------------------------------------------------------------------
+
+-- | This is used to ensure the comments are updated into the right
+-- place for makeDeltaAst.
+data HsModuleImpDecls
+    = HsModuleImpDecls {
+        id_cs     :: [LEpaComment],
+        id_imps   :: [LImportDecl GhcPs],
+        id_decls  :: [LHsDecl GhcPs]
+    } deriving Data
+
+instance ExactPrint HsModuleImpDecls where
+  -- Use an UnhelpfulSpan for the anchor, we are only interested in the comments
+  getAnnotationEntry mid = mkEntry (EpaSpan (UnhelpfulSpan UnhelpfulNoLocationInfo)) [] (EpaComments (id_cs mid))
+  setAnnotationAnchor mid _anc _ cs = mid { id_cs = priorComments cs ++ getFollowingComments cs }
+     `debug` ("HsModuleImpDecls.setAnnotationAnchor:cs=" ++ showAst cs)
+  exact (HsModuleImpDecls cs imports decls) = do
+    imports' <- markTopLevelList imports
+    decls' <- markTopLevelList (filter notDocDecl decls)
+    return (HsModuleImpDecls cs imports' decls')
 
-removeDocDecl :: LHsDecl GhcPs -> Bool
-removeDocDecl (L _ DocD{}) = False
-removeDocDecl _ = True
 
 -- ---------------------------------------------------------------------
 
@@ -1737,8 +1778,8 @@ instance ExactPrint InWarningCategory where
 
   exact (InWarningCategory tkIn source (L l wc)) = do
       tkIn' <- markEpToken tkIn
-      L _ (_,wc') <- markAnnotated (L l (source, wc))
-      return (InWarningCategory tkIn' source (L l wc'))
+      L l' (_,wc') <- markAnnotated (L l (source, wc))
+      return (InWarningCategory tkIn' source (L l' wc'))
 
 instance ExactPrint (SourceText, WarningCategory) where
   getAnnotationEntry _ = NoEntryVal
@@ -1943,14 +1984,14 @@ exactDataFamInstDecl an top_lvl
                            , feqn_pats   = pats
                            , feqn_fixity = fixity
                            , feqn_rhs    = defn })) = do
-    (an', an2', tycon', bndrs', _,  _mc, defn') <- exactDataDefn an2 pp_hdr defn
-                                                 -- See Note [an and an2 in exactDataFamInstDecl]
+    (an', an2', tycon', bndrs', pats', defn') <- exactDataDefn an2 pp_hdr defn
+                                          -- See Note [an and an2 in exactDataFamInstDecl]
     return
       (an',
        DataFamInstDecl ( FamEqn { feqn_ext    = an2'
                                 , feqn_tycon  = tycon'
                                 , feqn_bndrs  = bndrs'
-                                , feqn_pats   = pats
+                                , feqn_pats   = pats'
                                 , feqn_fixity = fixity
                                 , feqn_rhs    = defn' }))
                     `debug` ("exactDataFamInstDecl: defn' derivs:" ++ showAst (dd_derivs defn'))
@@ -2233,11 +2274,11 @@ instance ExactPrint (RoleAnnotDecl GhcPs) where
     an1 <- markEpAnnL an0 lidl AnnRole
     ltycon' <- markAnnotated ltycon
     let markRole (L l (Just r)) = do
-          (L _ r') <- markAnnotated (L l r)
-          return (L l (Just r'))
+          (L l' r') <- markAnnotated (L l r)
+          return (L l' (Just r'))
         markRole (L l Nothing) = do
-          printStringAtSs (locA l) "_"
-          return (L l Nothing)
+          e' <- printStringAtAA  (entry l) "_"
+          return (L (l { entry = e'}) Nothing)
     roles' <- mapM markRole roles
     return (RoleAnnotDecl an1 ltycon' roles')
 
@@ -2340,8 +2381,13 @@ instance (ExactPrint tm, ExactPrint ty, Outputable tm, Outputable ty)
   getAnnotationEntry = const NoEntryVal
   setAnnotationAnchor a _ _ _ = a
 
-  exact a@(HsValArg _ tm)   = markAnnotated tm >> return a
-  exact a@(HsTypeArg at ty) = markEpToken at >> markAnnotated ty >> return a
+  exact (HsValArg x tm) = do
+      tm' <- markAnnotated tm
+      return (HsValArg x tm')
+  exact (HsTypeArg at ty) = do
+      at' <- markEpToken at
+      ty' <- markAnnotated ty
+      return (HsTypeArg at' ty')
   exact x@(HsArgPar _sp)    = withPpr x -- Does not appear in original source
 
 -- ---------------------------------------------------------------------
@@ -2359,9 +2405,9 @@ instance ExactPrint (ClsInstDecl GhcPs) where
           (mbWarn', an0, mbOverlap', inst_ty') <- top_matter
           an1 <- markEpAnnL an0 lidl AnnOpenC
           an2 <- markEpAnnAllL' an1 lid AnnSemi
-          ds <- withSortKey sortKey
-                               [(ClsAtdTag, prepareListAnnotationA ats),
-                                (ClsAtdTag, prepareListAnnotationF an adts),
+          (sortKey', ds) <- withSortKey sortKey
+                               [(ClsAtTag, prepareListAnnotationA ats),
+                                (ClsAtdTag, prepareListAnnotationF adts),
                                 (ClsMethodTag, prepareListAnnotationA binds),
                                 (ClsSigTag, prepareListAnnotationA sigs)
                                ]
@@ -2371,7 +2417,7 @@ instance ExactPrint (ClsInstDecl GhcPs) where
             adts'  = undynamic ds
             binds' = undynamic ds
             sigs'  = undynamic ds
-          return (ClsInstDecl { cid_ext = (mbWarn', an3, sortKey)
+          return (ClsInstDecl { cid_ext = (mbWarn', an3, sortKey')
                               , cid_poly_ty = inst_ty', cid_binds = binds'
                               , cid_sigs = sigs', cid_tyfam_insts = ats'
                               , cid_overlap_mode = mbOverlap'
@@ -2452,15 +2498,29 @@ instance ExactPrint (HsBind GhcPs) where
     return (FunBind x fun_id' matches')
 
   exact (PatBind x pat q grhss) = do
+    q' <- markAnnotated q
     pat' <- markAnnotated pat
     grhss' <- markAnnotated grhss
-    return (PatBind x pat' q grhss')
+    return (PatBind x pat' q' grhss')
   exact (PatSynBind x bind) = do
     bind' <- markAnnotated bind
     return (PatSynBind x bind')
 
   exact x = error $ "HsBind: exact for " ++ showAst x
 
+instance ExactPrint (HsMultAnn GhcPs) where
+  getAnnotationEntry _ = NoEntryVal
+  setAnnotationAnchor a _ _ _ = a
+
+  exact (HsNoMultAnn x) = return (HsNoMultAnn x)
+  exact (HsPct1Ann tok) = do
+      tok' <- markEpToken tok
+      return (HsPct1Ann tok')
+  exact (HsMultAnn tok ty) = do
+      tok' <- markEpToken tok
+      ty' <- markAnnotated ty
+      return (HsMultAnn tok' ty')
+
 -- ---------------------------------------------------------------------
 
 instance ExactPrint (PatSynBind GhcPs GhcPs) where
@@ -2519,8 +2579,9 @@ instance ExactPrint (PatSynBind GhcPs GhcPs) where
 instance ExactPrint (RecordPatSynField GhcPs) where
   getAnnotationEntry = const NoEntryVal
   setAnnotationAnchor a _ _ _ = a
-  exact r@(RecordPatSynField { recordPatSynField = v }) = markAnnotated v
-        >> return r
+  exact (RecordPatSynField f v) = do
+      f' <- markAnnotated f
+      return (RecordPatSynField f' v)
 
 -- ---------------------------------------------------------------------
 
@@ -2648,15 +2709,20 @@ instance ExactPrint (HsLocalBinds GhcPs) where
 
     (an1, valbinds') <- markAnnList an0 $ markAnnotatedWithLayout valbinds
     debugM $ "exact HsValBinds: an1=" ++ showAst an1
-    return (HsValBinds an1 valbinds')
+    medr <- getExtraDPReturn
+    an2 <- case medr of
+             Nothing -> return an1
+             Just dp -> do
+                 setExtraDPReturn Nothing
+                 return $ an1 { anns = (anns an1) { al_anchor = Just (EpaDelta noSrcSpan dp []) }}
+    return (HsValBinds an2 valbinds')
 
   exact (HsIPBinds an bs) = do
-    (as, ipb) <- markAnnList an (markEpAnnL' an lal_rest AnnWhere
-                           >> markAnnotated bs
-                           >>= \bs' -> return (HsIPBinds an bs'::HsLocalBinds GhcPs))
-    case ipb of
-      HsIPBinds _ bs' -> return (HsIPBinds as bs'::HsLocalBinds GhcPs)
-      _ -> error "should not happen HsIPBinds"
+    (an2,bs') <- markAnnListA an $ \an0 -> do
+                           an1 <- markEpAnnL' an0 lal_rest AnnWhere
+                           bs' <- markAnnotated bs
+                           return (an1, bs')
+    return (HsIPBinds an2 bs')
   exact b@(EmptyLocalBinds _) = return b
 
 
@@ -2670,7 +2736,8 @@ instance ExactPrint (HsValBindsLR GhcPs GhcPs) where
     let
       binds' = concatMap decl2Bind decls
       sigs'  = concatMap decl2Sig decls
-    return (ValBinds sortKey binds' sigs')
+      sortKey' = captureOrderBinds decls
+    return (ValBinds sortKey' binds' sigs')
   exact (XValBindsLR _) = panic "XValBindsLR"
 
 undynamic :: Typeable a => [Dynamic] -> [a]
@@ -2682,7 +2749,9 @@ instance ExactPrint (HsIPBinds GhcPs) where
   getAnnotationEntry = const NoEntryVal
   setAnnotationAnchor a _ _ _ = a
 
-  exact b@(IPBinds _ binds) = setLayoutBoth $ markAnnotated binds >> return b
+  exact (IPBinds x binds) = setLayoutBoth $ do
+      binds' <- markAnnotated binds
+      return (IPBinds x binds')
 
 -- ---------------------------------------------------------------------
 
@@ -2703,18 +2772,18 @@ instance ExactPrint HsIPName where
   getAnnotationEntry = const NoEntryVal
   setAnnotationAnchor a _ _ _ = a
 
-  exact i@(HsIPName fs) = printStringAdvance ("?" ++ (unpackFS fs)) >> return i
+  exact i@(HsIPName fs) = printStringAdvanceA ("?" ++ (unpackFS fs)) >> return i
 
 -- ---------------------------------------------------------------------
 -- Managing lists which have been separated, e.g. Sigs and Binds
 
 prepareListAnnotationF :: (Monad m, Monoid w) =>
-  [AddEpAnn] -> [LDataFamInstDecl GhcPs] -> [(RealSrcSpan,EP w m Dynamic)]
-prepareListAnnotationF an ls = map (\b -> (realSrcSpan $ getLocA b, go b)) ls
+  [LDataFamInstDecl GhcPs] -> [(RealSrcSpan,EP w m Dynamic)]
+prepareListAnnotationF ls = map (\b -> (realSrcSpan $ getLocA b, go b)) ls
   where
     go (L l a) = do
-      d' <- markAnnotated (DataFamInstDeclWithContext an NotTopLevel a)
-      return (toDyn (L l (dc_d d')))
+      (L l' d') <- markAnnotated (L l (DataFamInstDeclWithContext noAnn NotTopLevel a))
+      return (toDyn (L l' (dc_d d')))
 
 prepareListAnnotationA :: (Monad m, Monoid w, ExactPrint (LocatedAn an a))
   => [LocatedAn an a] -> [(RealSrcSpan,EP w m Dynamic)]
@@ -2725,15 +2794,23 @@ prepareListAnnotationA ls = map (\b -> (realSrcSpan $ getLocA b,go b)) ls
       return (toDyn b')
 
 withSortKey :: (Monad m, Monoid w)
-  => AnnSortKey DeclTag -> [(DeclTag, [(RealSrcSpan, EP w m Dynamic)])] -> EP w m [Dynamic]
+  => AnnSortKey DeclTag -> [(DeclTag, [(RealSrcSpan, EP w m Dynamic)])]
+  -> EP w m (AnnSortKey DeclTag, [Dynamic])
 withSortKey annSortKey xs = do
   debugM $ "withSortKey:annSortKey=" ++ showAst annSortKey
-  let ordered = case annSortKey of
-                  NoAnnSortKey -> sortBy orderByFst $ concatMap snd xs
-                  AnnSortKey _keys -> orderedDecls annSortKey (Map.fromList xs)
-  mapM snd ordered
-orderByFst :: Ord a => (a, b1) -> (a, b2) -> Ordering
-orderByFst (a,_) (b,_) = compare a b
+  let (sk, ordered) = case annSortKey of
+                  NoAnnSortKey -> (annSortKey', map snd os)
+                    where
+                      doOne (tag, ds) = map (\d -> (tag, d)) ds
+                      xsExpanded = concatMap doOne xs
+                      os = sortBy orderByFst $ xsExpanded
+                      annSortKey' = AnnSortKey (map fst os)
+                  AnnSortKey _keys -> (annSortKey, orderedDecls annSortKey (Map.fromList xs))
+  ordered' <- mapM snd ordered
+  return (sk, ordered')
+
+orderByFst :: Ord a => (t, (a,b1)) -> (t, (a, b2)) -> Ordering
+orderByFst (_,(a,_)) (_,(b,_)) = compare a b
 
 -- ---------------------------------------------------------------------
 
@@ -2761,15 +2838,16 @@ instance ExactPrint (Sig GhcPs) where
         (an0, vars',ty') <- exactVarSig an vars ty
         return (ClassOpSig an0 is_deflt vars' ty')
 
-  exact (FixSig (an,src) (FixitySig x names (Fixity v fdir))) = do
+  exact (FixSig (an,src) (FixitySig ns names (Fixity v fdir))) = do
     let fixstr = case fdir of
          InfixL -> "infixl"
          InfixR -> "infixr"
          InfixN -> "infix"
     an0 <- markEpAnnLMS'' an  lidl AnnInfix (Just fixstr)
     an1 <- markEpAnnLMS'' an0 lidl AnnVal (Just (sourceTextToString src (show v)))
+    ns' <- markAnnotated ns
     names' <- markAnnotated names
-    return (FixSig (an1,src) (FixitySig x names' (Fixity v fdir)))
+    return (FixSig (an1,src) (FixitySig ns' names' (Fixity v fdir)))
 
   exact (InlineSig an ln inl) = do
     an0 <- markAnnOpen an (inl_src inl) "{-# INLINE"
@@ -2809,7 +2887,7 @@ instance ExactPrint (Sig GhcPs) where
 
   exact (CompleteMatchSig (an,src) cs mty) = do
     an0 <- markAnnOpen an src "{-# COMPLETE"
-    cs' <- markAnnotated cs
+    cs' <- mapM markAnnotated cs
     (an1, mty') <-
       case mty of
         Nothing -> return (an0, mty)
@@ -2822,6 +2900,20 @@ instance ExactPrint (Sig GhcPs) where
 
 -- ---------------------------------------------------------------------
 
+instance ExactPrint NamespaceSpecifier where
+  getAnnotationEntry _ = NoEntryVal
+  setAnnotationAnchor a _ _ _ = a
+
+  exact NoNamespaceSpecifier = return NoNamespaceSpecifier
+  exact (TypeNamespaceSpecifier typeTok) = do
+      typeTok' <- markEpToken typeTok
+      return (TypeNamespaceSpecifier typeTok')
+  exact (DataNamespaceSpecifier dataTok) = do
+      dataTok' <- markEpToken dataTok
+      return (DataNamespaceSpecifier dataTok')
+
+-- ---------------------------------------------------------------------
+
 exactVarSig :: (Monad m, Monoid w, ExactPrint a)
   => AnnSig -> [LocatedN RdrName] -> a -> EP w m (AnnSig, [LocatedN RdrName], a)
 exactVarSig an vars ty = do
@@ -2875,7 +2967,7 @@ instance ExactPrint (AnnDecl GhcPs) where
           n' <- markAnnotated n
           return (an1, TypeAnnProvenance n')
         ModuleAnnProvenance -> do
-          an1 <- markEpAnnL an lapr_rest AnnModule
+          an1 <- markEpAnnL an0 lapr_rest AnnModule
           return (an1, prov)
 
     e' <- markAnnotated e
@@ -2950,21 +3042,21 @@ instance ExactPrint (HsExpr GhcPs) where
       then markAnnotated n
       else return n
     return (HsVar x n')
-  exact x@(HsUnboundVar an _) = do
+  exact (HsUnboundVar an n) = do
     case an of
       Just (EpAnnUnboundVar (ob,cb) l) -> do
-        printStringAtAA ob "`" >> return ()
-        printStringAtAA l  "_" >> return ()
-        printStringAtAA cb "`" >> return ()
-        return x
+        ob' <-  printStringAtAA ob "`"
+        l' <- printStringAtAA l  "_"
+        cb' <- printStringAtAA cb "`"
+        return (HsUnboundVar (Just (EpAnnUnboundVar (ob',cb') l')) n)
       _ -> do
-        printStringAtLsDelta (SameLine 0) "_"
-        return x
+        printStringAdvanceA "_" >> return ()
+        return (HsUnboundVar an n)
   exact x@(HsOverLabel src l) = do
-    printStringAtLsDelta (SameLine 0) "#"
+    printStringAdvanceA "#" >> return ()
     case src of
-      NoSourceText   -> printStringAtLsDelta (SameLine 0) (unpackFS l)
-      SourceText txt -> printStringAtLsDelta (SameLine 0) (unpackFS txt)
+      NoSourceText   -> printStringAdvanceA (unpackFS l)  >> return ()
+      SourceText txt -> printStringAdvanceA (unpackFS txt) >> return ()
     return x
 
   exact x@(HsIPVar _ (HsIPName n))
@@ -3204,11 +3296,11 @@ instance ExactPrint (HsExpr GhcPs) where
 
   exact (HsTypedSplice an s)   = do
     an0 <- markEpAnnL an lidl AnnDollarDollar
-    s' <- exact s
+    s' <- markAnnotated s
     return (HsTypedSplice an0 s')
 
   exact (HsUntypedSplice an s) = do
-    s' <- exact s
+    s' <- markAnnotated s
     return (HsUntypedSplice an s')
 
   exact (HsProc an p c) = do
@@ -3274,12 +3366,15 @@ exactMdo an (Just module_name) kw = markEpAnnLMS'' an lal_rest kw (Just n)
 markMaybeDodgyStmts :: (Monad m, Monoid w, ExactPrint (LocatedAn an a))
   => AnnList -> LocatedAn an a -> EP w m (AnnList, LocatedAn an a)
 markMaybeDodgyStmts an stmts =
-  if isGoodSrcSpan (getLocA stmts)
+  if notDodgy stmts
     then do
       r <- markAnnotatedWithLayout stmts
       return (an, r)
     else return (an, stmts)
 
+notDodgy :: GenLocated (EpAnn ann) a -> Bool
+notDodgy (L (EpAnn anc _ _) _) = notDodgyE anc
+
 notDodgyE :: EpaLocation -> Bool
 notDodgyE anc =
   case anc of
@@ -3341,7 +3436,7 @@ instance ExactPrint (MatchGroup GhcPs (LocatedA (HsCmd GhcPs))) where
   setAnnotationAnchor a _ _ _ = a
   exact (MG x matches) = do
     -- TODO:AZ use SortKey, in MG ann.
-    matches' <- if isGoodSrcSpan (getLocA matches)
+    matches' <- if notDodgy matches
       then markAnnotated matches
       else return matches
     return (MG x matches')
@@ -3661,6 +3756,7 @@ instance ExactPrint (TyClDecl GhcPs) where
     -- There may be arbitrary parens around parts of the constructor
     -- that are infix.  Turn these into comments so that they feed
     -- into the right place automatically
+    -- TODO: no longer sorting on insert. What now?
     an0 <- annotationsToComments an lidl [AnnOpenP,AnnCloseP]
     an1 <- markEpAnnL an0 lidl AnnType
 
@@ -3674,7 +3770,7 @@ instance ExactPrint (TyClDecl GhcPs) where
   -- TODO: add a workaround for https://gitlab.haskell.org/ghc/ghc/-/issues/20452
   exact (DataDecl { tcdDExt = an, tcdLName = ltycon, tcdTyVars = tyvars
                   , tcdFixity = fixity, tcdDataDefn = defn }) = do
-    (_, an', ltycon', tyvars', _, _mctxt', defn') <-
+    (_, an', ltycon', tyvars', _, defn') <-
       exactDataDefn an (exactVanillaDeclHead ltycon tyvars fixity) defn
     return (DataDecl { tcdDExt = an', tcdLName = ltycon', tcdTyVars = tyvars'
                      , tcdFixity = fixity, tcdDataDefn = defn' })
@@ -3707,7 +3803,7 @@ instance ExactPrint (TyClDecl GhcPs) where
           (an0, fds', lclas', tyvars',context') <- top_matter
           an1 <- markEpAnnL    an0 lidl AnnOpenC
           an2 <- markEpAnnAllL' an1 lidl AnnSemi
-          ds <- withSortKey sortKey
+          (sortKey', ds) <- withSortKey sortKey
                                [(ClsSigTag, prepareListAnnotationA sigs),
                                 (ClsMethodTag, prepareListAnnotationA methods),
                                 (ClsAtTag, prepareListAnnotationA ats),
@@ -3720,7 +3816,7 @@ instance ExactPrint (TyClDecl GhcPs) where
             methods' = undynamic ds
             ats'     = undynamic ds
             at_defs' = undynamic ds
-          return (ClassDecl {tcdCExt = (an3, lo, sortKey),
+          return (ClassDecl {tcdCExt = (an3, lo, sortKey'),
                              tcdCtxt = context', tcdLName = lclas', tcdTyVars = tyvars',
                              tcdFixity = fixity,
                              tcdFDs  = fds',
@@ -3845,7 +3941,7 @@ exactDataDefn
   -> HsDataDefn GhcPs
   -> EP w m ( [AddEpAnn] -- ^ from exactHdr
             , [AddEpAnn] -- ^ updated one passed in
-            , LocatedN RdrName, a, b, Maybe (LHsContext GhcPs), HsDataDefn GhcPs)
+            , LocatedN RdrName, a, b, HsDataDefn GhcPs)
 exactDataDefn an exactHdr
                  (HsDataDefn { dd_ext = x, dd_ctxt = context
                              , dd_cType = mb_ct
@@ -3883,8 +3979,8 @@ exactDataDefn an exactHdr
           _ -> panic "exacprint NewTypeCon"
   an6 <- markEpAnnL an5 lidl AnnCloseC
   derivings' <- mapM markAnnotated derivings
-  return (anx, an6, ln', tvs', b, mctxt',
-                 (HsDataDefn { dd_ext = x, dd_ctxt = context
+  return (anx, an6, ln', tvs', b,
+                 (HsDataDefn { dd_ext = x, dd_ctxt = mctxt'
                              , dd_cType = mb_ct'
                              , dd_kindSig = mb_sig'
                              , dd_cons = condecls'', dd_derivs = derivings' }))
@@ -3941,22 +4037,23 @@ instance ExactPrint (InjectivityAnn GhcPs) where
 
 class Typeable flag => ExactPrintTVFlag flag where
   exactTVDelimiters :: (Monad m, Monoid w)
-    => [AddEpAnn] -> flag -> EP w m (HsTyVarBndr flag GhcPs)
-    -> EP w m ([AddEpAnn], (HsTyVarBndr flag GhcPs))
+    => [AddEpAnn] -> flag
+    -> ([AddEpAnn] -> EP w m ([AddEpAnn], HsTyVarBndr flag GhcPs))
+    -> EP w m ([AddEpAnn], flag, (HsTyVarBndr flag GhcPs))
 
 instance ExactPrintTVFlag () where
-  exactTVDelimiters an _ thing_inside = do
+  exactTVDelimiters an flag thing_inside = do
     an0 <- markEpAnnAllL' an lid AnnOpenP
-    r <- thing_inside
-    an1 <- markEpAnnAllL' an0 lid AnnCloseP
-    return (an1, r)
+    (an1, r) <- thing_inside an0
+    an2 <- markEpAnnAllL' an1 lid AnnCloseP
+    return (an2, flag, r)
 
 instance ExactPrintTVFlag Specificity where
   exactTVDelimiters an s thing_inside = do
     an0 <- markEpAnnAllL' an lid open
-    r <- thing_inside
-    an1 <- markEpAnnAllL' an0 lid close
-    return (an1, r)
+    (an1, r) <- thing_inside an0
+    an2 <- markEpAnnAllL' an1 lid close
+    return (an2, s, r)
     where
       (open, close) = case s of
         SpecifiedSpec -> (AnnOpenP, AnnCloseP)
@@ -3964,33 +4061,33 @@ instance ExactPrintTVFlag Specificity where
 
 instance ExactPrintTVFlag (HsBndrVis GhcPs) where
   exactTVDelimiters an0 bvis thing_inside = do
-    case bvis of
-      HsBndrRequired _ -> return ()
-      HsBndrInvisible at -> markEpToken at >> return ()
+    bvis' <- case bvis of
+      HsBndrRequired _ -> return bvis
+      HsBndrInvisible at -> HsBndrInvisible <$> markEpToken at
     an1 <- markEpAnnAllL' an0 lid AnnOpenP
-    r <- thing_inside
-    an2 <- markEpAnnAllL' an1 lid AnnCloseP
-    return (an2, r)
+    (an2, r) <- thing_inside an1
+    an3 <- markEpAnnAllL' an2 lid AnnCloseP
+    return (an3, bvis', r)
 
 instance ExactPrintTVFlag flag => ExactPrint (HsTyVarBndr flag GhcPs) where
   getAnnotationEntry _ = NoEntryVal
   setAnnotationAnchor a _ _ _ = a
 
   exact (UserTyVar an flag n) = do
-    r <- exactTVDelimiters an flag $ do
+    r <- exactTVDelimiters an flag $ \ani -> do
            n' <- markAnnotated n
-           return (UserTyVar an flag n')
+           return (ani, UserTyVar an flag n')
     case r of
-      (an', UserTyVar _ flag'' n'') -> return (UserTyVar an' flag'' n'')
+      (an', flag', UserTyVar _ _ n'') -> return (UserTyVar an' flag' n'')
       _ -> error "KindedTyVar should never happen here"
   exact (KindedTyVar an flag n k) = do
-    r <- exactTVDelimiters an flag $ do
+    r <- exactTVDelimiters an flag $ \ani -> do
           n' <- markAnnotated n
-          an0 <- markEpAnnL an lidl AnnDcolon
+          an0 <- markEpAnnL ani lidl AnnDcolon
           k' <- markAnnotated k
-          return (KindedTyVar an0 flag n' k')
+          return (an0, KindedTyVar an0 flag n' k')
     case r of
-      (an',KindedTyVar _ flag'' n'' k'') -> return (KindedTyVar an' flag'' n'' k'')
+      (an',flag', KindedTyVar _ _ n'' k'') -> return (KindedTyVar an' flag' n'' k'')
       _ -> error "UserTyVar should never happen here"
 
 -- ---------------------------------------------------------------------
@@ -4150,17 +4247,16 @@ instance ExactPrint (HsDerivingClause GhcPs) where
                           , deriv_clause_strategy = dcs
                           , deriv_clause_tys      = dct }) = do
     an0 <- markEpAnnL an lidl AnnDeriving
-    exact_strat_before
+    dcs0 <- case dcs of
+            Just (L _ ViaStrategy{}) -> return dcs
+            _ -> mapM markAnnotated dcs
     dct' <- markAnnotated dct
-    exact_strat_after
+    dcs1 <- case dcs0 of
+            Just (L _ ViaStrategy{}) -> mapM markAnnotated dcs0
+            _ -> return dcs0
     return (HsDerivingClause { deriv_clause_ext      = an0
-                             , deriv_clause_strategy = dcs
+                             , deriv_clause_strategy = dcs1
                              , deriv_clause_tys      = dct' })
-      where
-        (exact_strat_before, exact_strat_after) =
-          case dcs of
-            Just v@(L _ ViaStrategy{}) -> (pure (), markAnnotated v >> pure ())
-            _                          -> (mapM_ markAnnotated dcs, pure ())
 
 -- ---------------------------------------------------------------------
 
@@ -4467,7 +4563,9 @@ instance ExactPrint (ConDeclField GhcPs) where
 instance ExactPrint (FieldOcc GhcPs) where
   getAnnotationEntry = const NoEntryVal
   setAnnotationAnchor a _ _ _ = a
-  exact f@(FieldOcc _ n) = markAnnotated n >> return f
+  exact (FieldOcc x n) = do
+      n' <- markAnnotated n
+      return (FieldOcc x n')
 
 -- ---------------------------------------------------------------------
 
@@ -4535,7 +4633,7 @@ instance ExactPrint (LocatedL [LocatedA (IE GhcPs)]) where
     an0 <- markEpAnnL' an lal_rest AnnHiding
     p <- getPosP
     debugM $ "LocatedL [LIE:p=" ++ showPprUnsafe p
-    (an1, ies') <- markAnnList an0 (markAnnotated ies)
+    (an1, ies') <- markAnnList an0 (markAnnotated (filter notIEDoc ies))
     return (L an1 ies')
 
 instance (ExactPrint (Match GhcPs (LocatedA body)))
@@ -4985,6 +5083,14 @@ setExtraDP md = do
   debugM $ "setExtraDP:" ++ show md
   modify (\s -> s {uExtraDP = md})
 
+getExtraDPReturn :: (Monad m, Monoid w) => EP w m (Maybe DeltaPos)
+getExtraDPReturn = gets uExtraDPReturn
+
+setExtraDPReturn :: (Monad m, Monoid w) => Maybe DeltaPos -> EP w m ()
+setExtraDPReturn md = do
+  debugM $ "setExtraDPReturn:" ++ show md
+  modify (\s -> s {uExtraDPReturn = md})
+
 getPriorEndD :: (Monad m, Monoid w) => EP w m Pos
 getPriorEndD = gets dPriorEndPosition
 
@@ -5007,13 +5113,13 @@ setPriorEndNoLayoutD pe = do
   debugM $ "setPriorEndNoLayoutD:pe=" ++ show pe
   modify (\s -> s { dPriorEndPosition = pe })
 
-setPriorEndASTD :: (Monad m, Monoid w) => Bool -> RealSrcSpan -> EP w m ()
-setPriorEndASTD layout pe = setPriorEndASTPD layout (rs2range pe)
+setPriorEndASTD :: (Monad m, Monoid w) => RealSrcSpan -> EP w m ()
+setPriorEndASTD pe = setPriorEndASTPD (rs2range pe)
 
-setPriorEndASTPD :: (Monad m, Monoid w) => Bool -> (Pos,Pos) -> EP w m ()
-setPriorEndASTPD layout pe@(fm,to) = do
+setPriorEndASTPD :: (Monad m, Monoid w) => (Pos,Pos) -> EP w m ()
+setPriorEndASTPD pe@(fm,to) = do
   debugM $ "setPriorEndASTD:pe=" ++ show pe
-  when layout $ setLayoutStartD (snd fm)
+  setLayoutStartD (snd fm)
   modify (\s -> s { dPriorEndPosition = to } )
 
 setLayoutStartD :: (Monad m, Monoid w) => Int -> EP w m ()
@@ -5044,7 +5150,7 @@ getUnallocatedComments :: (Monad m, Monoid w) => EP w m [Comment]
 getUnallocatedComments = gets epComments
 
 putUnallocatedComments :: (Monad m, Monoid w) => [Comment] -> EP w m ()
-putUnallocatedComments cs = modify (\s -> s { epComments = cs } )
+putUnallocatedComments !cs = modify (\s -> s { epComments = cs } )
 
 -- | Push a fresh stack frame for the applied comments gatherer
 pushAppliedComments  :: (Monad m, Monoid w) => EP w m ()
@@ -5054,7 +5160,7 @@ pushAppliedComments = modify (\s -> s { epCommentsApplied = []:(epCommentsApplie
 -- takeAppliedComments, and clear them, not popping the stack
 takeAppliedComments :: (Monad m, Monoid w) => EP w m [Comment]
 takeAppliedComments = do
-  ccs <- gets epCommentsApplied
+  !ccs <- gets epCommentsApplied
   case ccs of
     [] -> do
       modify (\s -> s { epCommentsApplied = [] })
@@ -5067,7 +5173,7 @@ takeAppliedComments = do
 -- takeAppliedComments, and clear them, popping the stack
 takeAppliedCommentsPop :: (Monad m, Monoid w) => EP w m [Comment]
 takeAppliedCommentsPop = do
-  ccs <- gets epCommentsApplied
+  !ccs <- gets epCommentsApplied
   case ccs of
     [] -> do
       modify (\s -> s { epCommentsApplied = [] })
@@ -5080,7 +5186,7 @@ takeAppliedCommentsPop = do
 -- when doing delta processing
 applyComment :: (Monad m, Monoid w) => Comment -> EP w m ()
 applyComment c = do
-  ccs <- gets epCommentsApplied
+  !ccs <- gets epCommentsApplied
   case ccs of
     []    -> modify (\s -> s { epCommentsApplied = [[c]] } )
     (h:t) -> modify (\s -> s { epCommentsApplied = (c:h):t } )


=====================================
utils/check-exact/Main.hs
=====================================
@@ -470,7 +470,7 @@ changeAddDecl1 libdir top = do
   let (p',_,_) = runTransform doAddDecl
       doAddDecl = everywhereM (mkM replaceTopLevelDecls) top
       replaceTopLevelDecls :: ParsedSource -> Transform ParsedSource
-      replaceTopLevelDecls m = insertAtStart m decl'
+      replaceTopLevelDecls m = return $ insertAtStart m decl'
   return p'
 
 -- ---------------------------------------------------------------------
@@ -483,7 +483,7 @@ changeAddDecl2 libdir top = do
   let (p',_,_) = runTransform doAddDecl
       doAddDecl = everywhereM (mkM replaceTopLevelDecls) top
       replaceTopLevelDecls :: ParsedSource -> Transform ParsedSource
-      replaceTopLevelDecls m = insertAtEnd m decl'
+      replaceTopLevelDecls m = return $ insertAtEnd m decl'
   return p'
 
 -- ---------------------------------------------------------------------
@@ -500,7 +500,7 @@ changeAddDecl3 libdir top = do
           l2' = setEntryDP l2 (DifferentLine 2 0)
 
       replaceTopLevelDecls :: ParsedSource -> Transform ParsedSource
-      replaceTopLevelDecls m = insertAt f m decl'
+      replaceTopLevelDecls m = return $ insertAt f m decl'
   return p'
 
 -- ---------------------------------------------------------------------
@@ -571,8 +571,9 @@ changeLocalDecls2 libdir (L l p) = do
 changeWhereIn3a :: Changer
 changeWhereIn3a _libdir (L l p) = do
   let decls0 = hsmodDecls p
-      (decls,_,w) = runTransform (balanceCommentsList decls0)
-  debugM $ unlines w
+      decls = balanceCommentsList decls0
+      (_de0:_:de1:_d2:_) = decls
+  debugM $ "changeWhereIn3a:de1:" ++ showAst de1
   let p2 = p { hsmodDecls = decls}
   return (L l p2)
 
@@ -581,13 +582,12 @@ changeWhereIn3a _libdir (L l p) = do
 changeWhereIn3b :: Changer
 changeWhereIn3b _libdir (L l p) = do
   let decls0 = hsmodDecls p
-      (decls,_,w) = runTransform (balanceCommentsList decls0)
+      decls = balanceCommentsList decls0
       (de0:tdecls@(_:de1:d2:_)) = decls
       de0' = setEntryDP de0 (DifferentLine 2 0)
       de1' = setEntryDP de1 (DifferentLine 2 0)
       d2' = setEntryDP d2 (DifferentLine 2 0)
       decls' = d2':de1':de0':tdecls
-  debugM $ unlines w
   debugM $ "changeWhereIn3b:de1':" ++ showAst de1'
   let p2 = p { hsmodDecls = decls'}
   return (L l p2)
@@ -598,37 +598,37 @@ addLocaLDecl1 :: Changer
 addLocaLDecl1 libdir top = do
   Right (L ld (ValD _ decl)) <- withDynFlags libdir (\df -> parseDecl df "decl" "nn = 2")
   let decl' = setEntryDP (L ld decl) (DifferentLine 1 5)
-      doAddLocal = do
-        let lp = top
-        (de1:d2:d3:_) <- hsDecls lp
-        (de1'',d2') <- balanceComments de1 d2
-        (de1',_) <- modifyValD (getLocA de1'') de1'' $ \_m d -> do
-          return ((wrapDecl decl' : d),Nothing)
-        replaceDecls lp [de1', d2', d3]
-
-  (lp',_,w) <- runTransformT doAddLocal
-  debugM $ "addLocaLDecl1:" ++ intercalate "\n" w
+      doAddLocal :: ParsedSource
+      doAddLocal = replaceDecls lp [de1', d2', d3]
+        where
+          lp = top
+          (de1:d2:d3:_) = hsDecls lp
+          (de1'',d2') = balanceComments de1 d2
+          (de1',_) = modifyValD (getLocA de1'') de1'' $ \_m d -> ((wrapDecl decl' : d),Nothing)
+
+  let lp' = doAddLocal
   return lp'
 
 -- ---------------------------------------------------------------------
 
+
 addLocaLDecl2 :: Changer
 addLocaLDecl2 libdir lp = do
   Right newDecl <- withDynFlags libdir (\df -> parseDecl df "decl" "nn = 2")
   let
-      doAddLocal = do
-         (de1:d2:_) <- hsDecls lp
-         (de1'',d2') <- balanceComments de1 d2
+      doAddLocal = replaceDecls lp [parent',d2']
+        where
+         (de1:d2:_) = hsDecls lp
+         (de1'',d2') = balanceComments de1 d2
 
-         (parent',_) <- modifyValD (getLocA de1) de1'' $ \_m (d:ds) -> do
-           newDecl' <- transferEntryDP' d newDecl
-           let d' = setEntryDP d (DifferentLine 1 0)
-           return ((newDecl':d':ds),Nothing)
+         (parent',_) = modifyValD (getLocA de1) de1'' $ \_m (d:ds) ->
+             let
+               newDecl' = transferEntryDP' d (makeDeltaAst newDecl)
+               d' = setEntryDP d (DifferentLine 1 0)
+             in ((newDecl':d':ds),Nothing)
 
-         replaceDecls lp [parent',d2']
 
-  (lp',_,_w) <- runTransformT doAddLocal
-  debugM $ "log:[\n" ++ intercalate "\n" _w ++ "]log end\n"
+      lp' = doAddLocal
   return lp'
 
 -- ---------------------------------------------------------------------
@@ -637,19 +637,18 @@ addLocaLDecl3 :: Changer
 addLocaLDecl3 libdir top = do
   Right newDecl <- withDynFlags libdir (\df -> parseDecl df "decl" "nn = 2")
   let
-      doAddLocal = do
-         let lp = top
-         (de1:d2:_) <- hsDecls lp
-         (de1'',d2') <- balanceComments de1 d2
-
-         (parent',_) <- modifyValD (getLocA de1) de1'' $ \_m (d:ds) -> do
-           let newDecl' = setEntryDP newDecl (DifferentLine 1 0)
-           return (((d:ds) ++ [newDecl']),Nothing)
+      doAddLocal = replaceDecls (anchorEof lp) [parent',d2']
+        where
+         lp = top
+         (de1:d2:_) = hsDecls lp
+         (de1'',d2') = balanceComments de1 d2
 
-         replaceDecls (anchorEof lp) [parent',d2']
+         (parent',_) = modifyValD (getLocA de1) de1'' $ \_m (d:ds) ->
+           let
+             newDecl' = setEntryDP newDecl (DifferentLine 1 0)
+           in (((d:ds) ++ [newDecl']),Nothing)
 
-  (lp',_,_w) <- runTransformT doAddLocal
-  debugM $ "log:[\n" ++ intercalate "\n" _w ++ "]log end\n"
+      lp' = doAddLocal
   return lp'
 
 -- ---------------------------------------------------------------------
@@ -659,40 +658,38 @@ addLocaLDecl4 libdir lp = do
   Right newDecl <- withDynFlags libdir (\df -> parseDecl df "decl" "nn = 2")
   Right newSig  <- withDynFlags libdir (\df -> parseDecl df "sig"  "nn :: Int")
   let
-      doAddLocal = do
-         (parent:ds) <- hsDecls lp
+      doAddLocal = replaceDecls (anchorEof lp) (parent':ds)
+        where
+          (parent:ds) = hsDecls (makeDeltaAst lp)
 
-         let newDecl' = setEntryDP newDecl (DifferentLine 1 0)
-         let newSig'  = setEntryDP newSig  (DifferentLine 1 4)
+          newDecl' = setEntryDP (makeDeltaAst newDecl) (DifferentLine 1 0)
+          newSig'  = setEntryDP (makeDeltaAst newSig)  (DifferentLine 1 5)
 
-         (parent',_) <- modifyValD (getLocA parent) parent $ \_m decls -> do
-           return ((decls++[newSig',newDecl']),Nothing)
+          (parent',_) = modifyValD (getLocA parent) parent $ \_m decls ->
+                         ((decls++[newSig',newDecl']),Nothing)
 
-         replaceDecls (anchorEof lp) (parent':ds)
 
-  (lp',_,_w) <- runTransformT doAddLocal
-  debugM $ "log:[\n" ++ intercalate "\n" _w ++ "]log end\n"
+      lp' = doAddLocal
   return lp'
 
-
 -- ---------------------------------------------------------------------
 
 addLocaLDecl5 :: Changer
 addLocaLDecl5 _libdir lp = do
   let
-      doAddLocal = do
-         decls <- hsDecls lp
-         [s1,de1,d2,d3] <- balanceCommentsList decls
+      doAddLocal = replaceDecls lp [s1,de1',d3']
+        where
+          decls = hsDecls lp
+          [s1,de1,d2,d3] = balanceCommentsList decls
 
-         let d3' = setEntryDP d3 (DifferentLine 2 0)
+          d3' = setEntryDP d3 (DifferentLine 2 0)
 
-         (de1',_) <- modifyValD (getLocA de1) de1 $ \_m _decls -> do
-           let d2' = setEntryDP d2 (DifferentLine 1 0)
-           return ([d2'],Nothing)
-         replaceDecls lp [s1,de1',d3']
+          (de1',_) = modifyValD (getLocA de1) de1 $ \_m _decls ->
+                       let
+                         d2' = setEntryDP d2 (DifferentLine 1 0)
+                       in ([d2'],Nothing)
 
-  (lp',_,_w) <- runTransformT doAddLocal
-  debugM $ "log:[\n" ++ intercalate "\n" _w ++ "]log end\n"
+      lp' = doAddLocal
   return lp'
 
 -- ---------------------------------------------------------------------
@@ -701,39 +698,36 @@ addLocaLDecl6 :: Changer
 addLocaLDecl6 libdir lp = do
   Right newDecl <- withDynFlags libdir (\df -> parseDecl df "decl" "x = 3")
   let
-      newDecl' = setEntryDP newDecl (DifferentLine 1 4)
-      doAddLocal = do
-        decls0 <- hsDecls lp
-        [de1'',d2] <- balanceCommentsList decls0
+      newDecl' = setEntryDP (makeDeltaAst newDecl) (DifferentLine 1 5)
+      doAddLocal = replaceDecls lp [de1', d2]
+        where
+          decls0 = hsDecls lp
+          [de1'',d2] = balanceCommentsList decls0
 
-        let de1 = captureMatchLineSpacing de1''
-        let L _ (ValD _ (FunBind _ _ (MG _ (L _ ms)))) = de1
-        let [ma1,_ma2] = ms
+          de1 = captureMatchLineSpacing de1''
+          L _ (ValD _ (FunBind _ _ (MG _ (L _ ms)))) = de1
+          [ma1,_ma2] = ms
 
-        (de1',_) <- modifyValD (getLocA ma1) de1 $ \_m decls -> do
-           return ((newDecl' : decls),Nothing)
-        replaceDecls lp [de1', d2]
+          (de1',_) = modifyValD (getLocA ma1) de1 $ \_m decls ->
+                       ((newDecl' : decls),Nothing)
 
-  (lp',_,_w) <- runTransformT doAddLocal
-  debugM $ "log:[\n" ++ intercalate "\n" _w ++ "]log end\n"
+      lp' = doAddLocal
   return lp'
 
 -- ---------------------------------------------------------------------
 
 rmDecl1 :: Changer
-rmDecl1 _libdir top = do
-  let doRmDecl = do
-         let lp = top
-         tlDecs0 <- hsDecls lp
-         tlDecs' <- balanceCommentsList tlDecs0
-         let tlDecs = captureLineSpacing tlDecs'
-         let (de1:_s1:_d2:d3:ds) = tlDecs
-         let d3' = setEntryDP d3 (DifferentLine 2 0)
-
-         replaceDecls lp (de1:d3':ds)
-
-  (lp',_,_w) <- runTransformT doRmDecl
-  debugM $ "log:[\n" ++ intercalate "\n" _w ++ "]log end\n"
+rmDecl1 _libdir lp = do
+  let
+      doRmDecl = replaceDecls lp (de1:d3':ds)
+        where
+          tlDecs0 = hsDecls lp
+          tlDecs = balanceCommentsList tlDecs0
+          (de1:_s1:_d2:d3:ds) = tlDecs
+          d3' = setEntryDP d3 (DifferentLine 2 0)
+
+
+      lp' = doRmDecl
   return lp'
 
 -- ---------------------------------------------------------------------
@@ -745,13 +739,13 @@ rmDecl2 _libdir lp = do
         let
           go :: GHC.LHsExpr GhcPs -> Transform (GHC.LHsExpr GhcPs)
           go e@(GHC.L _ (GHC.HsLet{})) = do
-            decs0 <- hsDecls e
-            decs <- balanceCommentsList $ captureLineSpacing decs0
-            e' <- replaceDecls e (init decs)
+            let decs0 = hsDecls e
+            let decs = balanceCommentsList $ captureLineSpacing decs0
+            let e' = replaceDecls e (init decs)
             return e'
           go x = return x
 
-        everywhereM (mkM go) lp
+        everywhereM (mkM go) (makeDeltaAst lp)
 
   let (lp',_,_w) = runTransform doRmDecl
   debugM $ "log:[\n" ++ intercalate "\n" _w ++ "]log end\n"
@@ -762,17 +756,15 @@ rmDecl2 _libdir lp = do
 rmDecl3 :: Changer
 rmDecl3 _libdir lp = do
   let
-      doRmDecl = do
-         [de1,d2] <- hsDecls lp
-
-         (de1',Just sd1) <- modifyValD (getLocA de1) de1 $ \_m [sd1] -> do
-           let sd1' = setEntryDP sd1 (DifferentLine 2 0)
-           return ([],Just sd1')
-
-         replaceDecls lp [de1',sd1,d2]
+      doRmDecl = replaceDecls lp [de1',sd1,d2]
+        where
+          [de1,d2] = hsDecls lp
+          (de1',Just sd1) = modifyValD (getLocA de1) de1 $ \_m [sd1a] ->
+                       let
+                           sd1' = setEntryDP sd1a (DifferentLine 2 0)
+                       in ([],Just sd1')
 
-  (lp',_,_w) <- runTransformT doRmDecl
-  debugM $ "log:[\n" ++ intercalate "\n" _w ++ "]log end\n"
+      lp' = doRmDecl
   return lp'
 
 -- ---------------------------------------------------------------------
@@ -780,19 +772,15 @@ rmDecl3 _libdir lp = do
 rmDecl4 :: Changer
 rmDecl4 _libdir lp = do
   let
-      doRmDecl = do
-         [de1] <- hsDecls lp
-
-         (de1',Just sd1) <- modifyValD (getLocA de1) de1 $ \_m [sd1,sd2] -> do
-           sd2' <- transferEntryDP' sd1 sd2
-
-           let sd1' = setEntryDP sd1 (DifferentLine 2 0)
-           return ([sd2'],Just sd1')
-
-         replaceDecls (anchorEof lp) [de1',sd1]
-
-  (lp',_,_w) <- runTransformT doRmDecl
-  debugM $ "log:[\n" ++ intercalate "\n" _w ++ "]log end\n"
+      doRmDecl = replaceDecls (anchorEof lp) [de1',sd1]
+        where
+         [de1] = hsDecls lp
+         (de1',Just sd1) = modifyValD (getLocA de1) de1 $ \_m [sd1a,sd2] ->
+           let
+             sd2' = transferEntryDP' sd1a sd2
+             sd1' = setEntryDP sd1a (DifferentLine 2 0)
+           in ([sd2'],Just sd1')
+      lp' = doRmDecl
   return lp'
 
 -- ---------------------------------------------------------------------
@@ -805,10 +793,8 @@ rmDecl5 _libdir lp = do
           go :: HsExpr GhcPs -> Transform (HsExpr GhcPs)
           go (HsLet (tkLet, tkIn) lb expr) = do
             let decs = hsDeclsLocalBinds lb
-            let hdecs : _ = decs
             let dec = last decs
-            _ <- transferEntryDP hdecs dec
-            lb' <- replaceDeclsValbinds WithoutWhere lb [dec]
+            let lb' = replaceDeclsValbinds WithoutWhere lb [dec]
             return (HsLet (tkLet, tkIn) lb' expr)
           go x = return x
 
@@ -823,73 +809,61 @@ rmDecl5 _libdir lp = do
 rmDecl6 :: Changer
 rmDecl6 _libdir lp = do
   let
-      doRmDecl = do
-         [de1] <- hsDecls lp
-
-         (de1',_) <- modifyValD (getLocA de1) de1 $ \_m subDecs -> do
-           let subDecs' = captureLineSpacing subDecs
-           let (ss1:_sd1:sd2:sds) = subDecs'
-           sd2' <- transferEntryDP' ss1 sd2
-
-           return (sd2':sds,Nothing)
+      doRmDecl = replaceDecls lp [de1']
+        where
+          [de1] = hsDecls lp
 
-         replaceDecls lp [de1']
+          (de1',_) = modifyValD (getLocA de1) de1 $ \_m subDecs ->
+            let
+              subDecs' = captureLineSpacing subDecs
+              (ss1:_sd1:sd2:sds) = subDecs'
+              sd2' = transferEntryDP' ss1 sd2
+            in (sd2':sds,Nothing)
 
-  (lp',_,_w) <- runTransformT doRmDecl
-  debugM $ "log:[\n" ++ intercalate "\n" _w ++ "]log end\n"
+      lp' = doRmDecl
   return lp'
 
 -- ---------------------------------------------------------------------
 
 rmDecl7 :: Changer
-rmDecl7 _libdir top = do
+rmDecl7 _libdir lp = do
   let
-      doRmDecl = do
-         let lp = top
-         tlDecs <- hsDecls lp
-         [s1,de1,d2,d3] <- balanceCommentsList tlDecs
-
-         d3' <- transferEntryDP' d2 d3
-
-         replaceDecls lp [s1,de1,d3']
+      doRmDecl = replaceDecls lp [s1,de1,d3']
+        where
+          tlDecs = hsDecls lp
+          [s1,de1,d2,d3] = balanceCommentsList tlDecs
+          d3' = transferEntryDP' d2 d3
 
-  (lp',_,_w) <- runTransformT doRmDecl
-  debugM $ "log:[\n" ++ intercalate "\n" _w ++ "]log end\n"
+      lp' = doRmDecl
   return lp'
 
 -- ---------------------------------------------------------------------
 
 rmTypeSig1 :: Changer
 rmTypeSig1 _libdir lp = do
-  let doRmDecl = do
-         tlDecs <- hsDecls lp
-         let (s0:de1:d2) = tlDecs
-             s1 = captureTypeSigSpacing s0
-             (L l (SigD x1 (TypeSig x2 [n1,n2] typ))) = s1
-         L ln n2' <- transferEntryDP n1 n2
-         let s1' = (L l (SigD x1 (TypeSig x2 [L (noTrailingN ln) n2'] typ)))
-         replaceDecls lp (s1':de1:d2)
-
-  let (lp',_,_w) = runTransform doRmDecl
-  debugM $ "log:[\n" ++ intercalate "\n" _w ++ "]log end\n"
+  let doRmDecl = replaceDecls lp (s1':de1:d2)
+        where
+          tlDecs = hsDecls lp
+          (s0:de1:d2) = tlDecs
+          s1 = captureTypeSigSpacing s0
+          (L l (SigD x1 (TypeSig x2 [n1,n2] typ))) = s1
+          L ln n2' = transferEntryDP n1 n2
+          s1' = (L l (SigD x1 (TypeSig x2 [L (noTrailingN ln) n2'] typ)))
+
+      lp' = doRmDecl
   return lp'
 
 -- ---------------------------------------------------------------------
 
 rmTypeSig2 :: Changer
 rmTypeSig2 _libdir lp = do
-  let doRmDecl = do
-         tlDecs <- hsDecls lp
-         let [de1] = tlDecs
-
-         (de1',_) <- modifyValD (getLocA de1) de1 $ \_m [s,d] -> do
-           d' <- transferEntryDP' s d
-           return $ ([d'],Nothing)
-                  `debug` ("rmTypeSig2:(d,d')" ++ showAst (d,d'))
-         replaceDecls lp [de1']
+  let doRmDecl = replaceDecls lp [de1']
+        where
+          tlDecs = hsDecls lp
+          [de1] = tlDecs
+          (de1',_) = modifyValD (getLocA de1) de1 $ \_m [_s,d] -> ([d],Nothing)
 
-  let (lp',_,_w) = runTransform doRmDecl
-  debugM $ "log:[\n" ++ intercalate "\n" _w ++ "]log end\n"
+  let lp' = doRmDecl
   return lp'
 
 -- ---------------------------------------------------------------------
@@ -958,13 +932,15 @@ addClassMethod libdir lp = do
   let decl' = setEntryDP decl (DifferentLine 1 3)
   let  sig' = setEntryDP sig  (DifferentLine 2 3)
   let doAddMethod = do
-        [cd] <- hsDecls lp
-        (f1:f2s:f2d:_) <- hsDecls cd
-        let  f2s' = setEntryDP f2s  (DifferentLine 2 3)
-        cd' <- replaceDecls cd [f1, sig', decl', f2s', f2d]
-        replaceDecls lp [cd']
-
-  (lp',_,w) <- runTransformT doAddMethod
+        let
+          [cd] = hsDecls lp
+          (f1:f2s:f2d:_) = hsDecls cd
+          f2s' = setEntryDP f2s  (DifferentLine 2 3)
+          cd' = replaceDecls cd [f1, sig', decl', f2s', f2d]
+          lp' = replaceDecls lp [cd']
+        return lp'
+
+  let (lp',_,w) = runTransform doAddMethod
   debugM $ "addClassMethod:" ++ intercalate "\n" w
   return lp'
 


=====================================
utils/check-exact/Parsers.hs
=====================================
@@ -260,7 +260,7 @@ parseModuleEpAnnsWithCppInternal cppOptions dflags file = do
       GHC.PFailed pst
         -> Left (GHC.GhcPsMessage <$> GHC.getPsErrorMessages pst)
       GHC.POk _ pmod
-        -> Right $ (injectedComments, dflags', fixModuleTrailingComments pmod)
+        -> Right $ (injectedComments, dflags', fixModuleComments pmod)
 
 -- | Internal function. Exposed if you want to muck with DynFlags
 -- before parsing. Or after parsing.
@@ -269,8 +269,10 @@ postParseTransform
   -> Either a (GHC.ParsedSource)
 postParseTransform parseRes = fmap mkAnns parseRes
   where
-    -- TODO:AZ perhaps inject the comments into the parsedsource here already
-    mkAnns (_cs, _, m) = fixModuleTrailingComments m
+    mkAnns (_cs, _, m) = fixModuleComments m
+
+fixModuleComments :: GHC.ParsedSource -> GHC.ParsedSource
+fixModuleComments p = fixModuleHeaderComments $ fixModuleTrailingComments p
 
 fixModuleTrailingComments :: GHC.ParsedSource -> GHC.ParsedSource
 fixModuleTrailingComments (GHC.L l p) = GHC.L l p'
@@ -293,6 +295,47 @@ fixModuleTrailingComments (GHC.L l p) = GHC.L l p'
             in cs''
           _ -> cs
 
+-- Deal with https://gitlab.haskell.org/ghc/ghc/-/issues/23984
+-- The Lexer works bottom-up, so does not have module declaration info
+-- when the first top decl processed
+fixModuleHeaderComments :: GHC.ParsedSource -> GHC.ParsedSource
+fixModuleHeaderComments (GHC.L l p) = GHC.L l p'
+  where
+    moveComments :: GHC.EpaLocation -> GHC.LHsDecl GHC.GhcPs -> GHC.EpAnnComments
+                 -> (GHC.LHsDecl GHC.GhcPs, GHC.EpAnnComments)
+    moveComments GHC.EpaDelta{} dd cs = (dd,cs)
+    moveComments (GHC.EpaSpan (GHC.UnhelpfulSpan _)) dd cs = (dd,cs)
+    moveComments (GHC.EpaSpan (GHC.RealSrcSpan r _)) (GHC.L (GHC.EpAnn anc an csd) a) cs = (dd,css)
+      where
+        -- Move any comments on the decl that occur prior to the location
+        pc = GHC.priorComments csd
+        fc = GHC.getFollowingComments csd
+        bf (GHC.L anch _) = GHC.anchor anch > r
+        (move,keep) = break bf pc
+        csd' = GHC.EpaCommentsBalanced keep fc
+
+        dd = GHC.L (GHC.EpAnn anc an csd') a
+        css = cs <> GHC.EpaComments move
+
+    (ds',an') = rebalance (GHC.hsmodDecls p, GHC.hsmodAnn $ GHC.hsmodExt p)
+    p' = p { GHC.hsmodExt = (GHC.hsmodExt p){ GHC.hsmodAnn = an' },
+             GHC.hsmodDecls = ds'
+           }
+
+    rebalance :: ([GHC.LHsDecl GHC.GhcPs], GHC.EpAnn GHC.AnnsModule)
+              -> ([GHC.LHsDecl GHC.GhcPs], GHC.EpAnn GHC.AnnsModule)
+    rebalance (ds, GHC.EpAnn a an cs) = (ds1, GHC.EpAnn a an cs')
+      where
+        (ds1,cs') = case break (\(GHC.AddEpAnn k _) -> k == GHC.AnnWhere) (GHC.am_main an) of
+                     (_, (GHC.AddEpAnn _ whereLoc:_)) ->
+                           case GHC.hsmodDecls p of
+                               (d:ds0) -> (d':ds0, cs0)
+                                   where (d',cs0) = moveComments whereLoc d cs
+                               ds0 -> (ds0,cs)
+                     _ -> (ds,cs)
+
+
+
 -- | Internal function. Initializes DynFlags value for parsing.
 --
 -- Passes "-hide-all-packages" to the GHC API to prevent parsing of


=====================================
utils/check-exact/Transform.hs
=====================================
@@ -63,7 +63,7 @@ module Transform
         -- *** Low level operations used in 'HasDecls'
         , balanceComments
         , balanceCommentsList
-        , balanceCommentsList'
+        , balanceCommentsListA
         , anchorEof
 
         -- ** Managing lists, pure functions
@@ -92,6 +92,7 @@ import qualified Control.Monad.Fail as Fail
 
 import GHC  hiding (parseModule, parsedSource)
 import GHC.Data.FastString
+import GHC.Types.SrcLoc
 
 import Data.Data
 import Data.Maybe
@@ -154,6 +155,7 @@ logDataWithAnnsTr str ast = do
 
 -- |If we need to add new elements to the AST, they need their own
 -- 'SrcSpan' for this.
+-- This should no longer be needed, we use an @EpaDelta@ location instead.
 uniqueSrcSpanT :: (Monad m) => TransformT m SrcSpan
 uniqueSrcSpanT = do
   col <- get
@@ -171,15 +173,6 @@ srcSpanStartLine' _ = 0
 
 -- ---------------------------------------------------------------------
 
-captureOrderBinds :: [LHsDecl GhcPs] -> AnnSortKey BindTag
-captureOrderBinds ls = AnnSortKey $ map go ls
-  where
-    go (L _ (ValD _ _))       = BindTag
-    go (L _ (SigD _ _))       = SigDTag
-    go d      = error $ "captureOrderBinds:" ++ showGhc d
-
--- ---------------------------------------------------------------------
-
 captureMatchLineSpacing :: LHsDecl GhcPs -> LHsDecl GhcPs
 captureMatchLineSpacing (L l (ValD x (FunBind a b (MG c (L d ms )))))
                        = L l (ValD x (FunBind a b (MG c (L d ms'))))
@@ -253,7 +246,7 @@ setEntryDPDecl d dp = setEntryDP d dp
 
 -- |Set the true entry 'DeltaPos' from the annotation for a given AST
 -- element. This is the 'DeltaPos' ignoring any comments.
-setEntryDP :: NoAnn t => LocatedAn t a -> DeltaPos -> LocatedAn t a
+setEntryDP :: LocatedAn t a -> DeltaPos -> LocatedAn t a
 setEntryDP (L (EpAnn (EpaSpan ss@(UnhelpfulSpan _)) an cs) a) dp
   = L (EpAnn (EpaDelta ss dp []) an cs) a
 setEntryDP (L (EpAnn (EpaSpan ss) an (EpaComments [])) a) dp
@@ -293,7 +286,7 @@ setEntryDP (L (EpAnn (EpaSpan ss@(RealSrcSpan r _)) an cs) a) dp
         L (EpAnn (EpaDelta ss edp csd) an cs'') a
               where
                 cs'' = setPriorComments cs []
-                csd = L (EpaDelta ss dp NoComments) c:cs'
+                csd = L (EpaDelta ss dp NoComments) c:commentOrigDeltas cs'
                 lc = last $ (L ca c:cs')
                 delta = case getLoc lc of
                           EpaSpan (RealSrcSpan rr _) -> ss2delta (ss2pos rr) r
@@ -335,18 +328,15 @@ setEntryDPFromAnchor  off (EpaSpan (RealSrcSpan anc _)) ll@(L la _) = setEntryDP
 
 -- |Take the annEntryDelta associated with the first item and associate it with the second.
 -- Also transfer any comments occurring before it.
-transferEntryDP :: (Monad m, NoAnn t2, Typeable t1, Typeable t2)
-  => LocatedAn t1 a -> LocatedAn t2 b -> TransformT m (LocatedAn t2 b)
-transferEntryDP (L (EpAnn anc1 an1 cs1) _) (L (EpAnn _anc2 an2 cs2) b) = do
-  logTr $ "transferEntryDP': EpAnn,EpAnn"
+transferEntryDP :: (Typeable t1, Typeable t2)
+  => LocatedAn t1 a -> LocatedAn t2 b -> (LocatedAn t2 b)
+transferEntryDP (L (EpAnn anc1 an1 cs1) _) (L (EpAnn _anc2 an2 cs2) b) =
   -- Problem: if the original had preceding comments, blindly
   -- transferring the location is not correct
   case priorComments cs1 of
-    [] -> return (L (EpAnn anc1 (combine an1 an2) cs2) b)
+    [] -> (L (EpAnn anc1 (combine an1 an2) cs2) b)
     -- TODO: what happens if the receiving side already has comments?
-    (L anc _:_) -> do
-      logDataWithAnnsTr "transferEntryDP':priorComments anc=" anc
-      return (L (EpAnn anc1 (combine an1 an2) (cs1 <> cs2)) b)
+    (L _ _:_) -> (L (EpAnn anc1 (combine an1 an2) (cs1 <> cs2)) b)
 
 
 -- |If a and b are the same type return first arg, else return second
@@ -356,10 +346,11 @@ combine x y = fromMaybe y (cast x)
 -- |Take the annEntryDelta associated with the first item and associate it with the second.
 -- Also transfer any comments occurring before it.
 -- TODO: call transferEntryDP, and use pushDeclDP
-transferEntryDP' :: (Monad m) => LHsDecl GhcPs -> LHsDecl GhcPs -> TransformT m (LHsDecl GhcPs)
-transferEntryDP' la lb = do
-  (L l2 b) <- transferEntryDP la lb
-  return (L l2 (pushDeclDP b (SameLine 0)))
+transferEntryDP' :: LHsDecl GhcPs -> LHsDecl GhcPs -> (LHsDecl GhcPs)
+transferEntryDP' la lb =
+  let
+    (L l2 b) = transferEntryDP la lb
+  in (L l2 (pushDeclDP b (SameLine 0)))
 
 
 pushDeclDP :: HsDecl GhcPs -> DeltaPos -> HsDecl GhcPs
@@ -375,13 +366,24 @@ pushDeclDP d _dp = d
 
 -- ---------------------------------------------------------------------
 
-balanceCommentsList :: (Monad m) => [LHsDecl GhcPs] -> TransformT m [LHsDecl GhcPs]
-balanceCommentsList [] = return []
-balanceCommentsList [x] = return [x]
-balanceCommentsList (a:b:ls) = do
-  (a',b') <- balanceComments a b
-  r <- balanceCommentsList (b':ls)
-  return (a':r)
+-- | If we compile in haddock mode, the haddock processing inserts
+-- DocDecls to carry the Haddock Documentation. We ignore these in
+-- exact printing, as all the comments are also available in their
+-- normal location, and the haddock processing is lossy, in that it
+-- does not preserve all haddock-like comments. When we balance
+-- comments in a list, we migrate some to preceding or following
+-- declarations in the list. We must make sure we do not move any to
+-- these DocDecls, which are not printed.
+balanceCommentsList :: [LHsDecl GhcPs] -> [LHsDecl GhcPs]
+balanceCommentsList decls = balanceCommentsList' (filter notDocDecl decls)
+
+balanceCommentsList' :: [LHsDecl GhcPs] -> [LHsDecl GhcPs]
+balanceCommentsList' [] = []
+balanceCommentsList' [x] = [x]
+balanceCommentsList' (a:b:ls) = (a':r)
+  where
+    (a',b') = balanceComments a b
+    r = balanceCommentsList' (b':ls)
 
 -- |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.
@@ -389,28 +391,27 @@ balanceCommentsList (a:b:ls) = do
 -- from the second one to the 'annFollowingComments' of the first if they belong
 -- to it instead. This is typically required before deleting or duplicating
 -- either of the AST elements.
-balanceComments :: (Monad m)
-  => LHsDecl GhcPs -> LHsDecl GhcPs
-  -> TransformT m (LHsDecl GhcPs, LHsDecl GhcPs)
-balanceComments first second = do
+balanceComments :: LHsDecl GhcPs -> LHsDecl GhcPs
+                -> (LHsDecl GhcPs, LHsDecl GhcPs)
+balanceComments first second =
   case first of
-    (L l (ValD x fb@(FunBind{}))) -> do
-      (L l' fb',second') <- balanceCommentsFB (L l fb) second
-      return (L l' (ValD x fb'), second')
-    _ -> balanceComments' first second
+    (L l (ValD x fb@(FunBind{}))) ->
+      let
+        (L l' fb',second') = balanceCommentsFB (L l fb) second
+      in (L l' (ValD x fb'), second')
+    _ -> balanceCommentsA first second
 
--- |Once 'balanceComments' has been called to move trailing comments to a
+-- |Once 'balanceCommentsA has been called to move trailing comments to a
 -- 'FunBind', these need to be pushed down from the top level to the last
 -- 'Match' if that 'Match' needs to be manipulated.
-balanceCommentsFB :: (Monad m)
-  => LHsBind GhcPs -> LocatedA b -> TransformT m (LHsBind GhcPs, LocatedA b)
-balanceCommentsFB (L lf (FunBind x n (MG o (L lm matches)))) second = do
-  debugM $ "balanceCommentsFB entered: " ++ showGhc (ss2range $ locA lf)
+balanceCommentsFB :: LHsBind GhcPs -> LocatedA b -> (LHsBind GhcPs, LocatedA b)
+balanceCommentsFB (L lf (FunBind x n (MG o (L lm matches)))) second
+  = balanceCommentsA (packFunBind bind) second'
   -- 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
+  where
     (before,middle,after) = case entry lf of
         EpaSpan (RealSrcSpan ss _) ->
           let
@@ -426,40 +427,29 @@ balanceCommentsFB (L lf (FunBind x n (MG o (L lm matches)))) second = do
               getFollowingComments $ comments lf)
 
     lf' = setCommentsEpAnn lf (EpaComments before)
-  debugM $ "balanceCommentsFB (before, after): " ++ showAst (before, after)
-  debugM $ "balanceCommentsFB lf': " ++ showAst lf'
-  -- let matches' = case matches of
-  let matches' :: [LocatedA (Match GhcPs (LHsExpr GhcPs))]
-      matches' = case matches of
-                    (L lm' m':ms') ->
-                      (L (addCommentsToEpAnn lm' (EpaComments middle )) m':ms')
-                    _ -> error "balanceCommentsFB"
-  matches'' <- balanceCommentsList' matches'
-  let (m,ms) = case reverse matches'' of
-                 (L lm' m':ms') ->
-                   (L (addCommentsToEpAnn lm' (EpaCommentsBalanced [] after)) m',ms')
-                   -- (L (addCommentsToEpAnnS lm' (EpaCommentsBalanced [] after)) m',ms')
-                 _ -> error "balanceCommentsFB4"
-  debugM $ "balanceCommentsFB: (m,ms):" ++ showAst (m,ms)
-  (m',second') <- balanceComments' m second
-  m'' <- balanceCommentsMatch m'
-  let (m''',lf'') = case ms of
-        [] -> moveLeadingComments m'' lf'
-        _  -> (m'',lf')
-  debugM $ "balanceCommentsFB: (lf'', m'''):" ++ showAst (lf'',m''')
-  debugM $ "balanceCommentsFB done"
-  let bind = L lf'' (FunBind x n (MG o (L lm (reverse (m''':ms)))))
-  debugM $ "balanceCommentsFB returning:" ++ showAst bind
-  balanceComments' (packFunBind bind) second'
-balanceCommentsFB f s = balanceComments' f s
+    matches' :: [LocatedA (Match GhcPs (LHsExpr GhcPs))]
+    matches' = case matches of
+                  (L lm' m0:ms') ->
+                    (L (addCommentsToEpAnn lm' (EpaComments middle )) m0:ms')
+                  _ -> error "balanceCommentsFB"
+    matches'' = balanceCommentsListA matches'
+    (m,ms) = case reverse matches'' of
+               (L lm' m0:ms') ->
+                 (L (addCommentsToEpAnn lm' (EpaCommentsBalanced [] after)) m0,ms')
+               _ -> error "balanceCommentsFB4"
+    (m',second') = balanceCommentsA m second
+    m'' = balanceCommentsMatch m'
+    (m''',lf'') = case ms of
+      [] -> moveLeadingComments m'' lf'
+      _  -> (m'',lf')
+    bind = L lf'' (FunBind x n (MG o (L lm (reverse (m''':ms)))))
+balanceCommentsFB f s = balanceCommentsA f s
 
 -- | Move comments on the same line as the end of the match into the
 -- GRHS, prior to the binds
-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: (logInfo)=" ++ showAst (logInfo)
-  return (L l'' (Match am mctxt pats (GRHSs xg grhss' binds')))
+balanceCommentsMatch :: LMatch GhcPs (LHsExpr GhcPs) -> (LMatch GhcPs (LHsExpr GhcPs))
+balanceCommentsMatch (L l (Match am mctxt pats (GRHSs xg grhss binds)))
+  = (L l'' (Match am mctxt pats (GRHSs xg grhss' binds')))
   where
     simpleBreak (r,_) = r /= 0
     an1 = l
@@ -468,7 +458,7 @@ balanceCommentsMatch (L l (Match am mctxt pats (GRHSs xg grhss binds))) = do
     (move',stay') = break simpleBreak (trailingCommentsDeltas (anchorFromLocatedA (L l ())) cs1f)
     move = map snd move'
     stay = map snd stay'
-    (l'', grhss', binds', logInfo)
+    (l'', grhss', binds', _logInfo)
       = case reverse grhss of
           [] -> (l, [], binds,                 (EpaComments [], noSrcSpanA))
           (L lg (GRHS ag grs rhs):gs) ->
@@ -491,26 +481,24 @@ balanceCommentsMatch (L l (Match am mctxt pats (GRHSs xg grhss binds))) = do
 pushTrailingComments :: WithWhere -> EpAnnComments -> HsLocalBinds GhcPs -> (Bool, HsLocalBinds GhcPs)
 pushTrailingComments _ _cs b at EmptyLocalBinds{} = (False, b)
 pushTrailingComments _ _cs (HsIPBinds _ _) = error "TODO: pushTrailingComments:HsIPBinds"
-pushTrailingComments w cs lb@(HsValBinds an _)
-  = (True, HsValBinds an' vb)
+pushTrailingComments w cs lb@(HsValBinds an _) = (True, HsValBinds an' vb)
   where
     decls = hsDeclsLocalBinds lb
     (an', decls') = case reverse decls of
       [] -> (addCommentsToEpAnn an cs, decls)
       (L la d:ds) -> (an, L (addCommentsToEpAnn la cs) d:ds)
-    (vb,_ws2) = case runTransform (replaceDeclsValbinds w lb (reverse decls')) of
-      ((HsValBinds _ vb'), _, ws2') -> (vb', ws2')
-      _ -> (ValBinds NoAnnSortKey [] [], [])
+    vb = case replaceDeclsValbinds w lb (reverse decls') of
+      (HsValBinds _ vb') -> vb'
+      _ -> ValBinds NoAnnSortKey [] []
 
 
-balanceCommentsList' :: (Monad m) => [LocatedA a] -> TransformT m [LocatedA a]
-balanceCommentsList' [] = return []
-balanceCommentsList' [x] = return [x]
-balanceCommentsList' (a:b:ls) = do
-  logTr $ "balanceCommentsList' entered"
-  (a',b') <- balanceComments' a b
-  r <- balanceCommentsList' (b':ls)
-  return (a':r)
+balanceCommentsListA :: [LocatedA a] -> [LocatedA a]
+balanceCommentsListA [] = []
+balanceCommentsListA [x] = [x]
+balanceCommentsListA (a:b:ls) = (a':r)
+  where
+    (a',b') = balanceCommentsA a b
+    r = balanceCommentsListA (b':ls)
 
 -- |Prior to moving an AST element, make sure any trailing comments belonging to
 -- it are attached to it, and not the following element. Of necessity this is a
@@ -518,13 +506,8 @@ balanceCommentsList' (a:b:ls) = do
 -- with a passed-in decision function.
 -- The initial situation is that all comments for a given anchor appear as prior comments
 -- Many of these should in fact be following comments for the previous anchor
-balanceComments' :: (Monad m) => LocatedA a -> LocatedA b -> TransformT m (LocatedA a, LocatedA b)
-balanceComments' la1 la2 = do
-  debugM $ "balanceComments': (anc1)=" ++ showAst (anc1)
-  debugM $ "balanceComments': (cs1s)=" ++ showAst (cs1s)
-  debugM $ "balanceComments': (cs1stay,cs1move)=" ++ showAst (cs1stay,cs1move)
-  debugM $ "balanceComments': (an1',an2')=" ++ showAst (an1',an2')
-  return (la1', la2')
+balanceCommentsA :: LocatedA a -> LocatedA b -> (LocatedA a, LocatedA b)
+balanceCommentsA la1 la2 = (la1', la2')
   where
     simpleBreak n (r,_) = r > n
     L an1 f = la1
@@ -532,26 +515,31 @@ balanceComments' la1 la2 = do
     anc1 = comments an1
     anc2 = comments an2
 
-    cs1s = splitCommentsEnd (anchorFromLocatedA la1) anc1
-    cs1p = priorCommentsDeltas    (anchorFromLocatedA la1) (priorComments        cs1s)
-    cs1f = trailingCommentsDeltas (anchorFromLocatedA la1) (getFollowingComments cs1s)
+    (p1,m1,f1) = splitComments (anchorFromLocatedA la1) anc1
+    cs1p = priorCommentsDeltas    (anchorFromLocatedA la1) p1
 
-    cs2s = splitCommentsEnd (anchorFromLocatedA la2) anc2
-    cs2p = priorCommentsDeltas    (anchorFromLocatedA la2) (priorComments        cs2s)
-    cs2f = trailingCommentsDeltas (anchorFromLocatedA la2) (getFollowingComments cs2s)
+    -- Split cs1 following comments into those before any
+    -- TrailingAnn's on an1, and any after
+    cs1f = splitCommentsEnd (fullSpanFromLocatedA la1) $ EpaComments f1
+    cs1fp = priorCommentsDeltas    (anchorFromLocatedA la1) (priorComments        cs1f)
+    cs1ff = trailingCommentsDeltas (anchorFromLocatedA la1) (getFollowingComments cs1f)
 
-    -- Split cs1f into those that belong on an1 and ones that must move to an2
-    (cs1move,cs1stay) = break (simpleBreak 1) cs1f
+    -- Split cs1ff into those that belong on an1 and ones that must move to an2
+    (cs1move,cs1stay) = break (simpleBreak 1) cs1ff
+
+    (p2,m2,f2) = splitComments (anchorFromLocatedA la2) anc2
+    cs2p = priorCommentsDeltas    (anchorFromLocatedA la2) p2
+    cs2f = trailingCommentsDeltas (anchorFromLocatedA la2) f2
 
     (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 = sortEpaComments $ map snd (cs1move ++ move'' ++ move')
-    stay = sortEpaComments $ map snd (cs1stay ++ stay')
+    move = sortEpaComments $ map snd (cs1fp ++ cs1move ++ move'' ++ move')
+    stay = sortEpaComments $ m2 ++ map snd (cs1stay ++ stay')
 
-    an1' = setCommentsEpAnn (getLoc la1) (EpaCommentsBalanced (map snd cs1p) move)
-    an2' = setCommentsEpAnn (getLoc la2) (EpaCommentsBalanced stay (map snd cs2f))
+    an1' = setCommentsEpAnn (getLoc la1) (epaCommentsBalanced (m1 ++ map snd cs1p) move)
+    an2' = setCommentsEpAnn (getLoc la2) (epaCommentsBalanced stay (map snd cs2f))
     la1' = L an1' f
     la2' = L an2' s
 
@@ -569,10 +557,9 @@ trailingCommentsDeltas r (la@(L l _):las)
         (al,_) = ss2posEnd rs'
         (ll,_) = ss2pos (anchor loc)
 
--- AZ:TODO: this is identical to commentsDeltas
 priorCommentsDeltas :: RealSrcSpan -> [LEpaComment]
                     -> [(Int, LEpaComment)]
-priorCommentsDeltas r cs = go r (reverse $ sortEpaComments cs)
+priorCommentsDeltas r cs = go r (sortEpaComments cs)
   where
     go :: RealSrcSpan -> [LEpaComment] -> [(Int, LEpaComment)]
     go _   [] = []
@@ -588,6 +575,21 @@ priorCommentsDeltas r cs = go r (reverse $ sortEpaComments cs)
 
 -- ---------------------------------------------------------------------
 
+-- | Split comments into ones occurring before the end of the reference
+-- span, and those after it.
+splitComments :: RealSrcSpan -> EpAnnComments -> ([LEpaComment], [LEpaComment], [LEpaComment])
+splitComments p cs = (before, middle, after)
+  where
+    cmpe (L (EpaSpan (RealSrcSpan l _)) _) = ss2pos l > ss2posEnd p
+    cmpe (L _ _) = True
+
+    cmpb (L (EpaSpan (RealSrcSpan l _)) _) = ss2pos l > ss2pos p
+    cmpb (L _ _) = True
+
+    (beforeEnd, after) = break cmpe ((priorComments cs) ++ (getFollowingComments cs))
+    (before, middle) = break cmpb beforeEnd
+
+
 -- | Split comments into ones occurring before the end of the reference
 -- span, and those after it.
 splitCommentsEnd :: RealSrcSpan -> EpAnnComments -> EpAnnComments
@@ -598,8 +600,8 @@ splitCommentsEnd p (EpaComments cs) = cs'
     (before, after) = break cmp cs
     cs' = case after of
       [] -> EpaComments cs
-      _ -> EpaCommentsBalanced before after
-splitCommentsEnd p (EpaCommentsBalanced cs ts) = EpaCommentsBalanced cs' ts'
+      _ -> epaCommentsBalanced before after
+splitCommentsEnd p (EpaCommentsBalanced cs ts) = epaCommentsBalanced cs' ts'
   where
     cmp (L (EpaSpan (RealSrcSpan l _)) _) = ss2pos l > ss2posEnd p
     cmp (L _ _) = True
@@ -617,8 +619,8 @@ splitCommentsStart p (EpaComments cs) = cs'
     (before, after) = break cmp cs
     cs' = case after of
       [] -> EpaComments cs
-      _ -> EpaCommentsBalanced before after
-splitCommentsStart p (EpaCommentsBalanced cs ts) = EpaCommentsBalanced cs' ts'
+      _ -> epaCommentsBalanced before after
+splitCommentsStart p (EpaCommentsBalanced cs ts) = epaCommentsBalanced cs' ts'
   where
     cmp (L (EpaSpan (RealSrcSpan l _)) _) = ss2pos l > ss2posEnd p
     cmp (L _ _) = True
@@ -638,8 +640,8 @@ moveLeadingComments (L la a) lb = (L la' a, lb')
     -- TODO: need to set an entry delta on lb' to zero, and move the
     -- original spacing to the first comment.
 
-    la' = setCommentsEpAnn la (EpaCommentsBalanced [] after)
-    lb' = addCommentsToEpAnn lb (EpaCommentsBalanced before [])
+    la' = setCommentsEpAnn la (epaCommentsBalanced [] after)
+    lb' = addCommentsToEpAnn lb (epaCommentsBalanced before [])
 
 -- | A GHC comment includes the span of the preceding (non-comment)
 -- token.  Takes an original list of comments, and converts the
@@ -662,17 +664,27 @@ addCommentOrigDeltasAnn (EpAnn e a cs) = EpAnn e a (addCommentOrigDeltas cs)
 anchorFromLocatedA :: LocatedA a -> RealSrcSpan
 anchorFromLocatedA (L (EpAnn anc _ _) _) = anchor anc
 
+-- | Get the full span of interest for comments from a LocatedA.
+-- This extends up to the last TrailingAnn
+fullSpanFromLocatedA :: LocatedA a -> RealSrcSpan
+fullSpanFromLocatedA (L (EpAnn anc (AnnListItem tas)  _) _) = rr
+  where
+    r = anchor anc
+    trailing_loc ta = case ta_location ta of
+        EpaSpan (RealSrcSpan s _) -> [s]
+        _ -> []
+    rr = case reverse (concatMap trailing_loc tas) of
+        [] -> r
+        (s:_) -> combineRealSrcSpans r s
+
 -- ---------------------------------------------------------------------
 
-balanceSameLineComments :: (Monad m)
-  => LMatch GhcPs (LHsExpr GhcPs) -> TransformT m (LMatch GhcPs (LHsExpr GhcPs))
-balanceSameLineComments (L la (Match anm mctxt pats (GRHSs x grhss lb))) = do
-  logTr $ "balanceSameLineComments: (la)=" ++ showGhc (ss2range $ locA la)
-  logTr $ "balanceSameLineComments: [logInfo]=" ++ showAst logInfo
-  return (L la' (Match anm mctxt pats (GRHSs x grhss' lb)))
+balanceSameLineComments :: LMatch GhcPs (LHsExpr GhcPs) -> (LMatch GhcPs (LHsExpr GhcPs))
+balanceSameLineComments (L la (Match anm mctxt pats (GRHSs x grhss lb)))
+  = (L la' (Match anm mctxt pats (GRHSs x grhss' lb)))
   where
     simpleBreak n (r,_) = r > n
-    (la',grhss', logInfo) = case reverse grhss of
+    (la',grhss', _logInfo) = case reverse grhss of
       [] -> (la,grhss,[])
       (L lg (GRHS ga gs rhs):grs) -> (la'',reverse $ (L lg (GRHS ga' gs rhs)):grs,[(gac,(csp,csf))])
         where
@@ -684,7 +696,7 @@ balanceSameLineComments (L la (Match anm mctxt pats (GRHSs x grhss lb))) = do
           (move',stay') = break (simpleBreak 0) (trailingCommentsDeltas (anchor anc) csf)
           move = map snd move'
           stay = map snd stay'
-          cs1 = EpaCommentsBalanced csp stay
+          cs1 = epaCommentsBalanced csp stay
 
           gac = epAnnComments ga
           gfc = getFollowingComments gac
@@ -734,24 +746,21 @@ addComma (EpAnn anc (AnnListItem as) cs)
 -- | Insert a declaration into an AST element having sub-declarations
 -- (@HasDecls@) according to the given location function.
 insertAt :: (HasDecls ast)
-              => (LHsDecl GhcPs
-                  -> [LHsDecl GhcPs]
-                  -> [LHsDecl GhcPs])
-              -> ast
-              -> LHsDecl GhcPs
-              -> Transform ast
-insertAt f t decl = do
-  oldDecls <- hsDecls t
-  oldDeclsb <- balanceCommentsList oldDecls
-  let oldDecls' = oldDeclsb
-  replaceDecls t (f decl oldDecls')
+         => (LHsDecl GhcPs
+              -> [LHsDecl GhcPs]
+              -> [LHsDecl GhcPs])
+         -> ast
+         -> LHsDecl GhcPs
+         -> ast
+insertAt f t decl = replaceDecls t (f decl oldDecls')
+  where
+    oldDecls = hsDecls t
+    oldDeclsb = balanceCommentsList oldDecls
+    oldDecls' = oldDeclsb
 
 -- |Insert a declaration at the beginning or end of the subdecls of the given
 -- AST item
-insertAtStart, insertAtEnd :: (HasDecls ast)
-              => ast
-              -> LHsDecl GhcPs
-              -> Transform ast
+insertAtStart, insertAtEnd :: HasDecls ast => ast -> LHsDecl GhcPs -> ast
 
 insertAtEnd   = insertAt (\x xs -> xs ++ [x])
 
@@ -766,11 +775,11 @@ insertAtStart = insertAt insertFirst
 
 -- |Insert a declaration at a specific location in the subdecls of the given
 -- AST item
-insertAfter, insertBefore :: (HasDecls (LocatedA ast))
+insertAfter, insertBefore :: HasDecls (LocatedA ast)
                           => LocatedA old
                           -> LocatedA ast
                           -> LHsDecl GhcPs
-                          -> Transform (LocatedA ast)
+                          -> LocatedA ast
 insertAfter (getLocA -> k) = insertAt findAfter
   where
     findAfter x xs =
@@ -797,10 +806,10 @@ class (Data t) => HasDecls t where
     -- given syntax phrase. They are always returned in the wrapped 'HsDecl'
     -- form, even if orginating in local decls. This is safe, as annotations
     -- never attach to the wrapper, only to the wrapped item.
-    hsDecls :: (Monad m) => t -> TransformT m [LHsDecl GhcPs]
+    hsDecls :: t -> [LHsDecl GhcPs]
 
     -- | Replace the directly enclosed decl list by the given
-    --  decl list. Runs in the 'Transform' monad to be able to update list order
+    --  decl list. As part of replacing it will update list order
     --  annotations, and rebalance comments and other layout changes as needed.
     --
     -- For example, a call on replaceDecls for a wrapped 'FunBind' having no
@@ -818,96 +827,86 @@ class (Data t) => HasDecls t where
     --   where
     --     nn = 2
     -- @
-    replaceDecls :: (Monad m) => t -> [LHsDecl GhcPs] -> TransformT m t
+    replaceDecls :: t -> [LHsDecl GhcPs] -> t
 
 -- ---------------------------------------------------------------------
 
 instance HasDecls ParsedSource where
-  hsDecls (L _ (HsModule (XModulePs _ _lo _ _) _mn _exps _imps decls)) = return decls
+  hsDecls (L _ (HsModule (XModulePs _ _lo _ _) _mn _exps _imps decls)) = decls
 
   replaceDecls (L l (HsModule (XModulePs a lo deps haddocks) mname exps imps _decls)) decls
-    = do
-        logTr "replaceDecls LHsModule"
-        return (L l (HsModule (XModulePs a lo deps haddocks) mname exps imps decls))
+    = (L l (HsModule (XModulePs a lo deps haddocks) mname exps imps decls))
 
 -- ---------------------------------------------------------------------
 
 instance HasDecls (LocatedA (HsDecl GhcPs)) where
-  hsDecls (L _ (TyClD _ c at ClassDecl{}))  = return $ hsDeclsClassDecl c
-  hsDecls decl = do
-    error $ "hsDecls:decl=" ++ showAst decl
-  replaceDecls (L l (TyClD e dec at ClassDecl{})) decls = do
-    let decl' = replaceDeclsClassDecl dec decls
-    return (L l (TyClD e decl'))
-  replaceDecls decl _decls = do
-    error $ "replaceDecls:decl=" ++ showAst decl
+  hsDecls (L _ (TyClD _ c at ClassDecl{}))  = hsDeclsClassDecl c
+  hsDecls decl = error $ "hsDecls:decl=" ++ showAst decl
+  replaceDecls (L l (TyClD e dec at ClassDecl{})) decls =
+    let
+        decl' = replaceDeclsClassDecl dec decls
+    in (L l (TyClD e decl'))
+  replaceDecls decl _decls
+      = error $ "replaceDecls:decl=" ++ showAst decl
 
 -- ---------------------------------------------------------------------
 
 instance HasDecls (LocatedA (Match GhcPs (LocatedA (HsExpr GhcPs)))) where
-  hsDecls (L _ (Match _ _ _ (GRHSs _ _ lb))) = return $ hsDeclsLocalBinds lb
+  hsDecls (L _ (Match _ _ _ (GRHSs _ _ lb))) = hsDeclsLocalBinds lb
 
   replaceDecls (L l (Match xm c p (GRHSs xr rhs binds))) []
-    = do
-        logTr "replaceDecls LMatch empty decls"
-        binds'' <- replaceDeclsValbinds WithoutWhere binds []
-        return (L l (Match xm c p (GRHSs xr rhs binds'')))
+    = let
+        binds'' = replaceDeclsValbinds WithoutWhere binds []
+      in (L l (Match xm c p (GRHSs xr rhs binds'')))
 
   replaceDecls m@(L l (Match xm c p (GRHSs xr rhs binds))) newBinds
-    = do
-        logTr "replaceDecls LMatch nonempty decls"
+    = let
         -- Need to throw in a fresh where clause if the binds were empty,
         -- in the annotations.
-        (l', rhs') <- case binds of
-          EmptyLocalBinds{} -> do
-            logTr $ "replaceDecls LMatch empty binds"
-
-            logDataWithAnnsTr "Match.replaceDecls:balancing comments:m" m
-            L l' m' <- balanceSameLineComments m
-            logDataWithAnnsTr "Match.replaceDecls:(m1')" (L l' m')
-            return (l', grhssGRHSs $ m_grhss m')
-          _ -> return (l, rhs)
-        binds'' <- replaceDeclsValbinds WithWhere binds newBinds
-        logDataWithAnnsTr "Match.replaceDecls:binds'" binds''
-        return (L l' (Match xm c p (GRHSs xr rhs' binds'')))
+        (l', rhs') = case binds of
+          EmptyLocalBinds{} ->
+            let
+              L l0 m' = balanceSameLineComments m
+            in (l0, grhssGRHSs $ m_grhss m')
+          _ -> (l, rhs)
+        binds'' = replaceDeclsValbinds WithWhere binds newBinds
+      in (L l' (Match xm c p (GRHSs xr rhs' binds'')))
 
 -- ---------------------------------------------------------------------
 
 instance HasDecls (LocatedA (HsExpr GhcPs)) where
-  hsDecls (L _ (HsLet _ decls _ex)) = return $ hsDeclsLocalBinds decls
-  hsDecls _                         = return []
+  hsDecls (L _ (HsLet _ decls _ex)) = hsDeclsLocalBinds decls
+  hsDecls _                         = []
 
   replaceDecls (L ll (HsLet (tkLet, tkIn) binds ex)) newDecls
-    = do
-        logTr "replaceDecls HsLet"
-        let lastAnc = realSrcSpan $ spanHsLocaLBinds binds
+    = let
+        lastAnc = realSrcSpan $ spanHsLocaLBinds binds
         -- TODO: may be an intervening comment, take account for lastAnc
-        let (tkLet', tkIn', ex',newDecls') = case (tkLet, tkIn) of
-              (EpTok l, EpTok i) ->
-                let
-                  off = case l of
-                          (EpaSpan (RealSrcSpan r _)) -> LayoutStartCol $ snd $ ss2pos r
-                          (EpaSpan (UnhelpfulSpan _)) -> LayoutStartCol 0
-                          (EpaDelta _ (SameLine _) _) -> LayoutStartCol 0
-                          (EpaDelta _ (DifferentLine _ c) _) -> LayoutStartCol c
-                  ex'' = setEntryDPFromAnchor off i ex
-                  newDecls'' = case newDecls of
-                    [] -> newDecls
-                    (d:ds) -> setEntryDPDecl d (SameLine 0) : ds
-                in ( EpTok l
-                   , EpTok (addEpaLocationDelta off lastAnc i)
-                   , ex''
-                   , newDecls'')
-              (_,_) -> (tkLet, tkIn, ex, newDecls)
-        binds' <- replaceDeclsValbinds WithoutWhere binds newDecls'
-        return (L ll (HsLet (tkLet', tkIn') binds' ex'))
+        (tkLet', tkIn', ex',newDecls') = case (tkLet, tkIn) of
+          (EpTok l, EpTok i) ->
+            let
+              off = case l of
+                      (EpaSpan (RealSrcSpan r _)) -> LayoutStartCol $ snd $ ss2pos r
+                      (EpaSpan (UnhelpfulSpan _)) -> LayoutStartCol 0
+                      (EpaDelta _ (SameLine _) _) -> LayoutStartCol 0
+                      (EpaDelta _ (DifferentLine _ c) _) -> LayoutStartCol c
+              ex'' = setEntryDPFromAnchor off i ex
+              newDecls'' = case newDecls of
+                [] -> newDecls
+                (d:ds) -> setEntryDPDecl d (SameLine 0) : ds
+            in ( EpTok l
+               , EpTok (addEpaLocationDelta off lastAnc i)
+               , ex''
+               , newDecls'')
+          (_,_) -> (tkLet, tkIn, ex, newDecls)
+        binds' = replaceDeclsValbinds WithoutWhere binds newDecls'
+      in (L ll (HsLet (tkLet', tkIn') binds' ex'))
 
   -- TODO: does this make sense? Especially as no hsDecls for HsPar
   replaceDecls (L l (HsPar x e)) newDecls
-    = do
-        logTr "replaceDecls HsPar"
-        e' <- replaceDecls e newDecls
-        return (L l (HsPar x e'))
+    = let
+        e' = replaceDecls e newDecls
+      in (L l (HsPar x e'))
   replaceDecls old _new = error $ "replaceDecls (LHsExpr GhcPs) undefined for:" ++ showGhc old
 
 -- ---------------------------------------------------------------------
@@ -934,53 +933,51 @@ hsDeclsPatBind x = error $ "hsDeclsPatBind called for:" ++ showGhc x
 -- cannot be a member of 'HasDecls' because a 'FunBind' is not idempotent
 -- for 'hsDecls' \/ 'replaceDecls'. 'hsDeclsPatBindD' \/ 'replaceDeclsPatBindD' is
 -- idempotent.
-replaceDeclsPatBindD :: (Monad m) => LHsDecl GhcPs -> [LHsDecl GhcPs]
-                     -> TransformT m (LHsDecl GhcPs)
-replaceDeclsPatBindD (L l (ValD x d)) newDecls = do
-  (L _ d') <- replaceDeclsPatBind (L l d) newDecls
-  return (L l (ValD x d'))
+replaceDeclsPatBindD :: LHsDecl GhcPs -> [LHsDecl GhcPs] -> (LHsDecl GhcPs)
+replaceDeclsPatBindD (L l (ValD x d)) newDecls =
+  let
+    (L _ d') = replaceDeclsPatBind (L l d) newDecls
+  in (L l (ValD x d'))
 replaceDeclsPatBindD x _ = error $ "replaceDeclsPatBindD called for:" ++ showGhc x
 
 -- | Replace the immediate declarations for a 'PatBind'. This
 -- cannot be a member of 'HasDecls' because a 'FunBind' is not idempotent
 -- for 'hsDecls' \/ 'replaceDecls'. 'hsDeclsPatBind' \/ 'replaceDeclsPatBind' is
 -- idempotent.
-replaceDeclsPatBind :: (Monad m) => LHsBind GhcPs -> [LHsDecl GhcPs]
-                    -> TransformT m (LHsBind GhcPs)
+replaceDeclsPatBind :: LHsBind GhcPs -> [LHsDecl GhcPs] -> (LHsBind GhcPs)
 replaceDeclsPatBind (L l (PatBind x a p (GRHSs xr rhss binds))) newDecls
-    = do
-        logTr "replaceDecls PatBind"
-        binds'' <- replaceDeclsValbinds WithWhere binds newDecls
-        return (L l (PatBind x a p (GRHSs xr rhss binds'')))
+  =  (L l (PatBind x a p (GRHSs xr rhss binds'')))
+  where
+    binds'' = replaceDeclsValbinds WithWhere binds newDecls
 replaceDeclsPatBind x _ = error $ "replaceDeclsPatBind called for:" ++ showGhc x
 
 -- ---------------------------------------------------------------------
 
 instance HasDecls (LocatedA (Stmt GhcPs (LocatedA (HsExpr GhcPs)))) where
-  hsDecls (L _ (LetStmt _ lb))      = return $ hsDeclsLocalBinds lb
+  hsDecls (L _ (LetStmt _ lb))      = hsDeclsLocalBinds lb
   hsDecls (L _ (LastStmt _ e _ _))  = hsDecls e
   hsDecls (L _ (BindStmt _ _pat e)) = hsDecls e
   hsDecls (L _ (BodyStmt _ e _ _))  = hsDecls e
-  hsDecls _                         = return []
+  hsDecls _                         = []
 
   replaceDecls (L l (LetStmt x lb)) newDecls
-    = do
-        lb'' <- replaceDeclsValbinds WithWhere lb newDecls
-        return (L l (LetStmt x lb''))
+    = let
+        lb'' = replaceDeclsValbinds WithWhere lb newDecls
+      in (L l (LetStmt x lb''))
   replaceDecls (L l (LastStmt x e d se)) newDecls
-    = do
-        e' <- replaceDecls e newDecls
-        return (L l (LastStmt x e' d se))
+    = let
+        e' = replaceDecls e newDecls
+      in (L l (LastStmt x e' d se))
   replaceDecls (L l (BindStmt x pat e)) newDecls
-    = do
-      e' <- replaceDecls e newDecls
-      return (L l (BindStmt x pat e'))
+    = let
+        e' = replaceDecls e newDecls
+      in (L l (BindStmt x pat e'))
 
   replaceDecls (L l (BodyStmt x e a b)) newDecls
-    = do
-      e' <- replaceDecls e newDecls
-      return (L l (BodyStmt x e' a b))
-  replaceDecls x _newDecls = return x
+    = let
+        e' = replaceDecls e newDecls
+      in (L l (BodyStmt x e' a b))
+  replaceDecls x _newDecls = x
 
 -- =====================================================================
 -- end of HasDecls instances
@@ -1062,61 +1059,55 @@ data WithWhere = WithWhere
 -- care, as this does not manage the declaration order, the
 -- ordering should be done by the calling function from the 'HsLocalBinds'
 -- context in the AST.
-replaceDeclsValbinds :: (Monad m)
-                     => WithWhere
+replaceDeclsValbinds :: WithWhere
                      -> HsLocalBinds GhcPs -> [LHsDecl GhcPs]
-                     -> TransformT m (HsLocalBinds GhcPs)
-replaceDeclsValbinds _ _ [] = do
-  return (EmptyLocalBinds NoExtField)
+                     -> HsLocalBinds GhcPs
+replaceDeclsValbinds _ _ [] = EmptyLocalBinds NoExtField
 replaceDeclsValbinds w b@(HsValBinds a _) new
-    = do
-        logTr "replaceDeclsValbinds"
-        let oldSpan = spanHsLocaLBinds b
-        an <- oldWhereAnnotation a w (realSrcSpan oldSpan)
-        let decs = concatMap decl2Bind new
-        let sigs = concatMap decl2Sig new
-        let sortKey = captureOrderBinds new
-        return (HsValBinds an (ValBinds sortKey decs sigs))
+    = let
+        oldSpan = spanHsLocaLBinds b
+        an = oldWhereAnnotation a w (realSrcSpan oldSpan)
+        decs = concatMap decl2Bind new
+        sigs = concatMap decl2Sig new
+        sortKey = captureOrderBinds new
+      in (HsValBinds an (ValBinds sortKey decs sigs))
 replaceDeclsValbinds _ (HsIPBinds {}) _new    = error "undefined replaceDecls HsIPBinds"
 replaceDeclsValbinds w (EmptyLocalBinds _) new
-    = do
-        logTr "replaceDecls HsLocalBinds"
-        an <- newWhereAnnotation w
-        let newBinds = concatMap decl2Bind new
-            newSigs  = concatMap decl2Sig  new
-        let decs = newBinds
-        let sigs = newSigs
-        let sortKey = captureOrderBinds new
-        return (HsValBinds an (ValBinds sortKey decs sigs))
-
-oldWhereAnnotation :: (Monad m)
-  => EpAnn AnnList -> WithWhere -> RealSrcSpan -> TransformT m (EpAnn AnnList)
-oldWhereAnnotation (EpAnn anc an cs) ww _oldSpan = do
-  -- TODO: when we set DP (0,0) for the HsValBinds EpEpaLocation, change the AnnList anchor to have the correct DP too
-  let (AnnList ancl o c _r t) = an
-  let w = case ww of
-        WithWhere -> [AddEpAnn AnnWhere (EpaDelta noSrcSpan (SameLine 0) [])]
-        WithoutWhere -> []
-  (anc', ancl') <- do
-        case ww of
-          WithWhere -> return (anc, ancl)
-          WithoutWhere -> return (anc, ancl)
-  let an' = EpAnn anc'
-                  (AnnList ancl' o c w t)
-                  cs
-  return an'
-
-newWhereAnnotation :: (Monad m) => WithWhere -> TransformT m (EpAnn AnnList)
-newWhereAnnotation ww = do
-  let anc  = EpaDelta noSrcSpan (DifferentLine 1 3) []
-  let anc2 = EpaDelta noSrcSpan (DifferentLine 1 5) []
-  let w = case ww of
-        WithWhere -> [AddEpAnn AnnWhere (EpaDelta noSrcSpan (SameLine 0) [])]
-        WithoutWhere -> []
-  let an = EpAnn anc
-                  (AnnList (Just anc2) Nothing Nothing w [])
-                  emptyComments
-  return an
+    = let
+        an = newWhereAnnotation w
+        decs = concatMap decl2Bind new
+        sigs = concatMap decl2Sig  new
+        sortKey = captureOrderBinds new
+      in (HsValBinds an (ValBinds sortKey decs sigs))
+
+oldWhereAnnotation :: EpAnn AnnList -> WithWhere -> RealSrcSpan -> (EpAnn AnnList)
+oldWhereAnnotation (EpAnn anc an cs) ww _oldSpan = an'
+  -- TODO: when we set DP (0,0) for the HsValBinds EpEpaLocation,
+  -- change the AnnList anchor to have the correct DP too
+  where
+    (AnnList ancl o c _r t) = an
+    w = case ww of
+      WithWhere -> [AddEpAnn AnnWhere (EpaDelta noSrcSpan (SameLine 0) [])]
+      WithoutWhere -> []
+    (anc', ancl') =
+          case ww of
+            WithWhere -> (anc, ancl)
+            WithoutWhere -> (anc, ancl)
+    an' = EpAnn anc'
+                (AnnList ancl' o c w t)
+                cs
+
+newWhereAnnotation :: WithWhere -> (EpAnn AnnList)
+newWhereAnnotation ww = an
+  where
+  anc  = EpaDelta noSrcSpan (DifferentLine 1 3) []
+  anc2 = EpaDelta noSrcSpan (DifferentLine 1 5) []
+  w = case ww of
+    WithWhere -> [AddEpAnn AnnWhere (EpaDelta noSrcSpan (SameLine 0) [])]
+    WithoutWhere -> []
+  an = EpAnn anc
+              (AnnList (Just anc2) Nothing Nothing w [])
+              emptyComments
 
 -- ---------------------------------------------------------------------
 
@@ -1127,32 +1118,32 @@ type PMatch = LMatch GhcPs (LHsExpr GhcPs)
 -- declarations are extracted and returned after modification. For a
 -- 'FunBind' the supplied 'SrcSpan' is used to identify the specific
 -- 'Match' to be transformed, for when there are multiple of them.
-modifyValD :: forall m t. (HasTransform m)
-                => SrcSpan
+modifyValD :: forall t.
+                   SrcSpan
                 -> Decl
-                -> (PMatch -> [Decl] -> m ([Decl], Maybe t))
-                -> m (Decl,Maybe t)
+                -> (PMatch -> [Decl] -> ([Decl], Maybe t))
+                -> (Decl,Maybe t)
 modifyValD p pb@(L ss (ValD _ (PatBind {} ))) f =
   if (locA ss) == p
-     then do
-       let ds = hsDeclsPatBindD pb
-       (ds',r) <- f (error "modifyValD.PatBind should not touch Match") ds
-       pb' <- liftT $ replaceDeclsPatBindD pb ds'
-       return (pb',r)
-     else return (pb,Nothing)
-modifyValD p decl f = do
-  (decl',r) <- runStateT (everywhereM (mkM doModLocal) (unpackFunDecl decl)) Nothing
-  return (packFunDecl decl',r)
+     then
+       let
+           ds = hsDeclsPatBindD pb
+           (ds',r) = f (error "modifyValD.PatBind should not touch Match") ds
+           pb' = replaceDeclsPatBindD pb ds'
+       in (pb',r)
+     else (pb,Nothing)
+modifyValD p decl f = (packFunDecl decl', r)
   where
-    doModLocal :: PMatch -> StateT (Maybe t) m PMatch
+    (decl',r) = runState (everywhereM (mkM doModLocal) (unpackFunDecl decl)) Nothing
+    doModLocal :: PMatch -> State (Maybe t) PMatch
     doModLocal  (match@(L ss _) :: PMatch) = do
          if (locA ss) == p
            then do
-             ds <- lift $ liftT $ hsDecls match
-                `debug` ("modifyValD: match=" ++ showAst match)
-             (ds',r) <- lift $ f match ds
-             put r
-             match' <- lift $ liftT $ replaceDecls match ds'
+             let
+               ds = hsDecls match
+               (ds',r0) = f match ds
+             put r0
+             let match' = replaceDecls match ds'
              return match'
            else return match
 
@@ -1172,6 +1163,6 @@ modifyDeclsT :: (HasDecls t,HasTransform m)
              => ([LHsDecl GhcPs] -> m [LHsDecl GhcPs])
              -> t -> m t
 modifyDeclsT action t = do
-  decls <- liftT $ hsDecls t
+  let decls = hsDecls t
   decls' <- action decls
-  liftT $ replaceDecls t decls'
+  return $ replaceDecls t decls'


=====================================
utils/check-exact/Types.hs
=====================================
@@ -21,10 +21,6 @@ type Pos = (Int,Int)
 
 -- ---------------------------------------------------------------------
 
-data Rigidity = NormalLayout | RigidLayout deriving (Eq, Ord, Show)
-
--- ---------------------------------------------------------------------
-
 -- | A Haskell comment. The @AnnKeywordId@ is present if it has been converted
 -- from an @AnnKeywordId@ because the annotation must be interleaved into the
 -- stream and does not have a well-defined position


=====================================
utils/check-exact/Utils.hs
=====================================
@@ -20,9 +20,8 @@ module Utils
   where
 
 import Control.Monad (when)
+import GHC.Utils.Monad.State.Strict
 import Data.Function
-import Data.Maybe (isJust)
-import Data.Ord (comparing)
 
 import GHC.Hs.Dump
 import Lookup
@@ -36,9 +35,10 @@ import GHC.Driver.Ppr
 import GHC.Data.FastString
 import qualified GHC.Data.Strict as Strict
 import GHC.Base (NonEmpty(..))
+import GHC.Parser.Lexer (allocateComments)
 
 import Data.Data hiding ( Fixity )
-import Data.List (sortBy, elemIndex)
+import Data.List (sortBy, partition)
 import qualified Data.Map.Strict as Map
 
 import Debug.Trace
@@ -60,12 +60,32 @@ debug c s = if debugEnabledFlag
 debugM :: Monad m => String -> m ()
 debugM s = when debugEnabledFlag $ traceM s
 
--- ---------------------------------------------------------------------
-
 warn :: c -> String -> c
 -- warn = flip trace
 warn c _ = c
 
+-- ---------------------------------------------------------------------
+
+captureOrderBinds :: [LHsDecl GhcPs] -> AnnSortKey BindTag
+captureOrderBinds ls = AnnSortKey $ map go ls
+  where
+    go (L _ (ValD _ _))       = BindTag
+    go (L _ (SigD _ _))       = SigDTag
+    go d      = error $ "captureOrderBinds:" ++ showGhc d
+
+-- ---------------------------------------------------------------------
+
+notDocDecl :: LHsDecl GhcPs -> Bool
+notDocDecl (L _ DocD{}) = False
+notDocDecl _ = True
+
+notIEDoc :: LIE GhcPs -> Bool
+notIEDoc (L _ IEGroup {})    = False
+notIEDoc (L _ IEDoc {})      = False
+notIEDoc (L _ IEDocNamed {}) = False
+notIEDoc _ = True
+
+-- ---------------------------------------------------------------------
 -- | A good delta has no negative values.
 isGoodDelta :: DeltaPos -> Bool
 isGoodDelta (SameLine co) = co >= 0
@@ -108,7 +128,6 @@ pos2delta (refl,refc) (l,c) = deltaPos lo co
     lo = l - refl
     co = if lo == 0 then c - refc
                     else c
-                    -- else c - 1
 
 -- | Apply the delta to the current position, taking into account the
 -- current column offset if advancing to a new line
@@ -200,23 +219,6 @@ origDelta pos pp = ss2delta (ss2posEnd pp) pos
 
 -- ---------------------------------------------------------------------
 
--- |Given a list of items and a list of keys, returns a list of items
--- ordered by their position in the list of keys.
-orderByKey :: [(DeclTag,a)] -> [DeclTag] -> [(DeclTag,a)]
-orderByKey keys order
-    -- AZ:TODO: if performance becomes a problem, consider a Map of the order
-    -- SrcSpan to an index, and do a lookup instead of elemIndex.
-
-    -- Items not in the ordering are placed to the start
- = sortBy (comparing (flip elemIndex order . fst)) keys
-
--- ---------------------------------------------------------------------
-
-isListComp :: HsDoFlavour -> Bool
-isListComp = isDoComprehensionContext
-
--- ---------------------------------------------------------------------
-
 needsWhere :: DataDefnCons (LConDecl (GhcPass p)) -> Bool
 needsWhere (NewTypeCon _) = True
 needsWhere (DataTypeCons _ []) = True
@@ -225,21 +227,214 @@ needsWhere _ = False
 
 -- ---------------------------------------------------------------------
 
+-- | Insert the comments at the appropriate places in the AST
 insertCppComments ::  ParsedSource -> [LEpaComment] -> ParsedSource
-insertCppComments (L l p) cs = L l p'
+-- insertCppComments p [] = p
+insertCppComments (L l p) cs0 = insertRemainingCppComments (L l p2) remaining
+  where
+    (EpAnn anct ant cst) = hsmodAnn $ hsmodExt p
+    cs = sortEpaComments $ priorComments cst ++ getFollowingComments cst ++ cs0
+    p0 = p { hsmodExt = (hsmodExt p) { hsmodAnn = EpAnn anct ant emptyComments }}
+    -- Comments embedded within spans
+    -- everywhereM is a bottom-up traversal
+    (p1, toplevel) = runState (everywhereM (mkM   addCommentsListItem
+                                           `extM` addCommentsGrhs
+                                           `extM` addCommentsList) p0) cs
+    (p2, remaining) = insertTopLevelCppComments p1 toplevel
+
+    addCommentsListItem :: EpAnn AnnListItem -> State [LEpaComment] (EpAnn AnnListItem)
+    addCommentsListItem = addComments
+
+    addCommentsList :: EpAnn AnnList -> State [LEpaComment] (EpAnn AnnList)
+    addCommentsList = addComments
+
+    addCommentsGrhs :: EpAnn GrhsAnn -> State [LEpaComment] (EpAnn GrhsAnn)
+    addCommentsGrhs = addComments
+
+    addComments :: forall ann. EpAnn ann -> State [LEpaComment] (EpAnn ann)
+    addComments (EpAnn anc an ocs) = do
+      case anc of
+        EpaSpan (RealSrcSpan s _) -> do
+          unAllocated <- get
+          let
+            (rest, these) = GHC.Parser.Lexer.allocateComments s unAllocated
+            cs' = workInComments ocs these
+          put rest
+          return $ EpAnn anc an cs'
+
+        _ -> return $ EpAnn anc an ocs
+
+workInComments :: EpAnnComments -> [LEpaComment] -> EpAnnComments
+workInComments ocs [] = ocs
+workInComments ocs new = cs'
+  where
+    pc = priorComments ocs
+    fc = getFollowingComments ocs
+    cs' = case fc of
+      [] -> EpaComments $ sortEpaComments $ pc ++ fc ++ new
+      (L ac _:_) -> epaCommentsBalanced (sortEpaComments $ pc ++ cs_before)
+                                        (sortEpaComments $ fc ++ cs_after)
+             where
+               (cs_before,cs_after)
+                   = break (\(L ll _) -> (ss2pos $ anchor ll) > (ss2pos $ anchor ac) )
+                           new
+
+insertTopLevelCppComments ::  HsModule GhcPs -> [LEpaComment] -> (HsModule GhcPs, [LEpaComment])
+insertTopLevelCppComments (HsModule (XModulePs an lo mdeprec mbDoc) mmn mexports imports decls) cs
+  = (HsModule (XModulePs an4 lo mdeprec mbDoc) mmn mexports' imports' decls', cs3)
+    -- `debug` ("insertTopLevelCppComments: (cs2,cs3,hc0,hc1,hc_cs)" ++ showAst (cs2,cs3,hc0,hc1,hc_cs))
+    -- `debug` ("insertTopLevelCppComments: (cs2,cs3,hc0i,hc0,hc1,hc_cs)" ++ showAst (cs2,cs3,hc0i,hc0,hc1,hc_cs))
   where
-    an' = case GHC.hsmodAnn $ GHC.hsmodExt p of
-      (EpAnn a an ocs) -> EpAnn a an cs'
-        where
-          pc = priorComments ocs
-          fc = getFollowingComments ocs
-          cs' = case fc of
-            [] -> EpaComments $ sortEpaComments $ pc ++ fc ++ cs
-            (L ac _:_) -> EpaCommentsBalanced (sortEpaComments $ pc ++ cs_before)
-                                              (sortEpaComments $ fc ++ cs_after)
-                   where
-                     (cs_before,cs_after) = break (\(L ll _) ->   (ss2pos $ anchor ll) < (ss2pos $ anchor ac) ) cs
+    -- Comments at the top level.
+    (an0, cs0) =
+      case mmn of
+        Nothing -> (an, cs)
+        Just _ ->
+            -- We have a module name. Capture all comments up to the `where`
+            let
+              (these, remaining) = splitOnWhere Before (am_main $ anns an) cs
+              (EpAnn a anno ocs) = an :: EpAnn AnnsModule
+              anm = EpAnn a anno (workInComments ocs these)
+            in
+              (anm, remaining)
+    (an1,cs0a) = case lo of
+        EpExplicitBraces (EpTok (EpaSpan (RealSrcSpan s _))) _close ->
+            let
+                (stay,cs0a') = break (\(L ll _) -> (ss2pos $ anchor ll) > (ss2pos $ s)) cs0
+                cs' = workInComments (comments an0) stay
+            in (an0 { comments = cs' }, cs0a')
+        _ -> (an0,cs0)
+    -- Deal with possible leading semis
+    (an2, cs0b) = case am_decls $ anns an1 of
+        (AddSemiAnn (EpaSpan (RealSrcSpan s _)):_) -> (an1 {comments = cs'}, cs0b')
+          where
+            (stay,cs0b') = break (\(L ll _) -> (ss2pos $ anchor ll) > (ss2pos $ s)) cs0a
+            cs' = workInComments (comments an1) stay
+        _ -> (an1,cs0a)
+
+    (mexports', an3, cs1) =
+      case mexports of
+        Nothing -> (Nothing, an2, cs0b)
+        Just (L l exports) -> (Just (L l exports'), an3', cse)
+                         where
+                           hc1' = workInComments (comments an2) csh'
+                           an3' = an2 { comments = hc1' }
+                           (csh', cs0b') = case al_open $ anns l of
+                               Just (AddEpAnn _ (EpaSpan (RealSrcSpan s _))) ->(h, n)
+                                 where
+                                   (h,n) = break (\(L ll _) -> (ss2pos $ anchor ll) > (ss2pos s) )
+                                       cs0b
+
+                               _ -> ([], cs0b)
+                           (exports', cse) = allocPreceding exports cs0b'
+    (imports0, cs2) = allocPreceding imports cs1
+    (imports', hc0i) = balanceFirstLocatedAComments imports0
+
+    (decls0, cs3) = allocPreceding decls cs2
+    (decls', hc0d) = balanceFirstLocatedAComments decls0
+
+    -- Either hc0i or hc0d should have comments. Combine them
+    hc0 = hc0i ++ hc0d
+
+    (hc1,hc_cs) = if null ( am_main $ anns an3)
+        then (hc0,[])
+        else splitOnWhere After (am_main $ anns an3)  hc0
+    hc2 = workInComments (comments an3) hc1
+    an4 = an3 { anns = (anns an3) {am_cs = hc_cs}, comments = hc2 }
+
+    allocPreceding :: [LocatedA a] -> [LEpaComment] -> ([LocatedA a], [LEpaComment])
+    allocPreceding [] cs' = ([], cs')
+    allocPreceding (L (EpAnn anc4 an5 cs4) a:xs) cs' = ((L (EpAnn anc4 an5 cs4') a:xs'), rest')
+      where
+        (rest, these) =
+          case anc4 of
+            EpaSpan (RealSrcSpan s _) ->
+                allocatePriorComments (ss2pos s) cs'
+            _ -> (cs', [])
+        cs4' = workInComments cs4 these
+        (xs',rest') = allocPreceding xs rest
+
+data SplitWhere = Before | After
+splitOnWhere :: SplitWhere -> [AddEpAnn] -> [LEpaComment] -> ([LEpaComment], [LEpaComment])
+splitOnWhere _ [] csIn = (csIn,[])
+splitOnWhere w (AddEpAnn AnnWhere (EpaSpan (RealSrcSpan s _)):_) csIn = (hc, fc)
+  where
+    splitFunc Before anc_pos c_pos = c_pos < anc_pos
+    splitFunc After  anc_pos c_pos = anc_pos < c_pos
+    (hc,fc) = break (\(L ll _) ->  splitFunc w (ss2pos $ anchor ll) (ss2pos s)) csIn
+splitOnWhere _ (AddEpAnn AnnWhere _:_) csIn = (csIn, [])
+splitOnWhere f (_:as) csIn = splitOnWhere f as csIn
+
+balanceFirstLocatedAComments :: [LocatedA a] -> ([LocatedA a], [LEpaComment])
+balanceFirstLocatedAComments [] = ([],[])
+balanceFirstLocatedAComments ((L (EpAnn anc an csd) a):ds) = (L (EpAnn anc an csd0) a:ds, hc')
+  where
+    (csd0, hc') = case anc of
+        EpaSpan (RealSrcSpan s _) -> (csd', hc)
+               `debug` ("balanceFirstLocatedAComments: (csd,csd',attached,header)=" ++ showAst (csd,csd',attached,header))
+          where
+            (priors, inners) =  break (\(L ll _) -> (ss2pos $ anchor ll) > (ss2pos s) )
+                                       (priorComments csd)
+            pcds = priorCommentsDeltas' s priors
+            (attached, header) = break (\(d,_c) -> d /= 1) pcds
+            csd' = setPriorComments csd (reverse (map snd attached) ++ inners)
+            hc = reverse (map snd header)
+        _ -> (csd, [])
+
+
+
+priorCommentsDeltas' :: RealSrcSpan -> [LEpaComment]
+                    -> [(Int, LEpaComment)]
+priorCommentsDeltas' r cs = go r (reverse cs)
+  where
+    go :: RealSrcSpan -> [LEpaComment] -> [(Int, LEpaComment)]
+    go _   [] = []
+    go _   (la@(L l@(EpaDelta _ dp _) _):las) = (deltaLine dp, la) : go (anchor l) las
+    go rs' (la@(L l _):las) = deltaComment rs' la : go (anchor l) las
+
+    deltaComment :: RealSrcSpan -> LEpaComment -> (Int, LEpaComment)
+    deltaComment rs' (L loc c) = (abs(ll - al), L loc c)
+      where
+        (al,_) = ss2pos rs'
+        (ll,_) = ss2pos (anchor loc)
+
+allocatePriorComments
+  :: Pos
+  -> [LEpaComment]
+  -> ([LEpaComment], [LEpaComment])
+allocatePriorComments ss_loc comment_q =
+  let
+    cmp (L l _) = ss2pos (anchor l) <= ss_loc
+    (newAnns,after) = partition cmp comment_q
+  in
+    (after, newAnns)
+
+insertRemainingCppComments ::  ParsedSource -> [LEpaComment] -> ParsedSource
+insertRemainingCppComments (L l p) cs = L l p'
+    -- `debug` ("insertRemainingCppComments: (cs,an')=" ++ showAst (cs,an'))
+  where
+    (EpAnn a an ocs) = GHC.hsmodAnn $ GHC.hsmodExt p
+    an' = EpAnn a an (addTrailingComments end_loc ocs cs)
     p' = p { GHC.hsmodExt = (GHC.hsmodExt p) { GHC.hsmodAnn = an' } }
+    end_loc = case GHC.hsmodLayout $ GHC.hsmodExt p of
+        EpExplicitBraces _open close -> case close of
+            EpTok (EpaSpan (RealSrcSpan s _)) -> ss2pos s
+            _ -> (1,1)
+        _ -> (1,1)
+    (new_before, new_after) = break (\(L ll _) -> (ss2pos $ anchor ll) > end_loc ) cs
+
+    addTrailingComments end_loc' cur new = epaCommentsBalanced pc' fc'
+      where
+        pc = priorComments cur
+        fc = getFollowingComments cur
+        (pc', fc') = case reverse pc of
+            [] -> (sortEpaComments $ pc ++ new_before, sortEpaComments $ fc ++ new_after)
+            (L ac _:_) -> (sortEpaComments $ pc ++ cs_before, sortEpaComments $ fc ++ cs_after)
+              where
+               (cs_before,cs_after)
+                   = if (ss2pos $ anchor ac) > end_loc'
+                       then break (\(L ll _) -> (ss2pos $ anchor ll) > (ss2pos $ anchor ac) ) new
+                       else (new_before, new_after)
 
 -- ---------------------------------------------------------------------
 
@@ -291,11 +486,16 @@ dedentDocChunkBy  dedent (L (RealSrcSpan l mb) c) = L (RealSrcSpan l' mb) c
 
 dedentDocChunkBy _ x = x
 
+
+epaCommentsBalanced :: [LEpaComment] -> [LEpaComment] -> EpAnnComments
+epaCommentsBalanced priorCs     [] = EpaComments priorCs
+epaCommentsBalanced priorCs postCs = EpaCommentsBalanced priorCs postCs
+
 mkEpaComments :: [Comment] -> [Comment] -> EpAnnComments
 mkEpaComments priorCs []
   = EpaComments (map comment2LEpaComment priorCs)
 mkEpaComments priorCs postCs
-  = EpaCommentsBalanced (map comment2LEpaComment priorCs) (map comment2LEpaComment postCs)
+  = epaCommentsBalanced (map comment2LEpaComment priorCs) (map comment2LEpaComment postCs)
 
 comment2LEpaComment :: Comment -> LEpaComment
 comment2LEpaComment (Comment s anc r _mk) = mkLEpaComment s anc r
@@ -330,18 +530,11 @@ sortEpaComments cs = sortBy cmp cs
 mkKWComment :: AnnKeywordId -> NoCommentsLocation -> Comment
 mkKWComment kw (EpaSpan (RealSrcSpan ss mb))
   = Comment (keywordToString kw) (EpaSpan (RealSrcSpan ss mb)) ss (Just kw)
-mkKWComment kw (EpaSpan ss@(UnhelpfulSpan _))
-  = Comment (keywordToString kw) (EpaDelta ss (SameLine 0) NoComments) placeholderRealSpan (Just kw)
+mkKWComment kw (EpaSpan (UnhelpfulSpan _))
+  = Comment (keywordToString kw) (EpaDelta noSrcSpan (SameLine 0) NoComments) placeholderRealSpan (Just kw)
 mkKWComment kw (EpaDelta ss dp cs)
   = Comment (keywordToString kw) (EpaDelta ss dp cs) placeholderRealSpan (Just kw)
 
--- | Detects a comment which originates from a specific keyword.
-isKWComment :: Comment -> Bool
-isKWComment c = isJust (commentOrigin c)
-
-noKWComments :: [Comment] -> [Comment]
-noKWComments = filter (\c -> not (isKWComment c))
-
 sortAnchorLocated :: [GenLocated EpaLocation a] -> [GenLocated EpaLocation a]
 sortAnchorLocated = sortBy (compare `on` (anchor . getLoc))
 
@@ -379,11 +572,6 @@ name2String = showPprUnsafe
 
 -- ---------------------------------------------------------------------
 
-locatedAnAnchor :: LocatedAn a t -> RealSrcSpan
-locatedAnAnchor (L (EpAnn a _ _) _) = anchor a
-
--- ---------------------------------------------------------------------
-
 trailingAnnLoc :: TrailingAnn -> EpaLocation
 trailingAnnLoc (AddSemiAnn ss)    = ss
 trailingAnnLoc (AddCommaAnn ss)   = ss
@@ -401,46 +589,6 @@ setTrailingAnnLoc (AddDarrowUAnn _) ss = (AddDarrowUAnn ss)
 addEpAnnLoc :: AddEpAnn -> EpaLocation
 addEpAnnLoc (AddEpAnn _ l) = l
 
--- ---------------------------------------------------------------------
--- Horrible hack for dealing with some things still having a SrcSpan,
--- not an Anchor.
-
-{-
-A SrcSpan is defined as
-
-data SrcSpan =
-    RealSrcSpan !RealSrcSpan !(Maybe BufSpan)  -- See Note [Why Maybe BufPos]
-  | UnhelpfulSpan !UnhelpfulSpanReason
-
-data BufSpan =
-  BufSpan { bufSpanStart, bufSpanEnd :: {-# UNPACK #-} !BufPos }
-  deriving (Eq, Ord, Show)
-
-newtype BufPos = BufPos { bufPos :: Int }
-
-
-We use the BufPos to encode a delta, using bufSpanStart for the line,
-and bufSpanEnd for the col.
-
-To be absolutely sure, we make the delta versions use -ve values.
-
--}
-
-hackSrcSpanToAnchor :: SrcSpan -> EpaLocation
-hackSrcSpanToAnchor (UnhelpfulSpan s) = error $ "hackSrcSpanToAnchor : UnhelpfulSpan:" ++ show s
-hackSrcSpanToAnchor ss@(RealSrcSpan r mb)
-  = case mb of
-    (Strict.Just (BufSpan (BufPos s) (BufPos e))) ->
-      if s <= 0 && e <= 0
-      then EpaDelta ss (deltaPos (-s) (-e)) []
-        `debug` ("hackSrcSpanToAnchor: (r,s,e)=" ++ showAst (r,s,e) )
-      else EpaSpan (RealSrcSpan r mb)
-    _ -> EpaSpan (RealSrcSpan r mb)
-
-hackAnchorToSrcSpan :: EpaLocation -> SrcSpan
-hackAnchorToSrcSpan (EpaSpan s) = s
-hackAnchorToSrcSpan _ = error $ "hackAnchorToSrcSpan"
-
 -- ---------------------------------------------------------------------
 
 type DeclsByTag a = Map.Map DeclTag [(RealSrcSpan, a)]



View it on GitLab: https://gitlab.haskell.org/ghc/ghc/-/compare/02733138052c08b8e5149f4ca8d21712d3d23f11...dda9c763e92ae2c27a103b226e501d2f708e902b

-- 
View it on GitLab: https://gitlab.haskell.org/ghc/ghc/-/compare/02733138052c08b8e5149f4ca8d21712d3d23f11...dda9c763e92ae2c27a103b226e501d2f708e902b
You're receiving this email because of your account on gitlab.haskell.org.


-------------- next part --------------
An HTML attachment was scrubbed...
URL: <http://mail.haskell.org/pipermail/ghc-commits/attachments/20240912/11bc9713/attachment-0001.html>


More information about the ghc-commits mailing list