Language Server Features for Top-Level Expressions in Notebooks (micr…

…osoft#1726)

Updates the Language Server to allow features to work properly for
top-level expressions outside of callables for the notebook experience.

Related to: microsoft#1714 

Fixes: microsoft#1718

language_service/src/definition/tests.rs

@@ -672,3 +672,16 @@ fn notebook_callable_defined_in_later_cell() {
         ("cell2", "operation Callee() : Unit {}"),
     ]);
 }
+
+#[test]
+fn notebook_local_from_same_cell() {
+    assert_definition_notebook(&[("cell1", "let ◉x◉ = 3; let y = ↘x + 1;")]);
+}
+
+#[test]
+fn notebook_local_from_later_cell() {
+    assert_definition_notebook(&[
+        ("cell1", "let ◉x◉ = 3; let y = x + 1;"),
+        ("cell2", "let z = ↘x + 2;"),
+    ]);
+}

language_service/src/hover/tests.rs

@@ -1609,3 +1609,16 @@ fn notebook_local_definition() {
         "#]],
     );
 }
+
+#[test]
+fn notebook_local_reference() {
+    check_notebook(
+        &[("cell1", "let x = 3;"), ("cell2", "let y = ◉↘x◉ + 1;")],
+        &expect![[r#"
+            local
+            ```qsharp
+            x : Int
+            ```
+        "#]],
+    );
+}

language_service/src/name_locator.rs

@@ -5,7 +5,6 @@ use std::mem::replace;
 use std::rc::Rc;
 
 use crate::compilation::Compilation;
-use crate::qsc_utils::find_ident;
 use qsc::ast::visit::{walk_expr, walk_namespace, walk_pat, walk_ty, walk_ty_def, Visitor};
 use qsc::display::Lookup;
 use qsc::{ast, hir, resolve};
@@ -134,7 +133,7 @@ impl<'inner, 'package, T: Handler<'package>> Locator<'inner, 'package, T> {
     }
 
     fn get_field_def<'other>(
-        &mut self,
+        &self,
         udt_res: &'package hir::Res,
         field_ref: &'other ast::Ident,
     ) -> Option<(hir::ItemId, &'package hir::ty::UdtField)> {
@@ -148,6 +147,21 @@ impl<'inner, 'package, T: Handler<'package>> Locator<'inner, 'package, T> {
         }
         None
     }
+
+    fn find_ident(&self, node_id: ast::NodeId) -> Option<&'package ast::Ident> {
+        let mut finder = AstIdentFinder {
+            node_id,
+            ident: None,
+        };
+        {
+            if let Some(callable) = self.context.current_callable {
+                finder.visit_callable_decl(callable);
+            } else {
+                finder.visit_package(&self.compilation.user_unit().ast.package);
+            }
+        }
+        finder.ident
+    }
 }
 
 impl<'inner, 'package, T: Handler<'package>> Visitor<'package> for Locator<'inner, 'package, T> {
@@ -389,11 +403,9 @@ impl<'inner, 'package, T: Handler<'package>> Visitor<'package> for Locator<'inne
                         .split_first()
                         .expect("paths should have at least one part");
                     if first.span.touches(self.offset) {
-                        if let Some(curr) = self.context.current_callable {
-                            if let Some(definition) = find_ident(node_id, curr) {
-                                self.inner
-                                    .at_local_ref(&self.context, first, node_id, definition);
-                            }
+                        if let Some(definition) = self.find_ident(node_id) {
+                            self.inner
+                                .at_local_ref(&self.context, first, node_id, definition);
                         }
                     } else {
                         // Loop through the parts of the path to find the first part that touches the offset
@@ -438,15 +450,13 @@ impl<'inner, 'package, T: Handler<'package>> Visitor<'package> for Locator<'inne
                         }
                     }
                     Some(resolve::Res::Local(node_id)) => {
-                        if let Some(curr) = self.context.current_callable {
-                            if let Some(definition) = find_ident(*node_id, curr) {
-                                self.inner.at_local_ref(
-                                    &self.context,
-                                    &path.name,
-                                    *node_id,
-                                    definition,
-                                );
-                            }
+                        if let Some(definition) = self.find_ident(*node_id) {
+                            self.inner.at_local_ref(
+                                &self.context,
+                                &path.name,
+                                *node_id,
+                                definition,
+                            );
                         }
                     }
                     _ => {}
@@ -455,3 +465,47 @@ impl<'inner, 'package, T: Handler<'package>> Visitor<'package> for Locator<'inne
         }
     }
 }
