Improved (but possibly expensive) reference-finding

lsp/nls/src/analysis.rs

@@ -2,8 +2,9 @@ use std::collections::HashMap;
 
 use codespan::FileId;
 use nickel_lang_core::{
+    identifier::Ident,
     position::RawSpan,
-    term::{BinaryOp, RichTerm, Term, Traverse, TraverseControl},
+    term::{BinaryOp, RichTerm, Term, Traverse, TraverseControl, UnaryOp},
     typ::{Type, TypeF},
     typecheck::{reporting::NameReg, TypeTables, TypecheckVisitor, UnifType},
 };
@@ -100,6 +101,20 @@ impl ParentLookup {
     }
 }
 
+fn find_static_accesses(rt: &RichTerm) -> HashMap<Ident, Vec<RichTerm>> {
+    let mut map: HashMap<Ident, Vec<RichTerm>> = HashMap::new();
+    rt.traverse_ref(
+        &mut |rt: &RichTerm, _scope: &()| {
+            if let Term::Op1(UnaryOp::StaticAccess(id), _) = rt.as_ref() {
+                map.entry(id.ident()).or_default().push(rt.clone());
+            }
+            TraverseControl::Continue::<_, ()>
+        },
+        &(),
+    );
+    map
+}
+
 /// Essentially an iterator over pairs of `(ancestor, reversed_path_to_the_original)`.
 ///
 /// For example, if we are iterating over the AST of `foo.bar.baz`, the iterator
@@ -212,6 +227,10 @@ pub struct Analysis {
     pub usage_lookup: UsageLookup,
     pub parent_lookup: ParentLookup,
     pub type_lookup: CollectedTypes<Type>,
+
+    /// A lookup table for static accesses, for looking up all occurrences of,
+    /// say, `.foo` in a file.
+    pub static_accesses: HashMap<Ident, Vec<RichTerm>>,
 }
 
 impl Analysis {
@@ -224,6 +243,7 @@ impl Analysis {
             position_lookup: PositionLookup::new(term),
             usage_lookup: UsageLookup::new(term, initial_env),
             parent_lookup: ParentLookup::new(term),
+            static_accesses: find_static_accesses(term),
             type_lookup,
         }
     }
@@ -307,6 +327,15 @@ impl AnalysisRegistry {
         let file = rt.pos.as_opt_ref()?.src_id;
         Some(self.analysis.get(&file)?.parent_lookup.parent_chain(rt))
     }
+
+    pub fn get_static_accesses(&self, id: Ident) -> Vec<RichTerm> {
+        self.analysis
+            .values()
+            .filter_map(|a| a.static_accesses.get(&id))
+            .flatten()
+            .cloned()
+            .collect()
+    }
 }
 
 #[derive(Debug, Default)]

lsp/nls/src/files.rs

@@ -49,10 +49,21 @@ pub fn handle_open(server: &mut Server, params: DidOpenTextDocumentParams) -> Re
 
     // Invalidate the cache of every file that tried, but failed, to import a file
     // with a name like this.
-    let invalid = path
+    let mut invalid = path
         .file_name()
         .and_then(|name| server.failed_imports.remove(name))
         .unwrap_or_default();
+
+    // Replace the path (as opposed to adding it): we may already have this file in the
+    // cache if it was imported by an already-open file.
+    let file_id = server
+        .cache
+        .replace_string(SourcePath::Path(path), params.text_document.text);
+
+    // Invalidate any cached inputs that imported the newly-opened file, so that any
+    // cross-file references are updated.
+    invalid.extend(server.cache.get_rev_imports_transitive(file_id));
+
     for rev_dep in &invalid {
         server.analysis.remove(*rev_dep);
         // Reset the cached state (Parsed is the earliest one) so that it will
@@ -62,9 +73,6 @@ pub fn handle_open(server: &mut Server, params: DidOpenTextDocumentParams) -> Re
             .update_state(*rev_dep, nickel_lang_core::cache::EntryState::Parsed);
     }
 
