Improve hls-fixity-plugin (#3205)

Author: wz1000

plugins/hls-explicit-fixity-plugin/hls-explicit-fixity-plugin.cabal

@@ -35,6 +35,7 @@ library
     , hls-plugin-api       ^>=1.5
     , lsp                   >=1.2.0.1
     , text
+    , transformers
 
   ghc-options:
     -Wall

plugins/hls-explicit-fixity-plugin/src/Ide/Plugin/ExplicitFixity.hs

@@ -1,32 +1,30 @@
 {-# LANGUAGE DeriveGeneric     #-}
+{-# LANGUAGE TupleSections     #-}
 {-# LANGUAGE LambdaCase        #-}
 {-# LANGUAGE NamedFieldPuns    #-}
 {-# LANGUAGE OverloadedStrings #-}
 {-# LANGUAGE TypeFamilies      #-}
-{-# OPTIONS_GHC -Wno-deprecations #-}
-{-# OPTIONS_GHC -Wno-unrecognised-pragmas #-}
-{-# HLINT ignore "Use nubOrdOn" #-}
+{-# OPTIONS_GHC -Wno-orphans #-}
+{-# OPTIONS_GHC -Wno-incomplete-uni-patterns #-}
 
 module Ide.Plugin.ExplicitFixity(descriptor) where
 
 import           Control.DeepSeq
-import           Control.Monad                        (forM)
+import           Control.Monad.Trans.Maybe
 import           Control.Monad.IO.Class               (MonadIO, liftIO)
-import           Data.Coerce                          (coerce)
 import           Data.Either.Extra
 import           Data.Hashable
-import           Data.List.Extra                      (nubOn)
-import qualified Data.Map                             as M
+import qualified Data.Map.Strict                      as M
+import qualified Data.Set                             as S
 import           Data.Maybe
-import           Data.Monoid
 import qualified Data.Text                            as T
 import           Development.IDE                      hiding (pluginHandlers,
                                                        pluginRules)
 import           Development.IDE.Core.PositionMapping (idDelta)
 import           Development.IDE.Core.Shake           (addPersistentRule)
 import qualified Development.IDE.Core.Shake           as Shake
+import           Development.IDE.Spans.AtPoint
 import           Development.IDE.GHC.Compat
-import           Development.IDE.GHC.Compat.Util      (FastString)
 import qualified Development.IDE.GHC.Compat.Util      as Util
 import           Development.IDE.LSP.Notifications    (ghcideNotificationsPluginPriority)
 import           GHC.Generics                         (Generic)
@@ -48,14 +46,14 @@ descriptor recorder pluginId = (defaultPluginDescriptor pluginId)
 hover :: PluginMethodHandler IdeState TextDocumentHover
 hover state _ (HoverParams (TextDocumentIdentifier uri) pos _) = pluginResponse $ do
     nfp <- getNormalizedFilePath uri
-    fixityTrees <- handleMaybeM "Unable to get fixity"
-        $ liftIO
-        $ runAction "ExplicitFixity.GetFixity" state
-        $ use GetFixity nfp
-    -- We don't have much fixities on one position, so `nubOn` is acceptable.
-    pure $ toHover $ nubOn snd $ findInTree fixityTrees pos fNodeFixty
+    handleMaybeM "ExplicitFixity: Unable to get fixity" $ liftIO $ runIdeAction "ExplicitFixity" (shakeExtras state) $ runMaybeT $ do
+      (FixityMap fixmap, _) <- useE GetFixity nfp
+      (HAR{hieAst}, mapping) <- useE GetHieAst nfp
+      let ns = getNamesAtPoint hieAst pos mapping
+          fs = mapMaybe (\n -> (n,) <$> M.lookup n fixmap) ns
+      pure $ toHover $ fs
     where
-        toHover :: [(T.Text, Fixity)] -> Maybe Hover
+        toHover :: [(Name, Fixity)] -> Maybe Hover
         toHover [] = Nothing
         toHover fixities =
             let -- Splicing fixity info
@@ -64,44 +62,19 @@ hover state _ (HoverParams (TextDocumentIdentifier uri) pos _) = pluginResponse
                 contents' = "\n" <> sectionSeparator <> contents
             in  Just $ Hover (HoverContents $ unmarkedUpContent contents') Nothing
 
-        fixityText :: (T.Text, Fixity) -> T.Text
+        fixityText :: (Name, Fixity) -> T.Text
         fixityText (name, Fixity _ precedence direction) =
-            printOutputable direction <> " " <> printOutputable precedence <> " `" <> name <> "`"
-
--- | Transferred from ghc `selectSmallestContaining`
-selectSmallestContainingForFixityTree :: Span -> FixityTree -> Maybe FixityTree
-selectSmallestContainingForFixityTree sp node
-    | sp `containsSpan` fNodeSpan node = Just node
-    | fNodeSpan node `containsSpan` sp = getFirst $ mconcat
-        [ foldMap (First . selectSmallestContainingForFixityTree sp) $ fNodeChildren node
-        , First (Just node)
-        ]
-    | otherwise = Nothing
-
--- | Transferred from ghcide `pointCommand`
-findInTree :: FixityTrees -> Position -> (FixityTree -> [a]) -> [a]
-findInTree tree pos k =
-    concat $ M.elems $ flip M.mapWithKey tree $ \fs ast ->
-        maybe [] k (selectSmallestContainingForFixityTree (sp fs) ast)
-    where
-        sloc fs = mkRealSrcLoc fs (fromIntegral $ line+1) (fromIntegral $ cha+1)
-        sp fs = mkRealSrcSpan (sloc fs) (sloc fs)
-        line = _line pos
-        cha = _character pos
-
-data FixityTree = FNode
-    { fNodeSpan     :: Span
-    , fNodeChildren :: [FixityTree]
-    , fNodeFixty    :: [(T.Text, Fixity)]
-    } deriving (Generic)
+            printOutputable direction <> " " <> printOutputable precedence <> " `" <> printOutputable name <> "`"
 
-instance NFData FixityTree where
-    rnf = rwhnf
+newtype FixityMap = FixityMap (M.Map Name Fixity)
+instance Show FixityMap where
+  show _ = "FixityMap"
 
-instance Show FixityTree where
-    show _ = "<FixityTree>"
+instance NFData FixityMap where
+  rnf (FixityMap xs) = rnf xs
 
-type FixityTrees = M.Map FastString FixityTree
+instance NFData Fixity where
+  rnf = rwhnf
 
 newtype Log = LogShake Shake.Log
 
@@ -114,53 +87,33 @@ data GetFixity = GetFixity deriving (Show, Eq, Generic)
 instance Hashable GetFixity
 instance NFData GetFixity
 
-type instance RuleResult GetFixity = FixityTrees
-
-fakeFixityTrees :: FixityTrees
-fakeFixityTrees = M.empty
-
--- | Convert a HieASTs to FixityTrees with fixity info gathered
-hieAstsToFixitTrees :: MonadIO m => HscEnv -> TcGblEnv -> HieASTs a -> m FixityTrees
-hieAstsToFixitTrees hscEnv tcGblEnv ast =
-    -- coerce to avoid compatibility issues.
-    M.mapKeysWith const coerce <$>
-        sequence (M.map (hieAstToFixtyTree hscEnv tcGblEnv) (getAsts ast))
+type instance RuleResult GetFixity = FixityMap
 
 -- | Convert a HieAST to FixityTree with fixity info gathered
-hieAstToFixtyTree :: MonadIO m => HscEnv -> TcGblEnv -> HieAST a -> m FixityTree
-hieAstToFixtyTree hscEnv tcGblEnv ast = case ast of
-    (Node _ span []) -> FNode span [] <$> getFixities
-    (Node _ span children) -> do
-        fixities <- getFixities
-        childrenFixities <- mapM (hieAstToFixtyTree hscEnv tcGblEnv) children
-        pure $ FNode span childrenFixities fixities
-    where
-        -- Names at the current ast node
-        names :: [Name]
-        names = mapMaybe eitherToMaybe $ M.keys $ getNodeIds ast
-
-        getFixities :: MonadIO m => m [(T.Text, Fixity)]
-        getFixities = liftIO
-            $ fmap (filter ((/= defaultFixity) . snd) . mapMaybe pickFixity)
-            $ forM names $ \name ->
-                (,) (printOutputable name)
-                . snd
-                <$> Util.handleGhcException
-                    (const $ pure (emptyMessages, Nothing))
-                    (initTcWithGbl hscEnv tcGblEnv (realSrcLocSpan $ mkRealSrcLoc "<dummy>" 1 1) (lookupFixityRn name))
-
-        pickFixity :: (T.Text, Maybe Fixity) -> Maybe (T.Text, Fixity)
-        pickFixity (_, Nothing)   = Nothing
-        pickFixity (name, Just f) = Just (name, f)
+lookupFixities :: MonadIO m => HscEnv -> TcGblEnv -> S.Set Name -> m (M.Map Name Fixity)
+lookupFixities hscEnv tcGblEnv names
+    = liftIO
+    $ fmap (fromMaybe M.empty . snd)
+    $ initTcWithGbl hscEnv tcGblEnv (realSrcLocSpan $ mkRealSrcLoc "<dummy>" 1 1) 
+    $ M.traverseMaybeWithKey (\_ v -> v)
+    $ M.fromSet lookupFixity names
+  where
+    lookupFixity name = do
+      f <- Util.handleGhcException
+        (const $ pure Nothing)
+        (Just <$> lookupFixityRn name)
+      if f == Just defaultFixity
+      then pure Nothing
+      else pure f
 
 fixityRule :: Recorder (WithPriority Log) -> Rules ()
 fixityRule recorder = do
     define (cmapWithPrio LogShake recorder) $ \GetFixity nfp -> do
-        HAR{hieAst} <- use_ GetHieAst nfp
-        env <- hscEnv <$> use_ GhcSession nfp
+        HAR{refMap} <- use_ GetHieAst nfp
+        env <- hscEnv <$> use_ GhcSessionDeps nfp -- deps necessary so that we can consult already loaded in ifaces instead of loading in duplicates
         tcGblEnv <- tmrTypechecked <$> use_ TypeCheck nfp
-        trees <- hieAstsToFixitTrees env tcGblEnv hieAst
-        pure ([], Just trees)
+        fs <- lookupFixities env tcGblEnv (S.mapMonotonic (\(Right n) -> n) $ S.filter isRight $ M.keysSet refMap)
+        pure ([], Just (FixityMap fs))
 
     -- Ensure that this plugin doesn't block on startup
-    addPersistentRule GetFixity $ \_ -> pure $ Just (fakeFixityTrees, idDelta, Nothing)
+    addPersistentRule GetFixity $ \_ -> pure $ Just (FixityMap M.empty, idDelta, Nothing)