refactor: Add macro to join futures

quickemu-project · Sep 30, 2024 · db70182 · db70182
1 parent 5e5c3ad
commit db70182
Show file tree

Hide file tree

Showing 17 changed files with 106 additions and 202 deletions.
diff --git a/.gitignore b/.gitignore
@@ -1 +1,2 @@
 /target
+/join_futures/target
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -20,3 +20,4 @@ serde = "1.0.202"
 serde_json = "1.0.117"
 tokio = { version = "1.37.0", features = ["macros", "rt-multi-thread", "sync"] }
 zstd = "0.13.1"
+join_futures = { path = "join_futures" }
diff --git a/join_futures/Cargo.lock b/join_futures/Cargo.lock
diff --git a/join_futures/Cargo.toml b/join_futures/Cargo.toml
@@ -0,0 +1,10 @@
+[package]
+name = "join_futures"
+version = "0.1.0"
+edition = "2021"
+
+[profile.dev]
+opt-level = 3
+
+[lib]
+proc-macro = true
diff --git a/join_futures/src/lib.rs b/join_futures/src/lib.rs
@@ -0,0 +1,30 @@
+use proc_macro::{TokenStream, TokenTree};
+
+#[proc_macro]
+pub fn join_futures(input: TokenStream) -> TokenStream {
+    let mut tokens = input.into_iter();
+
+    let variable = match tokens.find(|t| matches!(t, TokenTree::Ident(_))) {
+        Some(TokenTree::Ident(variable)) => variable.to_string(),
+        _ => panic!("You must provide a variable containing futures to join"),
+    };
+    let flatten_amount = match tokens.find(|t| matches!(t, TokenTree::Literal(_))) {
+        Some(TokenTree::Literal(flatten_amount)) => flatten_amount.to_string().trim().parse::<usize>().unwrap(),
+        _ => 0,
+    };
+
+    let tokens = tokens.skip_while(|t| matches!(t, TokenTree::Punct(_)));
+    let value_type = match tokens.map(|t| t.to_string()).collect::<String>() {
+        s if s.is_empty() => "Vec<Config>".to_string(),
+        s => s,
+    };
+
+    let mut value = format!("futures::future::join_all({variable}).await");
+    if flatten_amount > 0 {
+        value.push_str(&format!(
+            ".into_iter(){}.collect::<{value_type}>()",
+            ".flatten()".repeat(flatten_amount),
+        ));
+    }
+    value.parse().unwrap()
+}
diff --git a/src/bsd/mod.rs b/src/bsd/mod.rs
@@ -1,5 +1,6 @@
 use crate::store_data::{ArchiveFormat, ChecksumSeparation, Config, Disk, Distro, Source, WebSource};
 use crate::utils::capture_page;
+use join_futures::join_futures;
 use quickemu::config::{Arch, GuestOS};
 use regex::Regex;
 use std::sync::Arc;
@@ -84,22 +85,14 @@ impl Distro for FreeBSD {
                             [normal_editions, vm_image]
                         })
                         .collect::<Vec<_>>();
-                    Some(futures::future::join_all(futures).await)
+                    Some(join_futures!(futures))
                 } else {
                     log::warn!("Failed to fetch FreeBSD {arch} releases");
                     None
                 }
             }
         });
-        futures::future::join_all(futures)
-            .await
-            .into_iter()
-            .flatten()
-            .flatten()
-            .flatten()
-            .flatten()
-            .collect::<Vec<Config>>()
-            .into()
+        Some(join_futures!(futures, 4))
     }
 }
 
@@ -205,16 +198,10 @@ impl Distro for GhostBSD {
                         }
                     })
                     .collect::<Vec<_>>();
-                Some(futures::future::join_all(futures).await)
+                Some(join_futures!(futures))
             }
         });
 
-        futures::future::join_all(futures)
-            .await
-            .into_iter()
-            .flatten()
-            .flatten()
-            .collect::<Vec<Config>>()
-            .into()
+        Some(join_futures!(futures, 2))
     }
 }
diff --git a/src/linux/arch.rs b/src/linux/arch.rs
@@ -4,6 +4,7 @@ use crate::{
     store_data::{ChecksumSeparation, Config, Distro, Source, WebSource},
     utils::{capture_page, GatherData, GithubAPI},
 };
+use join_futures::join_futures;
 use regex::Regex;
 use serde::Deserialize;
 use std::sync::Arc;
@@ -39,12 +40,7 @@ impl Distro for Archcraft {
                 })
             }
         });
-        futures::future::join_all(futures)
-            .await
-            .into_iter()
-            .flatten()
-            .collect::<Vec<Config>>()
-            .into()
+        Some(join_futures!(futures, 1))
     }
 }
 
@@ -147,17 +143,11 @@ impl Distro for ArcoLinux {
                             }
                         })
                         .collect::<Vec<_>>();
-                    Some(futures::future::join_all(futures).await)
+                    Some(join_futures!(futures))
                 }
             })
             .collect::<Vec<_>>();
-        futures::future::join_all(futures)
-            .await
-            .into_iter()
-            .flatten()
-            .flatten()
-            .collect::<Vec<Config>>()
-            .into()
+        Some(join_futures!(futures, 2))
     }
 }
 
@@ -235,12 +225,7 @@ impl Distro for AthenaOS {
             })
         });
 
