{-# 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)

{-|
Extra arguments for 'preaProcessAST', meant to be used in a 'Reader' context. We use 'Reader' to combine
-}
newtype PreProcessEnv a = PreProcessEnv
    { 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 :: 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

{-|
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 :: 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

{-|
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 :: 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)
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
    -- For every two adjacent nodes, we try to combine them into one.
    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

-- | 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 :: 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 (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

{-|
Given the AST node of a type signature, tries to find the identifier of it.
-}
identifierForTypeSig :: forall a. HieAST a -> Maybe Identifier
identifierForTypeSig :: 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)
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)

-- | Determines if the given occurence of an identifier is a function/variable definition in the outer span
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

    -- Does 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 :: (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