-- | 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 =
  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 ((SrcSpan
spn, a
kid), b
_) =
      case SrcSpan
spn of
        RealSrcSpan RealSrcSpan
rspn -> (RealSrcSpan, [a]) -> Maybe (RealSrcSpan, [a])
forall a. a -> Maybe a
Just (RealSrcSpan
rspn, [a
kid])
        UnhelpfulSpan FastString
_ -> 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 RealSrcSpan
rspn) (Anns 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 FastString
_) Anns
_ = []