{-# LANGUAGE OverloadedStrings #-}

module Ide.Plugin.CodeRange.ASTPreProcess
    ( preProcessAST
    , PreProcessEnv(..)
    , isCustomNode
    , CustomNodeType(..)
    ) where

import           Control.Monad.Reader       (Reader, asks)
import           Data.Foldable
import           Data.Functor.Identity      (Identity (Identity, runIdentity))
import           Data.List                  (groupBy)
import           Data.List.NonEmpty         (NonEmpty)
import qualified Data.List.NonEmpty         as NonEmpty
import           Data.Map.Strict            (Map)
import qualified Data.Map.Strict            as Map
import           Data.Maybe                 (fromMaybe, mapMaybe)
import           Data.Semigroup             (First (First, getFirst))
import           Data.Semigroup.Foldable    (foldlM1)
import qualified Data.Set                   as Set
import           Development.IDE.GHC.Compat hiding (nodeInfo)
import           Prelude                    hiding (span)

{-|
Extra arguments for 'preProcessAST'. It's expected to be used in a 'Reader' context
-}
newtype PreProcessEnv a = PreProcessEnv
    { forall a. PreProcessEnv a -> RefMap a
preProcessEnvRefMap :: RefMap a
    }

{-|
Before converting the HieAST to selection range, we need to run some passes on it. Each pass potentially modifies
the AST to handle some special cases.

'preProcessAST' combines the passes. Refer to 'mergeImports' or 'mergeSignatureWithDefinition' as
a concrete example example.

Adding another manipulation to the AST is simple, just implement a function of type
`HieAST a -> Reader (PreProcessEnv a) (HieAST a)`, and append it to 'preProcessAST' with `>>=`.

If it goes more complex, it may be more appropriate to split different manipulations to different modules.
-}
preProcessAST :: HieAST a -> Reader (PreProcessEnv a) (HieAST a)
preProcessAST :: forall a. HieAST a -> Reader (PreProcessEnv a) (HieAST a)
preProcessAST HieAST a
node = HieAST a -> Reader (PreProcessEnv a) (HieAST a)
forall a. HieAST a -> Reader (PreProcessEnv a) (HieAST a)
mergeImports HieAST a
node Reader (PreProcessEnv a) (HieAST a)
-> (HieAST a -> Reader (PreProcessEnv a) (HieAST a))
-> Reader (PreProcessEnv a) (HieAST a)
forall a b.
ReaderT (PreProcessEnv a) Identity a
-> (a -> ReaderT (PreProcessEnv a) Identity b)
-> ReaderT (PreProcessEnv a) Identity b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= HieAST a -> Reader (PreProcessEnv a) (HieAST a)
forall a. HieAST a -> Reader (PreProcessEnv a) (HieAST a)
mergeSignatureWithDefinition

{-|
Create a custom node in 'HieAST'. By "custom", we mean this node doesn't actually exist in the original 'HieAST'
provided by GHC, but created to suite the needs of hls-code-range-plugin.
-}
createCustomNode :: CustomNodeType -> NonEmpty (HieAST a) -> HieAST a
createCustomNode :: forall a. CustomNodeType -> NonEmpty (HieAST a) -> HieAST a
createCustomNode CustomNodeType
customNodeType NonEmpty (HieAST a)
children = NodeInfo a -> Span -> [HieAST a] -> HieAST a
forall a. NodeInfo a -> Span -> [HieAST a] -> HieAST a
mkAstNode NodeInfo a
forall {a}. NodeInfo a
customNodeInfo Span
span' (NonEmpty (HieAST a) -> [HieAST a]
forall a. NonEmpty a -> [a]
NonEmpty.toList NonEmpty (HieAST a)
children)
  where
    span' :: RealSrcSpan
    span' :: Span
