feat: support document symbols requests in lsp

This allows requesting a list/tree of symbols in the current document.
The following items are currently exposed:

- `load` statements:
    - The module path
    - The loaded symbol names
- `def` statements:
    - The function node
    - Argument names
- Closures, when assigned to a named variable
- Structs (calls to `struct`), when assigned to a named variable
- Build system targets (function calls containing a `name` argument)

starlark_lsp/src/server.rs

@@ -43,6 +43,7 @@ use lsp_types::notification::DidOpenTextDocument;
 use lsp_types::notification::LogMessage;
 use lsp_types::notification::PublishDiagnostics;
 use lsp_types::request::Completion;
+use lsp_types::request::DocumentSymbolRequest;
 use lsp_types::request::GotoDefinition;
 use lsp_types::request::HoverRequest;
 use lsp_types::CompletionItem;
@@ -55,6 +56,8 @@ use lsp_types::Diagnostic;
 use lsp_types::DidChangeTextDocumentParams;
 use lsp_types::DidCloseTextDocumentParams;
 use lsp_types::DidOpenTextDocumentParams;
+use lsp_types::DocumentSymbolParams;
+use lsp_types::DocumentSymbolResponse;
 use lsp_types::Documentation;
 use lsp_types::GotoDefinitionParams;
 use lsp_types::GotoDefinitionResponse;
@@ -107,6 +110,7 @@ use crate::definition::IdentifierDefinition;
 use crate::definition::LspModule;
 use crate::inspect::AstModuleInspect;
 use crate::inspect::AutocompleteType;
+use crate::symbols;
 use crate::symbols::find_symbols_at_location;
 
 /// The request to get the file contents for a starlark: URI
@@ -408,6 +412,7 @@ impl<T: LspContext> Backend<T> {
             definition_provider,
             completion_provider: Some(CompletionOptions::default()),
             hover_provider: Some(HoverProviderCapability::Simple(true)),
+            document_symbol_provider: Some(OneOf::Left(true)),
             ..ServerCapabilities::default()
         }
     }
@@ -506,6 +511,11 @@ impl<T: LspContext> Backend<T> {
         self.send_response(new_response(id, self.hover_info(params, initialize_params)));
     }
 
+    /// Offer an overview of symbols in the current document.
+    fn document_symbols(&self, id: RequestId, params: DocumentSymbolParams) {
+        self.send_response(new_response(id, self.get_document_symbols(params)));
+    }
+
     /// Get the file contents of a starlark: URI.
     fn get_starlark_file_contents(&self, id: RequestId, params: StarlarkFileContentsParams) {
         let response: anyhow::Result<_> = match params.uri {
@@ -1166,6 +1176,23 @@ impl<T: LspContext> Backend<T> {
         })
     }
 
