{-# 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
    { 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 = forall a. HieAST a -> Reader (PreProcessEnv a) (HieAST a)
mergeImports HieAST a
node forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= 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 :: forall a. HieAST a -> Reader (PreProcessEnv a) (HieAST a)
mergeImports HieAST a
node = forall (f :: * -> *) a. Applicative f => a -> f a
pure 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 = forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe [HieAST a] -> Maybe (HieAST a)
merge
        forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. (a -> a -> Bool) -> [a] -> [[a]]
groupBy (\HieAST a
x HieAST a
y -> forall a. HieAST a -> Bool
nodeIsImport HieAST a
x Bool -> Bool -> Bool
&& forall a. HieAST a -> Bool
nodeIsImport HieAST a
y)
        forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. HieAST a -> [HieAST a]
nodeChildren forall a b. (a -> b) -> a -> b
$ HieAST a
node

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

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

createVirtualNode :: NonEmpty (HieAST a) -> HieAST a
createVirtualNode :: forall a. NonEmpty (HieAST a) -> HieAST a
createVirtualNode NonEmpty (HieAST a)
children = forall a. NodeInfo a -> Span -> [HieAST a] -> HieAST a
mkAstNode (forall a.
Set NodeAnnotation -> [a] -> NodeIdentifiers a -> NodeInfo a
NodeInfo forall a. Monoid a => a
mempty forall a. Monoid a => a
mempty forall a. Monoid a => a
mempty) Span
span' (forall a. NonEmpty a -> [a]
NonEmpty.toList NonEmpty (HieAST a)
children)
  where
    span' :: RealSrcSpan
    span' :: Span
span' = forall a. Identity a -> a
runIdentity forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) (m :: * -> *) a.
(Foldable1 t, Monad m) =>
(a -> a -> m a) -> t a -> m a
foldlM1 (\Span
x Span
y -> forall a. a -> Identity a
Identity (Span -> Span -> Span
combineRealSrcSpans Span
x Span
y)) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a. HieAST a -> Span
nodeSpan 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 :: forall a. HieAST a -> Reader (PreProcessEnv a) (HieAST a)
mergeSignatureWithDefinition HieAST a
node = do
    RefMap a
refMap <- forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks forall a. PreProcessEnv a -> RefMap a
preProcessEnvRefMap
    -- Do this recursively for children, so that non top level functions can be handled.
    [HieAST a]
children' <- forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse forall a. HieAST a -> Reader (PreProcessEnv a) (HieAST a)
mergeSignatureWithDefinition (forall a. HieAST a -> [HieAST a]
nodeChildren HieAST a
node)
    forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ HieAST a
node { nodeChildren :: [HieAST a]
nodeChildren = forall a. [a] -> [a]
reverse forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (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 :: 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 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'forall a. a -> [a] -> [a]
:HieAST a
prevforall a. a -> [a] -> [a]
:[HieAST a]
others
            Just HieAST a
comb -> HieAST a
combforall 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.
    Either ModuleName Name
typeSigId <- forall a. HieAST a -> Maybe (Either ModuleName Name)
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 <- forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Either ModuleName Name
typeSigId RefMap a
refMap
    if forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (forall a. Span -> (Span, IdentifierDetails a) -> Bool
isIdentADef (forall a. HieAST a -> Span
nodeSpan HieAST a
n2)) [(Span, IdentifierDetails a)]
refs
    then forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. NonEmpty (HieAST a) -> HieAST a
createVirtualNode forall a b. (a -> b) -> a -> b
$ HieAST a
n1 forall a. a -> [a] -> NonEmpty a
NonEmpty.:| [HieAST a
n2]
    else forall a. Maybe a
Nothing
  where
    checkAnnotation :: Maybe ()
    checkAnnotation :: Maybe ()
checkAnnotation =
      if (FastString
"TypeSig", FastString
"Sig") forall a. (FastString, FastString) -> HieAST a -> Bool
`isAnnotationInAstNode` HieAST a
n1 Bool -> Bool -> Bool
&&
         ((FastString
"FunBind", FastString
"HsBindLR") forall a. (FastString, FastString) -> HieAST a -> Bool
`isAnnotationInAstNode` HieAST a
n2 Bool -> Bool -> Bool
|| (FastString
"VarBind", FastString
"HsBindLR") forall a. (FastString, FastString) -> HieAST a -> Bool
`isAnnotationInAstNode` HieAST a
n2)
      then forall a. a -> Maybe a
Just ()
      else 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 (Either ModuleName Name)
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 forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe HieAST a -> Maybe (Either ModuleName Name)
extractIdentifier [HieAST a]
nodes of
      []        -> forall a. Maybe a
Nothing
      (Either ModuleName Name
ident:[Either ModuleName Name]
_) -> forall a. a -> Maybe a
Just Either ModuleName Name
ident
  where
    nodes :: [HieAST a]
nodes = forall a. HieAST a -> [HieAST a]
flattenAst HieAST a
node

    extractIdentifier :: HieAST a -> Maybe Identifier
    extractIdentifier :: HieAST a -> Maybe (Either ModuleName Name)
extractIdentifier HieAST a
node' = forall a. HieAST a -> Maybe (NodeInfo a)
nodeInfoFromSource HieAST a
node' forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>=
        (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a b. (a, b) -> a
fst forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find (\(Either ModuleName Name
_, IdentifierDetails a
detail) -> ContextInfo
TyDecl forall a. Ord a => a -> Set a -> Bool
`Set.member` forall a. IdentifierDetails a -> Set ContextInfo
identInfo IdentifierDetails a
detail)
        forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Map k a -> [(k, a)]
Map.toList forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 :: forall a. Span -> (Span, IdentifierDetails a) -> Bool
isIdentADef Span
outerSpan (Span
span, IdentifierDetails a
detail) =
    Span -> RealSrcLoc
realSrcSpanStart Span
span forall a. Ord a => a -> a -> Bool
>= Span -> RealSrcLoc
realSrcSpanStart Span
outerSpan Bool -> Bool -> Bool
&& Span -> RealSrcLoc
realSrcSpanEnd Span
span forall a. Ord a => a -> a -> Bool
<= Span -> RealSrcLoc
realSrcSpanEnd Span
outerSpan
    Bool -> Bool -> Bool
&& Bool
isDef
  where
    isDef :: Bool
    isDef :: Bool
isDef = forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any ContextInfo -> Bool
isContextInfoDef forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Set a -> [a]
Set.toList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. IdentifierDetails a -> Set ContextInfo
identInfo 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 :: forall a. (FastString, FastString) -> HieAST a -> Bool
isAnnotationInAstNode (FastString, FastString)
p = forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
False (forall a. (FastString, FastString) -> NodeInfo a -> Bool
isAnnotationInNodeInfo (FastString, FastString)
p) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. HieAST a -> Maybe (NodeInfo a)
nodeInfoFromSource