Refactor source root handling

lsp/src/convert.rs

@@ -1,4 +1,4 @@
-use crate::{LineMap, StateSnapshot, Vfs, VfsPath};
+use crate::{LineMap, StateSnapshot, Vfs};
 use lsp_types::{
     self as lsp, DiagnosticSeverity, DiagnosticTag, Location, Position, Range,
     TextDocumentPositionParams,
@@ -10,17 +10,17 @@ pub(crate) fn from_file_pos(
     snap: &StateSnapshot,
     params: &TextDocumentPositionParams,
 ) -> Option<FilePos> {
-    let path = VfsPath::try_from(&params.text_document.uri).ok()?;
     let vfs = snap.vfs.read().unwrap();
-    let (file, line_map) = vfs.get(&path)?;
+    let file = vfs.get_file_for_uri(&params.text_document.uri)?;
+    let line_map = vfs.get_line_map(file)?;
     let pos = line_map.pos(params.position.line, params.position.character);
     Some(FilePos::new(file, pos))
 }
 
 pub(crate) fn to_location(vfs: &Vfs, frange: InFile<TextRange>) -> Option<Location> {
-    let url = vfs.file_path(frange.file_id)?.try_into().ok()?;
-    let line_map = vfs.file_line_map(frange.file_id)?;
-    Some(Location::new(url, to_range(line_map, frange.value)))
+    let uri = vfs.get_uri_for_file(frange.file_id)?;
+    let line_map = vfs.get_line_map(frange.file_id)?;
+    Some(Location::new(uri, to_range(line_map, frange.value)))
 }
 
 pub(crate) fn to_range(line_map: &LineMap, range: TextRange) -> Range {

lsp/src/handler.rs

@@ -59,7 +59,7 @@ pub(crate) fn completion(
     let fpos = convert::from_file_pos(&snap, &params.text_document_position)?;
     let items = snap.analysis.completions(fpos).ok()??;
     let vfs = snap.vfs.read().unwrap();
-    let line_map = vfs.file_line_map(fpos.file_id)?;
+    let line_map = vfs.get_line_map(fpos.file_id)?;
     let items = items
         .into_iter()
         .filter_map(|item| convert::to_completion_item(line_map, item))

lsp/src/lib.rs

@@ -3,8 +3,12 @@ mod handler;
 mod state;
 mod vfs;
 
+use lsp_types::InitializeParams;
+use std::env;
+use std::path::PathBuf;
+
 pub(crate) use state::{State, StateSnapshot};
-pub(crate) use vfs::{LineMap, Vfs, VfsPath};
+pub(crate) use vfs::{LineMap, Vfs};
 
 use anyhow::Result;
 use lsp_server::Connection;
@@ -14,7 +18,18 @@ pub fn main_loop(conn: Connection) -> Result<()> {
         conn.initialize(serde_json::to_value(&handler::server_capabilities()).unwrap())?;
     log::info!("Init params: {}", init_params);
 
-    let mut state = State::new(conn.sender.clone());
+    let init_params = serde_json::from_value::<InitializeParams>(init_params)?;
+    let workspace_path = (|| -> Option<PathBuf> {
+        if let Some(folders) = &init_params.workspace_folders {
+            return folders.get(0)?.uri.to_file_path().ok();
+        }
+        if let Some(uri) = &init_params.root_uri {
+            return uri.to_file_path().ok();
+        }
+        Some(env::current_dir().expect("Cannot get current directory"))
+    })();
+
+    let mut state = State::new(conn.sender.clone(), workspace_path);
     state.run(conn.receiver)?;
 
     log::info!("Leaving main loop");

lsp/src/state.rs

@@ -1,10 +1,11 @@
-use crate::{convert, handler, Vfs, VfsPath};
+use crate::{convert, handler, Vfs};
 use anyhow::{bail, Result};
 use crossbeam_channel::{Receiver, Sender};
 use lsp_server::{ErrorCode, Message, Notification, Request, Response};
 use lsp_types::notification::Notification as _;
 use lsp_types::{notification as notif, request as req, PublishDiagnosticsParams, Url};
 use nil::{Analysis, AnalysisHost};
+use std::path::PathBuf;
 use std::sync::{Arc, RwLock};
 
 const MAX_DIAGNOSTICS_CNT: usize = 128;
@@ -17,10 +18,11 @@ pub struct State {
 }
 
 impl State {
-    pub fn new(responder: Sender<Message>) -> Self {
+    pub fn new(responder: Sender<Message>, workspace_root: Option<PathBuf>) -> Self {
+        let vfs = Vfs::new(workspace_root.unwrap_or_else(|| PathBuf::from("/")));
         Self {
             host: Default::default(),
-            vfs: Default::default(),
+            vfs: Arc::new(RwLock::new(vfs)),
             sender: responder,
             is_shutdown: false,
         }
@@ -93,36 +95,31 @@ impl State {
     }
 
     fn set_vfs_file_content(&mut self, uri: &Url, text: Option<String>) {
-        if let Ok(path) = VfsPath::try_from(uri) {
-            let mut vfs = self.vfs.write().unwrap();
-            vfs.set_file_content(path, text);
+        let mut vfs = self.vfs.write().unwrap();
+        let file = vfs.set_uri_content(uri, text);
 
-            let change = vfs.take_change();
-            log::debug!("Files changed: {:?}", change);
-            self.host.apply_change(change.clone());
+        let change = vfs.take_change();
+        log::debug!("Change: {:?}", change);
+        self.host.apply_change(change);
 
-            // Currently we push down changes immediately.
-            assert_eq!(change.file_changes.len(), 1);
-            let (file, text) = &change.file_changes[0];
-            let diagnostics = vfs
-                .file_line_map(*file)
-                .and_then(|line_map| {
-                    let _ = text.as_deref()?;
-                    let diags = self.host.snapshot().diagnostics(*file).ok()?;
-                    let diags = diags
-                        .into_iter()
-                        .take(MAX_DIAGNOSTICS_CNT)
-                        .filter_map(|diag| convert::to_diagnostic(line_map, diag))
-                        .collect::<Vec<_>>();
-                    Some(diags)
-                })
-                .unwrap_or_default();
-            self.send_notification::<notif::PublishDiagnostics>(PublishDiagnosticsParams {
-                uri: uri.clone(),
-                diagnostics,
-                version: None,
-            });
-        }
+        // Currently we push down changes immediately.
+        let diagnostics = file
+            .and_then(|file| {
+                let line_map = vfs.get_line_map(file)?;
+                let diags = self.host.snapshot().diagnostics(file).ok()?;
+                let diags = diags
+                    .into_iter()
+                    .take(MAX_DIAGNOSTICS_CNT)
+                    .filter_map(|diag| convert::to_diagnostic(line_map, diag))
+                    .collect::<Vec<_>>();
+                Some(diags)
+            })
+            .unwrap_or_default();
+        self.send_notification::<notif::PublishDiagnostics>(PublishDiagnosticsParams {
+            uri: uri.clone(),
+            diagnostics,
+            version: None,
+        });
     }
 }
 

lsp/src/vfs.rs

@@ -1,82 +1,101 @@
-use indexmap::IndexMap;
 use lsp_types::Url;
-use nil::{Change, FileId};
+use nil::{Change, FileId, FileSet, SourceRoot, VfsPath};
 use std::collections::HashMap;
 use std::path::PathBuf;
 use std::sync::Arc;
 use std::{fmt, mem};
 use text_size::TextSize;
 
-#[derive(Debug, Clone, PartialEq, Eq, Hash)]
-pub struct VfsPath(PathBuf);
-
-impl<'a> TryFrom<&'a Url> for VfsPath {
-    type Error = ();
-    fn try_from(url: &'a Url) -> Result<Self, Self::Error> {
-        url.to_file_path().map(Self)
-    }
-}
-
-impl<'a> From<&'a VfsPath> for Url {
-    fn from(path: &'a VfsPath) -> Self {
-        Url::from_file_path(&path.0).unwrap()
-    }
-}
-
-#[derive(Default)]
 pub struct Vfs {
-    files: IndexMap<VfsPath, Option<(Arc<str>, LineMap)>>,
+    // FIXME: Currently this list is append-only.
+    files: Vec<Option<(Arc<str>, LineMap)>>,
+    local_root: PathBuf,
+    local_file_set: FileSet,
+    root_changed: bool,
     change: Change,
 }
 
 impl fmt::Debug for Vfs {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        struct Files(Vec<(usize, VfsPath)>);
-        impl fmt::Debug for Files {
-            fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-                f.debug_map()
-                    .entries(self.0.iter().map(|(k, v)| (k, v)))
-                    .finish()
-            }
-        }
-
-        let files = Files(self.files.keys().cloned().enumerate().collect());
         f.debug_struct("Vfs")
-            .field("files", &files)
+            .field("file_cnt", &self.files.len())
+            .field("local_root", &self.local_root)
             .finish_non_exhaustive()
     }
 }
 
 impl Vfs {
-    pub fn file_path(&self, file_id: FileId) -> Option<&VfsPath> {
-        self.files.get_index(file_id.0 as _).map(|(path, _)| path)
+    pub fn new(local_root: PathBuf) -> Self {
+        Self {
+            files: Vec::new(),
+            local_root,
+            local_file_set: FileSet::default(),
+            root_changed: false,
+            change: Change::default(),
+        }
+    }
+
+    fn alloc_file_id(&mut self) -> FileId {
+        let id = u32::try_from(self.files.len()).expect("Length overflow");
+        self.files.push(None);
+        FileId(id)
+    }
+
+    fn uri_to_vpath(&self, uri: &Url) -> Option<VfsPath> {
+        let path = uri.to_file_path().ok()?;
+        let relative_path = path.strip_prefix(&self.local_root).ok()?;
+        VfsPath::from_path(relative_path)
+    }
+
+    pub fn set_uri_content(&mut self, uri: &Url, text: Option<String>) -> Option<FileId> {
+        let vpath = self.uri_to_vpath(uri)?;
+        let content = text.and_then(LineMap::normalize);
+        let (file, (text, line_map)) =
+            match (self.local_file_set.get_file_for_path(&vpath), content) {
+                (Some(file), None) => {
+                    self.local_file_set.remove_file(file);
+                    self.root_changed = true;
+                    self.files[file.0 as usize] = None;
+                    return None;
+                }
+                (None, None) => return None,
+                (Some(file), Some(content)) => (file, content),
+                (None, Some(content)) => {
+                    let file = self.alloc_file_id();
+                    self.local_file_set.insert(file, vpath);
+                    self.root_changed = true;
+                    (file, content)
+                }
+            };
+        let text = <Arc<str>>::from(text);
+        self.change.change_file(file, Some(text.clone()));
+        self.files[file.0 as usize] = Some((text, line_map));
+        Some(file)
+    }
+
+    pub fn get_file_for_uri(&self, uri: &Url) -> Option<FileId> {
+        let vpath = self.uri_to_vpath(uri)?;
+        self.local_file_set.get_file_for_path(&vpath)
+    }
+
+    pub fn get_uri_for_file(&self, file: FileId) -> Option<Url> {
+        let vpath = self.local_file_set.get_path_for_file(file)?.as_str();
+        assert!(!vpath.is_empty(), "Root is a directory");
+        let path = self.local_root.join(vpath.strip_prefix('/')?);
+        Url::from_file_path(path).ok()
     }
 
     pub fn take_change(&mut self) -> Change {
-        mem::take(&mut self.change)
+        let mut change = mem::take(&mut self.change);
+        if self.root_changed {
+            self.root_changed = false;
+            change.set_roots(vec![SourceRoot::new_local(self.local_file_set.clone())]);
+        }
+        change
     }
 
-    pub fn set_file_content(&mut self, path: VfsPath, content: Option<String>) -> FileId {
-        let text_with_map = content
-            .and_then(LineMap::normalize)
-            .map(|(text, map)| (text.into(), map));
-        let text = text_with_map.as_ref().map(|(text, _)| Arc::clone(text));
-        let id = self.files.insert_full(path, text_with_map).0;
-        let file_id = FileId(u32::try_from(id).unwrap());
-        self.change.change_file(file_id, text);
-        file_id
-    }
-
-    pub fn get(&self, path: &VfsPath) -> Option<(FileId, &LineMap)> {
-        let (id, _, inner) = self.files.get_full(path)?;
-        let (_, line_map) = inner.as_ref()?;
-        Some((FileId(id as u32), line_map))
-    }
-
-    pub fn file_line_map(&self, file_id: FileId) -> Option<&LineMap> {
-        let (_, inner) = self.files.get_index(file_id.0 as usize)?;
-        let (_, line_map) = inner.as_ref()?;
-        Some(line_map)
+    pub fn get_line_map(&self, file_id: FileId) -> Option<&LineMap> {
+        Some(&self.files.get(file_id.0 as usize)?.as_ref()?.1)
     }
 }
 

src/base.rs

@@ -1,12 +1,115 @@
 use rowan::{TextRange, TextSize};
 use salsa::Durability;
+use std::collections::HashMap;
 use std::fmt;
+use std::path::Path;
 use std::sync::Arc;
-use syntax::Parse;
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
 pub struct FileId(pub u32);
 
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+pub struct SourceRootId(pub u32);
+
+/// An absolute path in format `(/.+)*`
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+pub struct VfsPath(String);
+
+impl VfsPath {
+    pub fn from_path(p: &Path) -> Option<Self> {
+        Self::new(p.to_str()?)
+    }
+
+    pub fn new(s: impl Into<String>) -> Option<Self> {
+        let mut s: String = s.into();
+        if s.is_empty() || s == "/" {
+            return Some(Self(String::new()));
+        }
+        if s.ends_with('/') || s.as_bytes().windows(2).any(|w| w == b"//") {
+            return None;
+        }
+        if !s.starts_with('/') {
+            s.insert(0, '/');
+        }
+        Some(Self(s))
+    }
+
+    pub fn push(&mut self, relative: &Self) {
+        self.0.push_str(&relative.0);
+    }
+
+    pub fn pop(&mut self) -> Option<()> {
+        self.0.truncate(self.0.rsplit_once('/')?.0.len());
+        Some(())
+    }
+
+    pub fn as_str(&self) -> &str {
+        &self.0
+    }
+}
+
+/// A set of [`VfsPath`]s identified by [`FileId`]s.
+#[derive(Default, Clone, PartialEq, Eq)]
+pub struct FileSet {
+    files: HashMap<VfsPath, FileId>,
+    paths: HashMap<FileId, VfsPath>,
+}
+
+impl FileSet {
+    pub fn insert(&mut self, file: FileId, path: VfsPath) {
+        self.files.insert(path.clone(), file);
+        self.paths.insert(file, path);
+    }
+
+    pub fn remove_file(&mut self, file: FileId) {
+        if let Some(path) = self.paths.remove(&file) {
+            self.files.remove(&path);
+        }
+    }
+
+    pub fn get_file_for_path(&self, path: &VfsPath) -> Option<FileId> {
+        self.files.get(path).copied()
+    }
+
+    pub fn get_path_for_file(&self, file: FileId) -> Option<&VfsPath> {
+        self.paths.get(&file)
+    }
+
+    pub fn iter(&self) -> impl Iterator<Item = (FileId, &'_ VfsPath)> + '_ {
+        self.paths.iter().map(|(&file, path)| (file, path))
+    }
+}
+
+impl fmt::Debug for FileSet {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_map().entries(&self.paths).finish()
+    }
+}
+
+/// A workspace unit, typically a Flake package.
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub struct SourceRoot {
+    file_set: FileSet,
+}
+
+impl SourceRoot {
+    pub fn new_local(file_set: FileSet) -> Self {
+        Self { file_set }
+    }
+
+    pub fn get_file_for_path(&self, path: &VfsPath) -> Option<FileId> {
+        self.file_set.get_file_for_path(path)
+    }
+
+    pub fn get_path_for_file(&self, file: FileId) -> Option<&VfsPath> {
+        self.file_set.get_path_for_file(file)
+    }
+
+    pub fn iter(&self) -> impl Iterator<Item = (FileId, &'_ VfsPath)> + '_ {
+        self.file_set.iter()
+    }
+}
+
 #[derive(Debug, PartialEq, Eq, Clone, Copy, Hash)]
 pub struct InFile<T> {
     pub file_id: FileId,
@@ -34,16 +137,16 @@ pub trait SourceDatabase {
     #[salsa::input]
     fn file_content(&self, file_id: FileId) -> Arc<str>;
 
-    fn parse(&self, file_id: FileId) -> Parse;
-}
+    #[salsa::input]
+    fn source_root(&self, id: SourceRootId) -> Arc<SourceRoot>;
 
-fn parse(db: &dyn SourceDatabase, file_id: FileId) -> Parse {
-    let content = db.file_content(file_id);
-    syntax::parse_file(&content)
+    #[salsa::input]
+    fn file_source_root(&self, file_id: FileId) -> SourceRootId;
 }
 
 #[derive(Default, Clone, PartialEq, Eq)]
 pub struct Change {
+    pub roots: Option<Vec<SourceRoot>>,
     pub file_changes: Vec<(FileId, Option<Arc<str>>)>,
 }
 
@@ -53,7 +156,11 @@ impl Change {
     }
 
     pub fn is_empty(&self) -> bool {
-        self.file_changes.is_empty()
+        self.roots.is_none() && self.file_changes.is_empty()
+    }
+
+    pub fn set_roots(&mut self, roots: Vec<SourceRoot>) {
+        self.roots = Some(roots);
     }
 
     pub fn change_file(&mut self, file_id: FileId, content: Option<Arc<str>>) {
@@ -61,10 +168,18 @@ impl Change {
     }
 
     pub(crate) fn apply(self, db: &mut dyn SourceDatabase) {
+        if let Some(roots) = self.roots {
+            u32::try_from(roots.len()).expect("Length overflow");
+            for (sid, root) in (0u32..).map(SourceRootId).zip(roots) {
+                for (fid, _) in root.iter() {
+                    db.set_file_source_root_with_durability(fid, sid, Durability::HIGH);
+                }
+                db.set_source_root_with_durability(sid, Arc::new(root), Durability::HIGH);
+            }
+        }
         for (file_id, content) in self.file_changes {
             let content = content.unwrap_or_else(|| String::new().into());
-            // TODO: Better guess of durability?
-            db.set_file_content_with_durability(file_id, content, Durability::HIGH);
+            db.set_file_content_with_durability(file_id, content, Durability::LOW);
         }
     }
 }
@@ -78,6 +193,7 @@ impl fmt::Debug for Change {
             .count();
         let cleared = self.file_changes.len() - modified;
         f.debug_struct("Change")
+            .field("roots", &self.roots.as_ref().map(|roots| roots.len()))
             .field("modified", &modified)
             .field("cleared", &cleared)
             .finish_non_exhaustive()

src/def/liveness.rs

@@ -178,8 +178,8 @@ impl<'a> Traversal<'a> {
 
 #[cfg(test)]
 mod tests {
-    use crate::def::DefDatabase;
     use crate::tests::TestDB;
+    use crate::DefDatabase;
     use expect_test::{expect, Expect};
 
     #[track_caller]

src/def/lower.rs

@@ -612,9 +612,8 @@ impl MergingEntry {
 #[cfg(test)]
 mod tests {
     use super::lower;
-    use crate::base::FileId;
-    use crate::def::DefDatabase;
     use crate::tests::TestDB;
+    use crate::{DefDatabase, FileId};
     use expect_test::{expect, Expect};
     use std::fmt::Write;
     use syntax::parse_file;

src/def/mod.rs

@@ -14,6 +14,7 @@ use smol_str::SmolStr;
 use std::collections::HashMap;
 use std::ops;
 use std::sync::Arc;
+use syntax::Parse;
 
 pub use self::liveness::LivenessCheck;
 pub use self::path::{Path, PathAnchor, PathData};
@@ -23,7 +24,9 @@ pub use syntax::ast::{BinaryOpKind as BinaryOp, UnaryOpKind as UnaryOp};
 #[salsa::query_group(DefDatabaseStorage)]
 pub trait DefDatabase: SourceDatabase {
     #[salsa::interned]
-    fn intern_path(&self, path: PathData) -> Path;
+    fn intern_path(&self, path_data: PathData) -> Path;
+
+    fn parse(&self, file_id: FileId) -> Parse;
 
     fn module_with_source_map(&self, file_id: FileId) -> (Arc<Module>, Arc<ModuleSourceMap>);
 
@@ -47,6 +50,11 @@ pub trait DefDatabase: SourceDatabase {
     fn liveness_check(&self, file_id: FileId) -> Arc<LivenessCheck>;
 }
 
+fn parse(db: &dyn DefDatabase, file_id: FileId) -> Parse {
+    let content = db.file_content(file_id);
+    syntax::parse_file(&content)
+}
+
 fn module_with_source_map(
     db: &dyn DefDatabase,
     file_id: FileId,

src/def/scope.rs

@@ -263,7 +263,7 @@ impl NameReferenceMap {
 #[cfg(test)]
 mod tests {
     use super::ScopeKind;
-    use crate::def::{AstPtr, DefDatabase, ResolveResult, SourceDatabase};
+    use crate::def::{AstPtr, DefDatabase, ResolveResult};
     use crate::tests::TestDB;
     use expect_test::{expect, Expect};
     use rowan::ast::AstNode;

src/ide/completion.rs

@@ -1,6 +1,6 @@
 use crate::builtin::BuiltinKind;
-use crate::def::{AstPtr, DefDatabase, NameDefKind};
-use crate::{builtin, FileId};
+use crate::def::{AstPtr, NameDefKind};
+use crate::{builtin, DefDatabase, FileId};
 use either::Either::{Left, Right};
 use rowan::ast::AstNode;
 use smol_str::SmolStr;

src/ide/diagnostics.rs

@@ -1,5 +1,4 @@
-use crate::def::DefDatabase;
-use crate::{Diagnostic, DiagnosticKind, FileId};
+use crate::{DefDatabase, Diagnostic, DiagnosticKind, FileId};
 use rowan::ast::AstNode;
 use syntax::ast;
 

src/ide/goto_definition.rs

@@ -1,6 +1,6 @@
 use super::NavigationTarget;
-use crate::def::{AstPtr, DefDatabase, ResolveResult};
-use crate::FileId;
+use crate::def::{AstPtr, ResolveResult};
+use crate::{DefDatabase, FileId};
 use rowan::ast::AstNode;
 use rowan::TextSize;
 use syntax::{ast, match_ast, SyntaxKind, T};

src/ide/references.rs

@@ -1,5 +1,5 @@
-use crate::def::{AstPtr, DefDatabase};
-use crate::{FileId, FileRange, InFile};
+use crate::def::AstPtr;
+use crate::{DefDatabase, FileId, FileRange, InFile};
 use rowan::ast::AstNode;
 use rowan::TextSize;
 use syntax::{ast, match_ast, SyntaxKind, T};

src/lib.rs

@@ -7,7 +7,11 @@ mod ide;
 #[cfg(test)]
 mod tests;
 
-pub use base::{Change, FileId, FilePos, FileRange, InFile};
+pub use base::{
+    Change, FileId, FilePos, FileRange, FileSet, InFile, SourceDatabase, SourceRoot, SourceRootId,
+    VfsPath,
+};
+pub use def::{DefDatabase, Module, ModuleSourceMap};
 pub use diagnostic::{Diagnostic, DiagnosticKind, Severity};
 pub use ide::{
     Analysis, AnalysisHost, CompletionItem, CompletionItemKind, NavigationTarget, RootDatabase,

src/tests.rs

@@ -1,6 +1,6 @@
-use crate::base::{SourceDatabase, SourceDatabaseStorage};
+use crate::base::SourceDatabaseStorage;
 use crate::def::DefDatabaseStorage;
-use crate::{Change, FileId};
+use crate::{Change, DefDatabase, FileId};
 use rowan::ast::AstNode;
 use rowan::TextSize;
 use std::ops;