fix(lsp): refactor, fix issues and add benchmark for code lens (#10841)

12 files changed, 780 insertions, 331 deletions

diff --git a/cli/bench/fixtures/code_lens.ts b/cli/bench/fixtures/code_lens.ts
new file mode 100644
index 000000000..0822e31c4
--- /dev/null
+++ b/cli/bench/fixtures/code_lens.ts
@@ -0,0 +1,56 @@
+interface A {
+  a: string;
+}
+
+interface B {
+  b: string;
+}
+
+interface C {
+  c: string;
+}
+
+interface D {
+  d: string;
+}
+
+interface E {
+  e: string;
+}
+
+interface F {
+  f: string;
+}
+
+interface G {
+  g: string;
+}
+
+interface H {
+  h: string;
+}
+
+class AB implements A, B {
+  a = "a";
+  b = "b";
+}
+
+class CD implements C, D {
+  c = "c";
+  d = "d";
+}
+
+class EF implements E, F {
+  e = "e";
+  f = "f";
+}
+
+class GH implements G, H {
+  g = "g";
+  h = "h";
+}
+
+new AB().a;
+new CD().c;
+new EF().e;
+new GH().g;
diff --git a/cli/bench/lsp.rs b/cli/bench/lsp.rs
index a227e5a99..3ec914d9e 100644
--- a/cli/bench/lsp.rs
+++ b/cli/bench/lsp.rs
@@ -13,6 +13,7 @@ use std::time::Duration;
 use test_util::lsp::LspClient;
 use test_util::lsp::LspResponseError;
 
+static FIXTURE_CODE_LENS_TS: &str = include_str!("fixtures/code_lens.ts");
 static FIXTURE_DB_TS: &str = include_str!("fixtures/db.ts");
 static FIXTURE_DB_MESSAGES: &[u8] = include_bytes!("fixtures/db_messages.json");
 static FIXTURE_INIT_JSON: &[u8] =
@@ -123,6 +124,70 @@ fn bench_big_file_edits(deno_exe: &Path) -> Result<Duration, AnyError> {
   Ok(client.duration())
 }
 
+fn bench_code_lens(deno_exe: &Path) -> Result<Duration, AnyError> {
+  let mut client = LspClient::new(deno_exe)?;
+
+  let params: Value = serde_json::from_slice(FIXTURE_INIT_JSON)?;
+  let (_, maybe_err) =
+    client.write_request::<_, _, Value>("initialize", params)?;
+  assert!(maybe_err.is_none());
+  client.write_notification("initialized", json!({}))?;
+
+  client.write_notification(
+    "textDocument/didOpen",
+    json!({
+      "textDocument": {
+        "uri": "file:///fixtures/code_lens.ts",
+        "languageId": "typescript",
+        "version": 1,
+        "text": FIXTURE_CODE_LENS_TS
+      }
+    }),
+  )?;
+
+  let (id, method, _): (u64, String, Option<Value>) = client.read_request()?;
+  assert_eq!(method, "workspace/configuration");
+
+  client.write_response(
+    id,
+    json!({
+      "enable": true
+    }),
+  )?;
+
+  let (method, _): (String, Option<Value>) = client.read_notification()?;
+  assert_eq!(method, "textDocument/publishDiagnostics");
+  let (method, _): (String, Option<Value>) = client.read_notification()?;
+  assert_eq!(method, "textDocument/publishDiagnostics");
+  let (method, _): (String, Option<Value>) = client.read_notification()?;
+  assert_eq!(method, "textDocument/publishDiagnostics");
+
+  let (maybe_res, maybe_err) = client
+    .write_request::<_, _, Vec<lsp::CodeLens>>(
+      "textDocument/codeLens",
+      json!({
+        "textDocument": {
+          "uri": "file:///fixtures/code_lens.ts"
+        }
+      }),
+    )
+    .unwrap();
+  assert!(maybe_err.is_none());
+  assert!(maybe_res.is_some());
+  let res = maybe_res.unwrap();
+  assert!(!res.is_empty());
+
+  for code_lens in res {
+    let (maybe_res, maybe_err) = client
+      .write_request::<_, _, lsp::CodeLens>("codeLens/resolve", code_lens)
+      .unwrap();
+    assert!(maybe_err.is_none());
+    assert!(maybe_res.is_some());
+  }
+
+  Ok(client.duration())
+}
+
 fn bench_find_replace(deno_exe: &Path) -> Result<Duration, AnyError> {
   let mut client = LspClient::new(deno_exe)?;
 
@@ -304,6 +369,16 @@ pub(crate) fn benchmarks(
   println!("      ({} runs, mean: {}ms)", times.len(), mean);
   exec_times.insert("find_replace".to_string(), mean);
 
+  println!("   - Code Lens");
+  let mut times = Vec::new();
+  for _ in 0..10 {
+    times.push(bench_code_lens(deno_exe)?);
+  }
+  let mean =
+    (times.iter().sum::<Duration>() / times.len() as u32).as_millis() as u64;
+  println!("      ({} runs, mean: {}ms)", times.len(), mean);
+  exec_times.insert("code_lens".to_string(), mean);
+
   println!("<- End benchmarking lsp");
 
   Ok(exec_times)
diff --git a/cli/lsp/analysis.rs b/cli/lsp/analysis.rs
index bd3ce799a..2235a2c83 100644
--- a/cli/lsp/analysis.rs
+++ b/cli/lsp/analysis.rs
@@ -15,7 +15,6 @@ use deno_core::error::anyhow;
 use deno_core::error::custom_error;
 use deno_core::error::AnyError;
 use deno_core::serde::Deserialize;
-use deno_core::serde::Serialize;
 use deno_core::serde_json;
 use deno_core::serde_json::json;
 use deno_core::ModuleResolutionError;
@@ -394,21 +393,6 @@ pub fn analyze_dependencies(
   (dependencies, maybe_type)
 }
 
-#[derive(Debug, Deserialize, Serialize)]
-pub enum CodeLensSource {
-  #[serde(rename = "implementations")]
-  Implementations,
-  #[serde(rename = "references")]
-  References,
-}
-
-#[derive(Debug, Deserialize)]
-#[serde(rename_all = "camelCase")]
-pub struct CodeLensData {
-  pub source: CodeLensSource,
-  pub specifier: ModuleSpecifier,
-}
-
 fn code_as_string(code: &Option<lsp::NumberOrString>) -> String {
   match code {
     Some(lsp::NumberOrString::String(str)) => str.clone(),
diff --git a/cli/lsp/code_lens.rs b/cli/lsp/code_lens.rs
new file mode 100644
index 000000000..9be5bed2c
--- /dev/null
+++ b/cli/lsp/code_lens.rs
@@ -0,0 +1,284 @@
+// Copyright 2018-2021 the Deno authors. All rights reserved. MIT license.
+
+use super::language_server;
+use super::tsc;
+
+use deno_core::error::anyhow;
+use deno_core::error::AnyError;
+use deno_core::resolve_url;
+use deno_core::serde::Deserialize;
+use deno_core::serde::Serialize;
+use deno_core::serde_json;
+use deno_core::serde_json::json;
+use deno_core::ModuleSpecifier;
+use lspower::lsp;
+use regex::Regex;
+use std::cell::RefCell;
+use std::rc::Rc;
+
+lazy_static::lazy_static! {
+  static ref ABSTRACT_MODIFIER: Regex = Regex::new(r"\babstract\b").unwrap();
+  static ref EXPORT_MODIFIER: Regex = Regex::new(r"\bexport\b").unwrap();
+}
+
+#[derive(Debug, Deserialize, Serialize)]
+pub enum CodeLensSource {
+  #[serde(rename = "implementations")]
+  Implementations,
+  #[serde(rename = "references")]
+  References,
+}
+
+#[derive(Debug, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub struct CodeLensData {
+  pub source: CodeLensSource,
+  pub specifier: ModuleSpecifier,
+}
+
+async fn resolve_implementation_code_lens(
+  code_lens: lsp::CodeLens,
+  data: CodeLensData,
+  language_server: &mut language_server::Inner,
+) -> Result<lsp::CodeLens, AnyError> {
+  let line_index = language_server
+    .get_line_index_sync(&data.specifier)
+    .unwrap();
+  let req = tsc::RequestMethod::GetImplementation((
+    data.specifier.clone(),
+    line_index.offset_tsc(code_lens.range.start)?,
+  ));
+  let snapshot = language_server.snapshot()?;
+  let maybe_implementations: Option<Vec<tsc::ImplementationLocation>> =
+    language_server.ts_server.request(snapshot, req).await?;
+  if let Some(implementations) = maybe_implementations {
+    let mut locations = Vec::new();
+    for implementation in implementations {
+      let implementation_specifier =
+        resolve_url(&implementation.document_span.file_name)?;
+      let implementation_location =
+        implementation.to_location(&line_index, language_server);
+      if !(implementation_specifier == data.specifier
+        && implementation_location.range.start == code_lens.range.start)
+      {
+        locations.push(implementation_location);
+      }
+    }
+    let command = if !locations.is_empty() {
+      let title = if locations.len() > 1 {
+        format!("{} implementations", locations.len())
+      } else {
+        "1 implementation".to_string()
+      };
+      lsp::Command {
+        title,
+        command: "deno.showReferences".to_string(),
+        arguments: Some(vec![
+          json!(data.specifier),
+          json!(code_lens.range.start),
+          json!(locations),
+        ]),
+      }
+    } else {
+      lsp::Command {
+        title: "0 implementations".to_string(),
+        command: "".to_string(),
+        arguments: None,
+      }
+    };
+    Ok(lsp::CodeLens {
+      range: code_lens.range,
+      command: Some(command),
+      data: None,
+    })
+  } else {
+    let command = Some(lsp::Command {
+      title: "0 implementations".to_string(),
+      command: "".to_string(),
+      arguments: None,
+    });
+    Ok(lsp::CodeLens {
+      range: code_lens.range,
+      command,
+      data: None,
+    })
+  }
+}
+
+async fn resolve_references_code_lens(
+  code_lens: lsp::CodeLens,
+  data: CodeLensData,
+  language_server: &mut language_server::Inner,
+) -> Result<lsp::CodeLens, AnyError> {
+  let line_index = language_server
+    .get_line_index_sync(&data.specifier)
+    .unwrap();
+  let req = tsc::RequestMethod::GetReferences((
+    data.specifier.clone(),
+    line_index.offset_tsc(code_lens.range.start)?,
+  ));
+  let snapshot = language_server.snapshot()?;
+  let maybe_references: Option<Vec<tsc::ReferenceEntry>> =
+    language_server.ts_server.request(snapshot, req).await?;
+  if let Some(references) = maybe_references {
+    let mut locations = Vec::new();
+    for reference in references {
+      if reference.is_definition {
+        continue;
+      }
+      let reference_specifier =
+        resolve_url(&reference.document_span.file_name)?;
+      let line_index =
+        language_server.get_line_index(reference_specifier).await?;
+      locations.push(reference.to_location(&line_index, language_server));
+    }
+    let command = if !locations.is_empty() {
+      let title = if locations.len() > 1 {
+        format!("{} references", locations.len())
+      } else {
+        "1 reference".to_string()
+      };
+      lsp::Command {
+        title,
+        command: "deno.showReferences".to_string(),
+        arguments: Some(vec![
+          json!(data.specifier),
+          json!(code_lens.range.start),
+          json!(locations),
+        ]),
+      }
+    } else {
+      lsp::Command {
+        title: "0 references".to_string(),
+        command: "".to_string(),
+        arguments: None,
+      }
+    };
+    Ok(lsp::CodeLens {
+      range: code_lens.range,
+      command: Some(command),
+      data: None,
+    })
+  } else {
+    let command = lsp::Command {
+      title: "0 references".to_string(),
+      command: "".to_string(),
+      arguments: None,
+    };
+    Ok(lsp::CodeLens {
+      range: code_lens.range,
+      command: Some(command),
+      data: None,
+    })
+  }
+}
+
+pub(crate) async fn resolve_code_lens(
+  code_lens: lsp::CodeLens,
+  language_server: &mut language_server::Inner,
+) -> Result<lsp::CodeLens, AnyError> {
+  let data: CodeLensData =
+    serde_json::from_value(code_lens.data.clone().unwrap())?;
+  match data.source {
+    CodeLensSource::Implementations => {
+      resolve_implementation_code_lens(code_lens, data, language_server).await
+    }
+    CodeLensSource::References => {
+      resolve_references_code_lens(code_lens, data, language_server).await
+    }
+  }
+}
+
+/// Return tsc navigation tree code lenses.
+pub(crate) async fn tsc_code_lenses(
+  specifier: &ModuleSpecifier,
+  language_server: &mut language_server::Inner,
+) -> Result<Vec<lsp::CodeLens>, AnyError> {
+  let workspace_settings = language_server.config.get_workspace_settings();
+  let line_index = language_server
+    .get_line_index_sync(&specifier)
+    .ok_or_else(|| anyhow!("Missing line index."))?;
+  let navigation_tree = language_server.get_navigation_tree(specifier).await?;
+  let code_lenses = Rc::new(RefCell::new(Vec::new()));
+  navigation_tree.walk(&|i, mp| {
+    let mut code_lenses = code_lenses.borrow_mut();
+
+    // TSC Implementations Code Lens
+    if workspace_settings.code_lens.implementations {
+      let source = CodeLensSource::Implementations;
+      match i.kind {
+        tsc::ScriptElementKind::InterfaceElement => {
+          code_lenses.push(i.to_code_lens(&line_index, specifier, &source));
+        }
+        tsc::ScriptElementKind::ClassElement
+        | tsc::ScriptElementKind::MemberFunctionElement
+        | tsc::ScriptElementKind::MemberVariableElement
+        | tsc::ScriptElementKind::MemberGetAccessorElement
+        | tsc::ScriptElementKind::MemberSetAccessorElement => {
+          if ABSTRACT_MODIFIER.is_match(&i.kind_modifiers) {
+            code_lenses.push(i.to_code_lens(&line_index, &specifier, &source));
+          }
+        }
+        _ => (),
+      }
+    }
+
+    // TSC References Code Lens
+    if workspace_settings.code_lens.references {
+      let source = CodeLensSource::References;
+      if let Some(parent) = &mp {
+        if parent.kind == tsc::ScriptElementKind::EnumElement {
+          code_lenses.push(i.to_code_lens(&line_index, &specifier, &source));
+        }
+      }
+      match i.kind {
+        tsc::ScriptElementKind::FunctionElement => {
+          if workspace_settings.code_lens.references_all_functions {
+            code_lenses.push(i.to_code_lens(&line_index, &specifier, &source));
+          }
+        }
+        tsc::ScriptElementKind::ConstElement
+        | tsc::ScriptElementKind::LetElement
+        | tsc::ScriptElementKind::VariableElement => {
+          if EXPORT_MODIFIER.is_match(&i.kind_modifiers) {
+            code_lenses.push(i.to_code_lens(&line_index, &specifier, &source));
+          }
+        }
+        tsc::ScriptElementKind::ClassElement => {
+          if i.text != "<class>" {
+            code_lenses.push(i.to_code_lens(&line_index, &specifier, &source));
+          }
+        }
+        tsc::ScriptElementKind::InterfaceElement
+        | tsc::ScriptElementKind::TypeElement
+        | tsc::ScriptElementKind::EnumElement => {
+          code_lenses.push(i.to_code_lens(&line_index, &specifier, &source));
+        }
+        tsc::ScriptElementKind::LocalFunctionElement
+        | tsc::ScriptElementKind::MemberGetAccessorElement
+        | tsc::ScriptElementKind::MemberSetAccessorElement
+        | tsc::ScriptElementKind::ConstructorImplementationElement
+        | tsc::ScriptElementKind::MemberVariableElement => {
+          if let Some(parent) = &mp {
+            if parent.spans[0].start != i.spans[0].start {
+              match parent.kind {
+                tsc::ScriptElementKind::ClassElement
+                | tsc::ScriptElementKind::InterfaceElement
+                | tsc::ScriptElementKind::TypeElement => {
+                  code_lenses.push(i.to_code_lens(
+                    &line_index,
+                    &specifier,
+                    &source,
+                  ));
+                }
+                _ => (),
+              }
+            }
+          }
+        }
+        _ => (),
+      }
+    }
+  });
+  Ok(Rc::try_unwrap(code_lenses).unwrap().into_inner())
+}
diff --git a/cli/lsp/documents.rs b/cli/lsp/documents.rs
index 5fbfcdf52..5e08245cb 100644
--- a/cli/lsp/documents.rs
+++ b/cli/lsp/documents.rs
@@ -2,6 +2,7 @@
 
 use super::analysis;
 use super::text::LineIndex;
+use super::tsc;
 
 use crate::media_type::MediaType;
 
@@ -66,6 +67,7 @@ pub struct DocumentData {
   bytes: Option<Vec<u8>>,
   language_id: LanguageId,
   line_index: Option<LineIndex>,
+  maybe_navigation_tree: Option<tsc::NavigationTree>,
   specifier: ModuleSpecifier,
   dependencies: Option<HashMap<String, analysis::Dependency>>,
   version: Option<i32>,
@@ -82,6 +84,7 @@ impl DocumentData {
       bytes: Some(source.as_bytes().to_owned()),
       language_id,
       line_index: Some(LineIndex::new(source)),
+      maybe_navigation_tree: None,
       specifier,
       dependencies: None,
       version: Some(version),
@@ -122,6 +125,7 @@ impl DocumentData {
     } else {
       Some(LineIndex::new(&content))
     };
+    self.maybe_navigation_tree = None;
     Ok(())
   }
 
@@ -236,6 +240,14 @@ impl DocumentCache {
     doc.dependencies.clone()
   }
 
+  pub fn get_navigation_tree(
+    &self,
+    specifier: &ModuleSpecifier,
+  ) -> Option<tsc::NavigationTree> {
+    let doc = self.docs.get(specifier)?;
+    doc.maybe_navigation_tree.clone()
+  }
+
   /// Determines if the specifier should be processed for diagnostics and other
   /// related language server features.
   pub fn is_diagnosable(&self, specifier: &ModuleSpecifier) -> bool {
@@ -340,6 +352,25 @@ impl DocumentCache {
     }
   }
 
+  pub fn set_navigation_tree(
+    &mut self,
+    specifier: &ModuleSpecifier,
+    navigation_tree: tsc::NavigationTree,
+  ) -> Result<(), AnyError> {
+    if let Some(doc) = self.docs.get_mut(specifier) {
+      doc.maybe_navigation_tree = Some(navigation_tree);
+      Ok(())
+    } else {
+      Err(custom_error(
+        "NotFound",
+        format!(
+          "The specifier (\"{}\") does not exist in the document cache.",
+          specifier
+        ),
+      ))
+    }
+  }
+
   pub fn specifiers(&self) -> Vec<ModuleSpecifier> {
     self.docs.keys().cloned().collect()
   }
diff --git a/cli/lsp/language_server.rs b/cli/lsp/language_server.rs
index cd8298573..10bb0478e 100644
--- a/cli/lsp/language_server.rs
+++ b/cli/lsp/language_server.rs
@@ -15,13 +15,9 @@ use lspower::jsonrpc::Result as LspResult;
 use lspower::lsp::request::*;
 use lspower::lsp::*;
 use lspower::Client;
-use regex::Regex;
 use serde_json::from_value;
-use std::cell::RefCell;
-use std::collections::HashMap;
 use std::env;
 use std::path::PathBuf;
-use std::rc::Rc;
 use std::sync::atomic::Ordering;
 use std::sync::Arc;
 use tokio::fs;
@@ -31,10 +27,9 @@ use super::analysis::fix_ts_import_changes;
 use super::analysis::ts_changes_to_edit;
 use super::analysis::CodeActionCollection;
 use super::analysis::CodeActionData;
-use super::analysis::CodeLensData;
-use super::analysis::CodeLensSource;
 use super::analysis::ResolvedDependency;
 use super::capabilities;
+use super::code_lens;
 use super::completions;
 use super::config::Config;
 use super::config::ConfigSnapshot;
@@ -67,11 +62,6 @@ use crate::tools::fmt::get_typescript_config;
 pub const REGISTRIES_PATH: &str = "registries";
 const SOURCES_PATH: &str = "deps";
 
-lazy_static::lazy_static! {
-  static ref ABSTRACT_MODIFIER: Regex = Regex::new(r"\babstract\b").unwrap();
-  static ref EXPORT_MODIFIER: Regex = Regex::new(r"\bexport\b").unwrap();
-}
-
 #[derive(Debug, Clone)]
 pub struct LanguageServer(Arc<tokio::sync::Mutex<Inner>>);
 
@@ -94,7 +84,7 @@ pub(crate) struct Inner {
   /// The LSP client that this LSP server is connected to.
   pub(crate) client: Client,
   /// Configuration information.
-  config: Config,
+  pub(crate) config: Config,
   diagnostics_server: diagnostics::DiagnosticsServer,
   /// The "in-memory" documents in the editor which can be updated and changed.
   documents: DocumentCache,
@@ -109,8 +99,6 @@ pub(crate) struct Inner {
   pub(crate) maybe_import_map: Option<ImportMap>,
   /// The URL for the import map which is used to determine relative imports.
   maybe_import_map_uri: Option<Url>,
-  /// A map of all the cached navigation trees.
-  navigation_trees: HashMap<ModuleSpecifier, tsc::NavigationTree>,
   /// A collection of measurements which instrument that performance of the LSP.
   performance: Performance,
   /// Cached sources that are read-only.
@@ -155,7 +143,6 @@ impl Inner {
       maybe_import_map_uri: Default::default(),
       module_registries,
       module_registries_location,
-      navigation_trees: Default::default(),
       performance,
       sources,
       ts_fixable_diagnostics: Default::default(),
@@ -224,7 +211,7 @@ impl Inner {
 
   /// Only searches already cached assets and documents for a line index.  If
   /// the line index cannot be found, `None` is returned.
-  fn get_line_index_sync(
+  pub fn get_line_index_sync(
     &self,
     specifier: &ModuleSpecifier,
   ) -> Option<LineIndex> {
@@ -269,7 +256,7 @@ impl Inner {
     }
   }
 
-  async fn get_navigation_tree(
+  pub(crate) async fn get_navigation_tree(
     &mut self,
     specifier: &ModuleSpecifier,
   ) -> Result<tsc::NavigationTree, AnyError> {
@@ -277,9 +264,19 @@ impl Inner {
       "get_navigation_tree",
       Some(json!({ "specifier": specifier })),
     );
-    if let Some(navigation_tree) = self.navigation_trees.get(specifier) {
-      self.performance.measure(mark);
-      Ok(navigation_tree.clone())
+    let maybe_navigation_tree = if specifier.scheme() == "asset" {
+      self
+        .assets
+        .get(specifier)
+        .map(|o| o.clone().map(|a| a.maybe_navigation_tree).flatten())
+        .flatten()
+    } else if self.documents.contains_key(specifier) {
+      self.documents.get_navigation_tree(specifier)
+    } else {
+      self.sources.get_navigation_tree(specifier)
+    };
+    let navigation_tree = if let Some(navigation_tree) = maybe_navigation_tree {
+      navigation_tree
     } else {
       let navigation_tree: tsc::NavigationTree = self
         .ts_server
@@ -288,12 +285,23 @@ impl Inner {
           tsc::RequestMethod::GetNavigationTree(specifier.clone()),
         )
         .await?;
-      self
-        .navigation_trees
-        .insert(specifier.clone(), navigation_tree.clone());
-      self.performance.measure(mark);
-      Ok(navigation_tree)
-    }
+      if specifier.scheme() == "asset" {
+        self
+          .assets
+          .set_navigation_tree(specifier, navigation_tree.clone())?;
+      } else if self.documents.contains_key(specifier) {
+        self
+          .documents
+          .set_navigation_tree(specifier, navigation_tree.clone())?;
+      } else {
+        self
+          .sources
+          .set_navigation_tree(specifier, navigation_tree.clone())?;
+      }
+      navigation_tree
+    };
+    self.performance.measure(mark);
+    Ok(navigation_tree)
   }
 
   pub(crate) fn snapshot(&self) -> LspResult<StateSnapshot> {
@@ -674,7 +682,6 @@ impl Inner {
     }
     let specifier = self.url_map.normalize_url(&params.text_document.uri);
     self.documents.close(&specifier);
-    self.navigation_trees.remove(&specifier);
 
     if self.documents.is_diagnosable(&specifier) {
       if let Err(err) = self.diagnostics_server.update() {
@@ -1090,302 +1097,36 @@ impl Inner {
     }
 
     let mark = self.performance.mark("code_lens", Some(&params));
-    let line_index = self.get_line_index_sync(&specifier).unwrap();
-    let navigation_tree =
-      self.get_navigation_tree(&specifier).await.map_err(|err| {
-        error!("Failed to retrieve nav tree: {}", err);
-        LspError::invalid_request()
+    let code_lenses = code_lens::tsc_code_lenses(&specifier, self)
+      .await
+      .map_err(|err| {
+        error!("Error getting code lenses for \"{}\": {}", specifier, err);
+        LspError::internal_error()
       })?;
-
-    // because we have to use this as a mutable in a closure, the compiler
-    // can't be sure when the vector will be mutated, and so a RefCell is
-    // required to "protect" the vector.
-    let cl = Rc::new(RefCell::new(Vec::new()));
-    navigation_tree.walk(&|i, mp| {
-      let mut code_lenses = cl.borrow_mut();
-      let workspace_settings = self.config.get_workspace_settings();
-
-      // TSC Implementations Code Lens
-      if workspace_settings.code_lens.implementations {
-        let source = CodeLensSource::Implementations;
-        match i.kind {
-          tsc::ScriptElementKind::InterfaceElement => {
-            code_lenses.push(i.to_code_lens(&line_index, &specifier, &source));
-          }
-          tsc::ScriptElementKind::ClassElement
-          | tsc::ScriptElementKind::MemberFunctionElement
-          | tsc::ScriptElementKind::MemberVariableElement
-          | tsc::ScriptElementKind::MemberGetAccessorElement
-          | tsc::ScriptElementKind::MemberSetAccessorElement => {
-            if ABSTRACT_MODIFIER.is_match(&i.kind_modifiers) {
-              code_lenses.push(i.to_code_lens(
-                &line_index,
-                &specifier,
-                &source,
-              ));
-            }
-          }
-          _ => (),
-        }
-      }
-
-      // TSC References Code Lens
-      if workspace_settings.code_lens.references {
-        let source = CodeLensSource::References;
-        if let Some(parent) = &mp {
-          if parent.kind == tsc::ScriptElementKind::EnumElement {
-            code_lenses.push(i.to_code_lens(&line_index, &specifier, &source));
-          }
-        }
-        match i.kind {
-          tsc::ScriptElementKind::FunctionElement => {
-            if workspace_settings.code_lens.references_all_functions {
-              code_lenses.push(i.to_code_lens(
-                &line_index,
-                &specifier,
-                &source,
-              ));
-            }
-          }
-          tsc::ScriptElementKind::ConstElement
-          | tsc::ScriptElementKind::LetElement
-          | tsc::ScriptElementKind::VariableElement => {
-            if EXPORT_MODIFIER.is_match(&i.kind_modifiers) {
-              code_lenses.push(i.to_code_lens(
-                &line_index,
-                &specifier,
-                &source,
-              ));
-            }
-          }
-          tsc::ScriptElementKind::ClassElement => {
-            if i.text != "<class>" {
-              code_lenses.push(i.to_code_lens(
-                &line_index,
-                &specifier,
-                &source,
-              ));
-            }
-          }
-          tsc::ScriptElementKind::InterfaceElement
-          | tsc::ScriptElementKind::TypeElement
-          | tsc::ScriptElementKind::EnumElement => {
-            code_lenses.push(i.to_code_lens(&line_index, &specifier, &source));
-          }
-          tsc::ScriptElementKind::LocalFunctionElement
-          | tsc::ScriptElementKind::MemberGetAccessorElement
-          | tsc::ScriptElementKind::MemberSetAccessorElement
-          | tsc::ScriptElementKind::ConstructorImplementationElement
-          | tsc::ScriptElementKind::MemberVariableElement => {
-            if let Some(parent) = &mp {
-              if parent.spans[0].start != i.spans[0].start {
-                match parent.kind {
-                  tsc::ScriptElementKind::ClassElement
-                  | tsc::ScriptElementKind::InterfaceElement
-                  | tsc::ScriptElementKind::TypeElement => {
-                    code_lenses.push(i.to_code_lens(
-                      &line_index,
-                      &specifier,
-                      &source,
-                    ));
-                  }
-                  _ => (),
-                }
-              }
-            }
-          }
-          _ => (),
-        }
-      }
-    });
-
     self.performance.measure(mark);
-    Ok(Some(Rc::try_unwrap(cl).unwrap().into_inner()))
+
+    Ok(Some(code_lenses))
   }
 
   async fn code_lens_resolve(
     &mut self,
-    params: CodeLens,
+    code_lens: CodeLens,
   ) -> LspResult<CodeLens> {
-    let mark = self.performance.mark("code_lens_resolve", Some(&params));
-    if let Some(data) = params.data.clone() {
-      let code_lens_data: CodeLensData = serde_json::from_value(data)
-        .map_err(|err| LspError::invalid_params(err.to_string()))?;
-      let code_lens = match code_lens_data.source {
-        CodeLensSource::Implementations => {
-          let line_index =
-            self.get_line_index_sync(&code_lens_data.specifier).unwrap();
-          let req = tsc::RequestMethod::GetImplementation((
-            code_lens_data.specifier.clone(),
-            line_index.offset_tsc(params.range.start)?,
-          ));
-          let maybe_implementations: Option<Vec<tsc::ImplementationLocation>> =
-            self
-              .ts_server
-              .request(self.snapshot()?, req)
-              .await
-              .map_err(|err| {
-                error!("Error processing TypeScript request: {}", err);
-                LspError::internal_error()
-              })?;
-          if let Some(implementations) = maybe_implementations {
-            let mut locations = Vec::new();
-            for implementation in implementations {
-              let implementation_specifier = resolve_url(
-                &implementation.document_span.file_name,
-              )
-              .map_err(|err| {
-                error!("Invalid specifier returned from TypeScript: {}", err);
-                LspError::internal_error()
-              })?;
-              let implementation_location =
-                implementation.to_location(&line_index, self);
-              if !(implementation_specifier == code_lens_data.specifier
-                && implementation_location.range.start == params.range.start)
-              {
-                locations.push(implementation_location);
-              }
-            }
-            let command = if !locations.is_empty() {
-              let title = if locations.len() > 1 {
-                format!("{} implementations", locations.len())
-              } else {
-                "1 implementation".to_string()
-              };
-              let url = self
-                .url_map
-                .normalize_specifier(&code_lens_data.specifier)
-                .map_err(|err| {
-                  error!("{}", err);
-                  LspError::internal_error()
-                })?;
-              Command {
-                title,
-                command: "deno.showReferences".to_string(),
-                arguments: Some(vec![
-                  serde_json::to_value(url).unwrap(),
-                  serde_json::to_value(params.range.start).unwrap(),
-                  serde_json::to_value(locations).unwrap(),
-                ]),
-              }
-            } else {
-              Command {
-                title: "0 implementations".to_string(),
-                command: "".to_string(),
-                arguments: None,
-              }
-            };
-            CodeLens {
-              range: params.range,
-              command: Some(command),
-              data: None,
-            }
-          } else {
-            let command = Command {
-              title: "0 implementations".to_string(),
-              command: "".to_string(),
-              arguments: None,
-            };
-            CodeLens {
-              range: params.range,
-              command: Some(command),
-              data: None,
-            }
-          }
-        }
-        CodeLensSource::References => {
-          let line_index =
-            self.get_line_index_sync(&code_lens_data.specifier).unwrap();
-          let req = tsc::RequestMethod::GetReferences((
-            code_lens_data.specifier.clone(),
-            line_index.offset_tsc(params.range.start)?,
-          ));
-          let maybe_references: Option<Vec<tsc::ReferenceEntry>> = self
-            .ts_server
-            .request(self.snapshot()?, req)
-            .await
-            .map_err(|err| {
-              error!("Error processing TypeScript request: {}", err);
-              LspError::internal_error()
-            })?;
-          if let Some(references) = maybe_references {
-            let mut locations = Vec::new();
-            for reference in references {
-              if reference.is_definition {
-                continue;
-              }
-              let reference_specifier = resolve_url(
-                &reference.document_span.file_name,
-              )
-              .map_err(|err| {
-                error!("Invalid specifier returned from TypeScript: {}", err);
-                LspError::internal_error()
-              })?;
-              let line_index = self
-                .get_line_index(reference_specifier)
-                .await
-                .map_err(|err| {
-                error!("Unable to get line index: {}", err);
-                LspError::internal_error()
-              })?;
-              locations.push(reference.to_location(&line_index, self));
-            }
-            let command = if !locations.is_empty() {
-              let title = if locations.len() > 1 {
-                format!("{} references", locations.len())
-              } else {
-                "1 reference".to_string()
-              };
-              let url = self
-                .url_map
-                .normalize_specifier(&code_lens_data.specifier)
-                .map_err(|err| {
-                  error!("{}", err);
-                  LspError::internal_error()
-                })?;
-              Command {
-                title,
-                command: "deno.showReferences".to_string(),
-                arguments: Some(vec![
-                  serde_json::to_value(url).unwrap(),
-                  serde_json::to_value(params.range.start).unwrap(),
-                  serde_json::to_value(locations).unwrap(),
-                ]),
-              }
-            } else {
-              Command {
-                title: "0 references".to_string(),
-                command: "".to_string(),
-                arguments: None,
-              }
-            };
-            CodeLens {
-              range: params.range,
-              command: Some(command),
-              data: None,
-            }
-          } else {
-            let command = Command {
-              title: "0 references".to_string(),
-              command: "".to_string(),
-              arguments: None,
-            };
-            CodeLens {
-              range: params.range,
-              command: Some(command),
-              data: None,
-            }
-          }
-        }
-      };
-      self.performance.measure(mark);
-      Ok(code_lens)
+    let mark = self.performance.mark("code_lens_resolve", Some(&code_lens));
+    let result = if code_lens.data.is_some() {
+      code_lens::resolve_code_lens(code_lens, self)
+        .await
+        .map_err(|err| {
+          error!("Error resolving code lens: {}", err);
+          LspError::internal_error()
+        })
     } else {
-      self.performance.measure(mark);
       Err(LspError::invalid_params(
         "Code lens is missing the \"data\" property.",
       ))
-    }
+    };
+    self.performance.measure(mark);
+    result
   }
 
   async fn document_highlight(
diff --git a/cli/lsp/mod.rs b/cli/lsp/mod.rs
index 488507a5b..367257911 100644
--- a/cli/lsp/mod.rs
+++ b/cli/lsp/mod.rs
@@ -6,6 +6,7 @@ use lspower::Server;
 
 mod analysis;
 mod capabilities;
+mod code_lens;
 mod completions;
 mod config;
 mod diagnostics;
diff --git a/cli/lsp/sources.rs b/cli/lsp/sources.rs
index c5586a440..74e1a6a60 100644
--- a/cli/lsp/sources.rs
+++ b/cli/lsp/sources.rs
@@ -2,6 +2,7 @@
 
 use super::analysis;
 use super::text::LineIndex;
+use super::tsc;
 
 use crate::file_fetcher::get_source_from_bytes;
 use crate::file_fetcher::map_content_type;
@@ -15,6 +16,7 @@ use crate::program_state::ProgramState;
 use crate::specifier_handler::FetchHandler;
 use crate::text_encoding;
 
+use deno_core::error::anyhow;
 use deno_core::error::AnyError;
 use deno_core::serde_json;
 use deno_core::ModuleSpecifier;
@@ -26,6 +28,7 @@ use std::path::PathBuf;
 use std::sync::Arc;
 use std::sync::Mutex;
 use std::time::SystemTime;
+use tsc::NavigationTree;
 
 pub async fn cache(
   specifier: &ModuleSpecifier,
@@ -104,6 +107,7 @@ struct Metadata {
   dependencies: Option<HashMap<String, analysis::Dependency>>,
   length_utf16: usize,
   line_index: LineIndex,
+  maybe_navigation_tree: Option<tsc::NavigationTree>,
   maybe_types: Option<analysis::ResolvedDependency>,
   maybe_warning: Option<String>,
   media_type: MediaType,
@@ -139,6 +143,7 @@ impl Metadata {
       dependencies,
       length_utf16: source.encode_utf16().count(),
       line_index,
+      maybe_navigation_tree: None,
       maybe_types,
       maybe_warning,
       media_type: media_type.to_owned(),
@@ -197,6 +202,13 @@ impl Sources {
     self.0.lock().unwrap().get_media_type(specifier)
   }
 
+  pub fn get_navigation_tree(
+    &self,
+    specifier: &ModuleSpecifier,
+  ) -> Option<tsc::NavigationTree> {
+    self.0.lock().unwrap().get_navigation_tree(specifier)
+  }
+
   pub fn get_script_version(
     &self,
     specifier: &ModuleSpecifier,
@@ -223,6 +235,18 @@ impl Sources {
   pub fn specifiers(&self) -> Vec<ModuleSpecifier> {
     self.0.lock().unwrap().metadata.keys().cloned().collect()
   }
+
+  pub fn set_navigation_tree(
+    &self,
+    specifier: &ModuleSpecifier,
+    navigation_tree: tsc::NavigationTree,
+  ) -> Result<(), AnyError> {
+    self
+      .0
+      .lock()
+      .unwrap()
+      .set_navigation_tree(specifier, navigation_tree)
+  }
 }
 
 impl Inner {
@@ -343,6 +367,16 @@ impl Inner {
     Some(metadata)
   }
 
+  fn get_navigation_tree(
+    &mut self,
+    specifier: &ModuleSpecifier,
+  ) -> Option<tsc::NavigationTree> {
+    let specifier =
+      resolve_specifier(specifier, &mut self.redirects, &self.http_cache)?;
+    let metadata = self.get_metadata(&specifier)?;
+    metadata.maybe_navigation_tree
+  }
+
   fn get_path(&mut self, specifier: &ModuleSpecifier) -> Option<PathBuf> {
     if specifier.scheme() == "file" {
       specifier.to_file_path().ok()
@@ -461,6 +495,19 @@ impl Inner {
       }
     }
   }
+
+  fn set_navigation_tree(
+    &mut self,
+    specifier: &ModuleSpecifier,
+    navigation_tree: NavigationTree,
+  ) -> Result<(), AnyError> {
+    let mut metadata = self
+      .metadata
+      .get_mut(specifier)
+      .ok_or_else(|| anyhow!("Specifier not found {}"))?;
+    metadata.maybe_navigation_tree = Some(navigation_tree);
+    Ok(())
+  }
 }
 
 #[cfg(test)]
diff --git a/cli/lsp/tsc.rs b/cli/lsp/tsc.rs
index 6a476054f..834ac0a0c 100644
--- a/cli/lsp/tsc.rs
+++ b/cli/lsp/tsc.rs
@@ -1,8 +1,8 @@
 // Copyright 2018-2021 the Deno authors. All rights reserved. MIT license.
 
-use super::analysis::CodeLensSource;
 use super::analysis::ResolvedDependency;
 use super::analysis::ResolvedDependencyErr;
+use super::code_lens;
 use super::config;
 use super::language_server;
 use super::language_server::StateSnapshot;
@@ -103,6 +103,7 @@ pub struct AssetDocument {
   pub text: String,
   pub length: usize,
   pub line_index: LineIndex,
+  pub maybe_navigation_tree: Option<NavigationTree>,
 }
 
 impl AssetDocument {
@@ -112,6 +113,7 @@ impl AssetDocument {
       text: text.to_string(),
       length: text.encode_utf16().count(),
       line_index: LineIndex::new(text),
+      maybe_navigation_tree: None,
     }
   }
 }
@@ -150,6 +152,22 @@ impl Assets {
   ) -> Option<Option<AssetDocument>> {
     self.0.insert(k, v)
   }
+
+  pub fn set_navigation_tree(
+    &mut self,
+    specifier: &ModuleSpecifier,
+    navigation_tree: NavigationTree,
+  ) -> Result<(), AnyError> {
+    let maybe_doc = self
+      .0
+      .get_mut(specifier)
+      .ok_or_else(|| anyhow!("Missing asset."))?;
+    let doc = maybe_doc
+      .as_mut()
+      .ok_or_else(|| anyhow!("Cannot get doc mutable"))?;
+    doc.maybe_navigation_tree = Some(navigation_tree);
+    Ok(())
+  }
 }
 
 /// Optionally returns an internal asset, first checking for any static assets
@@ -610,7 +628,7 @@ impl NavigationTree {
     &self,
     line_index: &LineIndex,
     specifier: &ModuleSpecifier,
-    source: &CodeLensSource,
+    source: &code_lens::CodeLensSource,
   ) -> lsp::CodeLens {
     let range = if let Some(name_span) = &self.name_span {
       name_span.to_range(line_index)
diff --git a/cli/tests/integration_tests_lsp.rs b/cli/tests/integration_tests_lsp.rs
index cb9eae5a8..10797a326 100644
--- a/cli/tests/integration_tests_lsp.rs
+++ b/cli/tests/integration_tests_lsp.rs
@@ -1092,7 +1092,7 @@ fn lsp_code_lens_impl() {
         "uri": "file:///a/file.ts",
         "languageId": "typescript",
         "version": 1,
-        "text": "interface A {\n  b(): void;\n}\n\nclass B implements A {\n  b() {\n    console.log(\"b\");\n  }\n}\n"
+        "text": "interface A {\n  b(): void;\n}\n\nclass B implements A {\n  b() {\n    console.log(\"b\");\n  }\n}\n\ninterface C {\n  c: string;\n}\n"
       }
     }),
   );
@@ -1137,6 +1137,47 @@ fn lsp_code_lens_impl() {
     maybe_res,
     Some(load_fixture("code_lens_resolve_response_impl.json"))
   );
+  let (maybe_res, maybe_err) = client
+    .write_request::<_, _, Value>(
+      "codeLens/resolve",
+      json!({
+        "range": {
+          "start": {
+            "line": 10,
+            "character": 10
+          },
+          "end": {
+            "line": 10,
+            "character": 11
+          }
+        },
+        "data": {
+          "specifier": "file:///a/file.ts",
+          "source": "implementations"
+        }
+      }),
+    )
+    .unwrap();
+  assert!(maybe_err.is_none());
+  assert_eq!(
+    maybe_res,
+    Some(json!({
+      "range": {
+        "start": {
+          "line": 10,
+          "character": 10
+        },
+        "end": {
+          "line": 10,
+          "character": 11
+        }
+      },
+      "command": {
+        "title": "0 implementations",
+        "command": ""
+      }
+    }))
+  );
   shutdown(&mut client);
 }
 
@@ -1226,6 +1267,79 @@ fn lsp_code_lens_non_doc_nav_tree() {
 }
 
 #[test]
+fn lsp_nav_tree_updates() {
+  let mut client = init("initialize_params.json");
+  did_open(
+    &mut client,
+    json!({
+      "textDocument": {
+        "uri": "file:///a/file.ts",
+        "languageId": "typescript",
+        "version": 1,
+        "text": "interface A {\n  b(): void;\n}\n\nclass B implements A {\n  b() {\n    console.log(\"b\");\n  }\n}\n\ninterface C {\n  c: string;\n}\n"
+      }
+    }),
+  );
+  let (maybe_res, maybe_err) = client
+    .write_request(
+      "textDocument/codeLens",
+      json!({
+        "textDocument": {
+          "uri": "file:///a/file.ts"
+        }
+      }),
+    )
+    .unwrap();
+  assert!(maybe_err.is_none());
+  assert_eq!(
+    maybe_res,
+    Some(load_fixture("code_lens_response_impl.json"))
+  );
+  client
+    .write_notification(
+      "textDocument/didChange",
+      json!({
+        "textDocument": {
+          "uri": "file:///a/file.ts",
+          "version": 2
+        },
+        "contentChanges": [
+          {
+            "range": {
+              "start": {
+                "line": 10,
+                "character": 0
+              },
+              "end": {
+                "line": 13,
+                "character": 0
+              }
+            },
+            "text": ""
+          }
+        ]
+      }),
+    )
+    .unwrap();
+  let (maybe_res, maybe_err) = client
+    .write_request(
+      "textDocument/codeLens",
+      json!({
+        "textDocument": {
+          "uri": "file:///a/file.ts"
+        }
+      }),
+    )
+    .unwrap();
+  assert!(maybe_err.is_none());
+  assert_eq!(
+    maybe_res,
+    Some(load_fixture("code_lens_response_changed.json"))
+  );
+  shutdown(&mut client);
+}
+
+#[test]
 fn lsp_signature_help() {
   let mut client = init("initialize_params.json");
   did_open(
diff --git a/cli/tests/lsp/code_lens_response_changed.json b/cli/tests/lsp/code_lens_response_changed.json
new file mode 100644
index 000000000..b0073a23f
--- /dev/null
+++ b/cli/tests/lsp/code_lens_response_changed.json
@@ -0,0 +1,50 @@
+[
+  {
+    "range": {
+      "start": {
+        "line": 0,
+        "character": 10
+      },
+      "end": {
+        "line": 0,
+        "character": 11
+      }
+    },
+    "data": {
+      "specifier": "file:///a/file.ts",
+      "source": "implementations"
+    }
+  },
+  {
+    "range": {
+      "start": {
+        "line": 0,
+        "character": 10
+      },
+      "end": {
+        "line": 0,
+        "character": 11
+      }
+    },
+    "data": {
+      "specifier": "file:///a/file.ts",
+      "source": "references"
+    }
+  },
+  {
+    "range": {
+      "start": {
+        "line": 4,
+        "character": 6
+      },
+      "end": {
+        "line": 4,
+        "character": 7
+      }
+    },
+    "data": {
+      "specifier": "file:///a/file.ts",
+      "source": "references"
+    }
+  }
+]
diff --git a/cli/tests/lsp/code_lens_response_impl.json b/cli/tests/lsp/code_lens_response_impl.json
index b0073a23f..c6e5bd92d 100644
--- a/cli/tests/lsp/code_lens_response_impl.json
+++ b/cli/tests/lsp/code_lens_response_impl.json
@@ -46,5 +46,53 @@
       "specifier": "file:///a/file.ts",
       "source": "references"
     }
+  },
+  {
+    "range": {
+      "start": {
+        "line": 10,
+        "character": 10
+      },
+      "end": {
+        "line": 10,
+        "character": 11
+      }
+    },
+    "data": {
+      "specifier": "file:///a/file.ts",
+      "source": "implementations"
+    }
+  },
+  {
+    "range": {
+      "start": {
+        "line": 10,
+        "character": 10
+      },
+      "end": {
+        "line": 10,
+        "character": 11
+      }
+    },
+    "data": {
+      "specifier": "file:///a/file.ts",
+      "source": "references"
+    }
+  },
+  {
+    "range": {
+      "start": {
+        "line": 11,
+        "character": 2
+      },
+      "end": {
+        "line": 11,
+        "character": 3
+      }
+    },
+    "data": {
+      "specifier": "file:///a/file.ts",
+      "source": "references"
+    }
   }
 ]