    -- | Ormolu-specific representation of GHC annotations.
    module Ormolu.Parser.Anns
      ( Anns (..),
        emptyAnns,
        mkAnns,
        lookupAnns,
      )
    where

    import Data.Map.Strict (Map)
    import qualified Data.Map.Strict as M
    import Data.Maybe (mapMaybe)
    import qualified GHC
    import qualified Lexer as GHC
    import SrcLoc

    -- | Ormolu-specific representation of GHC annotations.
    newtype Anns = Anns (Map RealSrcSpan [GHC.AnnKeywordId])
      deriving (Anns -> Anns -> Bool
    (Anns -> Anns -> Bool) -> (Anns -> Anns -> Bool) -> Eq Anns
    forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
    /= :: Anns -> Anns -> Bool
    $c/= :: Anns -> Anns -> Bool
    == :: Anns -> Anns -> Bool
    $c== :: Anns -> Anns -> Bool
    Eq)

    -- | Empty 'Anns'.
    emptyAnns :: Anns
    emptyAnns :: Anns
    emptyAnns = Map RealSrcSpan [AnnKeywordId] -> Anns
    Anns Map RealSrcSpan [AnnKeywordId]
    forall k a. Map k a
    M.empty

    -- | Create 'Anns' from 'GHC.PState'.
    mkAnns ::
      GHC.PState ->
      Anns
    mkAnns :: PState -> Anns
    mkAnns pstate :: PState
    pstate =
      Map RealSrcSpan [AnnKeywordId] -> Anns
    Anns (Map RealSrcSpan [AnnKeywordId] -> Anns)
    -> Map RealSrcSpan [AnnKeywordId] -> Anns
    forall a b. (a -> b) -> a -> b
    $
        ([AnnKeywordId] -> [AnnKeywordId] -> [AnnKeywordId])
    -> [(RealSrcSpan, [AnnKeywordId])]
    -> Map RealSrcSpan [AnnKeywordId]
    forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
    M.fromListWith [AnnKeywordId] -> [AnnKeywordId] -> [AnnKeywordId]
    forall a. [a] -> [a] -> [a]
    (++) ((((SrcSpan, AnnKeywordId), [SrcSpan])
     -> Maybe (RealSrcSpan, [AnnKeywordId]))
    -> [((SrcSpan, AnnKeywordId), [SrcSpan])]
    -> [(RealSrcSpan, [AnnKeywordId])]
    forall a b. (a -> Maybe b) -> [a] -> [b]
    mapMaybe ((SrcSpan, AnnKeywordId), [SrcSpan])
    -> Maybe (RealSrcSpan, [AnnKeywordId])
    forall a b. ((SrcSpan, a), b) -> Maybe (RealSrcSpan, [a])
    f (PState -> [((SrcSpan, AnnKeywordId), [SrcSpan])]
    GHC.annotations PState
    pstate))
      where
        f :: ((SrcSpan, a), b) -> Maybe (RealSrcSpan, [a])
    f ((spn :: SrcSpan
    spn, kid :: a
    kid), _) =
          case SrcSpan
    spn of
            RealSrcSpan rspn :: RealSrcSpan
    rspn -> (RealSrcSpan, [a]) -> Maybe (RealSrcSpan, [a])
    forall a. a -> Maybe a
    Just (RealSrcSpan
    rspn, [a
    kid])
            UnhelpfulSpan _ -> Maybe (RealSrcSpan, [a])
    forall a. Maybe a
    Nothing

    -- | Lookup 'GHC.AnnKeywordId's corresponding to a given 'SrcSpan'.
    lookupAnns ::
      -- | Span to lookup with
      SrcSpan ->
      -- | Collection of annotations
      Anns ->
      [GHC.AnnKeywordId]
    lookupAnns :: SrcSpan -> Anns -> [AnnKeywordId]
    lookupAnns (RealSrcSpan rspn :: RealSrcSpan
    rspn) (Anns m :: Map RealSrcSpan [AnnKeywordId]
    m) = [AnnKeywordId]
    -> RealSrcSpan -> Map RealSrcSpan [AnnKeywordId] -> [AnnKeywordId]
    forall k a. Ord k => a -> k -> Map k a -> a
    M.findWithDefault [] RealSrcSpan
    rspn Map RealSrcSpan [AnnKeywordId]
    m
    lookupAnns (UnhelpfulSpan _) _ = []