Merge pull request #450 from DropDemBits/placement-new-jars

Initialize jars in-place

book/src/plumbing/jars_and_ingredients.md

@@ -195,16 +195,16 @@ The `Default` implementation for `Storage<DB>` does the work:
 
 First, it creates an empty `Routes` instance.
 Then it invokes the `DB::create_jars` method.
-The implementation of this method is defined by the `#[salsa::db]` macro; it simply invokes the `Jar::create_jar` method on each of the jars:
+The implementation of this method is defined by the `#[salsa::db]` macro; it invokes `salsa::plumbing::create_jars_inplace` to allocate memory for the jars, and then invokes the `Jar::init_jar` method on each of the jars to initialize them:
 
 ```rust,ignore
 {{#include ../../../components/salsa-2022-macros/src/db.rs:create_jars}}
 ```
 
-This implementation for `create_jar` is geneated by the `#[salsa::jar]` macro, and simply walks over the representative type for each salsa item and asks *it* to create its ingredients
+This implementation for `init_jar` is generated by the `#[salsa::jar]` macro, and simply walks over the representative type for each salsa item and asks *it* to create its ingredients
 
 ```rust,ignore
-{{#include ../../../components/salsa-2022-macros/src/jar.rs:create_jar}}
+{{#include ../../../components/salsa-2022-macros/src/jar.rs:init_jar}}
 ```
 
 The code to create the ingredients for any particular item is generated by their associated macros (e.g., `#[salsa::tracked]`, `#[salsa::input]`), but it always follows a particular structure.

components/salsa-2022-macros/src/db.rs

@@ -1,5 +1,5 @@
 use proc_macro2::Literal;
-use syn::Token;
+use syn::{spanned::Spanned, Token};
 
 // Source:
 //
@@ -84,6 +84,12 @@ fn as_salsa_database_impl(input: &syn::ItemStruct) -> syn::ItemImpl {
 
 fn has_jars_impl(args: &Args, input: &syn::ItemStruct, storage: &syn::Ident) -> syn::ItemImpl {
     let jar_paths: Vec<&syn::Path> = args.jar_paths.iter().collect();
+    let jar_field_names: Vec<_> = args
+        .jar_paths
+        .iter()
+        .zip(0..)
+        .map(|(p, i)| syn::LitInt::new(&format!("{}", i), p.span()))
+        .collect();
     let db = &input.ident;
     parse_quote! {
         // ANCHOR: HasJars
@@ -100,12 +106,17 @@ fn has_jars_impl(args: &Args, input: &syn::ItemStruct, storage: &syn::Ident) ->
             }
 
             // ANCHOR: create_jars
-            fn create_jars(routes: &mut salsa::routes::Routes<Self>) -> Self::Jars {
-                (
-                    #(
-                        <#jar_paths as salsa::jar::Jar>::create_jar(routes),
-                    )*
-                )
+            fn create_jars(routes: &mut salsa::routes::Routes<Self>) -> Box<Self::Jars> {
+                unsafe {
+                    salsa::plumbing::create_jars_inplace::<#db>(|jars| {
+                        #(
+                            unsafe {
+                                let place = std::ptr::addr_of_mut!((*jars).#jar_field_names);
+                                <#jar_paths as salsa::jar::Jar>::init_jar(place, routes);
+                            }
+                        )*
+                    })
+                }
             }
             // ANCHOR_END: create_jars
         }

components/salsa-2022-macros/src/jar.rs

