Repository: astral-sh/ruff
File: crates/ruff_python_ast/src/nodes.rs (Rust)

Code before:
/// We rely on the same invariants outlined in the comment above `Parameters::len()`
/// in order to implement `ExactSizeIterator` here
impl ExactSizeIterator for ParametersIterator<'_> {}

impl<'a> IntoIterator for &'a Parameters {
    type IntoIter = ParametersIterator<'a>;
    type Item = AnyParameterRef<'a>;
    fn into_iter(self) -> Self::IntoIter {
        self.iter()
    }
}

impl<'a> IntoIterator for &'a Box<Parameters> {
    type IntoIter = ParametersIterator<'a>;
    type Item = AnyParameterRef<'a>;
    fn into_iter(self) -> Self::IntoIter {
        (&**self).into_iter()
    }
}

type NonVariadicParametersSourceOrderIterator<'a> = std::iter::Peekable<
    std::iter::Chain<
        std::iter::Chain<Iter<'a, ParameterWithDefault>, Iter<'a, ParameterWithDefault>>,
        Iter<'a, ParameterWithDefault>,
    >,
>;

/// The iterator returned by [`Parameters::iter_source_order`].
pub struct ParametersSourceOrderIterator<'a> {
    non_variadic: NonVariadicParametersSourceOrderIterator<'a>,
    variadics: [Option<&'a Parameter>; 2],
    next_variadic: usize,
}

impl<'a> ParametersSourceOrderIterator<'a> {
    /// Returns the next variadic parameter that appears before `before`, if any.
    fn next_variadic_before(&mut self, before: TextSize) -> Option<&'a Parameter> {
        let param = self.variadics.get(self.next_variadic).copied().flatten()?;
        if param.start() < before {
            self.next_variadic += 1;
            Some(param)
        } else {
            None
        }
    }
}

impl<'a> Iterator for ParametersSourceOrderIterator<'a> {
    type Item = AnyParameterRef<'a>;

