Reduce memory usage by deduplicating type information

We were storing the type information, 3 words wide, for each memo in each slot, while it is always constant wrt. the ingredient (different slots of the same ingredients will always have the same memos in the same order). This introduces some more unsafety, and the result wasn't as fast so I also had to use some lock-free structures, but the result is worth it: this shaves off 230mb from rust-analyzer with new Salsa.

src/accumulator.rs

@@ -14,6 +14,7 @@ use crate::{
     cycle::CycleRecoveryStrategy,
     ingredient::{fmt_index, Ingredient, Jar},
     plumbing::JarAux,
+    table::memo::MemoTableTypes,
     zalsa::IngredientIndex,
     zalsa_local::QueryOrigin,
     Database, DatabaseKeyIndex, Id, Revision,
@@ -61,6 +62,7 @@ impl<A: Accumulator> Jar for JarImpl<A> {
 
 pub struct IngredientImpl<A: Accumulator> {
     index: IngredientIndex,
+    memo_table_types: MemoTableTypes,
     phantom: PhantomData<Accumulated<A>>,
 }
 
@@ -80,6 +82,7 @@ impl<A: Accumulator> IngredientImpl<A> {
     pub fn new(index: IngredientIndex) -> Self {
         Self {
             index,
+            memo_table_types: MemoTableTypes::default(),
             phantom: PhantomData,
         }
     }
@@ -153,6 +156,10 @@ impl<A: Accumulator> Ingredient for IngredientImpl<A> {
     fn accumulated(&self, _db: &dyn Database, _key_index: Id) -> Option<&AccumulatedMap> {
         None
     }
+
+    fn memo_table_types(&self) -> &MemoTableTypes {
+        &self.memo_table_types
+    }
 }
 
 impl<A> std::fmt::Debug for IngredientImpl<A>

src/function.rs

@@ -7,6 +7,7 @@ use crate::{
     key::DatabaseKeyIndex,
     plumbing::JarAux,
     salsa_struct::SalsaStructInDb,
+    table::memo::MemoTableTypes,
     zalsa::{IngredientIndex, MemoIngredientIndex, Zalsa},
     zalsa_local::QueryOrigin,
     Cycle, Database, Id, Revision,
@@ -113,6 +114,9 @@ pub struct IngredientImpl<C: Configuration> {
     /// we don't know that we can trust the database to give us the same runtime
     /// everytime and so forth.
     deleted_entries: DeletedEntries<C>,
+
+    /// This is empty, but we need this for the trait and it doesn't consume a lot of memory anyway.
+    _memo_table_types: MemoTableTypes,
 }
 
 /// True if `old_value == new_value`. Invoked by the generated
@@ -132,6 +136,7 @@ where
             memo_ingredient_index: aux.next_memo_ingredient_index(struct_index, index),
             lru: Default::default(),
             deleted_entries: Default::default(),
+            _memo_table_types: MemoTableTypes::default(),
         }
     }
 
@@ -245,6 +250,10 @@ where
         C::DEBUG_NAME
     }
 
+    fn memo_table_types(&self) -> &MemoTableTypes {
+        &self._memo_table_types
+    }
+
     fn accumulated<'db>(
         &'db self,
         db: &'db dyn Database,

src/ingredient.rs

@@ -6,6 +6,7 @@ use std::{
 use crate::{
     accumulator::accumulated_map::AccumulatedMap,
     cycle::CycleRecoveryStrategy,
+    table::memo::MemoTableTypes,
     zalsa::{IngredientIndex, MemoIngredientIndex},
     zalsa_local::QueryOrigin,
     Database, DatabaseKeyIndex, Id,
@@ -121,6 +122,8 @@ pub trait Ingredient: Any + std::fmt::Debug + Send + Sync {
     /// [`IngredientRequiresReset::RESET_ON_NEW_REVISION`] to true.
     fn reset_for_new_revision(&mut self);
 
+    fn memo_table_types(&self) -> &MemoTableTypes;
+
     fn fmt_index(&self, index: Option<crate::Id>, fmt: &mut fmt::Formatter<'_>) -> fmt::Result;
 }
 

src/input.rs

@@ -18,7 +18,11 @@ use crate::{
     ingredient::{fmt_index, Ingredient},
     key::{DatabaseKeyIndex, DependencyIndex},
     plumbing::{Jar, JarAux, Stamp},
-    table::{memo::MemoTable, sync::SyncTable, Slot, Table},
+    table::{
+        memo::{MemoTable, MemoTableTypes},
+        sync::SyncTable,
+        Slot, Table,
+    },
     zalsa::{IngredientIndex, Zalsa},
     zalsa_local::QueryOrigin,
     Database, Durability, Id, Revision, Runtime,
@@ -75,6 +79,7 @@ pub struct IngredientImpl<C: Configuration> {
     ingredient_index: IngredientIndex,
     singleton_index: AtomicCell<Option<Id>>,
     singleton_lock: Mutex<()>,
+    memo_table_types: MemoTableTypes,
     _phantom: std::marker::PhantomData<C::Struct>,
 }
 
@@ -84,6 +89,7 @@ impl<C: Configuration> IngredientImpl<C> {
             ingredient_index: index,
             singleton_index: AtomicCell::new(None),
             singleton_lock: Default::default(),
+            memo_table_types: MemoTableTypes::default(),
             _phantom: std::marker::PhantomData,
         }
     }
@@ -283,6 +289,10 @@ impl<C: Configuration> Ingredient for IngredientImpl<C> {
     ) -> Option<&'db crate::accumulator::accumulated_map::AccumulatedMap> {
         None
     }
+
+    fn memo_table_types(&self) -> &MemoTableTypes {
+        &self.memo_table_types
+    }
 }
 
 impl<C: Configuration> std::fmt::Debug for IngredientImpl<C> {
@@ -327,4 +337,8 @@ where
     unsafe fn syncs(&self, _current_revision: Revision) -> &SyncTable {
         &self.syncs
     }
+
+    unsafe fn drop_memos(&mut self, types: &MemoTableTypes) {
+        self.memos.drop(types);
+    }
 }

src/input/input_field.rs

@@ -1,6 +1,6 @@
 use crate::cycle::CycleRecoveryStrategy;
 use crate::ingredient::{fmt_index, Ingredient};
-use crate::input::Configuration;
+use crate::input::{Configuration, MemoTableTypes};
 use crate::zalsa::IngredientIndex;
 use crate::zalsa_local::QueryOrigin;
 use crate::{Database, DatabaseKeyIndex, Id, Revision};
@@ -21,6 +21,7 @@ use super::{IngredientImpl, Value};
 pub struct FieldIngredientImpl<C: Configuration> {
     index: IngredientIndex,
     field_index: usize,
+    memo_table_types: MemoTableTypes,
     phantom: PhantomData<fn() -> Value<C>>,
 }
 
@@ -32,6 +33,7 @@ where
         Self {
             index: struct_index.successor(field_index),
             field_index,
+            memo_table_types: MemoTableTypes::default(),
             phantom: PhantomData,
         }
     }
@@ -104,6 +106,10 @@ where
     ) -> Option<&'db crate::accumulator::accumulated_map::AccumulatedMap> {
         None
     }