@@ -106,25 +106,27 @@ pub(crate) fn jar_impl(
         .fields
         .iter()
         .zip(0..)
-        .map(|(f, i)| Ident::new(&format!("i{}", i), f.ty.span()))
+        .map(|(f, i)| syn::LitInt::new(&format!("{}", i), f.ty.span()))
         .collect();
-    // ANCHOR: create_jar
+    // ANCHOR: init_jar
     quote! {
-        impl<'salsa_db> salsa::jar::Jar<'salsa_db> for #jar_struct {
+        unsafe impl<'salsa_db> salsa::jar::Jar<'salsa_db> for #jar_struct {
             type DynDb = dyn #jar_trait + 'salsa_db;
 
-            fn create_jar<DB>(routes: &mut salsa::routes::Routes<DB>) -> Self
+            unsafe fn init_jar<DB>(place: *mut Self, routes: &mut salsa::routes::Routes<DB>)
             where
                 DB: salsa::storage::JarFromJars<Self> + salsa::storage::DbWithJar<Self>,
             {
                 #(
-                    let #field_var_names = <#field_tys as salsa::storage::IngredientsFor>::create_ingredients(routes);
+                    unsafe {
+                        std::ptr::addr_of_mut!((*place).#field_var_names)
+                            .write(<#field_tys as salsa::storage::IngredientsFor>::create_ingredients(routes));
+                    }
                 )*
-                Self(#(#field_var_names),*)
             }
         }
     }
-    // ANCHOR_END: create_jar
+    // ANCHOR_END: init_jar
 }
 
 pub(crate) fn jar_struct(input: &ItemStruct) -> ItemStruct {

components/salsa-2022/src/jar.rs

@@ -5,10 +5,20 @@ use crate::{
 
 use super::routes::Routes;
 
-pub trait Jar<'db>: Sized {
+/// Representative trait of a salsa jar
+///
+/// # Safety
+///
+/// `init_jar` must fully initialize the jar
+pub unsafe trait Jar<'db>: Sized {
     type DynDb: ?Sized + HasJar<Self> + Database + 'db;
 
-    fn create_jar<DB>(routes: &mut Routes<DB>) -> Self
+    /// Initializes the jar at `place`
+    ///
+    /// # Safety
+    ///
+    /// `place` must be a valid pointer to this jar
+    unsafe fn init_jar<DB>(place: *mut Self, routes: &mut Routes<DB>)
     where
         DB: JarFromJars<Self> + DbWithJar<Self>;
 }

components/salsa-2022/src/plumbing.rs

@@ -1 +1,28 @@
+use std::{alloc, ptr};
 
+use crate::storage::HasJars;
+
+/// Initializes the `DB`'s jars in-place
+///
+/// # Safety:
+///
+/// `init` must fully initialize all of jars fields
+pub unsafe fn create_jars_inplace<DB: HasJars>(init: impl FnOnce(*mut DB::Jars)) -> Box<DB::Jars> {
+    let layout = alloc::Layout::new::<DB::Jars>();
+
+    if layout.size() == 0 {
+        // SAFETY: This is the recommended way of creating a Box
+        // to a ZST in the std docs
+        unsafe { Box::from_raw(ptr::NonNull::dangling().as_ptr()) }
+    } else {
+        // SAFETY: We've checked that the size isn't 0
+        let place = unsafe { alloc::alloc_zeroed(layout) };
+        let place = place.cast::<DB::Jars>();
+
+        init(place);
+
+        // SAFETY: Caller invariant requires that `init` must've
+        // initialized all of the fields
+        unsafe { Box::from_raw(place) }
+    }
+}

components/salsa-2022/src/storage.rs

@@ -60,7 +60,7 @@ where
         let jars = DB::create_jars(&mut routes);
         Self {
             shared: Shared {
-                jars: Some(Arc::new(jars)),
+                jars: Some(Arc::from(jars)),
                 cvar: Arc::new(Default::default()),
             },
             routes: Arc::new(routes),
@@ -192,7 +192,7 @@ pub trait HasJars: HasJarsDyn + Sized {
     /// and it will also cancel any ongoing work in the current revision.
     fn jars_mut(&mut self) -> (&mut Self::Jars, &mut Runtime);
 
-    fn create_jars(routes: &mut Routes<Self>) -> Self::Jars;
+    fn create_jars(routes: &mut Routes<Self>) -> Box<Self::Jars>;
 }
 
 pub trait DbWithJar<J>: HasJar<J> + Database {

salsa-2022-tests/tests/create-empty-database.rs

@@ -0,0 +1,35 @@
+//! Tests that we can create a database of only 0-size jars without invoking UB
+
+use salsa::storage::HasJars;
+
+#[salsa::jar(db = Db)]
+struct Jar();
+
+trait Db: salsa::DbWithJar<Jar> {}
+
+#[salsa::db(Jar)]
+#[derive(Default)]
+struct Database {
+    storage: salsa::Storage<Self>,
+}
+
+impl salsa::Database for Database {}
+
+impl Db for Database {}
+
+#[test]
+fn execute() {
+    let db = Database::default();
+    let jars = db.storage.jars().0;
+
+    ensure_init(jars);
+}
+
+fn ensure_init(place: *const <Database as HasJars>::Jars) {
+    use std::mem::forget;
+    use std::ptr::addr_of;
+
+    // SAFETY: Intentionally tries to access potentially uninitialized memory,
+    // so that miri can catch if we accidentally forget to initialize the memory.
+    forget(unsafe { addr_of!((*place).0).read() });
+}

salsa-2022-tests/tests/create-large-jar-database.rs

@@ -0,0 +1,154 @@
+//! Tests that we can create a database with very large jars without invoking UB
+
+use salsa::storage::HasJars;
+
+#[salsa::db(jar1::Jar1, jar2::Jar2, jar3::Jar3, jar4::Jar4)]
+#[derive(Default)]
+struct Database {
+    storage: salsa::Storage<Self>,
+}
+
+impl salsa::Database for Database {}
+
+#[test]
+fn execute() {
+    let db = Database::default();
+    let jars = db.storage.jars().0;
+
+    ensure_init(jars);
+}
+
+fn ensure_init(place: *const <Database as HasJars>::Jars) {
+    use std::mem::forget;
+    use std::ptr::addr_of;
+
+    // SAFETY: Intentionally tries to access potentially uninitialized memory,
+    // so that miri can catch if we accidentally forget to initialize the memory.
+    forget(unsafe { addr_of!((*place).0).read() });
+    forget(unsafe { addr_of!((*place).1).read() });
+    forget(unsafe { addr_of!((*place).2).read() });
+    forget(unsafe { addr_of!((*place).3).read() });
+}
+
+macro_rules! make_jarX {
+    ($jarX:ident, $JarX:ident) => {
+        mod $jarX {
+            #[salsa::jar(db = Db)]
+            pub(crate) struct $JarX(T1);
+
+            pub(crate) trait Db: salsa::DbWithJar<$JarX> {}
+
+            impl<DB> Db for DB where DB: salsa::DbWithJar<$JarX> {}
+
+            #[salsa::tracked(jar = $JarX)]
+            struct T1 {
+                a0: String,
+                a1: String,
+                a2: String,
+                a3: String,
+                a4: String,
+                a5: String,
+                a6: String,
+                a7: String,
+                a8: String,
+                a9: String,
+                a10: String,
+                a11: String,
+                a12: String,
+                a13: String,
+                a14: String,
+                a15: String,
+                a16: String,
+                a17: String,
+                a18: String,
+                a19: String,
+                a20: String,
+                a21: String,
+                a22: String,
+                a23: String,
+                a24: String,
+                a25: String,
+                a26: String,
+                a27: String,
+                a28: String,
+                a29: String,
+                a30: String,
+                a31: String,
+                a32: String,
+                a33: String,
+                a34: String,
+                a35: String,
+                a36: String,
+                a37: String,
+                a38: String,
+                a39: String,
+                a40: String,
+                a41: String,
+                a42: String,
+                a43: String,
+                a44: String,
+                a45: String,
+                a46: String,
+                a47: String,
+                a48: String,
+                a49: String,
+                a50: String,
+                a51: String,
+                a52: String,
+                a53: String,
+                a54: String,
+                a55: String,
+                a56: String,
+                a57: String,
+                a58: String,
+                a59: String,
+                a60: String,
+                a61: String,
+                a62: String,
+                a63: String,
+                a64: String,
+                a65: String,
+                a66: String,
+                a67: String,
+                a68: String,
+                a69: String,
+                a70: String,
+                a71: String,
+                a72: String,
+                a73: String,
+                a74: String,
+                a75: String,
+                a76: String,
+                a77: String,
+                a78: String,
+                a79: String,
+                a80: String,
+                a81: String,
+                a82: String,
+                a83: String,
+                a84: String,
+                a85: String,
+                a86: String,
+                a87: String,
+                a88: String,
+                a89: String,
+                a90: String,
+                a91: String,
+                a92: String,
+                a93: String,
+                a94: String,
+                a95: String,
+                a96: String,
+                a97: String,
+                a98: String,
+                a99: String,
+                a100: String,
+            }
+        }
+    };
+}
+
+make_jarX!(jar1, Jar1);
+make_jarX!(jar2, Jar2);
+make_jarX!(jar3, Jar3);
+make_jarX!(jar4, Jar4);