{-# LANGUAGE OverloadedStrings   #-}
{-# LANGUAGE ScopedTypeVariables #-}

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 = 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

{-|
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 = forall a. NodeInfo a -> Span -> [HieAST a] -> HieAST a
mkAstNode forall {a}. NodeInfo a
customNodeInfo 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

    customNodeInfo :: NodeInfo a
customNodeInfo = 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 <- forall a. HieAST a -> Maybe (NodeInfo a)
generatedNodeInfo HieAST a
node
    forall a. First a -> a
getFirst forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (FastStringCompat, FastStringCompat)
-> Maybe (First CustomNodeType)
go (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 forall a. Eq a => a -> a -> Bool
== FastStringCompat
"HlsCustom", Just CustomNodeType
v' <- Map FastStringCompat CustomNodeType
revCustomNodeTypeMapping forall k a. Ord k => Map k a -> k -> Maybe a
Map.!? FastStringCompat
v = forall a. a -> Maybe a
Just (forall a. a -> First a
First CustomNodeType
v')
        | Bool
otherwise = 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
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [CustomNodeType] -> ShowS
$cshowList :: [CustomNodeType] -> ShowS
show :: CustomNodeType -> String
$cshow :: CustomNodeType -> String
showsPrec :: Int -> CustomNodeType -> ShowS
$cshowsPrec :: Int -> CustomNodeType -> ShowS
Show, CustomNodeType -> CustomNodeType -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: CustomNodeType -> CustomNodeType -> Bool
$c/= :: CustomNodeType -> CustomNodeType -> Bool
== :: CustomNodeType -> CustomNodeType -> Bool
$c== :: CustomNodeType -> CustomNodeType -> Bool
Eq, Eq 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
min :: CustomNodeType -> CustomNodeType -> CustomNodeType
$cmin :: CustomNodeType -> CustomNodeType -> CustomNodeType
max :: CustomNodeType -> CustomNodeType -> CustomNodeType
$cmax :: CustomNodeType -> CustomNodeType -> CustomNodeType
>= :: CustomNodeType -> CustomNodeType -> Bool
$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
compare :: CustomNodeType -> CustomNodeType -> Ordering
$ccompare :: CustomNodeType -> CustomNodeType -> Ordering
Ord)

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

customNodeTypeToFastString :: CustomNodeType -> FastStringCompat
customNodeTypeToFastString :: CustomNodeType -> FastStringCompat
customNodeTypeToFastString CustomNodeType
k = forall a. a -> Maybe a -> a
fromMaybe FastStringCompat
"" (Map CustomNodeType FastStringCompat
customNodeTypeMapping 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 = 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. CustomNodeType -> NonEmpty (HieAST a) -> HieAST a
createCustomNode CustomNodeType
CustomNodeImportsGroup (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. (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 <- 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.
    Identifier
typeSigId <- 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 <- forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Identifier
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. CustomNodeType -> NonEmpty (HieAST a) -> HieAST a
createCustomNode CustomNodeType
CustomNodeAdjacentSignatureDefinition 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 (FastStringCompat
"TypeSig", FastStringCompat
"Sig") forall a. (FastStringCompat, FastStringCompat) -> HieAST a -> Bool
`isAnnotationInAstNode` HieAST a
n1 Bool -> Bool -> Bool
&&
         ((FastStringCompat
"FunBind", FastStringCompat
"HsBindLR") forall a. (FastStringCompat, FastStringCompat) -> HieAST a -> Bool
`isAnnotationInAstNode` HieAST a
n2 Bool -> Bool -> Bool
|| (FastStringCompat
"VarBind", FastStringCompat
"HsBindLR") forall a. (FastStringCompat, FastStringCompat) -> 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 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 forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe HieAST a -> Maybe Identifier
extractIdentifier [HieAST a]
nodes of
      []        -> forall a. Maybe a
Nothing
      (Identifier
ident:[Identifier]
_) -> forall a. a -> Maybe a
Just Identifier
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 Identifier
extractIdentifier HieAST a
node' = forall a. HieAST a -> Maybe (NodeInfo a)
sourceNodeInfo 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 (\(Identifier
_, 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 (t :: * -> *) a. Foldable t => t a -> [a]
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

    -- 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 = forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
False (forall a.
(FastStringCompat, FastStringCompat) -> NodeInfo a -> Bool
isAnnotationInNodeInfo (FastStringCompat, FastStringCompat)
p) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. HieAST a -> Maybe (NodeInfo a)
sourceNodeInfo