added bazel workspace lsp support

DanielEliad · DanielEliad · commit bc34622c50ce · 2022-08-21T20:45:17.000+03:00
diff --git a/starlark/bin/eval.rs b/starlark/bin/eval.rs
@@ -25,6 +25,7 @@ use std::path::PathBuf;
 use gazebo::prelude::*;
 use itertools::Either;
 use lsp_types::Diagnostic;
+use lsp_types::Range;
 use lsp_types::Url;
 use starlark::environment::FrozenModule;
 use starlark::environment::Globals;
@@ -50,6 +51,13 @@ pub(crate) enum ContextMode {
     Run,
 }
 
+#[derive(Debug, Clone)]
+pub(crate) struct BazelInfo {
+    pub(crate) workspace_root: PathBuf,
+    pub(crate) output_base: PathBuf,
+    pub(crate) execroot: PathBuf,
+}
+
 #[derive(Debug)]
 pub(crate) struct Context {
     pub(crate) mode: ContextMode,
@@ -58,6 +66,7 @@ pub(crate) struct Context {
     pub(crate) module: Option<Module>,
     pub(crate) builtin_docs: HashMap<LspUrl, String>,
     pub(crate) builtin_symbols: HashMap<String, LspUrl>,
+    pub(crate) bazel_info: Option<BazelInfo>,
 }
 
 /// The outcome of evaluating (checking, parsing or running) given starlark code.
@@ -110,6 +119,7 @@ impl Context {
             module,
             builtin_docs,
             builtin_symbols,
+            bazel_info: None,
         })
     }
 
@@ -241,6 +251,147 @@ impl Context {
     }
 }
 
