-- | 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 GHC.Parser.Annotation
import GHC.Parser.Lexer
import GHC.Types.SrcLoc

-- | Ormolu-specific representation of GHC annotations.
newtype Anns = Anns (Map RealSrcSpan [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 'PState'.
mkAnns ::
  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]
(++) (((RealSrcSpan, AnnKeywordId), [RealSrcSpan])
-> (RealSrcSpan, [AnnKeywordId])
forall a a b. ((a, a), b) -> (a, [a])
f (((RealSrcSpan, AnnKeywordId), [RealSrcSpan])
 -> (RealSrcSpan, [AnnKeywordId]))
-> [((RealSrcSpan, AnnKeywordId), [RealSrcSpan])]
-> [(RealSrcSpan, [AnnKeywordId])]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> PState -> [((RealSrcSpan, AnnKeywordId), [RealSrcSpan])]
annotations PState
pstate)
  where
    f :: ((a, a), b) -> (a, [a])
f ((a
rspn, a
kid), b
_) = (a
rspn, [a
kid])

-- | Lookup 'AnnKeywordId's corresponding to a given 'SrcSpan'.
lookupAnns ::
  -- | Span to lookup with
  SrcSpan ->
  -- | Collection of annotations
  Anns ->
  [AnnKeywordId]
lookupAnns :: SrcSpan -> Anns -> [AnnKeywordId]
lookupAnns (RealSrcSpan RealSrcSpan
rspn Maybe BufSpan
_) (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 UnhelpfulSpanReason
_) Anns
_ = []