From 299b0d3ce6f193caf61693efab47409f7c2a4c7a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jan=20Hr=C4=8Dek?= <honza.hrk@gmail.com>
Date: Wed, 21 Feb 2024 05:49:15 +0100
Subject: [PATCH] Fix it

---
 .../src/Ide/Plugin/Rename.hs                  | 20 +++++++++++++------
 1 file changed, 14 insertions(+), 6 deletions(-)

diff --git a/plugins/hls-rename-plugin/src/Ide/Plugin/Rename.hs b/plugins/hls-rename-plugin/src/Ide/Plugin/Rename.hs
index c25da1bd46..359cd676ef 100644
--- a/plugins/hls-rename-plugin/src/Ide/Plugin/Rename.hs
+++ b/plugins/hls-rename-plugin/src/Ide/Plugin/Rename.hs
@@ -25,7 +25,6 @@ import           Data.List.NonEmpty                    (NonEmpty ((:|)),
 import qualified Data.Map                              as M
 import           Data.Maybe
 import           Data.Mod.Word
-import qualified Data.Set                              as S
 import qualified Data.Text                             as T
 import           Development.IDE                       (Recorder, WithPriority,
                                                         usePropertyAction)
@@ -34,6 +33,7 @@ import           Development.IDE.Core.PositionMapping
 import           Development.IDE.Core.RuleTypes
 import           Development.IDE.Core.Service
 import           Development.IDE.Core.Shake
+import qualified Development.IDE.GHC.Compat            as Compat
 import           Development.IDE.GHC.Compat.Core
 import           Development.IDE.GHC.Compat.ExactPrint
 import           Development.IDE.GHC.Compat.Parser
@@ -212,12 +212,20 @@ handleGetHieAst state nfp =
 -- | We don't want to rename in code generated by GHC as this gives false positives.
 -- So we restrict the HIE file to remove all the generated code.
 removeGenerated :: HieAstResult -> HieAstResult
-removeGenerated HAR{..} = HAR{hieAst = go hieAst,..}
+removeGenerated HAR{..} = HAR{hieAst = sourceOnlyAsts, refMap = sourceOnlyRefMap, ..}
   where
-    go :: HieASTs a -> HieASTs a
-    go hf =
-      HieASTs (fmap goAst (getAsts hf))
-    goAst (Node nsi sp xs) = Node (SourcedNodeInfo $ M.restrictKeys (getSourcedNodeInfo nsi) (S.singleton SourceInfo)) sp (map goAst xs)
+    goAsts :: HieASTs a -> HieASTs a
+    goAsts (HieASTs asts) = HieASTs (fmap goAst asts)
+
+    goAst :: HieAST a -> HieAST a
+    goAst (Node (SourcedNodeInfo sniMap) sp children) =
+        let sourceOnlyNodeInfos = SourcedNodeInfo $ M.delete GeneratedInfo sniMap
+        in Node sourceOnlyNodeInfos sp $ map goAst children
+
+    sourceOnlyAsts = goAsts hieAst
+    -- Also need to regenerate the RefMap, because the one in HAR
+    -- is generated from HieASTs containing GeneratedInfo
+    sourceOnlyRefMap = Compat.generateReferencesMap $ Compat.getAsts sourceOnlyAsts
 
 collectWith :: (Hashable a, Eq b) => (a -> b) -> HashSet a -> [(b, HashSet a)]
 collectWith f = map (\(a :| as) -> (f a, HS.fromList (a:as))) . groupWith f . HS.toList