summaryrefslogtreecommitdiff
path: root/src/Language/Haskell/GHC/ExactPrint/Delta.hs
blob: 829d9884a3ec8541565836ba4b7058aeab3b6df1 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
{-# LANGUAGE CPP #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE ViewPatterns #-}

-- |  This module converts 'GHC.ApiAnns' into 'Anns' by traversing a
-- structure created by the "Annotate" module.
--
-- == Structure of an Annotation
--
-- As a rule of thumb, every located element in the GHC AST will have
-- a corresponding entry in 'Anns'. An 'Annotation' contains 6 fields which
-- can be modifed to change how the AST is printed.
--
-- == Layout Calculation
--
-- In order to properly place syntax nodes and comments properly after
-- refactoring them (in such a way that the indentation level changes), their
-- position (encoded in the 'addEntryDelta' field) is not expressed as absolute
-- but relative to their context. As further motivation, consider the simple
-- let-into-where-block refactoring, from:
--
-- @
-- foo = do
--   let bar = do
--         x
--         -- comment
--         y
--   bar
-- @
--
-- to
--
-- @
-- foo = do
--   bar
--  where
--   bar = do
--     x
--     -- comment
--     y
-- @
--
-- Notice how the column of @x@, @y@ and the comment change due to this
-- refactoring but certain relative positions (e.g. the comment starting at the
-- same column as @x@) remain unchanged.
--
-- Now, what does "context" mean exactly? Here we reference the
-- "indentation level" as used in the haskell report (see chapter 2.7:
-- <https://www.haskell.org/onlinereport/haskell2010/haskellch2.html#x7-210002.7>):
-- 'addEntryDelta' is mostly relative to the current (inner-most) indentation
-- level. But in order to get better results, for the purpose of defining
-- relative positions a the offside-rule is modified slightly: Normally it
-- fires (only) at the first elements after where/let/do/of, introducing a new
-- indentation level. In addition, the rule here fires also at the "@let@"
-- keyword (when it is part of a "@let-in@" construct) and at the "@if@" keyword.
--
-- The effect of this additional applications of the offside-rule is that any
-- elements (more or less directly) following the "@let@" ("@if@"")
-- keyword have a position relative to the "@let@" ("@if@")
-- keyword position, even when the regular offside-rule does apply not yet/not
-- anymore. This affects two concrete things: Comments directly following
-- "@let@"/"@if@", and the respective follow-up keywords: "@in@" or
-- "@then@"/"@else@".
--
-- Due to this additional indentation level, it is possible to observe/obtain
-- negative delta-positions; consider:
--
-- @
-- foo = let x = 1
--   in x
-- @
--
-- Here, the @in@ keyword has an 'annEntryDelta' of @DP (1, -4)@ as it appears
-- one line below the previous elements and 4 columns /left/ relative to the
-- start of the @let@ keyword.
--
-- In general, the element that defines such an indentation level (i.e. the
-- first element after a where/let/do/of) will have an 'annEntryDelta' relative
-- to the previous inner-most indentation level; in other words: a new
-- indentation level becomes relevant only after the construct introducing the
-- element received its 'annEntryDelta' position. (Otherwise these elements
-- always would have a zero horizontal position - relative to itself.)
--
-- (This affects comments, too: A comment preceding the first element of a
-- layout block will have a position relative to the outer block, not of the
-- newly introduced layout block.)
--
-- For example, in the following expression the statement corresponding to
-- @baz@ will be given a 'annEntryDelta' of @DP (1, 2)@ as it appears
-- 1 line and 2 columns after the @do@ keyword. On the other hand, @bar@
-- will be given a 'annEntryDelta' of @DP (1,0)@ as it appears 1 line
-- further than @baz@ but in the same column as the start of the layout
-- block.
--
-- @
-- foo = do
--   baz
--   bar
-- @
--
-- A useful way to think of these rules is that the 'DeltaPos' is relative
-- to the further left an expression could have been placed. In the
-- previous example, we could have placed @baz@ anywhere on the line as its
-- position determines where the other statements must be. @bar@ could have
-- not been placed any further left without resulting in a syntax error
-- which is why the relative column is 0.
--
-- === annTrueEntryDelta
-- A very useful function is 'annTrueEntryDelta' which calculates the
-- offset from the last syntactic element (ignoring comments). This is
-- different to 'annEntryDelta' which does not ignore comments.
--
--
--
module Language.Haskell.GHC.ExactPrint.Delta
  ( relativiseApiAnns
  , relativiseApiAnnsWithComments
  , relativiseApiAnnsWithOptions

  -- * Configuration
  , DeltaOptions(drRigidity)
  , deltaOptions
  , normalLayout
  ) where

-- import Control.Exception
import Control.Monad.RWS
import Control.Monad.Trans.Free

import Data.Data (Data)
import Data.List (sort, nub, partition, sortBy)

import Data.Ord

import Language.Haskell.GHC.ExactPrint.Utils
#if __GLASGOW_HASKELL__ <= 710
import Language.Haskell.GHC.ExactPrint.Lookup
#endif
import Language.Haskell.GHC.ExactPrint.Types
import Language.Haskell.GHC.ExactPrint.Annotate

import qualified GHC

import qualified Data.Map as Map
import qualified Data.Set as Set

-- import Debug.Trace

{-# ANN module "HLint: ignore Eta reduce" #-}
{-# ANN module "HLint: ignore Redundant do" #-}
{-# ANN module "HLint: ignore Reduce duplication" #-}

-- ---------------------------------------------------------------------
-- | Transform concrete annotations into relative annotations which are
-- more useful when transforming an AST.
#if __GLASGOW_HASKELL__ > 806
relativiseApiAnns :: (Data (GHC.SrcSpanLess ast), Annotate ast, GHC.HasSrcSpan ast)
                  => ast
#else
relativiseApiAnns :: Annotate ast
                  => GHC.Located ast
#endif
                  -> GHC.ApiAnns
                  -> Anns
relativiseApiAnns = relativiseApiAnnsWithComments []

-- | Exactly the same as 'relativiseApiAnns' but with the possibilty to
-- inject comments. This is typically used if the source has been preprocessed
-- by e.g. CPP, and the parts stripped out of the original source are re-added
-- as comments so they are not lost for round tripping.
relativiseApiAnnsWithComments ::
#if __GLASGOW_HASKELL__ > 806
                     (Data (GHC.SrcSpanLess ast), Annotate ast, GHC.HasSrcSpan ast)
                  => [Comment]
                  -> ast
#else
                     Annotate ast
                  => [Comment]
                  -> GHC.Located ast
#endif
                  -> GHC.ApiAnns
                  -> Anns
relativiseApiAnnsWithComments =
    relativiseApiAnnsWithOptions normalLayout

relativiseApiAnnsWithOptions ::
#if __GLASGOW_HASKELL__ > 806
                     (Data (GHC.SrcSpanLess ast), Annotate ast, GHC.HasSrcSpan ast)
                  => DeltaOptions
                  -> [Comment]
                  -> ast
#else
                     Annotate ast
                  => DeltaOptions
                  -> [Comment]
                  -> GHC.Located ast
#endif
                  -> GHC.ApiAnns
                  -> Anns
relativiseApiAnnsWithOptions opts cs modu ghcAnns
   = runDeltaWithComments
      opts cs (annotate modu) ghcAnns
      (ss2pos $ GHC.getLoc modu)

-- ---------------------------------------------------------------------
--
-- | Type used in the Delta Monad.
type Delta a = RWS DeltaOptions DeltaWriter DeltaState a

runDeltaWithComments :: DeltaOptions -> [Comment] -> Annotated () -> GHC.ApiAnns -> Pos -> Anns
runDeltaWithComments opts cs action ga priorEnd =
  mkAnns . snd
  . (\next -> execRWS next opts (defaultDeltaState cs priorEnd ga))
  . deltaInterpret $ action
  where
    mkAnns :: DeltaWriter -> Anns
    mkAnns = f . dwAnns
    f :: Monoid a => Endo a -> a
    f = ($ mempty) . appEndo

-- ---------------------------------------------------------------------

-- TODO: rename this, it is the R part of the RWS
data DeltaOptions = DeltaOptions
       {
         -- | Current `SrcSpan, part of current AnnKey`
         curSrcSpan :: !GHC.SrcSpan

         -- | Constuctor of current AST element, part of current AnnKey
       , annConName :: !AnnConName

        -- | Whether to use rigid or normal layout rules
       , drRigidity :: !Rigidity

       -- | Current higher level context. e.g. whether a Match is part of a
       -- LambdaExpr or a FunBind
       , drContext :: !AstContextSet
       }

data DeltaWriter = DeltaWriter
       { -- | Final list of annotations, and sort keys
         dwAnns :: Endo (Map.Map AnnKey Annotation)

         -- | Used locally to pass Keywords, delta pairs relevant to a specific
         -- subtree to the parent.
       , annKds         :: ![(KeywordId, DeltaPos)]
       , sortKeys       :: !(Maybe [GHC.SrcSpan])
       , dwCapturedSpan :: !(First AnnKey)
       }

data DeltaState = DeltaState
       { -- | Position reached when processing the last element
         priorEndPosition    :: !Pos

         -- | Ordered list of comments still to be allocated
       , apComments :: ![Comment]

         -- | The original GHC Delta Annotations
       , apAnns :: !GHC.ApiAnns

       , apMarkLayout :: Bool
       , apLayoutStart :: LayoutStartCol

       }

-- ---------------------------------------------------------------------

deltaOptions :: Rigidity -> DeltaOptions
deltaOptions ridigity =
  DeltaOptions
    { curSrcSpan = GHC.noSrcSpan
    , annConName = annGetConstr ()
    , drRigidity = ridigity
    , drContext  = defaultACS
    }

normalLayout :: DeltaOptions
normalLayout = deltaOptions NormalLayout

defaultDeltaState :: [Comment] -> Pos -> GHC.ApiAnns -> DeltaState
defaultDeltaState injectedComments priorEnd ga =
    DeltaState
      { priorEndPosition    = priorEnd
      , apComments = cs ++ injectedComments
      , apAnns     = ga
      , apLayoutStart = 1
      , apMarkLayout = False
      }
  where
    cs :: [Comment]
    cs = extractComments ga


-- Writer helpers

tellFinalAnn :: (AnnKey, Annotation) -> Delta ()
tellFinalAnn (k, v) =
  -- tell (mempty { dwAnns = Endo (Map.insertWith (<>) k v) })
  tell (mempty { dwAnns = Endo (Map.insert k v) })

tellSortKey :: [GHC.SrcSpan] -> Delta ()
tellSortKey xs = tell (mempty { sortKeys = Just xs } )

tellCapturedSpan :: AnnKey -> Delta ()
tellCapturedSpan key = tell ( mempty { dwCapturedSpan = First $ Just key })

tellKd :: (KeywordId, DeltaPos) -> Delta ()
tellKd kd = tell (mempty { annKds = [kd] })

#if __GLASGOW_HASKELL__ >= 804
instance Semigroup DeltaWriter where
  (<>) = mappend
#endif

instance Monoid DeltaWriter where
  mempty = DeltaWriter mempty mempty mempty mempty
  (DeltaWriter a b e g) `mappend` (DeltaWriter c d f h)
    = DeltaWriter (a <> c) (b <> d) (e <> f) (g <> h)

-----------------------------------
-- Free Monad Interpretation code

deltaInterpret :: Annotated a -> Delta a
deltaInterpret = iterTM go
  where
    go :: AnnotationF (Delta a) -> Delta a
    go (MarkEOF next)                   = addEofAnnotation >> next
    go (MarkPrim kwid _ next)           = addDeltaAnnotation kwid >> next
    go (MarkPPOptional kwid _ next)     = addDeltaAnnotation kwid >> next
#if __GLASGOW_HASKELL__ >= 800
    go (MarkInstead akwid kwid next)    = addDeltaAnnotationInstead akwid kwid >> next
#endif
    go (MarkOutside akwid kwid next)    = addDeltaAnnotationsOutside akwid kwid >> next
    go (MarkInside akwid next)          = addDeltaAnnotationsInside akwid >> next
    go (MarkMany akwid next)            = addDeltaAnnotations akwid >> next
    go (MarkManyOptional akwid next)    = addDeltaAnnotations akwid >> next
    go (MarkOffsetPrim akwid n _ next)  = addDeltaAnnotationLs akwid n >> next
    go (MarkOffsetPrimOptional akwid n _ next) = addDeltaAnnotationLs akwid n >> next
    go (WithAST lss prog next)          = withAST lss (deltaInterpret prog) >> next
    go (CountAnns kwid next)            = countAnnsDelta kwid >>= next
    go (SetLayoutFlag r action next)    = do
      rigidity <- asks drRigidity
      (if r <= rigidity then setLayoutFlag else id) (deltaInterpret action)
      next
    go (MarkAnnBeforeAnn ann1 ann2 next) = deltaMarkAnnBeforeAnn ann1 ann2 >> next
    go (MarkExternal ss akwid _ next)    = addDeltaAnnotationExt ss akwid >> next
    go (StoreOriginalSrcSpan _ key next) = storeOriginalSrcSpanDelta key >>= next
    go (GetSrcSpanForKw ss kw next)      = getSrcSpanForKw ss kw >>= next
#if __GLASGOW_HASKELL__ <= 710
    go (StoreString s ss next)           = storeString s ss >> next
#endif
    go (AnnotationsToComments     kws next) = annotationsToCommentsDelta kws >> next
#if __GLASGOW_HASKELL__ <= 710
    go (AnnotationsToCommentsBF _ kws next) = annotationsToCommentsDelta kws >> next
    go (FinalizeBF _ next)                  = next
#endif
    go (WithSortKey             kws next) = withSortKey kws >> next
    go (WithSortKeyContexts ctx kws next) = withSortKeyContexts ctx kws >> next

    go (SetContextLevel ctxt lvl action next) = setContextDelta ctxt lvl (deltaInterpret action) >> next
    go (UnsetContext   _ctxt action next) = deltaInterpret action >> next
    go (IfInContext    ctxt ifAction elseAction next) = ifInContextDelta ctxt ifAction elseAction >> next
    go (TellContext _ next)                  = next

withSortKey :: [(GHC.SrcSpan, Annotated b)] -> Delta ()
withSortKey kws =
  let order = sortBy (comparing fst) kws
  in do
    tellSortKey (map fst order)
    mapM_ (deltaInterpret . snd) order


withSortKeyContexts :: ListContexts -> [(GHC.SrcSpan, Annotated ())] -> Delta ()
withSortKeyContexts ctxts kws = do
  tellSortKey (map fst order)
  withSortKeyContextsHelper deltaInterpret ctxts order
  where
    order = sortBy (comparing fst) kws


setLayoutFlag :: Delta () -> Delta ()
setLayoutFlag action = do
  oldLay <- gets apLayoutStart
  modify (\s -> s { apMarkLayout = True } )
  let reset = do
                modify (\s -> s { apMarkLayout = False
                                , apLayoutStart = oldLay })
  action <* reset

-- ---------------------------------------------------------------------

setContextDelta :: Set.Set AstContext -> Int -> Delta () -> Delta ()
setContextDelta ctxt lvl =
  local (\s -> s { drContext = setAcsWithLevel ctxt lvl (drContext s) } )

ifInContextDelta :: Set.Set AstContext -> Annotated () -> Annotated () -> Delta ()
ifInContextDelta ctxt ifAction elseAction = do
  cur <- asks drContext
  let inContext = inAcs ctxt cur
  if inContext
    then deltaInterpret ifAction
    else deltaInterpret elseAction

-- ---------------------------------------------------------------------

storeOriginalSrcSpanDelta :: AnnKey -> Delta AnnKey
storeOriginalSrcSpanDelta key = do
  tellCapturedSpan key
  return key

#if __GLASGOW_HASKELL__ <= 710
storeString :: String -> GHC.SrcSpan -> Delta ()
storeString s ss = addAnnotationWorker (AnnString s) ss
#endif

-- ---------------------------------------------------------------------

-- |In order to interleave annotations into the stream, we turn them into
-- comments.
annotationsToCommentsDelta :: [GHC.AnnKeywordId] -> Delta ()
annotationsToCommentsDelta kws = do
  ss <- getSrcSpan
  cs <- gets apComments
  let
    doOne :: GHC.AnnKeywordId -> Delta [Comment]
    doOne kw = do
      (spans,_) <- getAndRemoveAnnotationDelta ss kw
      return $ map (mkKWComment kw) spans
    -- TODO:AZ make sure these are sorted/merged properly when the invariant for
    -- allocateComments is re-established.
  newComments <- mapM doOne kws
  putUnallocatedComments (cs ++ concat newComments)

-- ---------------------------------------------------------------------

-- | This function exists to overcome a shortcoming in the GHC AST for 7.10.1
getSrcSpanForKw :: GHC.SrcSpan -> GHC.AnnKeywordId -> Delta GHC.SrcSpan
getSrcSpanForKw _ kw = do
    ga <- gets apAnns
    ss <- getSrcSpan
    case GHC.getAnnotation ga ss kw of
      []     -> return GHC.noSrcSpan
      (sp:_) -> return sp

-- ---------------------------------------------------------------------

getSrcSpan :: Delta GHC.SrcSpan
getSrcSpan = asks curSrcSpan

#if __GLASGOW_HASKELL__ > 806
withSrcSpanDelta :: (Data a, Data (GHC.SrcSpanLess a), GHC.HasSrcSpan a) => a -> Delta b -> Delta b
withSrcSpanDelta (GHC.dL->GHC.L l a) =
#else
withSrcSpanDelta :: Data a => GHC.Located a -> Delta b -> Delta b
withSrcSpanDelta (GHC.L l a) =
#endif
  local (\s -> s { curSrcSpan = l
                 , annConName = annGetConstr a
                 , drContext = pushAcs (drContext s)
                    `debug` ("withSrcSpanDelta: (l,annConName,drContext)=" ++ showGhc (l,annGetConstr a, pushAcs (drContext s)))
                 })


getUnallocatedComments :: Delta [Comment]
getUnallocatedComments = gets apComments

putUnallocatedComments :: [Comment] -> Delta ()
putUnallocatedComments cs = modify (\s -> s { apComments = cs } )

-- ---------------------------------------------------------------------

adjustDeltaForOffsetM :: DeltaPos -> Delta DeltaPos
adjustDeltaForOffsetM dp = do
  colOffset <- gets apLayoutStart
  return (adjustDeltaForOffset colOffset dp)

adjustDeltaForOffset :: LayoutStartCol -> DeltaPos -> DeltaPos
adjustDeltaForOffset _colOffset              dp@(DP (0,_)) = dp -- same line
adjustDeltaForOffset (LayoutStartCol colOffset) (DP (l,c)) = DP (l,c - colOffset)

-- ---------------------------------------------------------------------

getPriorEnd :: Delta Pos
getPriorEnd = gets priorEndPosition

setPriorEnd :: Pos -> Delta ()
setPriorEnd pe =
  modify (\s -> s { priorEndPosition = pe })

setPriorEndAST :: GHC.SrcSpan -> Delta ()
setPriorEndAST pe = do
  setLayoutStart (snd (ss2pos pe))
  modify (\s -> s { priorEndPosition    = ss2posEnd pe } )

setLayoutStart :: Int -> Delta ()
setLayoutStart p = do
  DeltaState{apMarkLayout} <- get
  when apMarkLayout (
                      modify (\s -> s { apMarkLayout = False
                                     , apLayoutStart = LayoutStartCol p}))


-- -------------------------------------

peekAnnotationDelta :: GHC.AnnKeywordId -> Delta [GHC.SrcSpan]
peekAnnotationDelta an = do
    ga <- gets apAnns
    ss <- getSrcSpan
#if __GLASGOW_HASKELL__ <= 710
    return $ GHC.getAnnotation ga ss an
#else
    let unicodeAnns = case unicodeEquivalent an of
          [] -> []
          [kw] -> GHC.getAnnotation ga ss kw
          (kw:_) -> GHC.getAnnotation ga ss kw -- Keep exhaustiveness checker happy
    return $ unicodeAnns ++ GHC.getAnnotation ga ss an
#endif

getAnnotationDelta :: GHC.AnnKeywordId -> Delta ([GHC.SrcSpan],GHC.AnnKeywordId)
getAnnotationDelta an = do
    ss <- getSrcSpan
    getAndRemoveAnnotationDelta ss an

getAndRemoveAnnotationDelta :: GHC.SrcSpan -> GHC.AnnKeywordId -> Delta ([GHC.SrcSpan],GHC.AnnKeywordId)
getAndRemoveAnnotationDelta sp an = do
    ga <- gets apAnns
#if __GLASGOW_HASKELL__ <= 710
    let (r,ga') = GHC.getAndRemoveAnnotation ga sp an
        kw = an
#else
    let (r,ga',kw) = case GHC.getAndRemoveAnnotation ga sp an of
                    ([],_) -> (ss,g,k)
                      where
                        k = GHC.unicodeAnn an
                        (ss,g) = GHC.getAndRemoveAnnotation ga sp k
                    (ss,g)  -> (ss,g,an)
#endif
    modify (\s -> s { apAnns = ga' })
    return (r,kw)

getOneAnnotationDelta :: GHC.AnnKeywordId -> Delta ([GHC.SrcSpan],GHC.AnnKeywordId)
getOneAnnotationDelta an = do
    ss <- getSrcSpan
    getAndRemoveOneAnnotationDelta ss an

getAndRemoveOneAnnotationDelta :: GHC.SrcSpan -> GHC.AnnKeywordId -> Delta ([GHC.SrcSpan],GHC.AnnKeywordId)
getAndRemoveOneAnnotationDelta sp an = do
    (anns,cs) <- gets apAnns
#if __GLASGOW_HASKELL__ <= 710
    let (r,ga',kw) = case Map.lookup (sp,an) anns of
                    Nothing -> ([],(anns,cs),an)
                    Just []     -> ([], (Map.delete (sp,an)    anns,cs),an)
                    Just (s:ss) -> ([s],(Map.insert (sp,an) ss anns,cs),an)
#else
    let getKw kw =
          case Map.lookup (sp,kw) anns of
            Nothing -> ([],(anns,cs),kw)
            Just []     -> ([], (Map.delete (sp,kw)    anns,cs),kw)
            Just (s:ss) -> ([s],(Map.insert (sp,kw) ss anns,cs),kw)

    let (r,ga',kw) =
          case getKw an of
            ([],_,_) -> getKw (GHC.unicodeAnn an)
            v        -> v
#endif
    modify (\s -> s { apAnns = ga' })
    return (r,kw)

-- ---------------------------------------------------------------------

-- |Add some annotation to the currently active SrcSpan
addAnnotationsDelta :: Annotation -> Delta ()
addAnnotationsDelta ann = do
    l <- ask
    tellFinalAnn (getAnnKey l,ann)

getAnnKey :: DeltaOptions -> AnnKey
getAnnKey DeltaOptions {curSrcSpan, annConName}
  = AnnKey curSrcSpan annConName

-- -------------------------------------

addAnnDeltaPos :: KeywordId -> DeltaPos -> Delta ()
addAnnDeltaPos kw dp = tellKd (kw, dp)

-- -------------------------------------

-- | Enter a new AST element. Maintain SrcSpan stack
#if __GLASGOW_HASKELL__ > 806
withAST :: (Data a, Data (GHC.SrcSpanLess a), GHC.HasSrcSpan a)
        => a
        -> Delta b -> Delta b
withAST lss@(GHC.dL->GHC.L ss _) action = do
#else
withAST :: Data a
        => GHC.Located a
        -> Delta b -> Delta b
withAST lss@(GHC.L ss _) action = do
#endif
  -- Calculate offset required to get to the start of the SrcSPan
  off <- gets apLayoutStart
  (resetAnns .  withSrcSpanDelta lss) (do

    let maskWriter s = s { annKds = []
                         , sortKeys = Nothing
                         , dwCapturedSpan = mempty }

    -- make sure all kds are relative to the start of the SrcSpan
    let spanStart = ss2pos ss

    cs <- do
      priorEndBeforeComments <- getPriorEnd
      if GHC.isGoodSrcSpan ss && priorEndBeforeComments < ss2pos ss
        then
          commentAllocation (priorComment spanStart) return
        else
          return []
    priorEndAfterComments <- getPriorEnd
    let edp = adjustDeltaForOffset
                -- Use the propagated offset if one is set
                -- Note that we need to use the new offset if it has
                -- changed.
                off (ss2delta priorEndAfterComments ss)
    -- Preparation complete, perform the action
    when (GHC.isGoodSrcSpan ss && priorEndAfterComments < ss2pos ss) (do
      modify (\s -> s { priorEndPosition    = ss2pos ss } ))
    (res, w) <- censor maskWriter (listen action)

    let kds = annKds w
        an = Ann
               { annEntryDelta        = edp
               , annPriorComments     = cs
               , annFollowingComments = [] -- only used in Transform and Print
               , annsDP               = kds
               , annSortKey           = sortKeys w
               , annCapturedSpan      = getFirst $ dwCapturedSpan w }

    addAnnotationsDelta an
     `debug` ("leaveAST:(annkey,an)=" ++ show (mkAnnKey lss,an))
    return res)

resetAnns :: Delta a -> Delta a
resetAnns action = do
  ans <- gets apAnns
  action <* modify (\s -> s { apAnns = ans })


-- ---------------------------------------------------------------------
-- |Split the ordered list of comments into ones that occur prior to
-- the give SrcSpan and the rest
priorComment :: Pos -> Comment -> Bool
priorComment start c = (ss2pos . commentIdentifier $ c) < start

-- TODO:AZ: We scan the entire comment list here. It may be better to impose an
-- invariant that the comments are sorted, and consume them as the pos
-- advances. It then becomes a process of using `takeWhile p` rather than a full
-- partition.
allocateComments :: (Comment -> Bool) -> [Comment] -> ([Comment], [Comment])
allocateComments = partition

-- ---------------------------------------------------------------------

addAnnotationWorker :: KeywordId -> GHC.SrcSpan -> Delta ()
addAnnotationWorker ann pa =
  -- Zero-width source spans are injected by the GHC Lexer when it puts virtual
  -- '{', ';' and '}' tokens in for layout
  unless (isPointSrcSpan pa) $
    do
      pe <- getPriorEnd
      ss <- getSrcSpan
      let p = ss2delta pe pa
      case (ann,isGoodDelta p) of
        (G GHC.AnnComma,False) -> return ()
        (G GHC.AnnSemi, False) -> return ()
        (G GHC.AnnOpen, False) -> return ()
        (G GHC.AnnClose,False) -> return ()
        _ -> do
          p' <- adjustDeltaForOffsetM p
          commentAllocation (priorComment (ss2pos pa)) (mapM_ (uncurry addDeltaComment))
#if __GLASGOW_HASKELL__ <= 710
          addAnnDeltaPos (checkUnicode ann pa) p'
#else
          addAnnDeltaPos ann p'
#endif
          setPriorEndAST pa
              `debug` ("addAnnotationWorker:(ss,ss,pe,pa,p,p',ann)=" ++ show (showGhc ss,showGhc ss,pe,showGhc pa,p,p',ann))

#if __GLASGOW_HASKELL__ <= 710
checkUnicode :: KeywordId -> GHC.SrcSpan -> KeywordId
checkUnicode gkw@(G kw) ss =
  if kw `Set.member` unicodeSyntax
    then
      let s = keywordToString gkw in
      if length s /= spanLength ss
        then AnnUnicode kw
        else gkw
  else
    gkw
  where
    unicodeSyntax = Set.fromList
      [ GHC.AnnDcolon
      , GHC.AnnDarrow
      , GHC.AnnForall
      , GHC.AnnRarrow
      , GHC.AnnLarrow
      , GHC.Annlarrowtail
      , GHC.Annrarrowtail
      , GHC.AnnLarrowtail
      , GHC.AnnRarrowtail]
checkUnicode kwid _ = kwid
#else

unicodeEquivalent :: GHC.AnnKeywordId -> [GHC.AnnKeywordId]
unicodeEquivalent kw =
  case Map.lookup kw unicodeSyntax of
    Nothing -> []
    Just kwu -> [kwu]
  where
    unicodeSyntax = Map.fromList
      [ (GHC.AnnDcolon,     GHC.AnnDcolonU)
      , (GHC.AnnDarrow,     GHC.AnnDarrowU)
      , (GHC.AnnForall,     GHC.AnnForallU)
      , (GHC.AnnRarrow,     GHC.AnnRarrowU)
      , (GHC.AnnLarrow,     GHC.AnnLarrowU)
      , (GHC.Annlarrowtail, GHC.AnnlarrowtailU)
      , (GHC.Annrarrowtail, GHC.AnnrarrowtailU)
      , (GHC.AnnLarrowtail, GHC.AnnLarrowtailU)
      , (GHC.AnnRarrowtail, GHC.AnnRarrowtailU)
#if __GLASGOW_HASKELL__ > 801
      , (GHC.AnnCloseB,     GHC.AnnCloseBU)
      , (GHC.AnnCloseQ,     GHC.AnnCloseQU)
      , (GHC.AnnOpenB,      GHC.AnnOpenBU)
      , (GHC.AnnOpenEQ,     GHC.AnnOpenEQU)
#endif
      ]
#endif


-- ---------------------------------------------------------------------

commentAllocation :: (Comment -> Bool)
                  -> ([(Comment, DeltaPos)] -> Delta a)
                  -> Delta a
commentAllocation p k = do
  cs <- getUnallocatedComments
  let (allocated,cs') = allocateComments p cs
  putUnallocatedComments cs'
  k =<< mapM makeDeltaComment (sortBy (comparing commentIdentifier) allocated)


makeDeltaComment :: Comment -> Delta (Comment, DeltaPos)
makeDeltaComment c = do
  let pa = commentIdentifier c
  pe <- getPriorEnd
  let p = ss2delta pe pa
  p' <- adjustDeltaForOffsetM p
  setPriorEnd (ss2posEnd pa)
  return (c, p')

addDeltaComment :: Comment -> DeltaPos -> Delta ()
addDeltaComment d p = do
  addAnnDeltaPos (AnnComment d) p

-- ---------------------------------------------------------------------

-- |If the first annotation has a smaller SrcSpan than the second, then mark it.
deltaMarkAnnBeforeAnn :: GHC.AnnKeywordId -> GHC.AnnKeywordId -> Delta ()
deltaMarkAnnBeforeAnn annBefore annAfter = do
  ss <- getSrcSpan
  mb <- peekAnnotationDelta annBefore
  ma <- peekAnnotationDelta annAfter
  let
    before = sort $ filter (\s -> GHC.isSubspanOf s ss) mb
    after  = sort $ filter (\s -> GHC.isSubspanOf s ss) ma
  case (before,after) of
    (b:_, a:_) -> when (b < a) $ addDeltaAnnotation annBefore
    _ -> return ()

-- ---------------------------------------------------------------------

-- | Look up and add a Delta annotation at the current position, and
-- advance the position to the end of the annotation
addDeltaAnnotation :: GHC.AnnKeywordId -> Delta ()
addDeltaAnnotation ann' = do
  ss <- getSrcSpan
  (ma,ann) <- getOneAnnotationDelta ann'
  case nub ma of -- ++AZ++ TODO: get rid of duplicates earlier
    []     -> return () `debug` ("addDeltaAnnotation empty ma for:" ++ show (ss,ann))
    [pa]   -> addAnnotationWorker (G ann) pa
    (pa:_) -> addAnnotationWorker (G ann) pa `warn` ("addDeltaAnnotation:(ss,ann,ma)=" ++ showGhc (ss,ann,ma))

-- | Look up and add a Delta annotation at the current position, and
-- advance the position to the end of the annotation
addDeltaAnnotationLs :: GHC.AnnKeywordId -> Int -> Delta ()
addDeltaAnnotationLs ann off = do
  ss <- getSrcSpan
  ma <- peekAnnotationDelta ann
  let ma' = filter (\s -> GHC.isSubspanOf s ss) ma
  case drop off ma' of
    [] -> return ()
        `debug` ("addDeltaAnnotationLs:missed:(off,ann,ma)=" ++ showGhc (off,ss,ann))
    (pa:_) -> addAnnotationWorker (G ann) pa

-- | Look up and add possibly multiple Delta annotation at the current
-- position, and advance the position to the end of the annotations
addDeltaAnnotations :: GHC.AnnKeywordId -> Delta ()
addDeltaAnnotations ann = do
  (ma,kw) <- getAnnotationDelta ann
  let do_one ap' = addAnnotationWorker (G kw) ap'
                    `debug` ("addDeltaAnnotations:do_one:(ap',ann)=" ++ showGhc (ap',ann))
  mapM_ do_one (sort ma)

-- | Look up and add possibly multiple Delta annotations enclosed by
-- the current SrcSpan at the current position, and advance the
-- position to the end of the annotations
addDeltaAnnotationsInside :: GHC.AnnKeywordId -> Delta ()
addDeltaAnnotationsInside ann = do
  ss <- getSrcSpan
  ma <- peekAnnotationDelta ann
  let do_one ap' = addAnnotationWorker (G ann) ap'
                    -- `debug` ("addDeltaAnnotations:do_one:(ap',ann)=" ++ showGhc (ap',ann))
  let filtered = sort $ filter (\s -> GHC.isSubspanOf s ss) ma
  mapM_ do_one filtered

-- ---------------------------------------------------------------------
#if __GLASGOW_HASKELL__ >= 800
addDeltaAnnotationInstead :: GHC.AnnKeywordId  -> KeywordId -> Delta ()
addDeltaAnnotationInstead ann' kw = do
  ss <- getSrcSpan
  (ma,ann) <- getOneAnnotationDelta ann'
  case nub ma of -- ++AZ++ TODO: get rid of duplicates earlier
    []     -> return () `debug` ("addDeltaAnnotation empty ma for:" ++ show (ss,ann))
    [pa]   -> addAnnotationWorker kw pa
    (pa:_) -> addAnnotationWorker kw pa `warn` ("addDeltaAnnotationInstead:(ss,ann,kw,ma)=" ++ showGhc (ss,ann,kw,ma))
#endif
-- ---------------------------------------------------------------------

-- | Look up and add possibly multiple Delta annotations not enclosed by
-- the current SrcSpan at the current position, and advance the
-- position to the end of the annotations
-- The first argument (gann) is the one to look up in the GHC annotations, the
-- second is the one to apply in the ghc-exactprint ones. These are different
-- for GHC.AnnSemi mapping to AnnSemiSep, to ensure that it reflects the ';'
-- outside the current span.
addDeltaAnnotationsOutside :: GHC.AnnKeywordId -> KeywordId -> Delta ()
addDeltaAnnotationsOutside gann ann = do
  ss <- getSrcSpan
  (ma,kw) <- getAndRemoveAnnotationDelta ss gann
  let do_one ap' = if ann == AnnSemiSep
                     then addAnnotationWorker ann    ap'
                     else addAnnotationWorker (G kw) ap'
  mapM_ do_one (sort $ filter (\s -> not (GHC.isSubspanOf s ss)) ma)

-- | Add a Delta annotation at the current position, and advance the
-- position to the end of the annotation
addDeltaAnnotationExt :: GHC.SrcSpan -> GHC.AnnKeywordId -> Delta ()
addDeltaAnnotationExt s ann = addAnnotationWorker (G ann) s

addEofAnnotation :: Delta ()
addEofAnnotation = do
  pe <- getPriorEnd
  (ma,_kw) <- withSrcSpanDelta (GHC.noLoc () :: GHC.GenLocated GHC.SrcSpan ()) (getAnnotationDelta GHC.AnnEofPos)
  case ma of
    [] -> return ()
    (pa:pss) -> do
      commentAllocation (const True) (mapM_ (uncurry addDeltaComment))
      let DP (r,c) = ss2delta pe pa
      addAnnDeltaPos (G GHC.AnnEofPos) (DP (r, c - 1))
      setPriorEndAST pa `warn` ("Trailing annotations after Eof: " ++ showGhc pss)

-- ---------------------------------------------------------------------

countAnnsDelta :: GHC.AnnKeywordId -> Delta Int
countAnnsDelta ann = do
  ma <- peekAnnotationDelta ann
  return (length ma)