Allow fetching starlark: uris with new message type

Natalie Jameson · facebook-github-bot · commit 04cb6c4444fc · 2022-08-11T10:29:10.000-07:00
Summary: Add a new JSONRPC message type to allow fetching file contents for Starlark `LspUrl`s from the LSP context.

Reviewed By: bobyangyf

Differential Revision: D38452373

fbshipit-source-id: c05fe6fc798a1bb823adf47a1aa848bcece74e28
diff --git a/starlark/bin/eval.rs b/starlark/bin/eval.rs
@@ -15,6 +15,7 @@
  * limitations under the License.
  */
 
+use std::collections::HashMap;
 use std::fs;
 use std::io;
 use std::iter;
@@ -38,6 +39,10 @@ use starlark::lsp::server::ResolveLoadError;
 use starlark::lsp::server::StringLiteralResult;
 use starlark::syntax::AstModule;
 use starlark::syntax::Dialect;
+use starlark::values::docs::get_registered_docs;
+use starlark::values::docs::render_docs_as_code;
+use starlark::values::docs::Doc;
+use starlark::values::docs::DocItem;
 
 #[derive(Debug)]
 pub(crate) enum ContextMode {
@@ -51,6 +56,7 @@ pub(crate) struct Context {
     pub(crate) print_non_none: bool,
     pub(crate) prelude: Vec<FrozenModule>,
     pub(crate) module: Option<Module>,
+    pub(crate) builtin_docs: HashMap<LspUrl, String>,
 }
 
 /// The outcome of evaluating (checking, parsing or running) given starlark code.
@@ -84,15 +90,36 @@ impl Context {
         } else {
             None
         };
+        let mut builtins: HashMap<LspUrl, Vec<Doc>> = HashMap::new();
+        for doc in get_registered_docs() {
+            let uri = Self::url_for_doc(&doc);
+            builtins.entry(uri).or_default().push(doc);
+        }
+        let builtin_docs = builtins
+            .into_iter()
+            .map(|(u, ds)| (u, render_docs_as_code(&ds).join("\n\n")))
+            .collect();
 
         Ok(Self {
             mode,
             print_non_none,
             prelude,
             module,
+            builtin_docs,
         })
     }
 
+    fn url_for_doc(doc: &Doc) -> LspUrl {
+        let url = match &doc.item {
+            DocItem::Module(_) => Url::parse("starlark:/native/builtins.bzl").unwrap(),
+            DocItem::Object(_) => {
+                Url::parse(&format!("starlark:/native/builtins/{}.bzl", doc.id.name)).unwrap()
+            }
+            DocItem::Function(_) => Url::parse("starlark:/native/builtins.bzl").unwrap(),
+        };
+        LspUrl::try_from(url).unwrap()
+    }
+
     fn new_module(prelude: &[FrozenModule]) -> Module {
         let module = Module::new();
         for p in prelude {
@@ -271,6 +298,7 @@ impl LspContext for Context {
                 },
                 false => Err(LoadContentsError::NotAbsolute(uri.clone()).into()),
             },
+            LspUrl::Starlark(_) => Ok(self.builtin_docs.get(uri).cloned()),
             _ => Err(LoadContentsError::WrongScheme("file://".to_owned(), uri.clone()).into()),
         }
     }
diff --git a/starlark/src/lsp/server.rs b/starlark/src/lsp/server.rs
@@ -67,8 +67,32 @@ use serde::Serializer;
 
 use crate::analysis::DefinitionLocation;
 use crate::analysis::LspModule;
+use crate::lsp::server::LoadContentsError::WrongScheme;
 use crate::syntax::AstModule;
 