-        futures::future::join_all(futures)
-            .await
-            .into_iter()
-            .flatten()
-            .collect::<Vec<Config>>()
-            .into()
+        Some(join_futures!(futures, 1))
     }
 }
 
@@ -302,17 +287,11 @@ impl Distro for CachyOS {
                     })
                     .collect::<Vec<_>>();
 
-                Some(futures::future::join_all(futures).await)
+                Some(join_futures!(futures))
             }
         });
 
-        futures::future::join_all(futures)
-            .await
-            .into_iter()
-            .flatten()
-            .flatten()
-            .collect::<Vec<Config>>()
-            .into()
+        Some(join_futures!(futures, 2))
     }
 }
 
@@ -341,7 +320,7 @@ impl Distro for EndeavourOS {
                 }
             }
         });
-        futures::future::join_all(futures).await.into()
+        Some(join_futures!(futures))
     }
 }
 
@@ -383,11 +362,6 @@ impl Distro for Garuda {
             }
         });
 
-        futures::future::join_all(futures)
-            .await
-            .into_iter()
-            .flatten()
-            .collect::<Vec<Config>>()
-            .into()
+        Some(join_futures!(futures, 1))
     }
 }
diff --git a/src/linux/arch/manjaro.rs b/src/linux/arch/manjaro.rs
@@ -2,6 +2,7 @@ use crate::{
     store_data::{Config, Distro, Source, WebSource},
     utils::capture_page,
 };
+use join_futures::join_futures;
 use regex::Regex;
 
 const BIGLINUX_MIRROR: &str = "https://iso.biglinux.com.br/";
@@ -38,6 +39,6 @@ impl Distro for BigLinux {
             }
         });
 
-        futures::future::join_all(futures).await.into()
+        Some(join_futures!(futures))
     }
 }
diff --git a/src/linux/debian.rs b/src/linux/debian.rs
@@ -2,6 +2,7 @@ use crate::{
     store_data::{ChecksumSeparation, Config, Disk, Distro, Source, WebSource},
     utils::{capture_page, GatherData, GithubAPI},
 };
+use join_futures::join_futures;
 use quickemu::config::{Arch, DiskFormat};
 use quickget_core::data_structures::ArchiveFormat;
 use regex::Regex;
@@ -68,13 +69,7 @@ impl Distro for Antix {
             }
         });
 
-        futures::future::join_all(futures)
-            .await
-            .into_iter()
-            .flatten()
-            .flatten()
-            .collect::<Vec<Config>>()
-            .into()
+        Some(join_futures!(futures, 2))
     }
 }
 
@@ -96,12 +91,7 @@ impl Distro for BunsenLabs {
             let url = format!("{BUNSENLABS_MIRROR}{}", &c[1]);
             async move { ChecksumSeparation::Whitespace.build(&url).await }
         });
-        let mut checksums = futures::future::join_all(checksum_futures)
-            .await
-            .into_iter()
-            .flatten()
-            .flatten()
-            .collect::<HashMap<String, String>>();
+        let mut checksums = join_futures!(checksum_futures, 2, HashMap<String, String>);
 
         release_regex
             .captures_iter(&html)
@@ -254,14 +244,7 @@ impl Distro for Debian {
             })
             .flatten();
 
-        futures::future::join_all(futures)
-            .await
-            .into_iter()
-            .flatten()
-            .flatten()
-            .flatten()
-            .collect::<Vec<Config>>()
-            .into()
+        Some(join_futures!(futures, 3))
     }
 }
 
@@ -310,13 +293,7 @@ impl Distro for Devuan {
                 )
             }
         });
-        futures::future::join_all(futures)
-            .await
-            .into_iter()
-            .flatten()
-            .flatten()
-            .collect::<Vec<Config>>()
-            .into()
+        Some(join_futures!(futures, 2))
     }
 }
 
@@ -360,22 +337,10 @@ impl Distro for EasyOS {
                     }
                 });
 
-                Some(
-                    futures::future::join_all(futures)
-                        .await
-                        .into_iter()
-                        .flatten()
-                        .collect::<Vec<_>>(),
-                )
+                Some(join_futures!(futures))
             }
         });
-        let mut releases = futures::future::join_all(release_futures)
-            .await
-            .into_iter()
-            .flatten()
-            .flatten()
-            .flatten()
-            .collect::<Vec<_>>();
+        let mut releases = join_futures!(release_futures, 4, Vec<(String, String)>);
 
         releases.sort_by(|(a, _), (b, _)| {
             if let (Ok(a), Ok(b)) = (
@@ -425,12 +390,7 @@ impl Distro for EasyOS {
                 })
             }
         });
-        futures::future::join_all(futures)
-            .await
-            .into_iter()
-            .flatten()
-            .collect::<Vec<Config>>()
-            .into()
+        Some(join_futures!(futures, 1))
     }
 }
 
@@ -478,17 +438,10 @@ impl Distro for EndlessOS {
                         })
                     }
                 });
-                Some(futures::future::join_all(futures).await)
+                Some(join_futures!(futures))
             }
         });
 
-        futures::future::join_all(futures)
-            .await
-            .into_iter()
-            .flatten()
-            .flatten()
-            .flatten()
-            .collect::<Vec<_>>()
-            .into()
+        Some(join_futures!(futures, 3))
     }
 }