+
+    fn memo_table_types(&self) -> &MemoTableTypes {
+        &self.memo_table_types
+    }
 }
 
 impl<C> std::fmt::Debug for FieldIngredientImpl<C>

src/interned.rs

@@ -4,7 +4,7 @@ use crate::id::AsId;
 use crate::ingredient::fmt_index;
 use crate::key::DependencyIndex;
 use crate::plumbing::{Jar, JarAux};
-use crate::table::memo::MemoTable;
+use crate::table::memo::{MemoTable, MemoTableTypes};
 use crate::table::sync::SyncTable;
 use crate::table::Slot;
 use crate::zalsa::IngredientIndex;
@@ -66,6 +66,8 @@ pub struct IngredientImpl<C: Configuration> {
     /// but that will make anything dependent on those entries dirty and in need
     /// of being recomputed.
     reset_at: Revision,
+
+    memo_table_types: MemoTableTypes,
 }
 
 /// Struct storing the interned fields.
@@ -109,6 +111,7 @@ where
             ingredient_index,
             key_map: Default::default(),
             reset_at: Revision::start(),
+            memo_table_types: MemoTableTypes::default(),
         }
     }
 
@@ -280,6 +283,10 @@ where
         C::DEBUG_NAME
     }
 
+    fn memo_table_types(&self) -> &MemoTableTypes {
+        &self.memo_table_types
+    }
+
     fn accumulated<'db>(
         &'db self,
         _db: &'db dyn Database,
