{-# LANGUAGE CPP #-}
module Utils
  ( nodeHasAnnotation
  , getNodeInfo
  , foldNodeChildren
  ) where

#if __GLASGOW_HASKELL__ >= 900
import qualified Data.Map as M
#endif
import qualified Data.Set as S
import           Data.String

import           GHC.Api

#if __GLASGOW_HASKELL__ >= 900
mergeNodeInfo :: NodeInfo a -> NodeInfo a -> NodeInfo a
mergeNodeInfo :: forall a. NodeInfo a -> NodeInfo a -> NodeInfo a
mergeNodeInfo (NodeInfo Set NodeAnnotation
as [a]
ai NodeIdentifiers a
ad) (NodeInfo Set NodeAnnotation
bs [a]
bi NodeIdentifiers a
bd) =
  forall a.
Set NodeAnnotation -> [a] -> NodeIdentifiers a -> NodeInfo a
NodeInfo (Set NodeAnnotation
as forall a. Semigroup a => a -> a -> a
<> Set NodeAnnotation
bs) ([a]
ai forall a. Semigroup a => a -> a -> a
<> [a]
bi) (forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
M.unionWith forall a. Semigroup a => a -> a -> a
(<>) NodeIdentifiers a
ad NodeIdentifiers a
bd)
#endif

-- | Extract node info for an AST. GHC 9 includes generated things that need to
-- be removed.
getNodeInfo :: HieAST a -> NodeInfo a
#if __GLASGOW_HASKELL__ >= 900
getNodeInfo :: forall a. HieAST a -> NodeInfo a
getNodeInfo = forall a b k. (a -> b -> a) -> a -> Map k b -> a
M.foldl' forall a. NodeInfo a -> NodeInfo a -> NodeInfo a
mergeNodeInfo forall a. NodeInfo a
emptyNodeInfo
            forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Ord k => k -> Map k a -> Map k a
M.delete NodeOrigin
GeneratedInfo -- removed ghc generated nodes
            forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. SourcedNodeInfo a -> Map NodeOrigin (NodeInfo a)
getSourcedNodeInfo forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. HieAST a -> SourcedNodeInfo a
sourcedNodeInfo
#else
getNodeInfo = nodeInfo
#endif

nodeHasAnnotation :: String -> String -> HieAST a -> Bool
nodeHasAnnotation :: forall a. String -> String -> HieAST a -> Bool
nodeHasAnnotation String
constructor String
ty =
    forall a. Ord a => a -> Set a -> Bool
S.member
#if MIN_VERSION_ghc(9,2,0)
      (FastString -> FastString -> NodeAnnotation
NodeAnnotation (forall a. IsString a => String -> a
fromString String
constructor) (forall a. IsString a => String -> a
fromString String
ty))
#else
      (fromString constructor, fromString ty)
#endif
  forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. NodeInfo a -> Set NodeAnnotation
nodeAnnotations
  forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. HieAST a -> NodeInfo a
getNodeInfo

foldNodeChildren :: Monoid m => (HieAST a -> m) -> HieAST a -> m
foldNodeChildren :: forall m a. Monoid m => (HieAST a -> m) -> HieAST a -> m
foldNodeChildren HieAST a -> m
f = forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap HieAST a -> m
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. HieAST a -> [HieAST a]
nodeChildren