+/// The request to get the file contents for a starlark: URI
+struct StarlarkFileContentsRequest {}
+
+impl lsp_types::request::Request for StarlarkFileContentsRequest {
+    type Params = StarlarkFileContentsParams;
+    type Result = StarlarkFileContentsResponse;
+    const METHOD: &'static str = "starlark/fileContents";
+}
+
+/// Params to get the file contents for a starlark: URI.
+#[derive(Debug, PartialEq, Serialize, Deserialize, Clone)]
+#[serde(rename_all = "camelCase")] // camelCase to match idioms in LSP spec / typescript land.
+struct StarlarkFileContentsParams {
+    uri: LspUrl,
+}
+
+/// The contents of a starlark: URI if available.
+#[derive(Debug, PartialEq, Serialize, Deserialize, Clone)]
+#[serde(rename_all = "camelCase")] // camelCase to match idioms in LSP spec / typescript land.
+struct StarlarkFileContentsResponse {
+    contents: Option<String>,
+}
+
 /// Errors that can happen when converting LspUrl and Url to/from each other.
 #[derive(thiserror::Error, Debug)]
 pub enum LspUrlError {
@@ -370,6 +394,18 @@ impl<T: LspContext> Backend<T> {
         self.send_response(new_response(id, self.find_definition(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 {
+            LspUrl::Starlark(_) => self
+                .context
+                .get_load_contents(&params.uri)
+                .map(|contents| StarlarkFileContentsResponse { contents }),
+            _ => Err(WrongScheme("starlark:".to_owned(), params.uri).into()),
+        };
+        self.send_response(new_response(id, response));
+    }
+
     fn resolve_load_path(&self, path: &str, current_uri: &LspUrl) -> anyhow::Result<LspUrl> {
         match current_uri {
             LspUrl::File(_) => self.context.resolve_load(path, current_uri),
@@ -528,6 +564,8 @@ impl<T: LspContext> Backend<T> {
                     //            be handled client side.
                     if let Some(params) = as_request::<GotoDefinition>(&req) {
                         self.goto_definition(req.id, params);
+                    } else if let Some(params) = as_request::<StarlarkFileContentsRequest>(&req) {
+                        self.get_starlark_file_contents(req.id, params);
                     } else if self.connection.handle_shutdown(&req)? {
                         return Ok(());
                     }
@@ -685,6 +723,10 @@ mod test {
     use crate::analysis::FixtureWithRanges;
     use crate::codemap::ResolvedSpan;
     use crate::lsp::server::LspServerSettings;
+    use crate::lsp::server::LspUrl;
+    use crate::lsp::server::StarlarkFileContentsParams;
+    use crate::lsp::server::StarlarkFileContentsRequest;
+    use crate::lsp::server::StarlarkFileContentsResponse;
     use crate::lsp::test::TestServer;
 
     fn goto_definition_request(
@@ -1488,4 +1530,29 @@ mod test {
         assert!(goto_definition_enabled);
         Ok(())
     }
+
+    #[test]
+    fn returns_starlark_file_contents() -> anyhow::Result<()> {
+        let mut server = TestServer::new()?;
+
+        let uri = LspUrl::try_from(Url::parse("starlark:/native/builtin.bzl")?)?;
+        let req = server.new_request::<StarlarkFileContentsRequest>(StarlarkFileContentsParams {
+            uri: uri.clone(),
+        });
+        let request_id = server.send_request(req)?;
+        let response = server.get_response::<StarlarkFileContentsResponse>(request_id)?;
+        assert_eq!(
+            server.docs_as_code(&uri).unwrap(),
+            response.contents.unwrap()
+        );
+
+        let req = server.new_request::<StarlarkFileContentsRequest>(StarlarkFileContentsParams {
+            uri: LspUrl::try_from(Url::parse("starlark:/native/not_builtin.bzl")?)?,
+        });
+        let request_id = server.send_request(req)?;
+        let response = server.get_response::<StarlarkFileContentsResponse>(request_id)?;
+        assert!(response.contents.is_none());
+
+        Ok(())
+    }
 }
diff --git a/starlark/src/lsp/test.rs b/starlark/src/lsp/test.rs
@@ -19,6 +19,7 @@ use std::collections::hash_map::Entry;
 use std::collections::HashMap;
 use std::collections::HashSet;
 use std::collections::VecDeque;
+use std::path::Path;
 use std::path::PathBuf;
 use std::sync::Arc;
 use std::sync::RwLock;
@@ -68,6 +69,12 @@ use crate::lsp::server::ResolveLoadError;
 use crate::lsp::server::StringLiteralResult;
 use crate::syntax::AstModule;
 use crate::syntax::Dialect;
+use crate::values::docs::render_docs_as_code;
+use crate::values::docs::Doc;
+use crate::values::docs::DocItem;
+use crate::values::docs::Function;
+use crate::values::docs::Identifier;
+use crate::values::docs::Location;
 
 /// Get the path from a URL, trimming off things like the leading slash that gets
 /// appended in some windows test environments.
@@ -99,13 +106,14 @@ pub(crate) enum TestServerError {
 struct TestServerContext {
     file_contents: Arc<RwLock<HashMap<PathBuf, String>>>,
     dirs: Arc<RwLock<HashSet<PathBuf>>>,
+    builtin_docs: Arc<HashMap<LspUrl, String>>,
 }
 
 impl LspContext for TestServerContext {
     fn parse_file_with_contents(&self, uri: &LspUrl, content: String) -> LspEvalResult {
         match uri {
-            LspUrl::File(uri) => {
-                match AstModule::parse(&uri.to_string_lossy(), content, &Dialect::Extended) {
+            LspUrl::File(path) | LspUrl::Starlark(path) => {
+                match AstModule::parse(&path.to_string_lossy(), content, &Dialect::Extended) {
                     Ok(ast) => {
                         let diagnostics = ast.lint(None).into_map(|l| EvalMessage::from(l).into());
                         LspEvalResult {
@@ -114,7 +122,7 @@ impl LspContext for TestServerContext {
                         }
                     }
                     Err(e) => {
-                        let diagnostics = vec![EvalMessage::from_anyhow(uri, &e).into()];
+                        let diagnostics = vec![EvalMessage::from_anyhow(path, &e).into()];
                         LspEvalResult {
                             diagnostics,
                             ast: None,
@@ -196,6 +204,7 @@ impl LspContext for TestServerContext {
                     (false, _) => Err(LoadContentsError::NotAbsolute(uri.clone()).into()),
                 }
             }
+            LspUrl::Starlark(_) => Ok(self.builtin_docs.get(uri).cloned()),
             _ => Ok(None),
         }
     }
@@ -223,6 +232,8 @@ pub struct TestServer {
     dirs: Arc<RwLock<HashSet<PathBuf>>>,
     /// If it's been received, the response payload for initialization.
     initialize_response: Option<InitializeResult>,
+    /// Documentation for built in symbols.
+    builtin_docs: Arc<HashMap<LspUrl, String>>,
 }
 
 impl Drop for TestServer {
@@ -255,6 +266,49 @@ impl Drop for TestServer {
 }
 
 impl TestServer {
+    pub(crate) fn docs_as_code(&self, uri: &LspUrl) -> Option<String> {
+        self.builtin_docs.get(uri).cloned()
+    }
+
+    /// A static set of "builtins" to use for testing
+    fn testing_builtins(root: &Path) -> anyhow::Result<HashMap<LspUrl, Vec<Doc>>> {
+        let prelude_path = root.join("dir/prelude.bzl");
+        let ret = hashmap! {
+            LspUrl::try_from(Url::parse("starlark:/native/builtin.bzl")?)? => vec![
+                Doc {
+                    id: Identifier {
+                        name: "native_function1".to_owned(),
+                        location: None,
+                    },
+                    item: DocItem::Function(Function::default()),
+                    custom_attrs: Default::default(),
+                },
+                Doc {
+                    id: Identifier {
+                        name: "native_function2".to_owned(),
+                        location: None,
+                    },
+                    item: DocItem::Function(Function::default()),
+                    custom_attrs: Default::default(),
+                },
+            ],
+            LspUrl::try_from(Url::from_file_path(prelude_path).unwrap())? => vec![
+                Doc {
+                    id: Identifier {
+                        name: "prelude_function".to_owned(),
+                        location: Some(Location {
+                            path: "//dir/prelude.bzl".to_owned(),
+                            position: None,
+                        }),
+                    },
+                    item: DocItem::Function(Function::default()),
+                    custom_attrs: Default::default(),
+                },
+            ]
+        };
+        Ok(ret)
+    }
+
     /// Generate a new request ID
     fn next_request_id(&mut self) -> RequestId {
         self.req_counter += 1;
@@ -281,11 +335,25 @@ impl TestServer {
     pub(crate) fn new_with_settings(settings: Option<LspServerSettings>) -> anyhow::Result<Self> {
         let (server_connection, client_connection) = Connection::memory();
 
-        let file_contents = Arc::new(RwLock::new(HashMap::new()));
+        let builtin_docs = Arc::new(
+            Self::testing_builtins(&std::env::current_dir()?)?
+                .into_iter()
+                .map(|(u, ds)| (u, render_docs_as_code(&ds).join("\n\n")))
+                .collect::<HashMap<_, _>>(),
+        );
+        let prelude_file_contents = builtin_docs
+            .iter()
+            .filter_map(|(u, d)| match u {
+                LspUrl::File(p) => Some((p.clone(), d.clone())),
+                _ => None,
+            })
+            .collect();
+        let file_contents = Arc::new(RwLock::new(prelude_file_contents));
         let dirs = Arc::new(RwLock::new(HashSet::new()));
         let ctx = TestServerContext {
             file_contents: file_contents.dupe(),
             dirs: dirs.dupe(),
+            builtin_docs: builtin_docs.dupe(),
         };
 
         let server_thread = std::thread::spawn(|| {
@@ -305,6 +373,7 @@ impl TestServer {
             file_contents,
             dirs,
             initialize_response: None,
+            builtin_docs,
         };
         ret.initialize(settings)
     }
