Compress file text using lz4 in salsa

Author: lnicola

Cargo.lock

@@ -12,6 +12,8 @@ rust-version.workspace = true
 doctest = false
 
 [dependencies]
+lz4_flex = { version = "0.11", default-features = false }
+
 la-arena.workspace = true
 salsa.workspace = true
 rustc-hash.workspace = true

crates/base-db/Cargo.toml

@@ -7,7 +7,7 @@ use salsa::Durability;
 use triomphe::Arc;
 use vfs::FileId;
 
-use crate::{CrateGraph, SourceDatabaseExt, SourceRoot, SourceRootId};
+use crate::{CrateGraph, SourceDatabaseExt, SourceDatabaseExt2, SourceRoot, SourceRootId};
 
 /// Encapsulate a bunch of raw `.set` calls on the database.
 #[derive(Default)]

crates/base-db/src/change.rs

@@ -7,6 +7,7 @@ mod input;
 
 use std::panic;
 
+use salsa::Durability;
 use syntax::{ast, Parse, SourceFile};
 use triomphe::Arc;
 
@@ -42,6 +43,7 @@ pub trait Upcast<T: ?Sized> {
     fn upcast(&self) -> &T;
 }
 
+pub const DEFAULT_FILE_TEXT_LRU_CAP: usize = 16;
 pub const DEFAULT_PARSE_LRU_CAP: usize = 128;
 pub const DEFAULT_BORROWCK_LRU_CAP: usize = 1024;
 
@@ -89,7 +91,10 @@ fn parse(db: &dyn SourceDatabase, file_id: FileId) -> Parse<ast::SourceFile> {
 #[salsa::query_group(SourceDatabaseExtStorage)]
 pub trait SourceDatabaseExt: SourceDatabase {
     #[salsa::input]
+    fn compressed_file_text(&self, file_id: FileId) -> Arc<[u8]>;
+
     fn file_text(&self, file_id: FileId) -> Arc<str>;
+
     /// Path to a file, relative to the root of its source root.
     /// Source root of the file.
     #[salsa::input]
@@ -101,6 +106,44 @@ pub trait SourceDatabaseExt: SourceDatabase {
     fn source_root_crates(&self, id: SourceRootId) -> Arc<[CrateId]>;
 }
 
+fn file_text(db: &dyn SourceDatabaseExt, file_id: FileId) -> Arc<str> {
+    let bytes = db.compressed_file_text(file_id);
+    let bytes =
+        lz4_flex::decompress_size_prepended(&bytes).expect("lz4 decompression should not fail");
+    let text = std::str::from_utf8(&bytes).expect("file contents should be valid UTF-8");
+    Arc::from(text)
+}
+
+pub trait SourceDatabaseExt2 {
+    fn set_file_text(&mut self, file_id: FileId, text: Arc<str>) {
+        self.set_file_text_with_durability(file_id, text, Durability::LOW);
+    }
+
+    fn set_file_text_with_durability(
+        &mut self,
+        file_id: FileId,
+        text: Arc<str>,
+        durability: Durability,
+    );
+}
+
+impl<Db: ?Sized + SourceDatabaseExt> SourceDatabaseExt2 for Db {
+    fn set_file_text_with_durability(
+        &mut self,
+        file_id: FileId,
+        text: Arc<str>,
+        durability: Durability,
+    ) {
+        let bytes = text.as_bytes();
+        let compressed = lz4_flex::compress_prepend_size(&bytes);
+        self.set_compressed_file_text_with_durability(
+            file_id,
+            Arc::from(compressed.as_slice()),
+            durability,
+        )
+    }
+}
+
 fn source_root_crates(db: &dyn SourceDatabaseExt, id: SourceRootId) -> Arc<[CrateId]> {
     let graph = db.crate_graph();
     let mut crates = graph

crates/base-db/src/lib.rs

@@ -1,4 +1,4 @@
-use base_db::{SourceDatabase, SourceDatabaseExt};
+use base_db::{SourceDatabase, SourceDatabaseExt2 as _};
 use test_fixture::WithFixture;
 use triomphe::Arc;
 

crates/hir-def/src/nameres/tests/incremental.rs

@@ -12,7 +12,7 @@ mod traits;
 
 use std::env;
 
-use base_db::{FileRange, SourceDatabaseExt};
+use base_db::{FileRange, SourceDatabaseExt2 as _};
 use expect_test::Expect;
 use hir_def::{
     body::{Body, BodySourceMap, SyntheticSyntax},

crates/hir-ty/src/tests.rs

@@ -1,4 +1,4 @@
-use base_db::SourceDatabaseExt;
+use base_db::SourceDatabaseExt2 as _;
 use test_fixture::WithFixture;
 use triomphe::Arc;
 

crates/hir-ty/src/tests/incremental.rs

@@ -205,6 +205,7 @@ impl RootDatabase {
 
             // SourceDatabaseExt
             base_db::FileTextQuery
+            base_db::CompressedFileTextQuery
             base_db::FileSourceRootQuery
             base_db::SourceRootQuery
             base_db::SourceRootCratesQuery

crates/ide-db/src/apply_change.rs

@@ -51,6 +51,7 @@ use std::{fmt, mem::ManuallyDrop};
 use base_db::{
     salsa::{self, Durability},
     AnchoredPath, CrateId, FileId, FileLoader, FileLoaderDelegate, SourceDatabase, Upcast,
+    DEFAULT_FILE_TEXT_LRU_CAP,
 };
 use hir::db::{DefDatabase, ExpandDatabase, HirDatabase};
 use triomphe::Arc;
@@ -157,6 +158,7 @@ impl RootDatabase {
 
     pub fn update_base_query_lru_capacities(&mut self, lru_capacity: Option<usize>) {
         let lru_capacity = lru_capacity.unwrap_or(base_db::DEFAULT_PARSE_LRU_CAP);
+        base_db::FileTextQuery.in_db_mut(self).set_lru_capacity(DEFAULT_FILE_TEXT_LRU_CAP);
         base_db::ParseQuery.in_db_mut(self).set_lru_capacity(lru_capacity);
         // macro expansions are usually rather small, so we can afford to keep more of them alive
         hir::db::ParseMacroExpansionQuery.in_db_mut(self).set_lru_capacity(4 * lru_capacity);
@@ -166,6 +168,7 @@ impl RootDatabase {
     pub fn update_lru_capacities(&mut self, lru_capacities: &FxHashMap<Box<str>, usize>) {
         use hir::db as hir_db;
 
+        base_db::FileTextQuery.in_db_mut(self).set_lru_capacity(DEFAULT_FILE_TEXT_LRU_CAP);
         base_db::ParseQuery.in_db_mut(self).set_lru_capacity(
             lru_capacities
                 .get(stringify!(ParseQuery))
@@ -199,7 +202,7 @@ impl RootDatabase {
             // base_db::ProcMacrosQuery
 
             // SourceDatabaseExt
-            // base_db::FileTextQuery
+            base_db::FileTextQuery
             // base_db::FileSourceRootQuery
             // base_db::SourceRootQuery
             base_db::SourceRootCratesQuery