4 files changed, 38 insertions, 3 deletions

diff --git a/code/src/extension.ts b/code/src/extension.ts
index 554ac8af3..81153b7e6 100644
--- a/code/src/extension.ts
+++ b/code/src/extension.ts
@@ -134,6 +134,14 @@ export function activate(context: vscode.ExtensionContext) {
             textDocumentContentProvider.eventEmitter.fire(uris.syntaxTree)
         })
     }, null, context.subscriptions)
+    vscode.window.onDidChangeActiveTextEditor(async (editor) => {
+        if (!editor || editor.document.languageId != 'rust') return
+        let params: lc.TextDocumentIdentifier = {
+            uri: editor.document.uri.toString()
+        }
+        let decorations = await client.sendRequest<Decoration[]>("m/decorationsRequest", params)
+        setHighlights(editor, decorations)
+    })
 }
 
 // We need to order this after LS updates, but there's no API for that.
diff --git a/crates/server/src/main_loop/handlers.rs b/crates/server/src/main_loop/handlers.rs
index ee4072084..ec5421f06 100644
--- a/crates/server/src/main_loop/handlers.rs
+++ b/crates/server/src/main_loop/handlers.rs
@@ -7,7 +7,7 @@ use languageserver_types::{
     CompletionItem,
 };
 use serde_json::{to_value, from_value};
-use libanalysis::{Query, QuickFix};
+use libanalysis::{Query, QuickFix, FileId};
 use libeditor;
 use libsyntax2::{
     TextUnit,
@@ -401,18 +401,33 @@ pub fn publish_diagnostics(
     Ok(req::PublishDiagnosticsParams { uri, diagnostics })
 }
 
+pub fn handle_decorations(
+    world: ServerWorld,
+    params: TextDocumentIdentifier,
+) -> Result<Vec<Decoration>> {
+    let file_id = params.try_conv_with(&world)?;
+    highlight(&world, file_id)
+}
+
 pub fn publish_decorations(
     world: ServerWorld,
     uri: Url
 ) -> Result<req::PublishDecorationsParams> {
     let file_id = world.uri_to_file_id(&uri)?;
+    Ok(req::PublishDecorationsParams {
+        uri,
+        decorations: highlight(&world, file_id)?
+    })
+}
+
+fn highlight(world: &ServerWorld, file_id: FileId) -> Result<Vec<Decoration>> {
     let file = world.analysis().file_syntax(file_id)?;
     let line_index = world.analysis().file_line_index(file_id)?;
-    let decorations = libeditor::highlight(&file)
+    let res = libeditor::highlight(&file)
         .into_iter()
         .map(|h| Decoration {
             range: h.range.conv_with(&line_index),
             tag: h.tag,
         }).collect();
-    Ok(req::PublishDecorationsParams { uri, decorations })
+    Ok(res)
 }
diff --git a/crates/server/src/main_loop/mod.rs b/crates/server/src/main_loop/mod.rs
index 6d6ca6ae9..da1121cb4 100644
--- a/crates/server/src/main_loop/mod.rs
+++ b/crates/server/src/main_loop/mod.rs
@@ -30,6 +30,7 @@ use {
         handle_join_lines,
         handle_completion,
         handle_runnables,
+        handle_decorations,
     },
 };
 
@@ -157,6 +158,9 @@ fn on_request(
     handle_request_on_threadpool::<req::JoinLines>(
         &mut req, pool, world, sender, handle_join_lines,
     )?;
+    handle_request_on_threadpool::<req::DecorationsRequest>(
+        &mut req, pool, world, sender, handle_decorations,
+    )?;
     dispatch::handle_request::<req::ExecuteCommand, _>(&mut req, |params, resp| {
         io.send(RawMsg::Response(resp.into_response(Ok(None))?));
 
diff --git a/crates/server/src/req.rs b/crates/server/src/req.rs
index 269246dff..999fdb7c2 100644
--- a/crates/server/src/req.rs
+++ b/crates/server/src/req.rs
@@ -84,6 +84,14 @@ pub struct FindMatchingBraceParams {
     pub offsets: Vec<Position>,
 }
 
+pub enum DecorationsRequest {}
+
+impl Request for DecorationsRequest {
+    type Params = TextDocumentIdentifier;
+    type Result = Vec<Decoration>;
+    const METHOD: &'static str = "m/decorationsRequest";
+}
+
 pub enum PublishDecorations {}
 
 impl Notification for PublishDecorations {