uv-pep508: add more methods for simplifying extra-related expressions

Cargo.toml

@@ -222,6 +222,7 @@ get_unwrap = "warn"
 rc_buffer = "warn"
 rc_mutex = "warn"
 rest_pat_in_fully_bound_structs = "warn"
+if_not_else = "allow"
 
 [profile.release]
 strip = true

crates/uv-pep508/src/marker/algebra.rs

@@ -412,7 +412,7 @@ impl InternerGuard<'_> {
                 // Restrict this variable to the given output by merging it
                 // with the relevant child.
                 let node = if value { high } else { low };
-                return node.negate(i);
+                return self.restrict(node.negate(i), f);
             }
         }
 
@@ -421,6 +421,70 @@ impl InternerGuard<'_> {
         self.create_node(node.var.clone(), children)
     }
 
+    /// Returns a new tree where the only nodes remaining are non-`extra`
+    /// nodes.
+    ///
+    /// If there are only `extra` nodes, then this returns a tree that is
+    /// always true.
+    ///
+    /// This works by assuming all `extra` nodes are always true.
+    ///
+    /// For example, given a marker like
+    /// `((os_name == ... and extra == foo) or (sys_platform == ... and extra != foo))`,
+    /// this would return a marker
+    /// `os_name == ... or sys_platform == ...`.
+    pub(crate) fn without_extras(&mut self, mut i: NodeId) -> NodeId {
+        if matches!(i, NodeId::TRUE | NodeId::FALSE) {
+            return i;
+        }
+
+        let parent = i;
+        let node = self.shared.node(i);
+        if matches!(node.var, Variable::Extra(_)) {
+            i = NodeId::FALSE;
+            for child in node.children.nodes() {
+                i = self.or(i, child.negate(parent));
+            }
+            if i.is_true() {
+                return NodeId::TRUE;
+            }
+            self.without_extras(i)
+        } else {
+            // Restrict all nodes recursively.
+            let children = node.children.map(i, |node| self.without_extras(node));
+            self.create_node(node.var.clone(), children)
+        }
+    }
+
+    /// Returns a new tree where the only nodes remaining are `extra` nodes.
+    ///
+    /// If there are no extra nodes, then this returns a tree that is always
+    /// true.
+    ///
+    /// This works by assuming all non-`extra` nodes are always true.
+    pub(crate) fn only_extras(&mut self, mut i: NodeId) -> NodeId {
+        if matches!(i, NodeId::TRUE | NodeId::FALSE) {
+            return i;
+        }
+
+        let parent = i;
+        let node = self.shared.node(i);
+        if !matches!(node.var, Variable::Extra(_)) {
+            i = NodeId::FALSE;
+            for child in node.children.nodes() {
+                i = self.or(i, child.negate(parent));
+            }
+            if i.is_true() {
+                return NodeId::TRUE;
+            }
+            self.only_extras(i)
+        } else {
+            // Restrict all nodes recursively.
+            let children = node.children.map(i, |node| self.only_extras(node));
+            self.create_node(node.var.clone(), children)
+        }
+    }
+
     /// Simplify this tree by *assuming* that the Python version range provided
     /// is true and that the complement of it is false.
     ///

crates/uv-pep508/src/marker/tree.rs

@@ -1063,6 +1063,21 @@ impl MarkerTree {
         self.simplify_extras_with(|name| extras.contains(name))
     }
 
+    /// Remove negated extras from a marker, returning `None` if the marker
+    /// tree evaluates to `true`.
+    ///
+    /// Any negated `extra` markers that are always `true` given the provided
+    /// extras will be removed. Any `extra` markers that are always `false`
+    /// given the provided extras will be left unchanged.
+    ///
+    /// For example, if `dev` is a provided extra, given `sys_platform
+    /// == 'linux' and extra != 'dev'`, the marker will be simplified to
+    /// `sys_platform == 'linux'`.
+    #[must_use]
+    pub fn simplify_not_extras(self, extras: &[ExtraName]) -> MarkerTree {
+        self.simplify_not_extras_with(|name| extras.contains(name))
+    }
+
     /// Remove the extras from a marker, returning `None` if the marker tree evaluates to `true`.
     ///
     /// Any `extra` markers that are always `true` given the provided predicate will be removed.