    fn next(&mut self) -> Option<Self::Item> {

Review comment (suggestion):
that's a very impressive type alias, but I think I'd find something like this more readable 😆

```diff
diff --git a/crates/ruff_python_ast/src/nodes.rs b/crates/ruff_python_ast/src/nodes.rs
index 00287f5121..c52a2a1e2d 100644
--- a/crates/ruff_python_ast/src/nodes.rs
+++ b/crates/ruff_python_ast/src/nodes.rs
@@ -3115,12 +3115,11 @@ impl Parameters {
         variadics.sort_by_key(|param| param.map_or(TextSize::new(u32::MAX), Ranged::start));
 
         ParametersSourceOrderIterator {
-            non_variadic: self
-                .posonlyargs
-                .iter()
-                .chain(&self.args)
-                .chain(&self.kwonlyargs)
-                .peekable(),
+            non_variadic: NonVariadicParametersSourceOrderIterator {
+                posonlyargs: self.posonlyargs.iter(),
+                args: self.args.iter(),
+                kwonlyargs: self.kwonlyargs.iter(),
+            },
             variadics,
             next_variadic: 0,
         }
@@ -3266,12 +3265,34 @@ impl<'a> IntoIterator for &'a Box<Parameters> {
     }
 }
 
-type NonVariadicParametersSourceOrderIterator<'a> = std::iter::Peekable<
-    std::iter::Chain<
-        std::iter::Chain<Iter<'a, ParameterWithDefault>, Iter<'a, ParameterWithDefault>>,
-        Iter<'a, ParameterWithDefault>,
-    >,
->;
+struct NonVariadicParametersSourceOrderIterator<'a> {
+    posonlyargs: Iter<'a, ParameterWithDefault>,
+    args: Iter<'a, ParameterWithDefault>,
+    kwonlyargs: Iter<'a, ParameterWithDefault>,
+}
+
+impl<'a> NonVariadicParametersSourceOrderIterator<'a> {
+    fn peek(&mut self) -> Option<&'a ParameterWithDefault> {
+        self.posonlyargs
+            .as_slice()
+            .first()
+            .or_else(|| self.args.as_slice().first())
+            .or_else(|| self.kwonlyargs.as_slice().first())
+    }
+}
+
+impl<'a> Iterator for NonVariadicParametersSourceOrderIterator<'a> {
+    type Item = &'a ParameterWithDefault;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        self.posonlyargs
+            .next()
+            .or_else(|| self.args.next())
+            .or_else(|| self.kwonlyargs.next())
+    }
+}
+
+impl FusedIterator for NonVariadicParametersSourceOrderIterator<'_> {}
```

Or just store all the necessary state directly on the `ParametersSourceOrderIterator` struct?

```diff
diff --git a/crates/ruff_python_ast/src/nodes.rs b/crates/ruff_python_ast/src/nodes.rs
index 00287f5121..2fbbb431d6 100644
--- a/crates/ruff_python_ast/src/nodes.rs
+++ b/crates/ruff_python_ast/src/nodes.rs
@@ -3115,12 +3115,10 @@ impl Parameters {
         variadics.sort_by_key(|param| param.map_or(TextSize::new(u32::MAX), Ranged::start));
 
         ParametersSourceOrderIterator {
-            non_variadic: self
-                .posonlyargs
-                .iter()
-                .chain(&self.args)
-                .chain(&self.kwonlyargs)
-                .peekable(),
+            next_non_variadic_peeked: None,
+            posonlyargs: self.posonlyargs.iter(),
+            args: self.args.iter(),
+            kwonlyargs: self.kwonlyargs.iter(),
             variadics,
             next_variadic: 0,
         }
@@ -3266,16 +3264,12 @@ impl<'a> IntoIterator for &'a Box<Parameters> {
     }
 }
 
-type NonVariadicParametersSourceOrderIterator<'a> = std::iter::Peekable<
-    std::iter::Chain<
-        std::iter::Chain<Iter<'a, ParameterWithDefault>, Iter<'a, ParameterWithDefault>>,
-        Iter<'a, ParameterWithDefault>,
-    >,
->;
-
 /// The iterator returned by [`Parameters::iter_source_order`].
 pub struct ParametersSourceOrderIterator<'a> {
-    non_variadic: NonVariadicParametersSourceOrderIterator<'a>,
+    next_non_variadic_peeked: Option<&'a ParameterWithDefault>,
+    posonlyargs: Iter<'a, ParameterWithDefault>,
+    args: Iter<'a, ParameterWithDefault>,
+    kwonlyargs: Iter<'a, ParameterWithDefault>,
     variadics: [Option<&'a Parameter>; 2],
     next_variadic: usize,
 }
@@ -3291,6 +3285,20 @@ impl<'a> ParametersSourceOrderIterator<'a> {
             None
         }
     }
+
+    fn next_non_variadic(&mut self) -> Option<&'a ParameterWithDefault> {
+        self.next_non_variadic_peeked
+            .take()
+            .or_else(|| self.posonlyargs.next())
+            .or_else(|| self.args.next())
+            .or_else(|| self.kwonlyargs.next())
+    }
+
+    fn peek_next_non_variadic(&mut self) -> Option<&'a ParameterWithDefault> {
+        let next = self.next_non_variadic()?;
+        self.next_non_variadic_peeked = Some(next);
+        Some(next)
+    }
 }
 
 impl<'a> Iterator for ParametersSourceOrderIterator<'a> {
@@ -3299,16 +3307,14 @@ impl<'a> Iterator for ParametersSourceOrderIterator<'a> {
     fn next(&mut self) -> Option<Self::Item> {
         // If there's a variadic parameter that comes before the next
         // non-variadic parameter, emit it first.
-        let next_non_variadic_start = self
-            .non_variadic
-            .peek()
+        let next_non_variadic_start = self.peek_next_non_variadic()
             .map_or(TextSize::new(u32::MAX), Ranged::start);
 
         if let Some(variadic) = self.next_variadic_before(next_non_variadic_start) {
             return Some(AnyParameterRef::Variadic(variadic));
         }
 
-        if let Some(non_variadic) = self.non_variadic.next() {
+        if let Some(non_variadic) = self.next_non_variadic() {
             return Some(AnyParameterRef::NonVariadic(non_variadic));
         }
```

Code after:
impl ExactSizeIterator for ParametersIterator<'_> {}

impl<'a> IntoIterator for &'a Parameters {
    type IntoIter = ParametersIterator<'a>;
    type Item = AnyParameterRef<'a>;
    fn into_iter(self) -> Self::IntoIter {
        self.iter()
    }
}

impl<'a> IntoIterator for &'a Box<Parameters> {
    type IntoIter = ParametersIterator<'a>;
    type Item = AnyParameterRef<'a>;
    fn into_iter(self) -> Self::IntoIter {
        (&**self).into_iter()
    }
}

/// The iterator returned by [`Parameters::iter_source_order`].
pub struct ParametersSourceOrderIterator<'a> {
    next_non_variadic_peeked: Option<&'a ParameterWithDefault>,
    posonlyargs: Iter<'a, ParameterWithDefault>,
    args: Iter<'a, ParameterWithDefault>,
    kwonlyargs: Iter<'a, ParameterWithDefault>,
    variadics: [Option<&'a Parameter>; 2],
    next_variadic: usize,
}

impl<'a> ParametersSourceOrderIterator<'a> {
    /// Returns the next variadic parameter that appears before `before`, if any.
    fn next_variadic_before(&mut self, before: TextSize) -> Option<&'a Parameter> {
        let param = self.variadics.get(self.next_variadic).copied().flatten()?;
        if param.start() < before {
            self.next_variadic += 1;
            Some(param)
        } else {
            None
        }
    }

    fn next_non_variadic(&mut self) -> Option<&'a ParameterWithDefault> {
        self.next_non_variadic_peeked
            .take()
            .or_else(|| self.posonlyargs.next())
            .or_else(|| self.args.next())
            .or_else(|| self.kwonlyargs.next())
    }

    fn peek_next_non_variadic(&mut self) -> Option<&'a ParameterWithDefault> {
        let next = self.next_non_variadic()?;
        self.next_non_variadic_peeked = Some(next);