From b48e6b98625352e31496fb061a653f90c1e001dc Mon Sep 17 00:00:00 2001
From: Andrew Gallant <andrew@astral.sh>
Date: Thu, 6 Feb 2025 13:47:34 -0500
Subject: [PATCH] uv-pep508: add routine for visiting all extras in a marker
 tree

We'll need this to be able to decode conflict markers back into
a series of inclusion/exclusion rules used during resolution.
---
 crates/uv-pep508/src/marker/tree.rs | 89 +++++++++++++++++++++++++++++
 1 file changed, 89 insertions(+)

diff --git a/crates/uv-pep508/src/marker/tree.rs b/crates/uv-pep508/src/marker/tree.rs
index 7cdfad6f2..447a77ee7 100644
--- a/crates/uv-pep508/src/marker/tree.rs
+++ b/crates/uv-pep508/src/marker/tree.rs
@@ -1162,6 +1162,49 @@ impl MarkerTree {
         MarkerTree(INTERNER.lock().only_extras(self.0))
     }
 
+    /// Calls the provided function on every `extra` in this tree.
+    ///
+    /// The operator provided to the function is guaranteed to be
+    /// `MarkerOperator::Equal` or `MarkerOperator::NotEqual`.
+    pub fn visit_extras(self, mut f: impl FnMut(MarkerOperator, &ExtraName)) {
+        fn imp(tree: MarkerTree, f: &mut impl FnMut(MarkerOperator, &ExtraName)) {
+            match tree.kind() {
+                MarkerTreeKind::True | MarkerTreeKind::False => {}
+                MarkerTreeKind::Version(kind) => {
+                    for (tree, _) in simplify::collect_edges(kind.edges()) {
+                        imp(tree, f);
+                    }
+                }
+                MarkerTreeKind::String(kind) => {
+                    for (tree, _) in simplify::collect_edges(kind.children()) {
+                        imp(tree, f);
+                    }
+                }
+                MarkerTreeKind::In(kind) => {
+                    for (_, tree) in kind.children() {
+                        imp(tree, f);
+                    }
+                }
+                MarkerTreeKind::Contains(kind) => {
+                    for (_, tree) in kind.children() {
+                        imp(tree, f);
+                    }
+                }
+                MarkerTreeKind::Extra(kind) => {
+                    if kind.low.is_false() {
+                        f(MarkerOperator::Equal, kind.name().extra());
+                    } else {
+                        f(MarkerOperator::NotEqual, kind.name().extra());
+                    }
+                    for (_, tree) in kind.children() {
+                        imp(tree, f);
+                    }
+                }
+            }
+        }
+        imp(self, &mut f);
+    }
+
     fn simplify_extras_with_impl(self, is_extra: &impl Fn(&ExtraName) -> bool) -> MarkerTree {
         MarkerTree(INTERNER.lock().restrict(self.0, &|var| match var {
             Variable::Extra(name) => is_extra(name.extra()).then_some(true),
@@ -3211,4 +3254,50 @@ mod test {
             m("os_name == 'Linux' or os_name != 'Linux'"),
         );
     }
+
+    #[test]
+    fn visit_extras() {
+        let collect = |m: MarkerTree| -> Vec<(&'static str, String)> {
+            let mut collected = vec![];
+            m.visit_extras(|op, extra| {
+                let op = match op {
+                    MarkerOperator::Equal => "==",
+                    MarkerOperator::NotEqual => "!=",
+                    _ => unreachable!(),
+                };
+                collected.push((op, extra.as_str().to_string()));
+            });
+            collected
+        };
+        assert_eq!(collect(m("os_name == 'Linux'")), vec![]);
+        assert_eq!(
+            collect(m("extra == 'foo'")),
+            vec![("==", "foo".to_string())]
+        );
+        assert_eq!(
+            collect(m("extra == 'Linux' and extra == 'foo'")),
+            vec![("==", "foo".to_string()), ("==", "linux".to_string())]
+        );
+        assert_eq!(
+            collect(m("os_name == 'Linux' and extra == 'foo'")),
+            vec![("==", "foo".to_string())]
+        );
+
+        let marker = "
+            python_full_version >= '3.12'
+            and sys_platform == 'darwin'
+            and extra == 'extra-27-resolution-markers-for-days-cpu'
+            and extra != 'extra-27-resolution-markers-for-days-cu124'
+        ";
+        assert_eq!(
+            collect(m(marker)),
+            vec![
+                ("==", "extra-27-resolution-markers-for-days-cpu".to_string()),
+                (
+                    "!=",
+                    "extra-27-resolution-markers-for-days-cu124".to_string()
+                ),
+            ]
+        );
+    }
 }