+
+/// Call `visit_callable_decl` if the ident is local to a callable.
+/// Otherwise call `visit_package`.
+struct AstIdentFinder<'a> {
+    pub node_id: ast::NodeId,
+    pub ident: Option<&'a ast::Ident>,
+}
+
+impl<'a> ast::visit::Visitor<'a> for AstIdentFinder<'a> {
+    // Locals don't cross namespace boundaries, so don't visit namespaces.
+    fn visit_package(&mut self, package: &'a ast::Package) {
+        package.nodes.iter().for_each(|n| {
+            if let ast::TopLevelNode::Stmt(stmt) = n {
+                self.visit_stmt(stmt);
+            }
+        });
+        package.entry.iter().for_each(|e| self.visit_expr(e));
+    }
+
+    // Locals don't cross item boundaries, so don't visit items.
+    fn visit_stmt(&mut self, stmt: &'a ast::Stmt) {
+        match &*stmt.kind {
+            ast::StmtKind::Item(_) => {}
+            _ => ast::visit::walk_stmt(self, stmt),
+        }
+    }
+
+    fn visit_pat(&mut self, pat: &'a ast::Pat) {
+        match &*pat.kind {
+            ast::PatKind::Bind(ident, _) => {
+                if ident.id == self.node_id {
+                    self.ident = Some(ident);
+                }
+            }
+            _ => ast::visit::walk_pat(self, pat),
+        }
+    }
+
+    fn visit_expr(&mut self, expr: &'a ast::Expr) {
+        if self.ident.is_none() {
+            ast::visit::walk_expr(self, expr);
+        }
+    }
+}

language_service/src/qsc_utils.rs

@@ -4,7 +4,7 @@
 use crate::compilation::Compilation;
 use qsc::line_column::{Encoding, Range};
 use qsc::location::Location;
-use qsc::{ast, hir::PackageId, SourceMap, Span};
+use qsc::{hir::PackageId, SourceMap, Span};
 
 pub(crate) fn into_range(encoding: Encoding, span: Span, source_map: &SourceMap) -> Range {
     let lo_source = source_map
@@ -37,42 +37,3 @@ pub(crate) fn into_location(
         position_encoding,
     )
 }
-
-pub(crate) fn find_ident(
-    node_id: ast::NodeId,
-    callable: &ast::CallableDecl,
-) -> Option<&ast::Ident> {
-    let mut finder = AstIdentFinder {
-        node_id,
-        ident: None,
-    };
-    {
-        use ast::visit::Visitor;
-        finder.visit_callable_decl(callable);
-    }
-    finder.ident
-}
-
-struct AstIdentFinder<'a> {
-    pub node_id: ast::NodeId,
-    pub ident: Option<&'a ast::Ident>,
-}
-
-impl<'a> ast::visit::Visitor<'a> for AstIdentFinder<'a> {
-    fn visit_pat(&mut self, pat: &'a ast::Pat) {
-        match &*pat.kind {
-            ast::PatKind::Bind(ident, _) => {
-                if ident.id == self.node_id {
-                    self.ident = Some(ident);
-                }
-            }
-            _ => ast::visit::walk_pat(self, pat),
-        }
-    }
-
-    fn visit_expr(&mut self, expr: &'a ast::Expr) {
-        if self.ident.is_none() {
-            ast::visit::walk_expr(self, expr);
-        }
-    }
-}

language_service/src/references.rs

@@ -162,9 +162,9 @@ impl<'a> Handler<'a> for NameHandler<'a> {
         ident: &'a ast::Ident,
         _: &'a ast::Pat,
     ) {
-        if let Some(curr) = context.current_callable {
-            self.references = self.reference_finder.for_local(ident.id, curr);
-        }
+        self.references = self
+            .reference_finder
+            .for_local(ident.id, context.current_callable);
     }
 
     fn at_local_ref(
@@ -174,9 +174,9 @@ impl<'a> Handler<'a> for NameHandler<'a> {
         _: ast::NodeId,
         definition: &'a ast::Ident,
     ) {
-        if let Some(curr) = context.current_callable {
-            self.references = self.reference_finder.for_local(definition.id, curr);
-        }
+        self.references = self
+            .reference_finder
+            .for_local(definition.id, context.current_callable);
     }
 }
 