+    fn get_document_symbols(
+        &self,
+        params: DocumentSymbolParams,
+    ) -> anyhow::Result<DocumentSymbolResponse> {
+        let uri = params.text_document.uri.try_into()?;
+
+        let document = match self.get_ast(&uri) {
+            Some(document) => document,
+            None => return Ok(DocumentSymbolResponse::Nested(vec![])),
+        };
+
+        let result =
+            symbols::get_document_symbols(document.ast.codemap(), document.ast.statement());
+
+        Ok(result.into())
+    }
+
     fn get_workspace_root(
         workspace_roots: Option<&Vec<WorkspaceFolder>>,
         target: &LspUrl,
@@ -1223,6 +1250,8 @@ impl<T: LspContext> Backend<T> {
                         self.completion(req.id, params, &initialize_params);
                     } else if let Some(params) = as_request::<HoverRequest>(&req) {
                         self.hover(req.id, params, &initialize_params);
+                    } else if let Some(params) = as_request::<DocumentSymbolRequest>(&req) {
+                        self.document_symbols(req.id, params);
                     } else if self.connection.handle_shutdown(&req)? {
                         return Ok(());
                     }

starlark_lsp/src/symbols.rs

@@ -18,12 +18,21 @@
 //! Find which symbols are in scope at a particular point.
 
 use std::collections::HashMap;
+use std::ops::Deref;
 
+use lsp_types::DocumentSymbol;
+use lsp_types::SymbolKind as LspSymbolKind;
 use starlark::codemap::CodeMap;
+use starlark::codemap::Span;
 use starlark::docs::DocItem;
 use starlark::docs::DocParam;
 use starlark_syntax::codemap::ResolvedPos;
+use starlark_syntax::syntax::ast::ArgumentP;
 use starlark_syntax::syntax::ast::AssignP;
+use starlark_syntax::syntax::ast::AssignTargetP;
+use starlark_syntax::syntax::ast::AstAssignIdentP;
+use starlark_syntax::syntax::ast::AstExprP;
+use starlark_syntax::syntax::ast::AstLiteral;
 use starlark_syntax::syntax::ast::AstPayload;
 use starlark_syntax::syntax::ast::AstStmtP;
 use starlark_syntax::syntax::ast::ExprP;
@@ -161,6 +170,233 @@ pub(crate) fn find_symbols_at_location<P: AstPayload>(
     symbols
 }
 
+pub fn get_document_symbols<P: AstPayload>(
+    codemap: &CodeMap,
+    ast: &AstStmtP<P>,
+) -> Vec<DocumentSymbol> {
+    let mut symbols = Vec::new();
+    match &ast.node {
+        StmtP::Expression(expr) => {
+            if let Some(symbol) = get_document_symbol_for_expr(codemap, None, expr, ast.span) {
+                symbols.push(symbol);
+            }
+        }
+        StmtP::Assign(assign) => {
+            if let Some(symbol) = get_document_symbol_for_expr(
+                codemap,
+                match &assign.lhs.node {
+                    AssignTargetP::Tuple(_)
+                    | AssignTargetP::Index(_)
+                    | AssignTargetP::Dot(_, _) => None,
+                    AssignTargetP::Identifier(ident) => Some(ident),
+                },
+                &assign.rhs,
+                ast.span,
+            ) {
+                symbols.push(symbol);
+            }
+        }
+        StmtP::Statements(statements) => {
+            for stmt in statements {
+                symbols.extend(get_document_symbols(codemap, stmt));
+            }
+        }
+        StmtP::If(_, body) => {
+            symbols.extend(get_document_symbols(codemap, body));
+        }
+        StmtP::IfElse(_, bodies) => {
+            let (if_body, else_body) = bodies.deref();
+            symbols.extend(get_document_symbols(codemap, if_body));
+            symbols.extend(get_document_symbols(codemap, else_body));
+        }
+        StmtP::For(for_) => {
+            symbols.extend(get_document_symbols(codemap, &for_.body));
+        }
+        StmtP::Def(def) => {
+            symbols.push(make_document_symbol(
+                def.name.ident.clone(),
+                LspSymbolKind::FUNCTION,
+                ast.span,
+                def.name.span,
+                codemap,
+                Some(
+                    def.params
+                        .iter()
+                        .filter_map(|param| get_document_symbol_for_parameter(codemap, param))
+                        .chain(get_document_symbols(codemap, &def.body))
+                        .collect(),
+                ),
+            ));
+        }
+        StmtP::Load(load) => {
+            symbols.push(make_document_symbol(
+                load.module.node.clone(),
+                LspSymbolKind::MODULE,
+                ast.span,
+                load.module.span,
+                codemap,
+                Some(
+                    load.args
+                        .iter()
+                        .map(|loaded_symbol| {
+                            make_document_symbol(
+                                loaded_symbol.local.ident.clone(),
+                                LspSymbolKind::METHOD,
+                                loaded_symbol.span(),
+                                loaded_symbol.local.span,
+                                codemap,
+                                None,
+                            )
+                        })
+                        .collect(),
+                ),
+            ));
+        }
+
+        // These don't produce any symbols.
+        StmtP::Break
+        | StmtP::Continue
+        | StmtP::Pass
+        | StmtP::Return(_)
+        | StmtP::AssignModify(_, _, _) => {}
+    }
+
+    symbols
+}
+
+fn get_document_symbol_for_parameter<P: AstPayload>(
+    codemap: &CodeMap,
+    param: &ParameterP<P>,
+) -> Option<DocumentSymbol> {
+    match param {
+        ParameterP::NoArgs => None,
+        ParameterP::Normal(p, _)
+        | ParameterP::WithDefaultValue(p, _, _)
+        | ParameterP::Args(p, _)
+        | ParameterP::KwArgs(p, _) => Some(make_document_symbol(
+            p.ident.clone(),
+            LspSymbolKind::VARIABLE,
+            p.span,
+            p.span,
+            codemap,
+            None,
+        )),
+    }
+}
+
+fn get_document_symbol_for_expr<P: AstPayload>(
+    codemap: &CodeMap,
+    name: Option<&AstAssignIdentP<P>>,
+    expr: &AstExprP<P>,
+    outer_range: Span,
+) -> Option<DocumentSymbol> {
+    match &expr.node {
+        ExprP::Call(call, args) => {
+            if let ExprP::Identifier(func_name) = &call.node {
+                // Look for a call to `struct`. We'll require passing in a name from the assignment
+                // expression. The outer range is the range of the entire assignment expression.
+                if &func_name.node.ident == "struct" {
+                    name.map(|name| {
+                        make_document_symbol(
+                            name.ident.clone(),
+                            LspSymbolKind::STRUCT,
+                            outer_range,
+                            name.span,
+                            codemap,
+                            Some(
+                                args.iter()
+                                    .filter_map(|arg| match &arg.node {
+                                        ArgumentP::Named(name, _) => Some(make_document_symbol(
+                                            name.node.clone(),
+                                            LspSymbolKind::FIELD,
+                                            arg.span,
+                                            name.span,
+                                            codemap,
+                                            None,
+                                        )),
+                                        _ => None,
+                                    })
+                                    .collect(),
+                            ),
+                        )
+                    })
+                } else {
+                    // Check if this call has a named argument called "name". If so, we'll assume
+                    // that this is a buildable target, and expose it.
+                    args.iter()
+                        .find_map(|arg| match &arg.node {
+                            ArgumentP::Named(name, value) => match (name, &value.node) {
+                                (name, ExprP::Literal(AstLiteral::String(value)))
+                                    if &name.node == "name" =>
+                                {
+                                    Some(value)
+                                }
+                                _ => None,
+                            },
+                            _ => None,
+                        })
+                        .map(|target_name| {
+                            make_document_symbol(
+                                target_name.node.clone(),
+                                LspSymbolKind::CONSTANT,
+                                expr.span,
+                                target_name.span,
+                                codemap,
+                                None,
+                            )
+                        })
+                }
+            } else {
+                None
+            }
+        }
+        ExprP::Lambda(lambda) => name.map(|name| {
+            make_document_symbol(
+                name.ident.clone(),
+                LspSymbolKind::FUNCTION,
+                expr.span,
+                expr.span,
+                codemap,
+                Some(
+                    lambda
+                        .params
+                        .iter()
+                        .filter_map(|param| get_document_symbol_for_parameter(codemap, param))
+                        .chain(get_document_symbol_for_expr(
+                            codemap,
+                            None,
+                            &lambda.body,
+                            lambda.body.span,
+                        ))
+                        .collect(),
+                ),
+            )
+        }),
+
+        _ => None,
+    }
+}
+
+fn make_document_symbol(
+    name: String,
+    kind: LspSymbolKind,
+    range: Span,
+    selection_range: Span,
+    codemap: &CodeMap,
+    children: Option<Vec<DocumentSymbol>>,
+) -> DocumentSymbol {
+    DocumentSymbol {
+        name,
+        detail: None,
+        kind,
+        tags: None,
+        deprecated: None,
+        range: codemap.resolve_span(range).into(),
+        selection_range: codemap.resolve_span(selection_range).into(),
+        children,
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use std::collections::HashMap;