diff --git a/Cargo.lock b/Cargo.lock
index bbc0a00..58c2164 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -220,7 +220,7 @@ dependencies = [
 [[package]]
 name = "rnix"
 version = "0.11.0-dev"
-source = "git+https://github.com/nix-community/rnix-parser.git#f759fce8ac005ed14d13219579e62c06bdfba678"
+source = "git+https://github.com/oxalica/rnix-parser.git#80fabb304defc3b488f21d225128bde0758d1eb3"
 dependencies = [
  "cbitset",
  "rowan",
@@ -247,9 +247,9 @@ checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2"
 
 [[package]]
 name = "salsa"
-version = "0.16.1"
+version = "0.17.0-pre.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "4b84d9f96071f3f3be0dc818eae3327625d8ebc95b58da37d6850724f31d3403"
+checksum = "9b223dccb46c32753144d0b51290da7230bb4aedcd8379d6b4c9a474c18bf17a"
 dependencies = [
  "crossbeam-utils",
  "indexmap",
@@ -264,9 +264,9 @@ dependencies = [
 
 [[package]]
 name = "salsa-macros"
-version = "0.16.0"
+version = "0.17.0-pre.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "cd3904a4ba0a9d0211816177fd34b04c7095443f8cdacd11175064fe541c8fe2"
+checksum = "ac6c2e352df550bf019da7b16164ed2f7fa107c39653d1311d1bba42d1582ff7"
 dependencies = [
  "heck",
  "proc-macro2",
diff --git a/Cargo.toml b/Cargo.toml
index 43b207e..9cf511b 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -5,8 +5,8 @@ edition = "2021"
 
 [dependencies]
 la-arena = "0.2.1"
-rnix = { git = "https://github.com/nix-community/rnix-parser.git" }
-salsa = "0.16.1"
+rnix = { git = "https://github.com/oxalica/rnix-parser.git" }
+salsa = "0.17.0-pre.2"
 
 [dev-dependencies]
 expect-test = "1.3.0"
diff --git a/src/def/mod.rs b/src/def/mod.rs
index ed44152..339611b 100644
--- a/src/def/mod.rs
+++ b/src/def/mod.rs
@@ -10,33 +10,29 @@ use std::sync::Arc;
 
 #[salsa::query_group(DefDatabaseStorage)]
 pub trait DefDatabase: SourceDatabase {
-    #[salsa::invoke(Module::module_with_source_map_query)]
     fn module_with_source_map(&self, file_id: FileId) -> (Arc<Module>, Arc<ModuleSourceMap>);
 
-    #[salsa::invoke(Module::module_query)]
     fn module(&self, file_id: FileId) -> Arc<Module>;
 }
 
+fn module_with_source_map(
+    db: &dyn DefDatabase,
+    file_id: FileId,
+) -> (Arc<Module>, Arc<ModuleSourceMap>) {
+    let root = db.parse(file_id);
+    let (module, source_map) = lower::lower(root.map(|ast| ast.root()));
+    (Arc::new(module), Arc::new(source_map))
+}
+
+fn module(db: &dyn DefDatabase, file_id: FileId) -> Arc<Module> {
+    db.module_with_source_map(file_id).0
+}
+
 #[derive(Default, Debug, Clone, PartialEq, Eq)]
 pub struct Module {
     exprs: Arena<Expr>,
 }
 
-impl Module {
-    fn module_with_source_map_query(
-        db: &dyn DefDatabase,
-        file_id: FileId,
-    ) -> (Arc<Module>, Arc<ModuleSourceMap>) {
-        let root = db.parse(file_id);
-        let (module, source_map) = lower::lower(root);
-        (Arc::new(module), Arc::new(source_map))
-    }
-
-    fn module_query(db: &dyn DefDatabase, file_id: FileId) -> Arc<Module> {
-        db.module_with_source_map(file_id).0
-    }
-}
-
 impl ops::Index<ExprId> for Module {
     type Output = Expr;
     fn index(&self, index: ExprId) -> &Self::Output {
diff --git a/src/ide.rs b/src/ide.rs
new file mode 100644
index 0000000..fbc77f4
--- /dev/null
+++ b/src/ide.rs
@@ -0,0 +1,24 @@
+use crate::{def::DefDatabaseStorage, source::SourceDatabaseStorage};
+use std::fmt;
+
+#[salsa::database(SourceDatabaseStorage, DefDatabaseStorage)]
+#[derive(Default)]
+pub struct RootDatabase {
+    storage: salsa::Storage<Self>,
+}
+
+impl salsa::Database for RootDatabase {}
+
+impl salsa::ParallelDatabase for RootDatabase {
+    fn snapshot(&self) -> salsa::Snapshot<Self> {
+        salsa::Snapshot::new(RootDatabase {
+            storage: self.storage.snapshot(),
+        })
+    }
+}
+
+impl fmt::Debug for RootDatabase {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("RootDatabase").finish_non_exhaustive()
+    }
+}
diff --git a/src/lib.rs b/src/lib.rs
index 480e650..e50ac16 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,15 +1,8 @@
 mod def;
+mod ide;
 mod source;
 
 #[cfg(test)]
 mod tests;
 
-use crate::{def::DefDatabaseStorage, source::SourceDatabaseStorage};
-
-#[salsa::database(SourceDatabaseStorage, DefDatabaseStorage)]
-#[derive(Default)]
-pub struct RootDatabase {
-    storage: salsa::Storage<Self>,
-}
-
-impl salsa::Database for RootDatabase {}
+pub use ide::RootDatabase;
diff --git a/src/source.rs b/src/source.rs
index f8fce37..dea52b5 100644
--- a/src/source.rs
+++ b/src/source.rs
@@ -1,4 +1,4 @@
-use rnix::types as ast;
+use rnix::AST;
 use std::sync::Arc;
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
@@ -14,6 +14,13 @@ impl<T> InFile<T> {
     pub fn new(file_id: FileId, value: T) -> Self {
         Self { file_id, value }
     }
+
+    pub fn map<U, F: FnOnce(T) -> U>(self, f: F) -> InFile<U> {
+        InFile {
+            file_id: self.file_id,
+            value: f(self.value),
+        }
+    }
 }
 
 #[salsa::query_group(SourceDatabaseStorage)]
@@ -24,12 +31,12 @@ pub trait SourceDatabase {
     #[salsa::input]
     fn file_content(&self, file_id: FileId) -> Arc<[u8]>;
 
-    fn parse(&self, file_id: FileId) -> InFile<ast::Root>;
+    fn parse(&self, file_id: FileId) -> InFile<AST>;
 }
 
-fn parse(db: &dyn SourceDatabase, file_id: FileId) -> InFile<ast::Root> {
+fn parse(db: &dyn SourceDatabase, file_id: FileId) -> InFile<AST> {
     let content = db.file_content(file_id);
     let content = std::str::from_utf8(&content).unwrap_or_default();
     let ast = rnix::parse(content);
-    InFile::new(file_id, ast.root())
+    InFile::new(file_id, ast)
 }