-    let file_id = server
-        .cache
-        .add_string(SourcePath::Path(path), params.text_document.text);
     server.file_uris.insert(file_id, params.text_document.uri);
 
     parse_and_typecheck(server, file_id)?;

lsp/nls/src/requests/goto.rs

@@ -1,75 +1,11 @@
+use std::collections::HashSet;
+
 use lsp_server::{RequestId, Response, ResponseError};
 use lsp_types::{GotoDefinitionParams, GotoDefinitionResponse, Location, ReferenceParams};
-use nickel_lang_core::{
-    position::RawSpan,
-    term::{record::FieldMetadata, RichTerm, Term, UnaryOp},
-};
+use nickel_lang_core::position::RawSpan;
 use serde_json::Value;
 
-use crate::{
-    cache::CacheExt,
-    diagnostic::LocationCompat,
-    field_walker::{Def, FieldResolver},
-    identifier::LocIdent,
-    server::Server,
-};
-
-fn get_defs(term: &RichTerm, ident: Option<LocIdent>, server: &Server) -> Option<Vec<RawSpan>> {
-    let resolver = FieldResolver::new(server);
-    let ret = match (term.as_ref(), ident) {
-        (Term::Var(id), _) => {
-            let id = LocIdent::from(*id);
-            let def = server.analysis.get_def(&id)?;
-            let cousins = resolver.cousin_defs(def);
-            if cousins.is_empty() {
-                vec![def.ident().pos.unwrap()]
-            } else {
-                cousins
-                    .into_iter()
-                    .filter_map(|(loc, _)| loc.pos.into_opt())
-                    .collect()
-            }
-        }
-        (Term::Op1(UnaryOp::StaticAccess(id), parent), _) => {
-            let parents = resolver.resolve_record(parent);
-            parents
-                .iter()
-                .filter_map(|parent| {
-                    parent
-                        .field_loc(id.ident())
-                        .and_then(|def| def.pos.into_opt())
-                })
-                .collect()
-        }
-        (Term::LetPattern(_, pat, value, _), Some(hovered_id)) => {
-            let (mut path, _, _) = pat
-                .matches
-                .iter()
-                .flat_map(|m| m.to_flattened_bindings())
-                .find(|(_path, bound_id, _)| bound_id.ident() == hovered_id.ident)?;
-            path.reverse();
-            let (last, path) = path.split_last()?;
-            let path: Vec<_> = path.iter().map(|id| id.ident()).collect();
-            let parents = resolver.resolve_path(value, path.iter().copied());
-            parents
-                .iter()
-                .filter_map(|parent| {
-                    parent
-                        .field_loc(last.ident())
-                        .and_then(|def| def.pos.into_opt())
-                })
-                .collect()
-        }
-        (Term::ResolvedImport(file), _) => {
-            let pos = server.cache.terms().get(file)?.term.pos;
-            vec![pos.into_opt()?]
-        }
-        _ => {
-            return None;
-        }
-    };
-    Some(ret)
-}
+use crate::{cache::CacheExt, diagnostic::LocationCompat, server::Server};
 
 fn ids_to_locations(ids: impl IntoIterator<Item = RawSpan>, server: &Server) -> Vec<Location> {
     let mut spans: Vec<_> = ids.into_iter().collect();
@@ -97,7 +33,7 @@ pub fn handle_to_definition(
 
     let locations = server
         .lookup_term_by_position(pos)?
-        .and_then(|term| get_defs(term, ident, server))
+        .map(|term| server.get_defs(term, ident))
         .map(|defs| ids_to_locations(defs, server))
         .unwrap_or_default();
 
@@ -125,40 +61,14 @@ pub fn handle_references(
     // so first find the definitions and then find their usages.
     let term = server.lookup_term_by_position(pos)?;
     let mut def_locs = term
-        .and_then(|term| get_defs(term, ident, server))
+        .map(|term| server.get_defs(term, ident))
         .unwrap_or_default();
 
     // Maybe the position is pointing straight at the definition already.
     // In that case, def_locs won't have the definition yet; so add it.
-    if let Some(id) = server.lookup_ident_by_position(pos)? {
-        if let Some(span) = id.pos.into_opt() {
-            def_locs.push(span);
-            if let Some(parent) = term {
-                // If `id` is a field name in a record, we can search through cousins
-                // to find more definitions.
-                if matches!(parent.as_ref(), Term::RecRecord(..) | Term::Record(_)) {
-                    let def = Def::Field {
-                        ident: id,
-                        value: None,
-                        record: parent.clone(),
-                        metadata: FieldMetadata::default(),
-                    };
-                    let resolver = FieldResolver::new(server);
-                    let cousins = resolver.cousin_defs(&def);
-                    def_locs.extend(
-                        cousins
-                            .into_iter()
-                            .filter_map(|(loc, _)| loc.pos.into_opt()),
-                    )
-                }
-            }
-        }
-    }
+    def_locs.extend(ident.and_then(|id| id.pos.into_opt()));
 
-    // TODO: This usage map is based only on static scoping, and not on our "extended"
-    // scopes that we build up dynamically based on merges. Improving this probably
-    // requires building the extended scopes at static analysis time.
-    let mut usages: Vec<_> = def_locs
+    let mut usages: HashSet<_> = def_locs
         .iter()
         .flat_map(|id| server.analysis.get_usages(id))
         .filter_map(|id| id.pos.into_opt())
@@ -168,6 +78,10 @@ pub fn handle_references(
         usages.extend(def_locs.iter().cloned());
     }
 
+    for span in def_locs {
+        usages.extend(server.get_field_refs(span));
+    }
+
     let locations = ids_to_locations(usages, server);
 
     if locations.is_empty() {

lsp/nls/src/server.rs

@@ -24,9 +24,9 @@ use lsp_types::{
 
 use nickel_lang_core::{
     cache::{Cache, ErrorTolerance},
-    position::{RawPos, TermPos},
+    position::{RawPos, RawSpan, TermPos},
     stdlib::StdlibModule,
-    term::RichTerm,
+    term::{record::FieldMetadata, RichTerm, Term, UnaryOp},
 };
 use nickel_lang_core::{stdlib, typecheck::Context};
 
@@ -36,7 +36,8 @@ use crate::{
     cache::CacheExt,
     command,
     diagnostic::DiagnosticCompat,
-    field_walker::Def,
+    field_walker::{Def, FieldResolver},
+    identifier::LocIdent,
     requests::{completion, formatting, goto, hover, symbols},
     trace::Trace,
 };
@@ -349,4 +350,130 @@ impl Server {
             },
         ));
     }
+
+    /// Finds all the locations at which a term (or possibly an ident within a term) is "defined".
+    ///
+    /// "Ident within a term" applies when the term is a record or a pattern binding, so that we
+    /// can refer to fields in a record, or specific idents in a pattern binding.
+    ///
+    /// The return value contains all the spans of all the definition locations. It's a span instead
+    /// of a `LocIdent` because when `term` is an import, the definition location is the whole
+    /// included file. In every other case, the definition location will be the span of a LocIdent.
+    pub fn get_defs(&self, term: &RichTerm, ident: Option<LocIdent>) -> Vec<RawSpan> {
+        // The inner function returning Option is just for ?-early-return convenience.
+        fn inner(
+            server: &Server,
+            term: &RichTerm,
+            ident: Option<LocIdent>,
+        ) -> Option<Vec<RawSpan>> {
+            let resolver = FieldResolver::new(server);
+            let ret = match (term.as_ref(), ident) {
+                (Term::Var(id), _) => {
+                    let id = LocIdent::from(*id);
+                    let def = server.analysis.get_def(&id)?;
+                    let cousins = resolver.cousin_defs(def);
+                    if cousins.is_empty() {
+                        vec![def.ident().pos.unwrap()]
+                    } else {
+                        cousins
+                            .into_iter()
+                            .filter_map(|(loc, _)| loc.pos.into_opt())
+                            .collect()
+                    }
+                }
+                (Term::Op1(UnaryOp::StaticAccess(id), parent), _) => {
+                    let parents = resolver.resolve_record(parent);
+                    parents
+                        .iter()
+                        .filter_map(|parent| {
+                            parent
+                                .field_loc(id.ident())
+                                .and_then(|def| def.pos.into_opt())
+                        })
+                        .collect()
+                }
+                (Term::LetPattern(_, pat, value, _), Some(hovered_id)) => {
+                    let (mut path, _, _) = pat
+                        .matches
+                        .iter()
+                        .flat_map(|m| m.to_flattened_bindings())
+                        .find(|(_path, bound_id, _)| bound_id.ident() == hovered_id.ident)?;
+                    path.reverse();
+                    let (last, path) = path.split_last()?;
+                    let path: Vec<_> = path.iter().map(|id| id.ident()).collect();
+                    let parents = resolver.resolve_path(value, path.iter().copied());
+                    parents
+                        .iter()
+                        .filter_map(|parent| {
+                            parent
+                                .field_loc(last.ident())
+                                .and_then(|def| def.pos.into_opt())
+                        })
+                        .collect()
+                }
+                (Term::ResolvedImport(file), _) => {
+                    let pos = server.cache.terms().get(file)?.term.pos;
+                    vec![pos.into_opt()?]
+                }
+                (Term::RecRecord(..) | Term::Record(_), Some(id)) => {
+                    let def = Def::Field {
+                        ident: id,
+                        value: None,
+                        record: term.clone(),
+                        metadata: FieldMetadata::default(),
+                    };
+                    let cousins = resolver.cousin_defs(&def);
+                    cousins
+                        .into_iter()
+                        .filter_map(|(loc, _)| loc.pos.into_opt())
+                        .collect()
+                }
+                _ => {
+                    return None;
+                }
+            };
+            Some(ret)
+        }
+
+        inner(self, term, ident).unwrap_or_default()
+    }
+
+    /// If `span` is pointing at the identifier binding a record field, returns
+    /// all the places that the record field is referenced.
+    ///
+    /// This is a sort of inverse of `get_defs`, at least when the argument to `get_defs`
+    /// is a static access: the spans returned by this function are exactly the static accesses
+    /// that, when passed to `get_defs`, return `span`.
+    ///
+    /// This function can be expensive, because it calls `get_defs` on every static access
+    /// that could potentially be referencing this field.
+    pub fn get_field_refs(&self, span: RawSpan) -> Vec<RawSpan> {
+        // The inner function returning Option is just for ?-early-return convenience.
+        fn inner(server: &Server, span: RawSpan) -> Option<Vec<RawSpan>> {
+            let ident = server.lookup_ident_by_position(span.start_pos()).ok()??;
+            let term = server.lookup_term_by_position(span.start_pos()).ok()??;
+
+            if let Term::RecRecord(..) | Term::Record(_) = term.as_ref() {
+                let accesses = server.analysis.get_static_accesses(ident.ident);
+                Some(
+                    accesses
+                        .into_iter()
+                        .filter_map(|access| {
+                            let Term::Op1(UnaryOp::StaticAccess(id), _) = access.as_ref() else {
+                                return None;
+                            };
+                            if server.get_defs(&access, None).contains(&span) {
+                                id.pos.into_opt()
+                            } else {
+                                None
+                            }
+                        })
+                        .collect(),
+                )
+            } else {
+                None
+            }
+        }
+        inner(self, span).unwrap_or_default()
+    }
 }