1 files changed, 17 insertions, 5 deletions

diff --git a/crates/ra_lsp_server/src/main_loop/handlers.rs b/crates/ra_lsp_server/src/main_loop/handlers.rs
index 10e271376..307082865 100644
--- a/crates/ra_lsp_server/src/main_loop/handlers.rs
+++ b/crates/ra_lsp_server/src/main_loop/handlers.rs
@@ -18,7 +18,7 @@ use serde_json::to_value;
 
 use crate::{
     cargo_target_spec::{runnable_args, CargoTargetSpec},
-    conv::{to_location, Conv, ConvWith, MapConvWith, TryConvWith, TryConvWithToVec},
+    conv::{to_location, Conv, ConvWith, FoldConvCtx, MapConvWith, TryConvWith, TryConvWithToVec},
     req::{self, Decoration, InlayHint, InlayHintsParams, InlayKind},
     world::WorldSnapshot,
     LspError, Result,
@@ -383,8 +383,14 @@ pub fn handle_folding_range(
 ) -> Result<Option<Vec<FoldingRange>>> {
     let file_id = params.text_document.try_conv_with(&world)?;
     let folds = world.analysis().folding_ranges(file_id)?;
+    let text = world.analysis().file_text(file_id)?;
     let line_index = world.analysis().file_line_index(file_id)?;
-    let res = Some(folds.into_iter().map_conv_with(&*line_index).collect());
+    let ctx = FoldConvCtx {
+        text: &text,
+        line_index: &line_index,
+        line_folding_only: world.options.line_folding_only,
+    };
+    let res = Some(folds.into_iter().map_conv_with(&ctx).collect());
     Ok(res)
 }
 
@@ -475,7 +481,6 @@ pub fn handle_references(
     params: req::ReferenceParams,
 ) -> Result<Option<Vec<Location>>> {
     let position = params.text_document_position.try_conv_with(&world)?;
-    let line_index = world.analysis().file_line_index(position.file_id)?;
 
     let refs = match world.analysis().find_all_refs(position)? {
         None => return Ok(None),
@@ -484,13 +489,19 @@ pub fn handle_references(
 
     let locations = if params.context.include_declaration {
         refs.into_iter()
-            .filter_map(|r| to_location(r.file_id, r.range, &world, &line_index).ok())
+            .filter_map(|r| {
+                let line_index = world.analysis().file_line_index(r.file_id).ok()?;
+                to_location(r.file_id, r.range, &world, &line_index).ok()
+            })
             .collect()
     } else {
         // Only iterate over the references if include_declaration was false
         refs.references()
             .iter()
-            .filter_map(|r| to_location(r.file_id, r.range, &world, &line_index).ok())
+            .filter_map(|r| {
+                let line_index = world.analysis().file_line_index(r.file_id).ok()?;
+                to_location(r.file_id, r.range, &world, &line_index).ok()
+            })
             .collect()
     };
 
@@ -740,6 +751,7 @@ pub fn handle_document_highlight(
 
     Ok(Some(
         refs.into_iter()
+            .filter(|r| r.file_id == file_id)
             .map(|r| DocumentHighlight { range: r.range.conv_with(&line_index), kind: None })
             .collect(),
     ))