kt_list_comprehensions/

lib.rs

//! Python-like list comprehensions for Rust.
//! 
//! Provides a macro for list comprehensions with the following features:
//! - One or more nested list comprehensions
//! - Support for multiple `if` clauses
//! - Supports Rust expressions and patterns
//! 
//! # Examples
//! ```rust
//! use kt_list_comprehensions::list_comprehension;
//!
//! let vec = vec![-1, 2, 2, 3, 4];
//!
//! let result: Vec<i32> = list_comprehension![x * 2 for x in vec if x > 0].collect();
//!
//! assert_eq!(result, [4, 4, 6, 8]);
//! ```
//!
//! ```rust
//! use kt_list_comprehensions::list_comprehension;
//!
//! let vec_of_vectors = vec![vec![1, 2, 3], vec![4, 5, 6]];
//!
//! let result: Vec<i32> = list_comprehension![x for vec in vec_of_vectors for x in vec].collect();
//!
//! assert_eq!(result, [1, 2, 3, 4, 5, 6]);
//! ```
//! 
//! ```rust
//! use kt_list_comprehensions::list_comprehension;
//! 
//! let vec = vec!["1", "2", "not a number", "-3"];
//! 
//! let parse_i32 = |string: &str| string.parse::<i32>();
//! 
//! let result: Vec<i32> = list_comprehension![parse_i32(number).unwrap() for number in vec if parse_i32(number).is_ok() if parse_i32(number).unwrap() > 0].collect();
//! 
//! assert_eq!(result, [1, 2]);
//! ```

use proc_macro::TokenStream;
use syn::parse::ParseStream;

/// Python-like list comprehensions in Rust.
/// 
/// # Examples
/// ```rust
/// use kt_list_comprehensions::list_comprehension;
/// 
/// let vec = vec![-1, 2, 2, 3, 4];
///
/// let result: Vec<i32> = list_comprehension![x * 2 for x in vec if x > 0].collect();
/// 
/// assert_eq!(result, [4, 4, 6, 8]);
/// ```
/// 
/// ```rust
/// use kt_list_comprehensions::list_comprehension;
/// 
/// let vec_of_vectors = vec![vec![1, 2, 3], vec![4, 5, 6]];
///
/// let result: Vec<i32> = list_comprehension![x for vec in vec_of_vectors for x in vec].collect();
/// 
/// assert_eq!(result, [1, 2, 3, 4, 5, 6]);
/// ```
///
/// ```rust
/// use kt_list_comprehensions::list_comprehension;
///
/// let vec = vec!["1", "2", "not a number", "-3"];
///
/// let parse_i32 = |string: &str| string.parse::<i32>();
///
/// let result: Vec<i32> = list_comprehension![parse_i32(number).unwrap() for number in vec if parse_i32(number).is_ok() if parse_i32(number).unwrap() > 0].collect();
///
/// assert_eq!(result, [1, 2]);
/// ```
#[proc_macro]
pub fn list_comprehension(token_stream: TokenStream) -> TokenStream {
	let comprehension = syn::parse_macro_input!(token_stream as Comprehension);
	
	let tokens = quote::quote! {
		#comprehension
	};
	
	tokens.into()
}

struct Comprehension {
	mapping: Mapping,
	for_if_clause: ForIfClause,
	additional_for_if_clauses: Vec<ForIfClause>,
}

impl syn::parse::Parse for Comprehension {
	fn parse(input: ParseStream) -> syn::Result<Self> {
		Ok(Self {
			mapping: input.parse::<Mapping>()?,
			for_if_clause: input.parse::<ForIfClause>()?,
			additional_for_if_clauses: parse_zero_or_more::<ForIfClause>(input)
		})
	}
}

impl quote::ToTokens for Comprehension {
	fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
		let all_clauses = core::iter::once(&self.for_if_clause)
			.chain(&self.additional_for_if_clauses);
		
		let mut clause_iterator = all_clauses.rev();
		
		let mut output = {
			let innermost_clause = clause_iterator
				.next()
				.expect("List comprehensions require at least one for-if clause");

			let ForIfClause {
				pattern,
				expression: sequence,
				conditions
			} = innermost_clause;
			
			let Mapping(mapping) = &self.mapping;
			
			quote::quote! {
				::core::iter::IntoIterator::into_iter(#sequence)
					.filter_map(move |#pattern| {
						(true #(&& (#conditions))*)
							.then(|| #mapping)
					})
			}
		};
		
		output = clause_iterator.fold(output, |current_output, next_clause| {
			let ForIfClause {
				pattern,
				expression: sequence,
				conditions,
			} = next_clause;
			
			quote::quote! {
				::core::iter::IntoIterator::into_iter(#sequence)
					.filter_map(move |#pattern| {
						(true #(&& (#conditions))*)
							.then(|| #current_output)
					})
					.flatten()
			}
		});
		
		tokens.extend(output);
	}
}

struct Mapping(syn::Expr);

impl syn::parse::Parse for Mapping {
	fn parse(input: ParseStream) -> syn::Result<Self> {
		input.parse::<syn::Expr>()
			.map(Self)
	}
}

impl quote::ToTokens for Mapping {
	fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
		self.0.to_tokens(tokens);
	}
}

struct ForIfClause {
	pattern: Pattern,
	expression: syn::Expr,
	conditions: Vec<Condition>,
}

impl syn::parse::Parse for ForIfClause {
	fn parse(input: ParseStream) -> syn::Result<Self> {
		input.parse::<syn::Token![for]>()?;
		
		let pattern = input.parse()?;
		
		input.parse::<syn::Token![in]>()?;
		
		let expression = input.parse()?;
		
		let conditions = parse_zero_or_more::<Condition>(input);
		
		Ok(Self {
			pattern,
			expression,
			conditions
		})
	}
}

struct Pattern(syn::Pat);

impl syn::parse::Parse for Pattern {
	fn parse(input: ParseStream) -> syn::Result<Self> {
		input.call(syn::Pat::parse_single)
			.map(Self)
	}
}

impl quote::ToTokens for Pattern {
	fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
		self.0.to_tokens(tokens);
	}
}

struct Condition(syn::Expr);

impl syn::parse::Parse for Condition {
	fn parse(input: ParseStream) -> syn::Result<Self> {
		input.parse::<syn::Token![if]>()?;
		
		input.parse::<syn::Expr>()
			.map(Self)
	}
}

impl quote::ToTokens for Condition {
	fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
		self.0.to_tokens(tokens);
	}
}

fn parse_zero_or_more<T>(input: ParseStream) -> Vec<T>
where
	T: syn::parse::Parse
{
	let mut out = Vec::new();
	
	while let Ok(item) = input.parse::<T>() {
		out.push(item);
	}
	
	out
}