@@ -1082,12 +1097,57 @@ impl MarkerTree {
         self.simplify_extras_with_impl(&is_extra)
     }
 
+    /// Remove negated extras from a marker, returning `None` if the marker tree evaluates to
+    /// `true`.
+    ///
+    /// Any negated `extra` markers that are always `true` given the provided
+    /// predicate will be removed. Any `extra` markers that are always `false`
+    /// given the provided predicate will be left unchanged.
+    ///
+    /// For example, if `is_extra('dev')` is true, given
+    /// `sys_platform == 'linux' and extra != 'dev'`, the marker will be simplified to
+    /// `sys_platform == 'linux'`.
+    #[must_use]
+    pub fn simplify_not_extras_with(self, is_extra: impl Fn(&ExtraName) -> bool) -> MarkerTree {
+        // Because `simplify_extras_with_impl` is recursive, and we need to use
+        // our predicate in recursive calls, we need the predicate itself to
+        // have some indirection (or else we'd have to clone it). To avoid a
+        // recursive type at codegen time, we just introduce the indirection
+        // here, but keep the calling API ergonomic.
+        self.simplify_not_extras_with_impl(&is_extra)
+    }
+
+    /// Returns a new `MarkerTree` where all `extra` expressions are removed.
+    ///
+    /// If the marker only consisted of `extra` expressions, then a marker that
+    /// is always true is returned.
+    #[must_use]
+    pub fn without_extras(self) -> MarkerTree {
+        MarkerTree(INTERNER.lock().without_extras(self.0))
+    }
+
+    /// Returns a new `MarkerTree` where only `extra` expressions are removed.
+    ///
+    /// If the marker did not contain any `extra` expressions, then a marker
+    /// that is always true is returned.
+    #[must_use]
+    pub fn only_extras(self) -> MarkerTree {
+        MarkerTree(INTERNER.lock().only_extras(self.0))
+    }
+
     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),
             _ => None,
         }))
     }
