1.2系列提交: 在lib中新增iter()调用判别函数, 并正确应用到集合推导式中(新增visit依赖)

Author: Natural-selection1

Cargo.toml

@@ -14,13 +14,13 @@ repository = "https://github.com/Natural-selection1/better-comprehension-in-rust
 readme = "README.md"
 version = "1.1.0"
 edition = "2024"
-exclude = ["src/main.rs", "target", ".vscode", "README-CN.md",".gitignore"]
+exclude = ["src/main.rs", "target", ".vscode", "README-CN.md", ".gitignore"]
 
 
 [lib]
 proc-macro = true
 
 [dependencies]
-syn = { version = "2.0.100", features = ["full", "extra-traits"] }
+syn = { version = "2.0.100", features = ["full", "visit", "extra-traits"] }
 quote = "1.0.40"
 proc-macro2 = "1.0.94"

src/eager_evaluation.rs

@@ -43,19 +43,18 @@ pub(crate) fn handle_nested_loops<'a>(
                 need_to_shadow.push(iterable);
                 quote! { &#iterable }
             }
-            Expr::MethodCall(method_call) => {
-                let method_name = &method_call.method;
-                if method_name == "iter" {
-                    quote! { #iterable }
-                } else {
-                    panic!("不支持的类型: {:?}", iterable);
-                }
-            }
+            Expr::MethodCall(_) => match crate::is_iter(iterable) {
+                true => quote! { #iterable },
+                _ => panic!(
+                    "please ensure the first method call is iter(): {:?}",
+                    iterable
+                ),
+            },
             Expr::Paren(expr) => {
                 let iterable = &*expr.expr;
                 quote! { #iterable }
             }
-            _ => panic!("不支持的类型: {:?}", iterable),
+            _ => panic!("type is not supported: {:?}", iterable),
         };
 
         // 根据是否有if条件生成循环代码

src/lib.rs

@@ -267,6 +267,7 @@
 //!
 //! b_tree_map! : insert() to add key-value pairs
 
+// #![allow(unused, dead_code)]
 use proc_macro::TokenStream as TS;
 
 mod eager_evaluation;
@@ -344,3 +345,43 @@ pub(crate) fn common_parse(
 
     (mapping, iter_clauses)
 }
+
+struct IterMethodCallFinder {
+    is_iter: bool,
+}
+
+use syn::{ExprMethodCall, visit::Visit};
+impl<'ast> Visit<'ast> for IterMethodCallFinder {
+    fn visit_expr_method_call(&mut self, node: &'ast ExprMethodCall) {
+        match *node.receiver {
+            syn::Expr::Path(_) => {
+                if node.method == "iter" {
+                    self.is_iter = true;
+                }
+            }
+            _ => syn::visit::visit_expr(&mut *self, &node.receiver),
+        }
+    }
+}
+
+pub(crate) fn is_iter(expr: &syn::Expr) -> bool {
+    let mut finder = IterMethodCallFinder { is_iter: false };
+    finder.visit_expr(expr);
+
+    finder.is_iter
+}
+
+#[test]
+fn test_is_iter() {
+    // 最右侧是iter方法
+    let expr = syn::parse_quote!(some.method_1().method_2().iter());
+    assert!(!is_iter(&expr));
+    eprintln!("--------------------------------");
+    // 最左侧是iter方法
+    let expr = syn::parse_quote!(some.iter().method_3().method_4());
+    assert!(is_iter(&expr));
+    eprintln!("--------------------------------");
+
+    let expr = syn::parse_quote!(some.method_5().iter().method_6());
+    assert!(!is_iter(&expr));
+}

src/main.rs

@@ -18,6 +18,7 @@ fn main() {
     test_pattern_matching();
     test_nested_comprehension();
     test_ownership_handling();
+    test_option();
 }
 
 fn test_vec() {
@@ -350,3 +351,19 @@ fn test_ref_iterator() {
     assert_eq!(vec_1.len(), 3);
     assert_eq!(vec_2.len(), 3);
 }
+
+fn test_option() {
+    let vec = [Some("1".to_string()), None, Some("3".to_string())];
+
+    let _result = vector![
+        x.clone()
+        for x in vec.iter().flatten()
+    ];
+
+    let mut result = Vec::new();
+    for x in vec.iter().flatten() {
+        result.push(x.clone());
+    }
+
+    assert_eq!(result, vec!["1".to_string(), "3".to_string()]);
+}