@@ -277,14 +277,22 @@ impl<'a> ReferenceFinder<'a> {
         locations
     }
 
-    pub fn for_local(&self, node_id: ast::NodeId, callable: &ast::CallableDecl) -> Vec<Location> {
+    pub fn for_local(
+        &self,
+        node_id: ast::NodeId,
+        callable: Option<&ast::CallableDecl>,
+    ) -> Vec<Location> {
         let mut find_refs = FindLocalLocations {
             node_id,
             compilation: self.compilation,
             include_declaration: self.include_declaration,
             locations: vec![],
         };
-        find_refs.visit_callable_decl(callable);
+        if let Some(callable) = callable {
+            find_refs.visit_callable_decl(callable);
+        } else {
+            find_refs.visit_package(&self.compilation.user_unit().ast.package);
+        }
         find_refs
             .locations
             .into_iter()
@@ -440,7 +448,25 @@ struct FindLocalLocations<'a> {
     locations: Vec<Span>,
 }
 
-impl<'a> Visitor<'_> for FindLocalLocations<'a> {
+impl Visitor<'_> for FindLocalLocations<'_> {
+    // Locals don't cross namespace boundaries, so don't visit namespaces.
+    fn visit_package(&mut self, package: &ast::Package) {
+        package.nodes.iter().for_each(|n| {
+            if let ast::TopLevelNode::Stmt(stmt) = n {
+                self.visit_stmt(stmt);
+            }
+        });
+        package.entry.iter().for_each(|e| self.visit_expr(e));
+    }
+
+    // Locals don't cross item boundaries, so don't visit items.
+    fn visit_stmt(&mut self, stmt: &ast::Stmt) {
+        match &*stmt.kind {
+            ast::StmtKind::Item(_) => {}
+            _ => ast::visit::walk_stmt(self, stmt),
+        }
+    }
+
     fn visit_pat(&mut self, pat: &ast::Pat) {
         if self.include_declaration {
             match &*pat.kind {

language_service/src/references/tests.rs

@@ -917,3 +917,19 @@ fn notebook_defined_in_later_cell() {
         ("cell2", "operation Callee() : Unit {}"),
     ]);
 }
+
+#[test]
+fn notebook_local_definition() {
+    check_notebook_exclude_decl(&[
+        ("cell1", "let ↘x = 3; let y = ◉x◉ + 1;"),
+        ("cell2", "let z = ◉x◉ + 2;"),
+    ]);
+}
+
+#[test]
+fn notebook_local_reference() {
+    check_notebook_exclude_decl(&[
+        ("cell1", "let x = 3; let y = ◉x◉ + 1;"),
+        ("cell2", "let z = ◉↘x◉ + 2;"),
+    ]);
+}

language_service/src/rename.rs

@@ -141,7 +141,7 @@ impl<'a> Rename<'a> {
         &mut self,
         node_id: ast::NodeId,
         ast_name: &ast::Ident,
-        current_callable: &ast::CallableDecl,
+        current_callable: Option<&ast::CallableDecl>,
     ) {
         if self.is_prepare {
             self.prepare = Some((ast_name.span, ast_name.name.to_string()));
@@ -265,9 +265,7 @@ impl<'a> Handler<'a> for Rename<'a> {
         ident: &'a ast::Ident,
         _: &'a ast::Pat,
     ) {
-        if let Some(curr) = context.current_callable {
-            self.get_spans_for_local_rename(ident.id, ident, curr);
-        }
+        self.get_spans_for_local_rename(ident.id, ident, context.current_callable);
     }
 
     fn at_local_ref(
@@ -277,9 +275,7 @@ impl<'a> Handler<'a> for Rename<'a> {
         node_id: ast::NodeId,
         _: &'a ast::Ident,
     ) {
-        if let Some(curr) = context.current_callable {
-            self.get_spans_for_local_rename(node_id, name, curr);
-        }
+        self.get_spans_for_local_rename(node_id, name, context.current_callable);
     }
 }