From d190bebcacbd62fe9d3e56f397fae911beaad88a Mon Sep 17 00:00:00 2001
From: Niko Matsakis <niko@alum.mit.edu>
Date: Mon, 13 May 2024 05:30:49 -0400
Subject: [PATCH] introduce helper functions

We'll need these for use with tracked functions
---
 .../salsa-2022-macros/src/db_lifetime.rs      | 48 +++++++++++++++++++
 components/salsa-2022-macros/src/lib.rs       |  1 +
 .../salsa-2022-macros/src/salsa_struct.rs     | 36 ++------------
 3 files changed, 54 insertions(+), 31 deletions(-)
 create mode 100644 components/salsa-2022-macros/src/db_lifetime.rs

diff --git a/components/salsa-2022-macros/src/db_lifetime.rs b/components/salsa-2022-macros/src/db_lifetime.rs
new file mode 100644
index 00000000..7e3d9494
--- /dev/null
+++ b/components/salsa-2022-macros/src/db_lifetime.rs
@@ -0,0 +1,48 @@
+//! Helper functions for working with fns, structs, and other generic things
+//! that are allowed to have a `'db` lifetime.
+
+use proc_macro2::Span;
+use syn::spanned::Spanned;
+
+/// Normally we try to use whatever lifetime parameter the use gave us
+/// to represent `'db`; but if they didn't give us one, we need to use a default
+/// name. We choose `'db`.
+pub(crate) fn default_db_lifetime(span: Span) -> syn::Lifetime {
+    syn::Lifetime {
+        apostrophe: span,
+        ident: syn::Ident::new("db", span),
+    }
+}
+
+/// Require that either there are no generics or exactly one lifetime parameter.
+pub(crate) fn require_db_lifetime(generics: &syn::Generics) -> syn::Result<()> {
+    if generics.params.len() == 0 {
+        return Ok(());
+    }
+
+    for (param, index) in generics.params.iter().zip(0..) {
+        let error = match param {
+            syn::GenericParam::Lifetime(_) => index > 0,
+            syn::GenericParam::Type(_) | syn::GenericParam::Const(_) => true,
+        };
+
+        if error {
+            return Err(syn::Error::new_spanned(
+                param,
+                "only a single lifetime parameter is accepted",
+            ));
+        }
+    }
+
+    Ok(())
+}
+
+/// Return the `'db` lifetime given be the user, or a default.
+/// The generics ought to have been checked with `require_db_lifetime` already.
+pub(crate) fn db_lifetime(generics: &syn::Generics) -> syn::Lifetime {
+    if let Some(lt) = generics.lifetimes().next() {
+        lt.lifetime.clone()
+    } else {
+        default_db_lifetime(generics.span())
+    }
+}
diff --git a/components/salsa-2022-macros/src/lib.rs b/components/salsa-2022-macros/src/lib.rs
index 0fcdccaf..596cb18b 100644
--- a/components/salsa-2022-macros/src/lib.rs
+++ b/components/salsa-2022-macros/src/lib.rs
@@ -39,6 +39,7 @@ pub(crate) fn literal(ident: &proc_macro2::Ident) -> proc_macro2::Literal {
 mod accumulator;
 mod configuration;
 mod db;
+mod db_lifetime;
 mod debug;
 mod input;
 mod interned;
diff --git a/components/salsa-2022-macros/src/salsa_struct.rs b/components/salsa-2022-macros/src/salsa_struct.rs
index 1e3c2a7e..a96a56ce 100644
--- a/components/salsa-2022-macros/src/salsa_struct.rs
+++ b/components/salsa-2022-macros/src/salsa_struct.rs
@@ -26,6 +26,7 @@
 //!     * this could be optimized, particularly for interned fields
 
 use crate::{
+    db_lifetime::{self, db_lifetime, default_db_lifetime},
     options::{AllowedOptions, Options},
     xform::ChangeLt,
 };
@@ -99,28 +100,9 @@ impl<A: AllowedOptions> SalsaStruct<A> {
         Ok(())
     }
 
+    /// Require that either there are no generics or exactly one lifetime parameter.
     pub(crate) fn require_db_lifetime(&self) -> syn::Result<()> {
-        let generics = &self.struct_item.generics;
-
-        if generics.params.len() == 0 {
-            return Ok(());
-        }
-
-        for (param, index) in generics.params.iter().zip(0..) {
-            let error = match param {
-                syn::GenericParam::Lifetime(_) => index > 0,
-                syn::GenericParam::Type(_) | syn::GenericParam::Const(_) => true,
-            };
-
-            if error {
-                return Err(syn::Error::new_spanned(
-                    param,
-                    "only a single lifetime parameter is accepted",
-                ));
-            }
-        }
-
-        Ok(())
+        db_lifetime::require_db_lifetime(&self.struct_item.generics)
     }
 
     /// Some salsa structs require a "Configuration" struct
@@ -144,11 +126,7 @@ impl<A: AllowedOptions> SalsaStruct<A> {
         if self.struct_item.generics.params.is_empty() {
             TheStructKind::Id
         } else {
-            if let Some(lt) = self.struct_item.generics.lifetimes().next() {
-                TheStructKind::Pointer(lt.lifetime.clone())
-            } else {
-                TheStructKind::Pointer(self.default_db_lifetime())
-            }
+            TheStructKind::Pointer(db_lifetime(&self.struct_item.generics))
         }
     }
 
@@ -378,11 +356,7 @@ impl<A: AllowedOptions> SalsaStruct<A> {
     /// to represent `'db`; but if they didn't give us one, we need to use a default
     /// name. We choose `'db`.
     fn default_db_lifetime(&self) -> syn::Lifetime {
-        let span = self.struct_item.ident.span();
-        syn::Lifetime {
-            apostrophe: span,
-            ident: syn::Ident::new("db", span),
-        }
+        default_db_lifetime(self.struct_item.generics.span())
     }
 
     /// Generate `impl salsa::AsId for Foo`