@@ -311,6 +318,10 @@ where
     unsafe fn syncs(&self, _current_revision: Revision) -> &crate::table::sync::SyncTable {
         &self.syncs
     }
+
+    unsafe fn drop_memos(&mut self, types: &MemoTableTypes) {
+        self.memos.drop(types);
+    }
 }
 
 /// The `Lookup` trait is a more flexible variant on [`std::borrow::Borrow`]

src/runtime.rs

@@ -350,4 +350,8 @@ impl Runtime {
             .lock()
             .unblock_runtimes_blocked_on(database_key, wait_result);
     }
+
+    pub(crate) fn take_table(&mut self) -> Table {
+        std::mem::take(&mut self.table)
+    }
 }

src/table.rs

@@ -12,7 +12,11 @@ use memo::MemoTable;
 use parking_lot::Mutex;
 use sync::SyncTable;
 
-use crate::{zalsa::transmute_data_ptr, Id, IngredientIndex, Revision};
+use crate::{
+    table::memo::MemoTableTypes,
+    zalsa::{transmute_data_ptr, Zalsa},
+    Id, IngredientIndex, Revision,
+};
 
 pub(crate) mod memo;
 pub(crate) mod sync;
@@ -27,6 +31,17 @@ pub(crate) struct Table {
     pages: AppendOnlyVec<Box<dyn TablePage>>,
 }
 