+
+    fn simplify_not_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(false),
+            _ => None,
+        }))
+    }
 }
 
 impl fmt::Debug for MarkerTree {
@@ -2068,6 +2128,57 @@ mod test {
         assert_eq!(simplified, expected);
     }
 
+    #[test]
+    fn test_simplify_not_extras() {
+        // Given `os_name == "nt" and extra != "dev"`, simplify to `os_name == "nt"`.
+        let markers = MarkerTree::from_str(r#"os_name == "nt" and extra != "dev""#).unwrap();
+        let simplified = markers.simplify_not_extras(&[ExtraName::from_str("dev").unwrap()]);
+        let expected = MarkerTree::from_str(r#"os_name == "nt""#).unwrap();
+        assert_eq!(simplified, expected);
+
+        // Given `os_name == "nt" or extra != "dev"`, remove the marker entirely.
+        let markers = MarkerTree::from_str(r#"os_name == "nt" or extra != "dev""#).unwrap();
+        let simplified = markers.simplify_not_extras(&[ExtraName::from_str("dev").unwrap()]);
+        assert_eq!(simplified, MarkerTree::TRUE);
+
+        // Given `extra != "dev"`, remove the marker entirely.
+        let markers = MarkerTree::from_str(r#"extra != "dev""#).unwrap();
+        let simplified = markers.simplify_not_extras(&[ExtraName::from_str("dev").unwrap()]);
+        assert_eq!(simplified, MarkerTree::TRUE);
+
+        // Given `extra != "dev" and extra != "test"`, simplify to `extra != "test"`.
+        let markers = MarkerTree::from_str(r#"extra != "dev" and extra != "test""#).unwrap();
+        let simplified = markers.simplify_not_extras(&[ExtraName::from_str("dev").unwrap()]);
+        let expected = MarkerTree::from_str(r#"extra != "test""#).unwrap();
+        assert_eq!(simplified, expected);
+
+        // Given `os_name == "nt" and extra != "test"`, don't simplify.
+        let markers = MarkerTree::from_str(r#"os_name != "nt" and extra != "test""#).unwrap();
+        let simplified = markers.simplify_not_extras(&[ExtraName::from_str("dev").unwrap()]);
+        assert_eq!(simplified, markers);
+
+        // Given `os_name == "nt" and (python_version == "3.7" or extra != "dev")`, simplify to
+        // `os_name == "nt".
+        let markers = MarkerTree::from_str(
+            r#"os_name == "nt" and (python_version == "3.7" or extra != "dev")"#,
+        )
+        .unwrap();
+        let simplified = markers.simplify_not_extras(&[ExtraName::from_str("dev").unwrap()]);
+        let expected = MarkerTree::from_str(r#"os_name == "nt""#).unwrap();
+        assert_eq!(simplified, expected);
+
+        // Given `os_name == "nt" or (python_version == "3.7" and extra != "dev")`, simplify to
+        // `os_name == "nt" or python_version == "3.7"`.
+        let markers = MarkerTree::from_str(
+            r#"os_name == "nt" or (python_version == "3.7" and extra != "dev")"#,
+        )
+        .unwrap();
+        let simplified = markers.simplify_not_extras(&[ExtraName::from_str("dev").unwrap()]);
+        let expected =
+            MarkerTree::from_str(r#"os_name == "nt" or python_version == "3.7""#).unwrap();
+        assert_eq!(simplified, expected);
+    }
+
     #[test]
     fn test_marker_simplification() {
         assert_false("python_version == '3.9.1'");
@@ -3008,4 +3119,71 @@ mod test {
             m("python_full_version >= '3.9'"),
         );
     }
+
+    #[test]
+    fn without_extras() {
+        assert_eq!(
+            m("os_name == 'Linux'").without_extras(),
+            m("os_name == 'Linux'"),
+        );
+        assert!(m("extra == 'foo'").without_extras().is_true());
+        assert_eq!(
+            m("os_name == 'Linux' and extra == 'foo'").without_extras(),
+            m("os_name == 'Linux'"),
+        );
+
+        assert!(m("
+                (os_name == 'Linux' and extra == 'foo')
+                or (os_name != 'Linux' and extra == 'bar')",)
+        .without_extras()
+        .is_true());
+
+        assert_eq!(
+            m("os_name == 'Linux' and extra != 'foo'").without_extras(),
+            m("os_name == 'Linux'"),
+        );
+
+        assert!(
+            m("extra != 'extra-project-bar' and extra == 'extra-project-foo'")
+                .without_extras()
+                .is_true()
+        );
+
+        assert_eq!(
+            m("(os_name == 'Darwin' and extra == 'foo') \
+               or (sys_platform == 'Linux' and extra != 'foo')")
+            .without_extras(),
+            m("os_name == 'Darwin' or sys_platform == 'Linux'"),
+        );
+    }
+
+    #[test]
+    fn only_extras() {
+        assert!(m("os_name == 'Linux'").only_extras().is_true());
+        assert_eq!(m("extra == 'foo'").only_extras(), m("extra == 'foo'"));
+        assert_eq!(
+            m("os_name == 'Linux' and extra == 'foo'").only_extras(),
+            m("extra == 'foo'"),
+        );
+        assert!(m("
+                (os_name == 'foo' and extra == 'foo')
+                or (os_name == 'bar' and extra != 'foo')",)
+        .only_extras()
+        .is_true());
+        assert_eq!(
+            m("
+                (os_name == 'Linux' and extra == 'foo')
+                or (os_name != 'Linux' and extra == 'bar')")
+            .only_extras(),
+            m("extra == 'foo' or extra == 'bar'"),
+        );
+
+        assert_eq!(
+            m("
+                (implementation_name != 'pypy' and os_name == 'nt' and sys_platform == 'darwin')
+                or (os_name == 'nt' and sys_platform == 'win32')")
+            .only_extras(),
+            m("os_name == 'Linux' or os_name != 'Linux'"),
+        );
+    }
 }