return a NonNull instead of a &'db

nikomatsakis · nikomatsakis · commit 07d0ead9f40d · 2024-05-30T01:59:49.000-04:00
In old code, we converted to a `&amp;'db` when
creating a new tracked struct or interning,
but this value in fact persisted beyond the end
of `'db` (i.e., into the new revision).

We now refactor so that we create the `Foo&lt;'db&gt;`
from a `NonNull&lt;T&gt;` instead of a `&amp;'db T`, and
then only create safe references when users
access fields.

This makes miri happy.
diff --git a/components/salsa-2022-macros/src/interned.rs b/components/salsa-2022-macros/src/interned.rs
@@ -160,13 +160,24 @@ impl InternedStruct {
         data_ident: &syn::Ident,
         config_ident: &syn::Ident,
     ) -> syn::ItemImpl {
+        let the_ident = self.the_ident();
         let lt_db = &self.named_db_lifetime();
         let (_, _, _, type_generics, _) = self.the_ident_and_generics();
         parse_quote_spanned!(
             config_ident.span() =>
 
             impl salsa::interned::Configuration for #config_ident {
                 type Data<#lt_db> = #data_ident #type_generics;
+
+                type Struct<#lt_db> = #the_ident < #lt_db >;
+
+                unsafe fn struct_from_raw<'db>(ptr: std::ptr::NonNull<salsa::interned::ValueStruct<Self>>) -> Self::Struct<'db> {
+                    #the_ident(ptr, std::marker::PhantomData)
+                }
+
+                fn deref_struct<'db>(s: Self::Struct<'db>) -> &'db salsa::interned::ValueStruct<Self> {
+                    unsafe { s.0.as_ref() }
+                }
             }
         )
     }