+impl Table {
+    pub(crate) fn drop(self, zalsa: &Zalsa) {
+        for mut page in self.pages.into_vec() {
+            // SAFETY: We own `self` and don't use it after.
+            unsafe {
+                TablePage::drop(&mut *page, zalsa);
+            }
+        }
+    }
+}
+
 pub(crate) trait TablePage: Any + Send + Sync {
     fn hidden_type_name(&self) -> &'static str;
 
@@ -35,19 +50,27 @@ pub(crate) trait TablePage: Any + Send + Sync {
     /// # Safety condition
     ///
     /// The `current_revision` MUST be the current revision of the database owning this table page.
-    unsafe fn memos(&self, slot: SlotIndex, current_revision: Revision) -> &MemoTable;
+    unsafe fn memos_and_ingredient(
+        &self,
+        slot: SlotIndex,
+        current_revision: Revision,
+    ) -> (&MemoTable, IngredientIndex);
 
     /// Access the syncs attached to `slot`.
     ///
     /// # Safety condition
     ///
     /// The `current_revision` MUST be the current revision of the database owning this table page.
     unsafe fn syncs(&self, slot: SlotIndex, current_revision: Revision) -> &SyncTable;
+
+    /// # Safety
+    ///
+    /// You cannot use `self` after calling this method.
+    unsafe fn drop(&mut self, zalsa: &Zalsa);
 }
 
 pub(crate) struct Page<T: Slot> {
     /// The ingredient for elements on this page.
-    #[allow(dead_code)] // pretty sure we'll need this
     ingredient: IngredientIndex,
 
     /// Number of elements of `data` that are initialized.
@@ -65,6 +88,9 @@ pub(crate) struct Page<T: Slot> {
 
     /// The potentially uninitialized data of this page. As we initialize new entries, we increment `allocated`.
     data: Box<[UnsafeCell<MaybeUninit<T>>; PAGE_LEN]>,
+
+    /// A helper to ensure we don't forget to call `drop()`. It will panic if we forget or if we call it twice.
+    dropped: bool,
 }
 
 pub(crate) trait Slot: Any + Send + Sync {
@@ -81,6 +107,11 @@ pub(crate) trait Slot: Any + Send + Sync {
     ///
     /// The current revision MUST be the current revision of the database containing this slot.
     unsafe fn syncs(&self, current_revision: Revision) -> &SyncTable;
+
+    /// # Safety
+    ///
+    /// `types` must be correct.
+    unsafe fn drop_memos(&mut self, types: &MemoTableTypes);
 }
 
 unsafe impl<T: Slot> Send for Page<T> {}
@@ -165,9 +196,14 @@ impl Table {
     ///
     /// The parameter `current_revision` MUST be the current revision
     /// of the owner of database owning this table.
-    pub unsafe fn memos(&self, id: Id, current_revision: Revision) -> &MemoTable {
+    pub unsafe fn memos_and_ingredient(
+        &self,
+        id: Id,
+        current_revision: Revision,
+    ) -> (&MemoTable, IngredientIndex) {
         let (page, slot) = split_id(id);
-        self.pages[page.0].memos(slot, current_revision)
+        let page = &self.pages[page.0];
+        page.memos_and_ingredient(slot, current_revision)
     }
 
     /// Get the sync table associated with `id`
@@ -190,6 +226,7 @@ impl<T: Slot> Page<T> {
             allocated: Default::default(),
             allocation_lock: Default::default(),
             data: Box::new([const { UnsafeCell::new(MaybeUninit::uninit()) }; PAGE_LEN]),
+            dropped: false,
         }
     }
 
@@ -253,28 +290,42 @@ impl<T: Slot> TablePage for Page<T> {
         std::any::type_name::<Self>()
     }
 
-    unsafe fn memos(&self, slot: SlotIndex, current_revision: Revision) -> &MemoTable {
-        self.get(slot).memos(current_revision)
+    unsafe fn memos_and_ingredient(
+        &self,
+        slot: SlotIndex,
+        current_revision: Revision,
+    ) -> (&MemoTable, IngredientIndex) {
+        (self.get(slot).memos(current_revision), self.ingredient)
     }
 
     unsafe fn syncs(&self, slot: SlotIndex, current_revision: Revision) -> &SyncTable {
         self.get(slot).syncs(current_revision)
     }
-}
 
-impl<T: Slot> Drop for Page<T> {
-    fn drop(&mut self) {
+    unsafe fn drop(&mut self, zalsa: &Zalsa) {
+        assert!(!self.dropped, "cannot drop a Page twice");
+        self.dropped = true;
+        let types = zalsa.lookup_ingredient(self.ingredient).memo_table_types();
         // Execute destructors for all initialized elements
-        let len = self.allocated.load(Ordering::Acquire);
+        let len = *self.allocated.get_mut();
         // SAFETY: self.data is initialized for T's up to len
         unsafe {
             // FIXME: Should be ptr::from_raw_parts_mut but that is unstable
             let to_drop = slice::from_raw_parts_mut(self.data.as_mut_ptr().cast::<T>(), len);
-            ptr::drop_in_place(to_drop)
+            for v in &mut *to_drop {
+                v.drop_memos(types);
+            }
+            ptr::drop_in_place(to_drop);
         }
     }
 }
 
+impl<T: Slot> Drop for Page<T> {
+    fn drop(&mut self) {
+        assert!(self.dropped, "Page not dropped");
+    }
+}
+
 impl dyn TablePage {
     fn assert_type<T: Any>(&self) -> &T {
         assert_eq!(