+fn handle_local_bazel_repository(
+    info: &Option<BazelInfo>,
+    path: &str,
+    path_buf: PathBuf,
+    current_file_dir: &PathBuf,
+) -> Result<PathBuf, ResolveLoadError> {
+    match info {
+        None => Err(ResolveLoadError::MissingBazelInfo(path_buf)),
+        Some(info) => {
+            let malformed_err = ResolveLoadError::PathMalformed(path_buf.clone());
+            let mut split_parts = path.trim_start_match("//").split(':');
+            let package = split_parts.next().ok_or(malformed_err.clone())?;
+            let target = split_parts.next().ok_or(malformed_err.clone())?;
+            match split_parts.next().is_some() {
+                true => Err(malformed_err.clone()),
+                false => {
+                    let file_path = PathBuf::from(package).join(target);
+                    let root_path = current_file_dir
+                        .ancestors()
+                        .find(|a| match a.read_dir() {
+                            Ok(mut entries) => entries
+                                .find(|f| match f {
+                                    Ok(f) => ["MODULE.bazel", "WORKSPACE", "WORKSPACE.bazel"]
+                                        .contains(&f.file_name().to_str().unwrap_or("")),
+                                    _ => false,
+                                })
+                                .is_some(),
+                            _ => false,
+                        })
+                        .unwrap_or(&info.workspace_root);
+                    Ok(root_path.join(file_path))
+                }
+            }
+        }
+    }
+}
+fn handle_remote_bazel_repository(
+    info: &Option<BazelInfo>,
+    path: &str,
+    path_buf: PathBuf,
+) -> Result<PathBuf, ResolveLoadError> {
+    match info {
+        None => Err(ResolveLoadError::MissingBazelInfo(path_buf)),
+        Some(info) => {
+            let malformed_err = ResolveLoadError::PathMalformed(path_buf.clone());
+            let mut split_parts = path.trim_start_match("@").split("//");
+            let repository = split_parts.next().ok_or(malformed_err.clone())?;
+            split_parts = split_parts.next().ok_or(malformed_err.clone())?.split(":");
+            let package = split_parts.next().ok_or(malformed_err.clone())?;
+            let target = split_parts.next().ok_or(malformed_err.clone())?;
+            match split_parts.next().is_some() {
+                true => Err(malformed_err.clone()),
+                false => {
+                    let execroot_dirname =
+                        info.execroot.file_name().ok_or(malformed_err.clone())?;
+
+                    if repository == execroot_dirname {
+                        Ok(info.workspace_root.join(package).join(target))
+                    } else {
+                        Ok(info
+                            .output_base
+                            .join("external")
+                            .join(repository)
+                            .join(package)
+                            .join(target))
+                    }
+                }
+            }
+        }
+    }
+}
+
+fn get_relative_file(
+    current_file_dir: &PathBuf,
+    path: &str,
+    pathbuf: PathBuf,
+) -> Result<PathBuf, ResolveLoadError> {
+    let malformed_err = ResolveLoadError::MissingCurrentFilePath(pathbuf.clone());
+    let mut split_parts = path.split(":");
+    let package = split_parts.next().ok_or(malformed_err.clone())?;
+    let target = split_parts.next().ok_or(malformed_err.clone())?;
+    match split_parts.next().is_some() {
+        true => Err(malformed_err.clone()),
+        false => Ok(current_file_dir.join(package).join(target)),
+    }
+}
+fn label_into_file(
+    bazel_info: &Option<BazelInfo>,
+    path: &str,
+    current_file_path: &PathBuf,
+) -> Result<PathBuf, ResolveLoadError> {
+    let current_file_dir = current_file_path.parent();
+    let path_buf = PathBuf::from(path);
+
+    if path.starts_with("@") {
+        handle_remote_bazel_repository(bazel_info, path, path_buf.clone())
+    } else if path.starts_with("//") {
+        handle_local_bazel_repository(bazel_info, path, path_buf.clone(), current_file_path)
+    } else if path.contains(":") {
+        match current_file_dir {
+            Some(dir) => get_relative_file(&dir.to_path_buf(), path, path_buf.clone()),
+            None => Err(ResolveLoadError::MissingCurrentFilePath(path_buf)),
+        }
+    } else {
+        match (current_file_dir, path_buf.is_absolute()) {
+            (_, true) => Ok(path_buf),
+            (Some(current_file_dir), false) => Ok(current_file_dir.join(&path_buf)),
+            (None, false) => Err(ResolveLoadError::MissingCurrentFilePath(path_buf)),
+        }
+    }
+}
+
+fn replace_fake_file_with_build_target(fake_file: PathBuf) -> Option<PathBuf> {
+    fake_file.parent().and_then(|p| {
+        let build = p.join("BUILD");
+        let build_bazel = p.join("BUILD.bazel");
+        if build.exists() {
+            Some(build)
+        } else if build_bazel.exists() {
+            Some(build_bazel)
+        } else {
+            None
+        }
+    })
+}
+
+fn find_location_in_build_file(
+    info: Option<BazelInfo>,
+    literal: String,
+    current_file_pathbuf: PathBuf,
+    ast: &AstModule,
+) -> anyhow::Result<Option<Range>> {
+    let resolved_file = label_into_file(&info, literal.as_str(), &current_file_pathbuf)?;
+    let basename = resolved_file.file_name().and_then(|f| f.to_str()).ok_or(
+        ResolveLoadError::ResolvedDoesNotExist(resolved_file.clone()),
+    )?;
+    let resolved_span = ast
+        .find_function_call_with_name(basename)
+        .and_then(|r| Some(Range::from(r)));
+    Ok(resolved_span)
+}
 impl LspContext for Context {
     fn parse_file_with_contents(&self, uri: &LspUrl, content: String) -> LspEvalResult {
         match uri {
@@ -257,16 +408,23 @@ impl LspContext for Context {
     }
 
     fn resolve_load(&self, path: &str, current_file: &LspUrl) -> anyhow::Result<LspUrl> {
-        let path = PathBuf::from(path);
         match current_file {
             LspUrl::File(current_file_path) => {
-                let current_file_dir = current_file_path.parent();
-                let absolute_path = match (current_file_dir, path.is_absolute()) {
-                    (_, true) => Ok(path),
-                    (Some(current_file_dir), false) => Ok(current_file_dir.join(&path)),
-                    (None, false) => Err(ResolveLoadError::MissingCurrentFilePath(path)),
+                let mut resolved_file = label_into_file(&self.bazel_info, path, current_file_path)?;
+                resolved_file = match resolved_file.canonicalize() {
+                    Ok(f) => {
+                        if f.exists() {
+                            Ok(f)
+                        } else {
+                            replace_fake_file_with_build_target(resolved_file.clone())
+                                .ok_or(ResolveLoadError::ResolvedDoesNotExist(resolved_file))
+                        }
+                    }
+                    _ => replace_fake_file_with_build_target(resolved_file.clone())
+                        .ok_or(ResolveLoadError::ResolvedDoesNotExist(resolved_file)),
                 }?;
-                Ok(Url::from_file_path(absolute_path).unwrap().try_into()?)
+
+                Ok(Url::from_file_path(resolved_file).unwrap().try_into()?)
             }
             _ => Err(
                 ResolveLoadError::WrongScheme("file://".to_owned(), current_file.clone()).into(),
@@ -279,11 +437,26 @@ impl LspContext for Context {
         literal: &str,
         current_file: &LspUrl,
     ) -> anyhow::Result<Option<StringLiteralResult>> {
+        let current_file_pathbuf = current_file.path().to_path_buf();
         self.resolve_load(literal, current_file).map(|url| {
-            Some(StringLiteralResult {
-                url,
-                location_finder: None,
-            })
+            let p = url.path();
+            // TODO: we can always give literal location finder
+            // TODO: but if its a file it will always try to resolve the location but won't be able to and an error will be printed
+            if p.ends_with("BUILD") || p.ends_with("BUILD.bazel") {
+                let info = self.bazel_info.clone();
+                let literal_copy = literal.to_owned();
+                Some(StringLiteralResult {
+                    url,
+                    location_finder: Some(box |ast: &AstModule, _url| {
+                        find_location_in_build_file(info, literal_copy, current_file_pathbuf, ast)
+                    }),
+                })
+            } else {
+                Some(StringLiteralResult {
+                    url,
+                    location_finder: None,
+                })
+            }
         })
     }
 
diff --git a/starlark/bin/main.rs b/starlark/bin/main.rs
@@ -30,9 +30,11 @@ use std::ffi::OsStr;
 use std::fmt;
 use std::fmt::Display;
 use std::path::PathBuf;
+use std::process::Command;
 use std::sync::Arc;
 
 use anyhow::anyhow;
+use eval::BazelInfo;
 use eval::Context;
 use gazebo::prelude::*;
 use itertools::Either;
@@ -214,6 +216,28 @@ fn interactive(ctx: &Context) -> anyhow::Result<()> {
     }
 }
 
+fn get_bazel_info_path(path_type: &str) -> anyhow::Result<PathBuf> {
+    let output = Command::new("bazel").arg("info").arg(path_type).output()?;
+
+    if !output.status.success() {
+        return Err(anyhow!("Failed running command bazel info"));
+    }
+    let s = std::str::from_utf8(output.stdout.as_slice())?;
+    Ok(PathBuf::from(s.trim()))
+}
+
+fn get_bazel_info() -> Option<BazelInfo> {
+    let workspace_root = get_bazel_info_path("workspace").ok()?;
+    let output_base = get_bazel_info_path("output_base").ok()?;
+    let execroot = get_bazel_info_path("execution_root").ok()?;
+
+    Some(BazelInfo {
+        workspace_root,
+        output_base,
+        execroot,
+    })
+}
+
 fn main() -> anyhow::Result<()> {
     gazebo::terminate_on_panic();
 
@@ -242,6 +266,7 @@ fn main() -> anyhow::Result<()> {
 
         if args.lsp {
             ctx.mode = ContextMode::Check;
+            ctx.bazel_info = get_bazel_info();
             lsp::server::stdio_server(ctx)?;
         } else if is_interactive {
             interactive(&ctx)?;
diff --git a/starlark/src/lsp/server.rs b/starlark/src/lsp/server.rs
@@ -292,15 +292,24 @@ pub trait LspContext {
 }
 
 /// Errors when [`LspContext::resolve_load()`] cannot resolve a given path.
-#[derive(thiserror::Error, Debug)]
+#[derive(thiserror::Error, Debug, Clone)]
 pub enum ResolveLoadError {
+    /// Attempted to resolve a load but the path was malformed
+    #[error("path `{}` provided, but was malformed", .0.display())]
+    PathMalformed(PathBuf),
+    /// Attempted to resolve a bazel dependent load but no bazel info could be found
+    #[error("path `{}` provided, but bazel info could not be determined", .0.display())]
+    MissingBazelInfo(PathBuf),
     /// Attempted to resolve a relative path, but no current_file_path was provided,
     /// so it is not known what to resolve the path against.
     #[error("Relative path `{}` provided, but current_file_path could not be determined", .0.display())]
     MissingCurrentFilePath(PathBuf),
     /// The scheme provided was not correct or supported.
     #[error("Url `{}` was expected to be of type `{}`", .1, .0)]
     WrongScheme(String, LspUrl),
+    /// Resolved Loaded file does not exist.
+    #[error("Resolved file `{}` did not exist", .0.display())]
+    ResolvedDoesNotExist(PathBuf),
 }
 
 /// Errors when loading contents of a starlark program.
@@ -494,8 +503,7 @@ impl<T: LspContext> Backend<T> {
                         }) => {
                             // If there's an error loading the file to parse it, at least
                             // try to get to the file.
-                            let target_range = self
-                                .get_ast_or_load_from_disk(&url)
+                            self.get_ast_or_load_from_disk(&url)
                                 .and_then(|ast| match ast {
                                     Some(module) => location_finder(&module.ast, &url),
                                     None => Ok(None),
@@ -504,13 +512,20 @@ impl<T: LspContext> Backend<T> {
                                     eprintln!("Error jumping to definition: {:#}", e);
                                 })
                                 .unwrap_or_default()
-                                .unwrap_or_default();
-                            Some(LocationLink {
-                                origin_selection_range: Some(source.into()),
-                                target_uri: url.try_into()?,
-                                target_range,
-                                target_selection_range: target_range,
-                            })
+                                .and_then(|target_range| {
+                                    Some(LocationLink {
+                                        origin_selection_range: Some(source.into()),
+                                        target_uri: url.clone().try_into().ok()?,
+                                        target_range,
+                                        target_selection_range: target_range,
+                                    })
+                                })
+                                .or(Some(LocationLink {
+                                    origin_selection_range: Some(source.into()),
+                                    target_uri: url.try_into()?,
+                                    target_range: Range::default(),
+                                    target_selection_range: Range::default(),
+                                }))
                         }
                         Some(StringLiteralResult {
                             url,