@@ -191,21 +202,21 @@ impl InternedStruct {
 
         let field_getters: Vec<syn::ImplItemFn> = self
             .all_fields()
-            .map(|field| {
+            .map(|field: &crate::salsa_struct::SalsaField| {
                 let field_name = field.name();
                 let field_ty = field.ty();
                 let field_vis = field.vis();
                 let field_get_name = field.get_name();
                 if field.is_clone_field() {
                     parse_quote_spanned! { field_get_name.span() =>
                         #field_vis fn #field_get_name(self, _db: & #db_lt #db_dyn_ty) -> #field_ty {
-                            std::clone::Clone::clone(&unsafe { &*self.0 }.data().#field_name)
+                            std::clone::Clone::clone(&unsafe { self.0.as_ref() }.data().#field_name)
                         }
                     }
                 } else {
                     parse_quote_spanned! { field_get_name.span() =>
                         #field_vis fn #field_get_name(self, _db: & #db_lt #db_dyn_ty) -> & #db_lt #field_ty {
-                            &unsafe { &*self.0 }.data().#field_name
+                            &unsafe { self.0.as_ref() }.data().#field_name
                         }
                     }
                 }
@@ -218,18 +229,15 @@ impl InternedStruct {
         let constructor_name = self.constructor_name();
         let new_method: syn::ImplItemFn = parse_quote_spanned! { constructor_name.span() =>
             #vis fn #constructor_name(
-                db: &#db_dyn_ty,
+                db: &#db_lt #db_dyn_ty,
                 #(#field_names: #field_tys,)*
             ) -> Self {
                 let (jar, runtime) = <_ as salsa::storage::HasJar<#jar_ty>>::jar(db);
                 let ingredients = <#jar_ty as salsa::storage::HasIngredientsFor< #the_ident #type_generics >>::ingredient(jar);
-                Self(
-                    ingredients.intern(runtime, #data_ident {
-                        #(#field_names,)*
-                        __phantom: std::marker::PhantomData,
-                    }),
-                    std::marker::PhantomData,
-                )
+                ingredients.intern(runtime, #data_ident {
+                    #(#field_names,)*
+                    __phantom: std::marker::PhantomData,
+                })
             }
         };
 
@@ -262,6 +270,7 @@ impl InternedStruct {
             self.the_ident_and_generics();
         let db_dyn_ty = self.db_dyn_ty();
         let jar_ty = self.jar_ty();
+        let db_lt = self.named_db_lifetime();
 
         let field_getters: Vec<syn::ImplItemFn> = self
             .all_fields()
@@ -296,7 +305,7 @@ impl InternedStruct {
         let constructor_name = self.constructor_name();
         let new_method: syn::ImplItemFn = parse_quote_spanned! { constructor_name.span() =>
             #vis fn #constructor_name(
-                db: &#db_dyn_ty,
+                db: & #db_lt #db_lt #db_dyn_ty,
                 #(#field_names: #field_tys,)*
             ) -> Self {
                 let (jar, runtime) = <_ as salsa::storage::HasJar<#jar_ty>>::jar(db);
@@ -384,7 +393,7 @@ impl InternedStruct {
                         fn lookup_id(id: salsa::Id, db: & #db_lt DB) -> Self {
                             let (jar, _) = <_ as salsa::storage::HasJar<#jar_ty>>::jar(db);
                             let ingredients = <#jar_ty as salsa::storage::HasIngredientsFor<#ident #type_generics>>::ingredient(jar);
-                            Self(ingredients.interned_value(id), std::marker::PhantomData)
+                            ingredients.interned_value(id)
                         }
                     }
                 })
diff --git a/components/salsa-2022-macros/src/salsa_struct.rs b/components/salsa-2022-macros/src/salsa_struct.rs
@@ -349,7 +349,7 @@ impl<A: AllowedOptions> SalsaStruct<A> {
                 #(#attrs)*
                 #[derive(Copy, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
                 #visibility struct #ident #generics (
-                    *const salsa::#module::ValueStruct < #config_ident >,
+                    std::ptr::NonNull<salsa::#module::ValueStruct < #config_ident >>,
                     std::marker::PhantomData < & #lifetime salsa::#module::ValueStruct < #config_ident > >
                 );
             })
@@ -360,7 +360,9 @@ impl<A: AllowedOptions> SalsaStruct<A> {
     pub(crate) fn access_salsa_id_from_self(&self) -> syn::Expr {
         match self.the_struct_kind() {
             TheStructKind::Id => parse_quote!(self.0),
-            TheStructKind::Pointer(_) => parse_quote!(salsa::id::AsId::as_id(unsafe { &*self.0 })),
+            TheStructKind::Pointer(_) => {
+                parse_quote!(salsa::id::AsId::as_id(unsafe { self.0.as_ref() }))
+            }
         }
     }
 
@@ -434,7 +436,7 @@ impl<A: AllowedOptions> SalsaStruct<A> {
                     #where_clause
                     {
                         fn as_id(&self) -> salsa::Id {
-                            salsa::id::AsId::as_id(unsafe { &*self.0 })
+                            salsa::id::AsId::as_id(unsafe { self.0.as_ref() })
                         }
                     }
 
diff --git a/components/salsa-2022-macros/src/tracked_fn.rs b/components/salsa-2022-macros/src/tracked_fn.rs
@@ -343,6 +343,16 @@ fn interned_configuration_impl(
     parse_quote!(
         impl salsa::interned::Configuration for #config_ty {
             type Data<#db_lt> = #intern_data_ty;
+
+            type Struct<#db_lt> = & #db_lt salsa::interned::ValueStruct<Self>;
+
+            unsafe fn struct_from_raw<'db>(ptr: std::ptr::NonNull<salsa::interned::ValueStruct<Self>>) -> Self::Struct<'db> {
+                unsafe { ptr.as_ref() }
+            }
+
+            fn deref_struct<'db>(s: Self::Struct<'db>) -> &'db salsa::interned::ValueStruct<Self> {
+                s
+            }
         }
     )
 }
diff --git a/components/salsa-2022-macros/src/tracked_struct.rs b/components/salsa-2022-macros/src/tracked_struct.rs
@@ -92,6 +92,7 @@ impl TrackedStruct {
         let field_tys: Vec<_> = self.all_fields().map(SalsaField::ty).collect();
         let id_field_indices = self.id_field_indices();
         let arity = self.all_field_count();
+        let the_ident = self.the_ident();
         let lt_db = &self.named_db_lifetime();
 
         // Create the function body that will update the revisions for each field.
@@ -132,8 +133,19 @@ impl TrackedStruct {
         parse_quote! {
             impl salsa::tracked_struct::Configuration for #config_ident {
                 type Fields<#lt_db> = ( #(#field_tys,)* );
+
+                type Struct<#lt_db> = #the_ident<#lt_db>;
+
                 type Revisions = [salsa::Revision; #arity];
 
+                unsafe fn struct_from_raw<'db>(ptr: std::ptr::NonNull<salsa::tracked_struct::ValueStruct<Self>>) -> Self::Struct<'db> {
+                    #the_ident(ptr, std::marker::PhantomData)
+                }
+
+                fn deref_struct<'db>(s: Self::Struct<'db>) -> &'db salsa::tracked_struct::ValueStruct<Self> {
+                    unsafe { s.0.as_ref() }
+                }
+
                 #[allow(clippy::unused_unit)]
                 fn id_fields(fields: &Self::Fields<'_>) -> impl std::hash::Hash {
                     ( #( &fields.#id_field_indices ),* )
@@ -205,7 +217,7 @@ impl TrackedStruct {
                             #field_vis fn #field_get_name(self, __db: & #lt_db #db_dyn_ty) -> & #lt_db #field_ty
                             {
                                 let (_, __runtime) = <_ as salsa::storage::HasJar<#jar_ty>>::jar(__db);
-                                let fields = unsafe { &*self.0 }.field(__runtime, #field_index);
+                                let fields = unsafe { self.0.as_ref() }.field(__runtime, #field_index);
                                 &fields.#field_index
                             }
                         }
@@ -214,7 +226,7 @@ impl TrackedStruct {
                             #field_vis fn #field_get_name(self, __db: & #lt_db #db_dyn_ty) -> #field_ty
                             {
                                 let (_, __runtime) = <_ as salsa::storage::HasJar<#jar_ty>>::jar(__db);
-                                let fields = unsafe { &*self.0 }.field(__runtime, #field_index);
+                                let fields = unsafe { self.0.as_ref() }.field(__runtime, #field_index);
                                 fields.#field_index.clone()
                             }
                         }
@@ -232,11 +244,6 @@ impl TrackedStruct {
 
         let salsa_id = self.access_salsa_id_from_self();
 
-        let ctor = match the_kind {
-            TheStructKind::Id => quote!(salsa::id::FromId::from_as_id(#data)),
-            TheStructKind::Pointer(_) => quote!(Self(#data, std::marker::PhantomData)),
-        };
-
         let lt_db = self.maybe_elided_db_lifetime();
         parse_quote! {
             #[allow(dead_code, clippy::pedantic, clippy::complexity, clippy::style)]
@@ -246,11 +253,10 @@ impl TrackedStruct {
                 {
                     let (__jar, __runtime) = <_ as salsa::storage::HasJar<#jar_ty>>::jar(__db);
                     let __ingredients = <#jar_ty as salsa::storage::HasIngredientsFor< Self >>::ingredient(__jar);
-                    let #data = __ingredients.0.new_struct(
+                    __ingredients.0.new_struct(
                         __runtime,
                         (#(#field_names,)*),
-                    );
-                    #ctor
+                    )
                 }
 
                 pub fn salsa_id(&self) -> salsa::Id {
@@ -354,7 +360,7 @@ impl TrackedStruct {
                         fn lookup_id(id: salsa::Id, db: & #db_lt DB) -> Self {
                             let (jar, runtime) = <_ as salsa::storage::HasJar<#jar_ty>>::jar(db);
                             let ingredients = <#jar_ty as salsa::storage::HasIngredientsFor<#ident #type_generics>>::ingredient(jar);
-                            Self(ingredients.#tracked_struct_ingredient.lookup_struct(runtime, id), std::marker::PhantomData)
+                            ingredients.#tracked_struct_ingredient.lookup_struct(runtime, id)
                         }
                     }
                 })
diff --git a/components/salsa-2022/src/alloc.rs b/components/salsa-2022/src/alloc.rs
@@ -13,30 +13,28 @@ impl<T> Alloc<T> {
             data: unsafe { NonNull::new_unchecked(data) },
         }
     }
-}
 
-impl<T> Drop for Alloc<T> {
-    fn drop(&mut self) {
-        let data: *mut T = self.data.as_ptr();
-        let data: Box<T> = unsafe { Box::from_raw(data) };
-        drop(data);
+    pub fn as_raw(&self) -> NonNull<T> {
+        self.data
     }
-}
 
-impl<T> std::ops::Deref for Alloc<T> {
-    type Target = T;
-
-    fn deref(&self) -> &Self::Target {
+    pub unsafe fn as_ref(&self) -> &T {
         unsafe { self.data.as_ref() }
     }
-}
 
-impl<T> std::ops::DerefMut for Alloc<T> {
-    fn deref_mut(&mut self) -> &mut Self::Target {
+    pub unsafe fn as_mut(&mut self) -> &mut T {
         unsafe { self.data.as_mut() }
     }
 }
 
+impl<T> Drop for Alloc<T> {
+    fn drop(&mut self) {
+        let data: *mut T = self.data.as_ptr();
+        let data: Box<T> = unsafe { Box::from_raw(data) };
+        drop(data);
+    }
+}
+
 unsafe impl<T> Send for Alloc<T> where T: Send {}
 
 unsafe impl<T> Sync for Alloc<T> where T: Sync {}
diff --git a/components/salsa-2022/src/interned.rs b/components/salsa-2022/src/interned.rs
@@ -2,6 +2,7 @@ use crossbeam::atomic::AtomicCell;
 use std::fmt;
 use std::hash::Hash;
 use std::marker::PhantomData;
+use std::ptr::NonNull;
 
 use crate::alloc::Alloc;
 use crate::durability::Durability;
@@ -18,8 +19,28 @@ use super::ingredient::Ingredient;
 use super::routes::IngredientIndex;
 use super::Revision;
 
-pub trait Configuration {
+pub trait Configuration: Sized {
     type Data<'db>: InternedData;
+    type Struct<'db>: Copy;
+
+    /// Create an end-user struct from the underlying raw pointer.
+    ///
+    /// This call is an "end-step" to the tracked struct lookup/creation
+    /// process in a given revision: it occurs only when the struct is newly
+    /// created or, if a struct is being reused, after we have updated its
+    /// fields (or confirmed it is green and no updates are required).
+    ///
+    /// # Unsafety
+    ///
+    /// Requires that `ptr` represents a "confirmed" value in this revision,
+    /// which means that it will remain valid and immutable for the remainder of this
+    /// revision, represented by the lifetime `'db`.
+    unsafe fn struct_from_raw<'db>(ptr: NonNull<ValueStruct<Self>>) -> Self::Struct<'db>;
+
+    /// Deref the struct to yield the underlying value struct.
+    /// Since we are still part of the `'db` lifetime in which the struct was created,
+    /// this deref is safe, and the value-struct fields are immutable and verified.
+    fn deref_struct<'db>(s: Self::Struct<'db>) -> &'db ValueStruct<Self>;
 }
 
 pub trait InternedData: Sized + Eq + Hash + Clone {}
@@ -83,15 +104,11 @@ where
     }
 
     pub fn intern_id<'db>(&'db self, runtime: &'db Runtime, data: C::Data<'db>) -> crate::Id {
-        self.intern(runtime, data).as_id()
+        C::deref_struct(self.intern(runtime, data)).as_id()
     }
 
     /// Intern data to a unique reference.
-    pub fn intern<'db>(
-        &'db self,
-        runtime: &'db Runtime,
-        data: C::Data<'db>,
-    ) -> &'db ValueStruct<C> {
+    pub fn intern<'db>(&'db self, runtime: &'db Runtime, data: C::Data<'db>) -> C::Struct<'db> {
         runtime.report_tracked_read(
             DependencyIndex::for_table(self.ingredient_index),
             Durability::MAX,
@@ -126,27 +143,27 @@ where
                         id: next_id,
                         fields: internal_data,
                     }));
-                // SAFETY: Items are only removed from the `value_map` with an `&mut self` reference.
-                let value_ref = unsafe { transmute_lifetime(self, &**value) };
+                let value_raw = value.as_raw();
                 drop(value);
                 entry.insert(next_id);
-                value_ref
+                // SAFETY: Items are only removed from the `value_map` with an `&mut self` reference.
+                unsafe { C::struct_from_raw(value_raw) }
             }
         }
     }
 
-    pub fn interned_value<'db>(&'db self, id: Id) -> &'db ValueStruct<C> {
+    pub fn interned_value<'db>(&'db self, id: Id) -> C::Struct<'db> {
         let r = self.value_map.get(&id).unwrap();
 
         // SAFETY: Items are only removed from the `value_map` with an `&mut self` reference.
-        unsafe { transmute_lifetime(self, &**r) }
+        unsafe { C::struct_from_raw(r.as_raw()) }
     }
 
     /// Lookup the data for an interned value based on its id.
     /// Rarely used since end-users generally carry a struct with a pointer directly
     /// to the interned item.
     pub fn data<'db>(&'db self, id: Id) -> &'db C::Data<'db> {
-        self.interned_value(id).data()
+        C::deref_struct(self.interned_value(id)).data()
     }
 
     /// Variant of `data` that takes a (unnecessary) database argument.
diff --git a/components/salsa-2022/src/tracked_struct.rs b/components/salsa-2022/src/tracked_struct.rs
diff --git a/components/salsa-2022/src/tracked_struct/struct_map.rs b/components/salsa-2022/src/tracked_struct/struct_map.rs
diff --git a/components/salsa-2022/src/tracked_struct/tracked_field.rs b/components/salsa-2022/src/tracked_struct/tracked_field.rs

Original file line number	Diff line number	Diff line change
`@@ -349,7 +349,7 @@ impl<A: AllowedOptions> SalsaStruct<A> {`
`349`	`349`	`#(#attrs)*`
`350`	`350`	`#[derive(Copy, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]`
`351`	`351`	`#visibility struct #ident #generics (`
`352`		`- *const salsa::#module::ValueStruct < #config_ident >,`
	`352`	`+ std::ptr::NonNull<salsa::#module::ValueStruct < #config_ident >>,`
`353`	`353`	`std::marker::PhantomData < & #lifetime salsa::#module::ValueStruct < #config_ident > >`
`354`	`354`	`);`
`355`	`355`	`})`
`@@ -360,7 +360,9 @@ impl<A: AllowedOptions> SalsaStruct<A> {`
`360`	`360`	`pub(crate) fn access_salsa_id_from_self(&self) -> syn::Expr {`
`361`	`361`	`match self.the_struct_kind() {`
`362`	`362`	`TheStructKind::Id => parse_quote!(self.0),`
`363`		`- TheStructKind::Pointer(_) => parse_quote!(salsa::id::AsId::as_id(unsafe { &*self.0 })),`
	`363`	`+ TheStructKind::Pointer(_) => {`
	`364`	`+ parse_quote!(salsa::id::AsId::as_id(unsafe { self.0.as_ref() }))`
	`365`	`+ }`
`364`	`366`	`}`
`365`	`367`	`}`
`366`	`368`
`@@ -434,7 +436,7 @@ impl<A: AllowedOptions> SalsaStruct<A> {`
`434`	`436`	`#where_clause`
`435`	`437`	`{`
`436`	`438`	`fn as_id(&self) -> salsa::Id {`
`437`		`- salsa::id::AsId::as_id(unsafe { &*self.0 })`
	`439`	`+ salsa::id::AsId::as_id(unsafe { self.0.as_ref() })`
`438`	`440`	`}`
`439`	`441`	`}`
`440`	`442`
Original file line number	Diff line number	Diff line change
`@@ -343,6 +343,16 @@ fn interned_configuration_impl(`
`343`	`343`	`parse_quote!(`
`344`	`344`	`impl salsa::interned::Configuration for #config_ty {`
`345`	`345`	`type Data<#db_lt> = #intern_data_ty;`
	`346`	`+`
	`347`	`+ type Struct<#db_lt> = & #db_lt salsa::interned::ValueStruct<Self>;`
	`348`	`+`
	`349`	`+ unsafe fn struct_from_raw<'db>(ptr: std::ptr::NonNull<salsa::interned::ValueStruct<Self>>) -> Self::Struct<'db> {`
	`350`	`+ unsafe { ptr.as_ref() }`
	`351`	`+ }`
	`352`	`+`
	`353`	`+ fn deref_struct<'db>(s: Self::Struct<'db>) -> &'db salsa::interned::ValueStruct<Self> {`
	`354`	`+ s`
	`355`	`+ }`
`346`	`356`	`}`
`347`	`357`	`)`
`348`	`358`	`}`