span' = Identity Span -> Span
forall a. Identity a -> a
runIdentity (Identity Span -> Span)
-> (NonEmpty (HieAST a) -> Identity Span)
-> NonEmpty (HieAST a)
-> Span
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Span -> Span -> Identity Span) -> NonEmpty Span -> Identity Span
forall (t :: * -> *) (m :: * -> *) a.
(Foldable1 t, Monad m) =>
(a -> a -> m a) -> t a -> m a
foldlM1 (\Span
x Span
y -> Span -> Identity Span
forall a. a -> Identity a
Identity (Span -> Span -> Span
combineRealSrcSpans Span
x Span
y)) (NonEmpty Span -> Identity Span)
-> (NonEmpty (HieAST a) -> NonEmpty Span)
-> NonEmpty (HieAST a)
-> Identity Span
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (HieAST a -> Span) -> NonEmpty (HieAST a) -> NonEmpty Span
forall a b. (a -> b) -> NonEmpty a -> NonEmpty b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap HieAST a -> Span
forall a. HieAST a -> Span
nodeSpan (NonEmpty (HieAST a) -> Span) -> NonEmpty (HieAST a) -> Span
forall a b. (a -> b) -> a -> b
$ NonEmpty (HieAST a)
children

    customNodeInfo :: NodeInfo a
customNodeInfo = FastStringCompat -> FastStringCompat -> NodeInfo a
forall a. FastStringCompat -> FastStringCompat -> NodeInfo a
simpleNodeInfoCompat FastStringCompat
"HlsCustom" (CustomNodeType -> FastStringCompat
customNodeTypeToFastString CustomNodeType
customNodeType)

isCustomNode :: HieAST a -> Maybe CustomNodeType
isCustomNode :: forall a. HieAST a -> Maybe CustomNodeType
isCustomNode HieAST a
node = do
    NodeInfo a
nodeInfo <- HieAST a -> Maybe (NodeInfo a)
forall a. HieAST a -> Maybe (NodeInfo a)
generatedNodeInfo HieAST a
node
    First CustomNodeType -> CustomNodeType
forall a. First a -> a
getFirst (First CustomNodeType -> CustomNodeType)
-> Maybe (First CustomNodeType) -> Maybe CustomNodeType
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((FastStringCompat, FastStringCompat)
 -> Maybe (First CustomNodeType))
-> Set (FastStringCompat, FastStringCompat)
-> Maybe (First CustomNodeType)
forall m a. Monoid m => (a -> m) -> Set a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (FastStringCompat, FastStringCompat)
-> Maybe (First CustomNodeType)
go (NodeInfo a -> Set (FastStringCompat, FastStringCompat)
forall a. NodeInfo a -> Set (FastStringCompat, FastStringCompat)
nodeAnnotations NodeInfo a
nodeInfo)
  where
    go :: (FastStringCompat, FastStringCompat) -> Maybe (First CustomNodeType)
    go :: (FastStringCompat, FastStringCompat)
-> Maybe (First CustomNodeType)
go (FastStringCompat
k, FastStringCompat
v)
        | FastStringCompat
