{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
module Ide.Plugin.SelectionRange.ASTPreProcess
( preProcessAST
, PreProcessEnv(..)
) where
import Control.Monad.Reader (Reader, asks)
import Data.Foldable (find, foldl')
import Data.Functor.Identity (Identity (Identity, runIdentity))
import Data.List (groupBy)
import Data.List.NonEmpty (NonEmpty)
import qualified Data.List.NonEmpty as NonEmpty
import qualified Data.Map.Strict as Map
import Data.Maybe (mapMaybe)
import Data.Semigroup.Foldable (foldlM1)
import qualified Data.Set as Set
import Development.IDE.GHC.Compat (ContextInfo (MatchBind, TyDecl, ValBind),
HieAST (..), Identifier,
IdentifierDetails (identInfo),
NodeInfo (NodeInfo, nodeIdentifiers),
RealSrcSpan, RefMap, Span,
combineRealSrcSpans,
flattenAst,
isAnnotationInNodeInfo,
mkAstNode, nodeInfoFromSource,
realSrcSpanEnd,
realSrcSpanStart)
import Development.IDE.GHC.Compat.Util (FastString)
import Prelude hiding (span)
newtype PreProcessEnv a = PreProcessEnv
{ PreProcessEnv a -> RefMap a
preProcessEnvRefMap :: RefMap a
}
preProcessAST :: HieAST a -> Reader (PreProcessEnv a) (HieAST a)
preProcessAST :: 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 (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
mergeImports :: forall a. HieAST a -> Reader (PreProcessEnv a) (HieAST a)
mergeImports :: HieAST a -> Reader (PreProcessEnv a) (HieAST a)
mergeImports HieAST a
node = HieAST a -> Reader (PreProcessEnv a) (HieAST 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 :: [HieAST a]
nodeChildren = [HieAST a]
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
$ NonEmpty (HieAST a) -> HieAST a
forall a. NonEmpty (HieAST a) -> HieAST a
createVirtualNode (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 :: HieAST a -> Bool
nodeIsImport = (FastString, FastString) -> HieAST a -> Bool
forall a. (FastString, FastString) -> HieAST a -> Bool
isAnnotationInAstNode (FastString
"ImportDecl", FastString
"ImportDecl")
createVirtualNode :: NonEmpty (HieAST a) -> HieAST a
createVirtualNode :: NonEmpty (HieAST a) -> HieAST a
createVirtualNode NonEmpty (HieAST a)
children = NodeInfo a -> Span -> [HieAST a] -> HieAST a
forall a. NodeInfo a -> Span -> [HieAST a] -> HieAST a
mkAstNode (Set (FastString, FastString)
-> [a] -> NodeIdentifiers a -> NodeInfo a
forall a.
Set (FastString, FastString)
-> [a] -> NodeIdentifiers a -> NodeInfo a
NodeInfo Set (FastString, FastString)
forall a. Monoid a => a
mempty [a]
forall a. Monoid a => a
mempty NodeIdentifiers a
forall a. Monoid a => a
mempty) 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 (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
mergeSignatureWithDefinition :: HieAST a -> Reader (PreProcessEnv a) (HieAST a)
mergeSignatureWithDefinition :: 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
[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)
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 (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 :: [HieAST a]
nodeChildren = [HieAST a] -> [HieAST a]
forall a. [a] -> [a]
reverse ([HieAST a] -> [HieAST a]) -> [HieAST a] -> [HieAST a]
forall a b. (a -> b) -> a -> b
$ ([HieAST a] -> HieAST a -> [HieAST a])
-> [HieAST a] -> [HieAST a] -> [HieAST a]
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (RefMap a -> [HieAST a] -> HieAST a -> [HieAST a]
forall a. RefMap a -> [HieAST a] -> HieAST a -> [HieAST a]
go RefMap a
refMap) [] [HieAST a]
children' }
where
go :: RefMap a -> [HieAST a] -> HieAST a -> [HieAST a]
go :: 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
mergeAdjacentSigDef :: RefMap a -> (HieAST a, HieAST a) -> Maybe (HieAST a)
mergeAdjacentSigDef :: RefMap a -> (HieAST a, HieAST a) -> Maybe (HieAST a)
mergeAdjacentSigDef RefMap a
refMap (HieAST a
n1, HieAST a
n2) = do
Maybe ()
checkAnnotation
Identifier
typeSigId <- HieAST a -> Maybe Identifier
forall a. HieAST a -> Maybe Identifier
identifierForTypeSig HieAST a
n1
[(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 (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
. NonEmpty (HieAST a) -> HieAST a
forall a. NonEmpty (HieAST a) -> HieAST a
createVirtualNode (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 (FastString
"TypeSig", FastString
"Sig") (FastString, FastString) -> HieAST a -> Bool
forall a. (FastString, FastString) -> HieAST a -> Bool
`isAnnotationInAstNode` HieAST a
n1 Bool -> Bool -> Bool
&&
((FastString
"FunBind", FastString
"HsBindLR") (FastString, FastString) -> HieAST a -> Bool
forall a. (FastString, FastString) -> HieAST a -> Bool
`isAnnotationInAstNode` HieAST a
n2 Bool -> Bool -> Bool
|| (FastString
"VarBind", FastString
"HsBindLR") (FastString, FastString) -> HieAST a -> Bool
forall a. (FastString, FastString) -> HieAST a -> Bool
`isAnnotationInAstNode` HieAST a
n2)
then () -> Maybe ()
forall a. a -> Maybe a
Just ()
else Maybe ()
forall a. Maybe a
Nothing
identifierForTypeSig :: forall a. HieAST a -> Maybe Identifier
identifierForTypeSig :: HieAST a -> Maybe Identifier
identifierForTypeSig HieAST a
node =
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)
nodeInfoFromSource HieAST a
node' Maybe (NodeInfo a)
-> (NodeInfo a -> Maybe Identifier) -> Maybe Identifier
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>=
(((Identifier, IdentifierDetails a) -> Identifier)
-> Maybe (Identifier, IdentifierDetails a) -> Maybe Identifier
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)
isIdentADef :: Span -> (Span, IdentifierDetails a) -> Bool
isIdentADef :: 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]
Set.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
isContextInfoDef :: ContextInfo -> Bool
isContextInfoDef :: ContextInfo -> Bool
isContextInfoDef ValBind{} = Bool
True
isContextInfoDef ContextInfo
MatchBind = Bool
True
isContextInfoDef ContextInfo
_ = Bool
False
isAnnotationInAstNode :: (FastString, FastString) -> HieAST a -> Bool
isAnnotationInAstNode :: (FastString, FastString) -> HieAST a -> Bool
isAnnotationInAstNode (FastString, FastString)
p = Bool -> (NodeInfo a -> Bool) -> Maybe (NodeInfo a) -> Bool
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
False ((FastString, FastString) -> NodeInfo a -> Bool
forall a. (FastString, FastString) -> NodeInfo a -> Bool
isAnnotationInNodeInfo (FastString, FastString)
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)
nodeInfoFromSource