cached: Merge lock_key into lock

jaemk · Sep 19, 2024 · a1d0f40 · a1d0f40
1 parent 32d1bd0
commit a1d0f40
Showing 1 changed file with 26 additions and 32 deletions.
diff --git a/cached_proc_macro/src/cached.rs b/cached_proc_macro/src/cached.rs
@@ -219,23 +219,24 @@ pub fn cached(args: TokenStream, input: TokenStream) -> TokenStream {
     let no_cache_fn_ident = Ident::new(&format!("{}_no_cache", &fn_ident), fn_ident.span());
 
     let lock;
-    let lock_key;
     let function_no_cache;
     let function_call;
     let ty;
     if asyncness.is_some() {
-        lock = quote! {
-            let mut cache = #cache_ident.lock().await;
-        };
-
-        lock_key = quote! {
-            let mut locks = #cache_ident_key.lock().await;
-            let lock = locks
-                .entry(key.clone())
-                .or_insert_with(|| std::sync::Arc::new(::cached::async_sync::Mutex::new(#cache_create)))
-                .clone();
-            drop(locks);
-            let mut cache = lock.lock().await;
+        lock = if args.sync_writes_by_key {
+            quote! {
+                let mut locks = #cache_ident_key.lock().await;
+                let lock = locks
+                    .entry(key.clone())
+                    .or_insert_with(|| std::sync::Arc::new(::cached::async_sync::Mutex::new(#cache_create)))
+                    .clone();
+                drop(locks);
+                let mut cache = lock.lock().await;
+            }
+        } else {
+            quote! {
+                let mut cache = #cache_ident.lock().await;
+            }
         };
 
         function_no_cache = quote! {
@@ -251,15 +252,17 @@ pub fn cached(args: TokenStream, input: TokenStream) -> TokenStream {
             #visibility static #cache_ident_key: ::cached::once_cell::sync::Lazy<::cached::async_sync::Mutex<std::collections::HashMap<#cache_key_ty, std::sync::Arc<::cached::async_sync::Mutex<#cache_ty>>>>> = ::cached::once_cell::sync::Lazy::new(|| ::cached::async_sync::Mutex::new(std::collections::HashMap::new()));
         };
     } else {
-        lock = quote! {
-            let mut cache = #cache_ident.lock().unwrap();
-        };
-
-        lock_key = quote! {
-            let mut locks = #cache_ident_key.lock().unwrap();
-            let lock = locks.entry(key.clone()).or_insert_with(|| std::sync::Arc::new(std::sync::Mutex::new(#cache_create))).clone();
-            drop(locks);
-            let mut cache = lock.lock().unwrap();
+        lock = if args.sync_writes_by_key {
+            quote! {
+                let mut locks = #cache_ident_key.lock().unwrap();
+                let lock = locks.entry(key.clone()).or_insert_with(|| std::sync::Arc::new(std::sync::Mutex::new(#cache_create))).clone();
+                drop(locks);
+                let mut cache = lock.lock().unwrap();
+            }
+        } else {
+            quote! {
+                let mut cache = #cache_ident.lock().unwrap();
+            }
         };
 
         function_no_cache = quote! {
@@ -284,16 +287,7 @@ pub fn cached(args: TokenStream, input: TokenStream) -> TokenStream {
         #set_cache_and_return
     };
 
-    let do_set_return_block = if args.sync_writes_by_key {
-        quote! {
-            #lock_key
-            if let Some(result) = cache.cache_get(&key) {
-                #return_cache_block
-            }
-            #function_call
-            #set_cache_and_return
-        }
-    } else if args.sync_writes {
+    let do_set_return_block = if args.sync_writes_by_key || args.sync_writes {
         quote! {
             #lock
             if let Some(result) = cache.cache_get(&key) {