k FastStringCompat -> FastStringCompat -> Bool
forall a. Eq a => a -> a -> Bool
== FastStringCompat
"HlsCustom", Just CustomNodeType
v' <- Map FastStringCompat CustomNodeType
revCustomNodeTypeMapping Map FastStringCompat CustomNodeType
-> FastStringCompat -> Maybe CustomNodeType
forall k a. Ord k => Map k a -> k -> Maybe a
Map.!? FastStringCompat
v = First CustomNodeType -> Maybe (First CustomNodeType)
forall a. a -> Maybe a
Just (CustomNodeType -> First CustomNodeType
forall a. a -> First a
First CustomNodeType
v')
        | Bool
otherwise = Maybe (First CustomNodeType)
forall a. Maybe a
Nothing

data CustomNodeType =
    -- | a group of imports
    CustomNodeImportsGroup
    -- | adjacent type signature and value definition are paired under a custom parent node
  | CustomNodeAdjacentSignatureDefinition
    deriving (Int -> CustomNodeType -> ShowS
[CustomNodeType] -> ShowS
CustomNodeType -> String
(Int -> CustomNodeType -> ShowS)
-> (CustomNodeType -> String)
-> ([CustomNodeType] -> ShowS)
-> Show CustomNodeType
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> CustomNodeType -> ShowS
showsPrec :: Int -> CustomNodeType -> ShowS
$cshow :: CustomNodeType -> String
show :: CustomNodeType -> String
$cshowList :: [CustomNodeType] -> ShowS
showList :: [CustomNodeType] -> ShowS
Show, CustomNodeType -> CustomNodeType -> Bool
(CustomNodeType -> CustomNodeType -> Bool)
-> (CustomNodeType -> CustomNodeType -> Bool) -> Eq CustomNodeType
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: CustomNodeType -> CustomNodeType -> Bool
== :: CustomNodeType -> CustomNodeType -> Bool
$c/= :: CustomNodeType -> CustomNodeType -> Bool
/= :: CustomNodeType -> CustomNodeType -> Bool
Eq, Eq CustomNodeType
Eq CustomNodeType =>
(CustomNodeType -> CustomNodeType -> Ordering)
-> (CustomNodeType -> CustomNodeType -> Bool)
-> (CustomNodeType -> CustomNodeType -> Bool)
-> (CustomNodeType -> CustomNodeType -> Bool)
-> (CustomNodeType -> CustomNodeType -> Bool)
-> (CustomNodeType -> CustomNodeType -> CustomNodeType)
-> (CustomNodeType -> CustomNodeType -> CustomNodeType)
-> Ord CustomNodeType
CustomNodeType -> CustomNodeType -> Bool
CustomNodeType -> CustomNodeType -> Ordering
CustomNodeType -> CustomNodeType -> CustomNodeType
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: CustomNodeType -> CustomNodeType -> Ordering
compare :: CustomNodeType -> CustomNodeType -> Ordering
$c< :: CustomNodeType -> CustomNodeType -> Bool
< :: CustomNodeType -> CustomNodeType -> Bool
$c<= :: CustomNodeType -> CustomNodeType -> Bool
<= :: CustomNodeType -> CustomNodeType -> Bool
$c> :: CustomNodeType -> CustomNodeType -> Bool
> :: CustomNodeType -> CustomNodeType -> Bool
$c>= :: CustomNodeType -> CustomNodeType -> Bool
>= :: CustomNodeType -> CustomNodeType -> Bool
$cmax :: CustomNodeType -> CustomNodeType -> CustomNodeType
max :: CustomNodeType -> CustomNodeType -> CustomNodeType
$cmin :: CustomNodeType -> CustomNodeType -> CustomNodeType
min :: CustomNodeType -> CustomNodeType -> CustomNodeType
Ord)

customNodeTypeMapping :: Map CustomNodeType FastStringCompat
customNodeTypeMapping :: Map CustomNodeType FastStringCompat
customNodeTypeMapping = [(CustomNodeType, FastStringCompat)]
-> Map CustomNodeType FastStringCompat
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList
    [ (CustomNodeType
CustomNodeImportsGroup, FastStringCompat
"Imports")
    , (CustomNodeType
CustomNodeAdjacentSignatureDefinition, FastStringCompat
"AdjacentSignatureDefinition")
    ]

revCustomNodeTypeMapping :: Map FastStringCompat CustomNodeType
revCustomNodeTypeMapping :: Map FastStringCompat CustomNodeType
revCustomNodeTypeMapping = [(FastStringCompat, CustomNodeType)]
-> Map FastStringCompat CustomNodeType
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList ([(FastStringCompat, CustomNodeType)]
 -> Map FastStringCompat CustomNodeType)
-> (Map CustomNodeType FastStringCompat
    -> [(FastStringCompat, CustomNodeType)])
-> Map CustomNodeType FastStringCompat
-> Map FastStringCompat CustomNodeType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((CustomNodeType, FastStringCompat)
 -> (FastStringCompat, CustomNodeType))
-> [(CustomNodeType, FastStringCompat)]
-> [(FastStringCompat, CustomNodeType)]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\(CustomNodeType
k, FastStringCompat
v) -> (FastStringCompat
v, CustomNodeType
k)) ([(CustomNodeType, FastStringCompat)]
 -> [(FastStringCompat, CustomNodeType)])
-> (Map CustomNodeType FastStringCompat
    -> [(CustomNodeType, FastStringCompat)])
