From 3dabb0ddc89e6ab1901704a5e0c43b86653cdee2 Mon Sep 17 00:00:00 2001
From: Niko Matsakis <niko@alum.mit.edu>
Date: Thu, 15 Aug 2024 17:38:16 +0300
Subject: [PATCH] make memos take read-lock (and infallible return)

We want to ensure that accessing the memos only occurs
in revision R after the struct is created.
---
 src/input.rs          |  4 +--
 src/interned.rs       |  4 +--
 src/table.rs          | 16 +++++++-----
 src/tracked_struct.rs | 58 ++++++++++++++++++++++++-------------------
 4 files changed, 46 insertions(+), 36 deletions(-)

diff --git a/src/input.rs b/src/input.rs
index 47ea6d76..3ee6c53b 100644
--- a/src/input.rs
+++ b/src/input.rs
@@ -308,7 +308,7 @@ impl<C> Slot for Value<C>
 where
     C: Configuration,
 {
-    fn memos(&self) -> Option<&crate::table::memo::MemoTable> {
-        Some(&self.memos)
+    fn memos(&self, _current_revision: Revision) -> &crate::table::memo::MemoTable {
+        &self.memos
     }
 }
diff --git a/src/interned.rs b/src/interned.rs
index 6a101720..c863ccbd 100644
--- a/src/interned.rs
+++ b/src/interned.rs
@@ -276,7 +276,7 @@ impl<C> Slot for Value<C>
 where
     C: Configuration,
 {
-    fn memos(&self) -> Option<&MemoTable> {
-        Some(&self.memos)
+    fn memos(&self, _current_revision: Revision) -> &MemoTable {
+        &self.memos
     }
 }
diff --git a/src/table.rs b/src/table.rs
index 20242eb7..a7b38c8f 100644
--- a/src/table.rs
+++ b/src/table.rs
@@ -24,7 +24,7 @@ pub(crate) struct Table {
 pub(crate) trait TablePage: Any + Send + Sync {
     fn hidden_type_name(&self) -> &'static str;
 
-    fn memos(&self, slot: SlotIndex) -> Option<&MemoTable>;
+    fn memos(&self, slot: SlotIndex, current_revision: Revision) -> &MemoTable;
 }
 
 pub(crate) struct Page<T: Slot> {
@@ -51,7 +51,7 @@ pub(crate) struct Page<T: Slot> {
 }
 
 pub(crate) trait Slot: Any + Send + Sync {
-    fn memos(&self) -> Option<&MemoTable>;
+    fn memos(&self, current_revision: Revision) -> &MemoTable;
 }
 
 unsafe impl<T: Slot> Send for Page<T> {}
@@ -117,9 +117,13 @@ impl Table {
     }
 
     /// Get the memo table associated with `id` (if any)
-    pub fn memos(&self, id: Id) -> Option<&MemoTable> {
+    ///
+    /// # Panics
+    ///
+    /// If the ingredient for `id` doesn't have associated memo-tables on its slots.
+    pub fn memos(&self, id: Id, current_revision: Revision) -> &MemoTable {
         let (page, slot) = split_id(id);
-        self.pages[page.0].memos(slot)
+        self.pages[page.0].memos(slot, current_revision)
     }
 }
 
@@ -194,8 +198,8 @@ impl<T: Slot> TablePage for Page<T> {
         std::any::type_name::<Self>()
     }
 
-    fn memos(&self, slot: SlotIndex) -> Option<&MemoTable> {
-        self.get(slot).memos()
+    fn memos(&self, slot: SlotIndex, current_revision: Revision) -> &MemoTable {
+        self.get(slot).memos(current_revision)
     }
 }
 
diff --git a/src/tracked_struct.rs b/src/tracked_struct.rs
index 21c9b120..9ec160c9 100644
--- a/src/tracked_struct.rs
+++ b/src/tracked_struct.rs
@@ -521,7 +521,7 @@ where
         let field_ingredient_index = self.ingredient_index.successor(field_index);
         let data = Self::data(zalsa.table(), id);
 
-        self.read_lock(data, zalsa.current_revision());
+        data.read_lock(zalsa.current_revision());
 
         let field_changed_at = data.revisions[field_index];
 
@@ -536,29 +536,6 @@ where
 
         unsafe { self.to_self_ref(&data.fields) }
     }
-
-    fn read_lock(&self, data: &Value<C>, current_revision: Revision) {
-        loop {
-            match data.updated_at.load() {
-                None => {
-                    panic!("access to field whilst the value is being initialized");
-                }
-                Some(r) => {
-                    if r == current_revision {
-                        return;
-                    }
-
-                    if data
-                        .updated_at
-                        .compare_exchange(Some(r), Some(current_revision))
-                        .is_ok()
-                    {
-                        break;
-                    }
-                }
-            }
-        }
-    }
 }
 
 impl<C> Ingredient for IngredientImpl<C>
@@ -640,11 +617,40 @@ where
     }
 }
 
+impl<C> Value<C>
+where
+    C: Configuration,
+{
+    fn read_lock(&self, current_revision: Revision) {
+        loop {
+            match self.updated_at.load() {
+                None => {
+                    panic!("access to field whilst the value is being initialized");
+                }
+                Some(r) => {
+                    if r == current_revision {
+                        return;
+                    }
+
+                    if self
+                        .updated_at
+                        .compare_exchange(Some(r), Some(current_revision))
+                        .is_ok()
+                    {
+                        break;
+                    }
+                }
+            }
+        }
+    }
+}
+
 impl<C> Slot for Value<C>
 where
     C: Configuration,
 {
-    fn memos(&self) -> Option<&crate::table::memo::MemoTable> {
-        Some(&self.memos)
+    fn memos(&self, current_revision: Revision) -> &crate::table::memo::MemoTable {
+        self.read_lock(current_revision);
+        &self.memos
     }
 }