aboutsummaryrefslogtreecommitdiff
path: root/src/select.rs
diff options
context:
space:
mode:
Diffstat (limited to 'src/select.rs')
-rw-r--r--src/select.rs303
1 files changed, 303 insertions, 0 deletions
diff --git a/src/select.rs b/src/select.rs
new file mode 100644
index 0000000..e309600
--- /dev/null
+++ b/src/select.rs
@@ -0,0 +1,303 @@
+use proc_macro::TokenStream;
+use proc_macro2::{Spacing, TokenStream as TokenStream2, TokenTree};
+use quote::{ToTokens, quote};
+use syn::parse::{Parse, ParseStream};
+use syn::punctuated::Punctuated;
+use syn::{Expr, LitStr, Token, parse_macro_input};
+
+// ─── SelectArm ──────────────────────────────────────────────────────────────────────────
+
+/// A single arm inside `select!`.
+pub enum SelectArm {
+ /// "feat_name" => expr
+ Explicit { feat: LitStr, body: Expr },
+ /// ! => expr
+ Not { body: Expr },
+ /// expr (no feature name — auto-detect by .await)
+ Implicit { body: Expr },
+}
+
+impl Parse for SelectArm {
+ fn parse(input: ParseStream) -> syn::Result<Self> {
+ if input.peek(LitStr) {
+ let feat: LitStr = input.parse()?;
+ input.parse::<Token![=>]>()?;
+ let body: Expr = input.parse()?;
+ Ok(SelectArm::Explicit { feat, body })
+ } else if input.peek(Token![!]) {
+ let _not: Token![!] = input.parse()?;
+ input.parse::<Token![=>]>()?;
+ let body: Expr = input.parse()?;
+ Ok(SelectArm::Not { body })
+ } else {
+ let body: Expr = input.parse()?;
+ Ok(SelectArm::Implicit { body })
+ }
+ }
+}
+
+impl ToTokens for SelectArm {
+ fn to_tokens(&self, tokens: &mut TokenStream2) {
+ match self {
+ SelectArm::Explicit { feat, body } => {
+ feat.to_tokens(tokens);
+ Token![=>].to_tokens(tokens);
+ body.to_tokens(tokens);
+ }
+ SelectArm::Not { body } => {
+ Token![!].to_tokens(tokens);
+ Token![=>].to_tokens(tokens);
+ body.to_tokens(tokens);
+ }
+ SelectArm::Implicit { body } => {
+ body.to_tokens(tokens);
+ }
+ }
+ }
+}
+
+// ─── SelectInput ────────────────────────────────────────────────────────────────────────
+
+struct SelectInput {
+ arms: Vec<SelectArm>,
+}
+
+impl Parse for SelectInput {
+ fn parse(input: ParseStream) -> syn::Result<Self> {
+ let punctuated: Punctuated<SelectArm, Token![,]> = Punctuated::parse_terminated(input)?;
+
+ if punctuated.len() != 2 {
+ return Err(syn::Error::new(
+ input.span(),
+ "select! requires exactly 2 arms",
+ ));
+ }
+
+ let arms: Vec<SelectArm> = punctuated.into_iter().collect();
+ Ok(SelectInput { arms })
+ }
+}
+
+// ─── Select macro ───────────────────────────────────────────────────────────────────────
+
+pub(crate) fn select(input: TokenStream) -> TokenStream {
+ let input = parse_macro_input!(input as SelectInput);
+ let expanded = input.expand();
+ TokenStream::from(expanded)
+}
+
+impl SelectInput {
+ fn expand(&self) -> TokenStream2 {
+ let arm0 = &self.arms[0];
+ let arm1 = &self.arms[1];
+
+ match (arm0, arm1) {
+ // Both explicit with feature names
+ (
+ SelectArm::Explicit { feat: f0, body: b0 },
+ SelectArm::Explicit { feat: f1, body: b1 },
+ ) => {
+ let f0_str = f0.value();
+ let f1_str = f1.value();
+
+ if has_not_prefix(&f0_str) {
+ let inner = &f0_str[1..];
+ quote! {{
+ #[cfg(not(feature = #inner))]
+ { #b0 }
+ #[cfg(feature = #inner)]
+ { #b1 }
+ }}
+ } else if has_not_prefix(&f1_str) {
+ let inner = &f1_str[1..];
+ quote! {{
+ #[cfg(feature = #f0)]
+ { #b0 }
+ #[cfg(not(feature = #f0))]
+ { #b1 }
+ }}
+ } else {
+ quote! {{
+ #[cfg(feature = #f0)]
+ { #b0 }
+ #[cfg(feature = #f1)]
+ { #b1 }
+ }}
+ }
+ }
+
+ // Explicit + Not
+ (SelectArm::Explicit { feat, body }, SelectArm::Not { body: not_body }) => {
+ let feat_str = feat.value();
+ if has_not_prefix(&feat_str) {
+ let inner = &feat_str[1..];
+ quote! {{
+ #[cfg(not(feature = #inner))]
+ { #body }
+ #[cfg(feature = #inner)]
+ { #not_body }
+ }}
+ } else {
+ quote! {{
+ #[cfg(feature = #feat_str)]
+ { #body }
+ #[cfg(not(feature = #feat_str))]
+ { #not_body }
+ }}
+ }
+ }
+ (SelectArm::Not { body: not_body }, SelectArm::Explicit { feat, body }) => {
+ let feat_str = feat.value();
+ if has_not_prefix(&feat_str) {
+ let inner = &feat_str[1..];
+ quote! {{
+ #[cfg(feature = #inner)]
+ { #body }
+ #[cfg(not(feature = #inner))]
+ { #not_body }
+ }}
+ } else {
+ quote! {{
+ #[cfg(not(feature = #feat_str))]
+ { #not_body }
+ #[cfg(feature = #feat_str)]
+ { #body }
+ }}
+ }
+ }
+
+ // Explicit + Implicit
+ (SelectArm::Explicit { feat, body }, SelectArm::Implicit { body: imp_body }) => {
+ let feat_str = feat.value();
+ if has_not_prefix(&feat_str) {
+ let inner = &feat_str[1..];
+ quote! {{
+ #[cfg(not(feature = #inner))]
+ { #body }
+ #[cfg(feature = #inner)]
+ { #imp_body }
+ }}
+ } else {
+ quote! {{
+ #[cfg(feature = #feat_str)]
+ { #body }
+ #[cfg(not(feature = #feat_str))]
+ { #imp_body }
+ }}
+ }
+ }
+ (SelectArm::Implicit { body: imp_body }, SelectArm::Explicit { feat, body }) => {
+ let feat_str = feat.value();
+ if has_not_prefix(&feat_str) {
+ let inner = &feat_str[1..];
+ quote! {{
+ #[cfg(feature = #inner)]
+ { #body }
+ #[cfg(not(feature = #inner))]
+ { #imp_body }
+ }}
+ } else {
+ quote! {{
+ #[cfg(not(feature = #feat_str))]
+ { #imp_body }
+ #[cfg(feature = #feat_str)]
+ { #body }
+ }}
+ }
+ }
+
+ // Both implicit — auto-detect .await
+ (SelectArm::Implicit { body: b0 }, SelectArm::Implicit { body: b1 }) => {
+ let b0_has_await = token_stream_has_await(&b0.to_token_stream());
+ let b1_has_await = token_stream_has_await(&b1.to_token_stream());
+
+ match (b0_has_await, b1_has_await) {
+ (true, false) => {
+ quote! {{
+ #[cfg(feature = "async")]
+ { #b0 }
+ #[cfg(not(feature = "async"))]
+ { #b1 }
+ }}
+ }
+ (false, true) => {
+ quote! {{
+ #[cfg(feature = "async")]
+ { #b1 }
+ #[cfg(not(feature = "async"))]
+ { #b0 }
+ }}
+ }
+ (true, true) => {
+ // Both have .await — use second as sync (strip .await)
+ let b1_stripped = strip_await_from_tokens(&b1.to_token_stream());
+ quote! {{
+ #[cfg(feature = "async")]
+ { #b0 }
+ #[cfg(not(feature = "async"))]
+ { #b1_stripped }
+ }}
+ }
+ (false, false) => {
+ // Neither has .await — use second as async (add .await)
+ quote! {{
+ #[cfg(feature = "async")]
+ { #b0 .await }
+ #[cfg(not(feature = "async"))]
+ { #b1 }
+ }}
+ }
+ }
+ }
+
+ // Two Not arms — doesn't make sense, fallback
+ (SelectArm::Not { body: b0 }, SelectArm::Not { body: b1 }) => {
+ quote! {{
+ #[cfg(feature = "async")]
+ { #b0 }
+ #[cfg(not(feature = "async"))]
+ { #b1 }
+ }}
+ }
+ }
+ }
+}
+
+// ─── Helpers ─────────────────────────────────────────────────────────────────────────────
+
+fn has_not_prefix(s: &str) -> bool {
+ s.starts_with('!')
+}
+
+fn token_stream_has_await(ts: &TokenStream2) -> bool {
+ let mut tokens = ts.clone().into_iter();
+ while let Some(token) = tokens.next() {
+ if let TokenTree::Punct(p) = &token {
+ if p.as_char() == '.' && p.spacing() == Spacing::Alone {
+ if let Some(TokenTree::Ident(ident)) = tokens.next() {
+ if ident == "await" {
+ return true;
+ }
+ }
+ }
+ }
+ }
+ false
+}
+
+fn strip_await_from_tokens(ts: &TokenStream2) -> TokenStream2 {
+ let tokens: Vec<_> = ts.clone().into_iter().collect();
+ let len = tokens.len();
+ if len >= 2 {
+ if let TokenTree::Punct(p) = &tokens[len - 2] {
+ if p.as_char() == '.' {
+ if let TokenTree::Ident(ident) = &tokens[len - 1] {
+ if ident == "await" {
+ return tokens[..len - 2].iter().cloned().collect();
+ }
+ }
+ }
+ }
+ }
+ ts.clone()
+}
generated by cgit (git 2.34.1) at 2026-08-17 17:39:50 +0000