-> Map CustomNodeType FastStringCompat
-> [(FastStringCompat, CustomNodeType)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map CustomNodeType FastStringCompat
-> [(CustomNodeType, FastStringCompat)]
forall k a. Map k a -> [(k, a)]
Map.toList (Map CustomNodeType FastStringCompat
 -> Map FastStringCompat CustomNodeType)
-> Map CustomNodeType FastStringCompat
-> Map FastStringCompat CustomNodeType
forall a b. (a -> b) -> a -> b
$ Map CustomNodeType FastStringCompat
customNodeTypeMapping

customNodeTypeToFastString :: CustomNodeType -> FastStringCompat
customNodeTypeToFastString :: CustomNodeType -> FastStringCompat
customNodeTypeToFastString CustomNodeType
k = FastStringCompat -> Maybe FastStringCompat -> FastStringCompat
forall a. a -> Maybe a -> a
fromMaybe FastStringCompat
"" (Map CustomNodeType FastStringCompat
customNodeTypeMapping Map CustomNodeType FastStringCompat
-> CustomNodeType -> Maybe FastStringCompat
forall k a. Ord k => Map k a -> k -> Maybe a
Map.!? CustomNodeType
k)

{-|
Combines adjacent import declarations under a new parent node, so that the user will have an extra step selecting
the whole import area while expanding/shrinking the selection range.
-}
mergeImports :: forall a. HieAST a -> Reader (PreProcessEnv a) (HieAST a)
mergeImports :: forall a. HieAST a -> Reader (PreProcessEnv a) (HieAST a)
mergeImports HieAST a
node = HieAST a -> ReaderT (PreProcessEnv a) Identity (HieAST a)
forall a. a -> ReaderT (PreProcessEnv a) Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (HieAST a -> ReaderT (PreProcessEnv a) Identity (HieAST a))
-> HieAST a -> ReaderT (PreProcessEnv a) Identity (HieAST a)
forall a b. (a -> b) -> a -> b
$ HieAST a
node { nodeChildren = children }
  where
    children :: [HieAST a]
    children :: [HieAST a]
children = ([HieAST a] -> Maybe (HieAST a)) -> [[HieAST a]] -> [HieAST a]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe [HieAST a] -> Maybe (HieAST a)
merge
        ([[HieAST a]] -> [HieAST a])
-> (HieAST a -> [[HieAST a]]) -> HieAST a -> [HieAST a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (HieAST a -> HieAST a -> Bool) -> [HieAST a] -> [[HieAST a]]
forall a. (a -> a -> Bool) -> [a] -> [[a]]
groupBy (\HieAST a
x HieAST a
y -> HieAST a -> Bool
forall a. HieAST a -> Bool
nodeIsImport HieAST a
x Bool -> Bool -> Bool
&& HieAST a -> Bool
forall a. HieAST a -> Bool
nodeIsImport HieAST a
y)
        ([HieAST a] -> [[HieAST a]])
-> (HieAST a -> [HieAST a]) -> HieAST a -> [[HieAST a]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. HieAST a -> [HieAST a]
forall a. HieAST a -> [HieAST a]
nodeChildren (HieAST a -> [HieAST a]) -> HieAST a -> [HieAST a]
forall a b. (a -> b) -> a -> b
$ HieAST a
node

    merge :: [HieAST a] -> Maybe (HieAST a)
    merge :: [HieAST a] -> Maybe (HieAST a)
merge []     = Maybe (HieAST a)
forall a. Maybe a
Nothing
    merge [HieAST a
x]    = HieAST a -> Maybe (HieAST a)
forall a. a -> Maybe a
Just HieAST a
x
    merge (HieAST a
x:[HieAST a]
xs) = HieAST a -> Maybe (HieAST a)
forall a. a -> Maybe a
Just (HieAST a -> Maybe (HieAST a)) -> HieAST a -> Maybe (HieAST a)
forall a b. (a -> b) -> a -> b
$ CustomNodeType -> NonEmpty (HieAST a) -> HieAST a
forall a. CustomNodeType -> NonEmpty (HieAST a) -> HieAST a
createCustomNode CustomNodeType
CustomNodeImportsGroup (HieAST a
x HieAST a -> [HieAST a] -> NonEmpty (HieAST a)
forall a. a -> [a] -> NonEmpty a
NonEmpty.:| [HieAST a]
xs)

nodeIsImport :: HieAST a -> Bool
nodeIsImport :: forall a. HieAST a -> Bool
nodeIsImport = (FastStringCompat, FastStringCompat) -> HieAST a -> Bool
forall a. (FastStringCompat, FastStringCompat) -> HieAST a -> Bool
isAnnotationInAstNode (FastStringCompat
"ImportDecl", FastStringCompat
"ImportDecl")

{-|
Combine type signature with variable definition under a new parent node, if the signature is placed right before the
definition. This allows the user to have a step selecting both type signature and its accompanying definition.
-}
mergeSignatureWithDefinition :: HieAST a -> Reader (PreProcessEnv a) (HieAST a)
mergeSignatureWithDefinition :: forall a. HieAST a -> Reader (PreProcessEnv a) (HieAST a)
mergeSignatureWithDefinition HieAST a
node = do
    RefMap a
refMap <- (PreProcessEnv a -> RefMap a)
-> ReaderT (PreProcessEnv a) Identity (RefMap a)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks PreProcessEnv a -> RefMap a
forall a. PreProcessEnv a -> RefMap a
preProcessEnvRefMap
    -- Do this recursively for children, so that non top level functions can be handled.
    [HieAST a]
children' <- (HieAST a -> Reader (PreProcessEnv a) (HieAST a))
-> [HieAST a] -> ReaderT (PreProcessEnv a) Identity [HieAST a]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> [a] -> f [b]
traverse HieAST a -> Reader (PreProcessEnv a) (HieAST a)
forall a. HieAST a -> Reader (PreProcessEnv a) (HieAST a)
mergeSignatureWithDefinition (HieAST a -> [HieAST a]
forall a. HieAST a -> [HieAST a]
nodeChildren HieAST a
node)
    HieAST a -> Reader (PreProcessEnv a) (HieAST a)
forall a. a -> ReaderT (PreProcessEnv a) Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (HieAST a -> Reader (PreProcessEnv a) (HieAST a))
-> HieAST a -> Reader (PreProcessEnv a) (HieAST a)
forall a b. (a -> b) -> a -> b
$ HieAST a
node { nodeChildren = reverse $ foldl' (go refMap) [] children' }
  where
    -- For every two adjacent nodes, we try to combine them into one.
    go :: RefMap a -> [HieAST a] -> HieAST a -> [HieAST a]
    go :: forall a. RefMap a -> [HieAST a] -> HieAST a -> [HieAST a]
go RefMap a
_ [] HieAST a
node' = [HieAST a
node']
    go RefMap a
refMap (HieAST a
prev:[HieAST a]
others) HieAST a
node' =
        case RefMap a -> (HieAST a, HieAST a) -> Maybe (HieAST a)
forall a. RefMap a -> (HieAST a, HieAST a) -> Maybe (HieAST a)
mergeAdjacentSigDef RefMap a
refMap (HieAST a
prev, HieAST a
node') of
            Maybe (HieAST a)
Nothing   -> HieAST a
node'HieAST a -> [HieAST a] -> [HieAST a]
forall a. a -> [a] -> [a]
:HieAST a
prevHieAST a -> [HieAST a] -> [HieAST a]
forall a. a -> [a] -> [a]
:[HieAST a]
others
            Just HieAST a
comb -> HieAST a
combHieAST a -> [HieAST a] -> [HieAST a]
forall a. a -> [a] -> [a]
:[HieAST a]
others

-- | Merge adjacent type signature and variable/function definition, if the type signature belongs to that variable or
-- function.
--
-- The implementation potentially has some corner cases not handled properly.
mergeAdjacentSigDef :: RefMap a -> (HieAST a, HieAST a) -> Maybe (HieAST a)
mergeAdjacentSigDef :: forall a. RefMap a -> (HieAST a, HieAST a) -> Maybe (HieAST a)
mergeAdjacentSigDef RefMap a
refMap (HieAST a
n1, HieAST a
n2) = do
    -- Let's check the node's annotation. There should be a function binding following its type signature.
    Maybe ()
checkAnnotation
    -- Find the identifier of the type signature.
    Identifier
typeSigId <- HieAST a -> Maybe Identifier
forall a. HieAST a -> Maybe Identifier
identifierForTypeSig HieAST a
n1
    -- Does that identifier appear in the second AST node as a definition? If so, we combines the two nodes.
    [(Span, IdentifierDetails a)]
refs <- Identifier -> RefMap a -> Maybe [(Span, IdentifierDetails a)]
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Identifier
typeSigId RefMap a
refMap
    if ((Span, IdentifierDetails a) -> Bool)
-> [(Span, IdentifierDetails a)] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Span -> (Span, IdentifierDetails a) -> Bool
forall a. Span -> (Span, IdentifierDetails a) -> Bool
isIdentADef (HieAST a -> Span
forall a. HieAST a -> Span
nodeSpan HieAST a
n2)) [(Span, IdentifierDetails a)]
refs
    then HieAST a -> Maybe (HieAST a)
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (HieAST a -> Maybe (HieAST a))
-> (NonEmpty (HieAST a) -> HieAST a)
-> NonEmpty (HieAST a)
-> Maybe (HieAST a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CustomNodeType -> NonEmpty (HieAST a) -> HieAST a
forall a. CustomNodeType -> NonEmpty (HieAST a) -> HieAST a
createCustomNode CustomNodeType
CustomNodeAdjacentSignatureDefinition (NonEmpty (HieAST a) -> Maybe (HieAST a))
-> NonEmpty (HieAST a) -> Maybe (HieAST a)
forall a b. (a -> b) -> a -> b
$ HieAST a
n1 HieAST a -> [HieAST a] -> NonEmpty (HieAST a)
forall a. a -> [a] -> NonEmpty a
NonEmpty.:| [HieAST a
n2]
    else Maybe (HieAST a)
forall a. Maybe a
Nothing
  where
    checkAnnotation :: Maybe ()
    checkAnnotation :: Maybe ()
checkAnnotation =
      if (FastStringCompat
"TypeSig", FastStringCompat
"Sig") (FastStringCompat, FastStringCompat) -> HieAST a -> Bool
forall a. (FastStringCompat, FastStringCompat) -> HieAST a -> Bool
`isAnnotationInAstNode` HieAST a
n1 Bool -> Bool -> Bool
&&
         ((FastStringCompat
"FunBind", FastStringCompat
"HsBindLR") (FastStringCompat, FastStringCompat) -> HieAST a -> Bool
forall a. (FastStringCompat, FastStringCompat) -> HieAST a -> Bool
`isAnnotationInAstNode` HieAST a
n2 Bool -> Bool -> Bool
|| (FastStringCompat
"VarBind", FastStringCompat
"HsBindLR") (FastStringCompat, FastStringCompat) -> HieAST a -> Bool
forall a. (FastStringCompat, FastStringCompat) -> HieAST a -> Bool
`isAnnotationInAstNode` HieAST a
n2)
      then () -> Maybe ()
forall a. a -> Maybe a
Just ()
      else Maybe ()
forall a. Maybe a
Nothing

{-|
Given the AST node of a type signature, tries to find the identifier of it.
-}
identifierForTypeSig :: forall a. HieAST a -> Maybe Identifier
identifierForTypeSig :: forall a. HieAST a -> Maybe Identifier
identifierForTypeSig HieAST a
node =
    {-
        It seems that the identifier lives in one of the children, so we search for the first 'TyDecl' node in
        its children recursively.
    -}
    case (HieAST a -> Maybe Identifier) -> [HieAST a] -> [Identifier]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe HieAST a -> Maybe Identifier
extractIdentifier [HieAST a]
nodes of
      []        -> Maybe Identifier
forall a. Maybe a
Nothing
      (Identifier
ident:[Identifier]
_) -> Identifier -> Maybe Identifier
forall a. a -> Maybe a
Just Identifier
ident
  where
    nodes :: [HieAST a]
nodes = HieAST a -> [HieAST a]
forall a. HieAST a -> [HieAST a]
flattenAst HieAST a
node

    extractIdentifier :: HieAST a -> Maybe Identifier
    extractIdentifier :: HieAST a -> Maybe Identifier
extractIdentifier HieAST a
node' = HieAST a -> Maybe (NodeInfo a)
forall a. HieAST a -> Maybe (NodeInfo a)
sourceNodeInfo HieAST a
node' Maybe (NodeInfo a)
-> (NodeInfo a -> Maybe Identifier) -> Maybe Identifier
forall a b. Maybe a -> (a -> Maybe b) -> Maybe b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>=
        (((Identifier, IdentifierDetails a) -> Identifier)
-> Maybe (Identifier, IdentifierDetails a) -> Maybe Identifier
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Identifier, IdentifierDetails a) -> Identifier
forall a b. (a, b) -> a
fst (Maybe (Identifier, IdentifierDetails a) -> Maybe Identifier)
-> (NodeInfo a -> Maybe (Identifier, IdentifierDetails a))
-> NodeInfo a
-> Maybe Identifier
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Identifier, IdentifierDetails a) -> Bool)
-> [(Identifier, IdentifierDetails a)]
-> Maybe (Identifier, IdentifierDetails a)
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find (\(Identifier
_, IdentifierDetails a
detail) -> ContextInfo
TyDecl ContextInfo -> Set ContextInfo -> Bool
forall a. Ord a => a -> Set a -> Bool
`Set.member` IdentifierDetails a -> Set ContextInfo
forall a. IdentifierDetails a -> Set ContextInfo
identInfo IdentifierDetails a
detail)
        ([(Identifier, IdentifierDetails a)]
 -> Maybe (Identifier, IdentifierDetails a))
-> (NodeInfo a -> [(Identifier, IdentifierDetails a)])
-> NodeInfo a
-> Maybe (Identifier, IdentifierDetails a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map Identifier (IdentifierDetails a)
-> [(Identifier, IdentifierDetails a)]
forall k a. Map k a -> [(k, a)]
Map.toList (Map Identifier (IdentifierDetails a)
 -> [(Identifier, IdentifierDetails a)])
-> (NodeInfo a -> Map Identifier (IdentifierDetails a))
-> NodeInfo a
-> [(Identifier, IdentifierDetails a)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. NodeInfo a -> Map Identifier (IdentifierDetails a)
forall a. NodeInfo a -> NodeIdentifiers a
nodeIdentifiers)

-- | Determines if the given occurrence of an identifier is a function/variable definition in the outer span
isIdentADef :: Span -> (Span, IdentifierDetails a) -> Bool
isIdentADef :: forall a. Span -> (Span, IdentifierDetails a) -> Bool
isIdentADef Span
outerSpan (Span
span, IdentifierDetails a
detail) =
    Span -> RealSrcLoc
realSrcSpanStart Span
span RealSrcLoc -> RealSrcLoc -> Bool
forall a. Ord a => a -> a -> Bool
>= Span -> RealSrcLoc
realSrcSpanStart Span
outerSpan Bool -> Bool -> Bool
&& Span -> RealSrcLoc
realSrcSpanEnd Span
span RealSrcLoc -> RealSrcLoc -> Bool
forall a. Ord a => a -> a -> Bool
<= Span -> RealSrcLoc
realSrcSpanEnd Span
outerSpan
    Bool -> Bool -> Bool
&& Bool
isDef
  where
    isDef :: Bool
    isDef :: Bool
isDef = (ContextInfo -> Bool) -> [ContextInfo] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any ContextInfo -> Bool
isContextInfoDef ([ContextInfo] -> Bool)
-> (IdentifierDetails a -> [ContextInfo])
-> IdentifierDetails a
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Set ContextInfo -> [ContextInfo]
forall a. Set a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList (Set ContextInfo -> [ContextInfo])
-> (IdentifierDetails a -> Set ContextInfo)
-> IdentifierDetails a
-> [ContextInfo]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IdentifierDetails a -> Set ContextInfo
forall a. IdentifierDetails a -> Set ContextInfo
identInfo (IdentifierDetails a -> Bool) -> IdentifierDetails a -> Bool
forall a b. (a -> b) -> a -> b
$ IdentifierDetails a
detail

    -- Determines if the 'ContextInfo' represents a variable/function definition
    isContextInfoDef :: ContextInfo -> Bool
    isContextInfoDef :: ContextInfo -> Bool
isContextInfoDef ValBind{} = Bool
True
    isContextInfoDef ContextInfo
MatchBind = Bool
True
    isContextInfoDef ContextInfo
_         = Bool
False

isAnnotationInAstNode :: (FastStringCompat, FastStringCompat) -> HieAST a -> Bool
isAnnotationInAstNode :: forall a. (FastStringCompat, FastStringCompat) -> HieAST a -> Bool
isAnnotationInAstNode (FastStringCompat, FastStringCompat)
p = Bool -> (NodeInfo a -> Bool) -> Maybe (NodeInfo a) -> Bool
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
False ((FastStringCompat, FastStringCompat) -> NodeInfo a -> Bool
forall a.
(FastStringCompat, FastStringCompat) -> NodeInfo a -> Bool
isAnnotationInNodeInfo (FastStringCompat, FastStringCompat)
p) (Maybe (NodeInfo a) -> Bool)
-> (HieAST a -> Maybe (NodeInfo a)) -> HieAST a -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. HieAST a -> Maybe (NodeInfo a)
forall a. HieAST a -> Maybe (NodeInfo a)
sourceNodeInfo