From c4e7f062472e90bef0c2d4e5a639b0910eedd2f3 Mon Sep 17 00:00:00 2001
From: fa-sharp <farshadkt@gmail.com>
Date: Thu, 19 Mar 2026 22:02:08 -0400
Subject: [PATCH] add AppState derive macro for convenient state extraction

---
 Cargo.lock                                | 10 ++++
 Cargo.toml                                |  5 ++
 crates/axum-app-wrapper-macros/Cargo.toml | 12 ++++
 crates/axum-app-wrapper-macros/src/lib.rs | 71 +++++++++++++++++++++++
 src/lib.rs                                | 17 +++++-
 5 files changed, 114 insertions(+), 1 deletion(-)
 create mode 100644 crates/axum-app-wrapper-macros/Cargo.toml
 create mode 100644 crates/axum-app-wrapper-macros/src/lib.rs

diff --git a/Cargo.lock b/Cargo.lock
index bd328d6..3e547df 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -53,10 +53,20 @@ version = "0.1.0"
 dependencies = [
  "anyhow",
  "axum",
+ "axum-app-wrapper-macros",
  "futures",
  "type-map",
 ]
 
+[[package]]
+name = "axum-app-wrapper-macros"
+version = "0.1.0"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn",
+]
+
 [[package]]
 name = "axum-core"
 version = "0.5.5"
diff --git a/Cargo.toml b/Cargo.toml
index 5011552..aef059b 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,3 +1,7 @@
+[workspace]
+members = [".", "crates/*"]
+resolver = "2"
+
 [package]
 name = "axum-app-wrapper"
 version = "0.1.0"
@@ -8,3 +12,4 @@ anyhow = "1.0"
 axum = "0.8"
 futures = "0.3"
 type-map = "0.5"
+axum-app-wrapper-macros = { path = "crates/axum-app-wrapper-macros" }
diff --git a/crates/axum-app-wrapper-macros/Cargo.toml b/crates/axum-app-wrapper-macros/Cargo.toml
new file mode 100644
index 0000000..0d0690d
--- /dev/null
+++ b/crates/axum-app-wrapper-macros/Cargo.toml
@@ -0,0 +1,12 @@
+[package]
+name = "axum-app-wrapper-macros"
+version = "0.1.0"
+edition = "2024"
+
+[lib]
+proc-macro = true
+
+[dependencies]
+syn = "2"
+quote = "1"
+proc-macro2 = "1"
diff --git a/crates/axum-app-wrapper-macros/src/lib.rs b/crates/axum-app-wrapper-macros/src/lib.rs
new file mode 100644
index 0000000..a553c09
--- /dev/null
+++ b/crates/axum-app-wrapper-macros/src/lib.rs
@@ -0,0 +1,71 @@
+use proc_macro::TokenStream;
+use quote::quote;
+use syn::{Data, DeriveInput, Fields, parse_macro_input};
+
+/// Derive macro that implements [`TryFrom<TypeMap>`] for `T`.
+///
+/// Apply this to your state struct and each named field will be extracted from the
+/// [`TypeMap`] by type automatically. Keep in mind each field must be of a
+/// different type — if you need to store multiple values of the same type, use
+/// a wrapper like `struct Foo(type)` or `enum Foo { A(type), B(type) }`.
+///
+/// How the state is stored and cloned is up to you —
+/// if you want cheap clones, wrap your fields in `Arc` individually or use
+/// `Arc<AppState>` as your axum state type.
+///
+/// ```ignore
+/// #[derive(AppState)]
+/// struct AppState {
+///     config: Config,
+///     pool: Pool,
+///     db: Db,
+/// }
+///
+/// // Plain state — axum will clone the whole struct on each request:
+/// App::<AppState>::new()
+///
+/// // Or wrap in Arc yourself for cheap clones:
+/// App::<Arc<AppState>>::new()
+/// ```
+#[proc_macro_derive(AppState)]
+pub fn derive_app_state(input: TokenStream) -> TokenStream {
+    let input = parse_macro_input!(input as DeriveInput);
+    let name = &input.ident;
+
+    let fields = match &input.data {
+        Data::Struct(s) => match &s.fields {
+            Fields::Named(f) => &f.named,
+            _ => {
+                return syn::Error::new_spanned(
+                    name,
+                    "#[derive(AppState)] only supports structs with named fields",
+                )
+                .to_compile_error()
+                .into();
+            }
+        },
+        _ => {
+            return syn::Error::new_spanned(
+                name,
+                "#[derive(AppState)] can only be applied to structs",
+            )
+            .to_compile_error()
+            .into();
+        }
+    };
+
+    let field_names = fields.iter().map(|f| &f.ident);
+
+    quote! {
+        impl ::std::convert::TryFrom<::axum_app_wrapper::TypeMap> for #name {
+            type Error = ::anyhow::Error;
+
+            fn try_from(mut map: ::axum_app_wrapper::TypeMap) -> ::std::result::Result<Self, Self::Error> {
+                Ok(#name {
+                    #(#field_names: ::axum_app_wrapper::extract_type_field(&mut map)?,)*
+                })
+            }
+        }
+    }
+    .into()
+}
diff --git a/src/lib.rs b/src/lib.rs
index bc12f9c..8af1288 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -8,7 +8,22 @@ use futures::{
     FutureExt,
     future::{BoxFuture, join_all},
 };
-use type_map::concurrent::TypeMap;
+
+// State extraction utilities
+
+pub use axum_app_wrapper_macros::AppState;
+pub use type_map::concurrent::TypeMap;
+
+/// Extracts a value of type `T` from a [`TypeMap`], returning an Anyhow error if the type is not present.
+///
+/// This is used internally by the [`AppState`] macro but is also available for manual
+/// [`TryFrom<TypeMap>`] implementations.
+pub fn extract_type_field<T: Send + Sync + 'static>(map: &mut TypeMap) -> anyhow::Result<T> {
+    map.remove::<T>()
+        .ok_or_else(|| anyhow!("Missing type in TypeMap: {}", std::any::type_name::<T>()))
+}
+
+// Plugin system
 
 /// A lightweight wrapper around axum that enables building plugins around the router
 pub struct App<S = Arc<TypeMap>> {