Factor supported environments out of lockfile markers

crates/uv-distribution-types/src/requires_python.rs

@@ -612,6 +612,11 @@ impl SimplifiedMarkerTree {
     pub fn as_simplified_marker_tree(self) -> MarkerTree {
         self.0
     }
+
+    /// Combine this simplified marker with another via a conjunction.
+    pub fn and(&mut self, marker: Self) {
+        self.0.and(marker.0);
+    }
 }
 
 #[cfg(test)]

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

@@ -524,12 +524,16 @@ impl InternerGuard<'_> {
         }
     }
 
-    // Restrict the output of a given boolean variable in the tree.
+    // Restrict the output of selected boolean variables in the tree.
     //
     // If the provided function `f` returns a `Some` boolean value, the tree will be simplified
-    // with the assumption that the given variable is restricted to that value. If the function
+    // with the assumption that each variable is restricted to that value. If the function
     // returns `None`, the variable will not be affected.
-    pub(crate) fn restrict(&mut self, i: NodeId, f: &impl Fn(&Variable) -> Option<bool>) -> NodeId {
+    pub(crate) fn restrict_by(
+        &mut self,
+        i: NodeId,
+        f: &impl Fn(&Variable) -> Option<bool>,
+    ) -> NodeId {
         if matches!(i, NodeId::TRUE | NodeId::FALSE) {
             return i;
         }
@@ -540,15 +544,108 @@ impl InternerGuard<'_> {
                 // Restrict this variable to the given output by merging it
                 // with the relevant child.
                 let node = if value { high } else { low };
-                return self.restrict(node.negate(i), f);
+                return self.restrict_by(node.negate(i), f);
             }
         }
 
         // Restrict all nodes recursively.
-        let children = node.children.map(i, |node| self.restrict(node, f));
+        let children = node.children.map(i, |node| self.restrict_by(node, f));
         self.create_node(node.var.clone(), children)
     }
 
+    /// Restrict a marker by assuming that another marker is true.
+    ///
+    /// The returned marker is equivalent to `value` wherever `assumption` is true. Its value
+    /// outside of `assumption` is unspecified, which lets us eliminate decisions that are only
+    /// needed to restate the assumption.
+    pub(crate) fn restrict(&mut self, value: NodeId, assumption: NodeId) -> NodeId {
+        let mut cache = FxHashMap::default();
+        self.restrict_cached(value, assumption, &mut cache)
+    }
+
+    fn restrict_cached(
+        &mut self,
+        value: NodeId,
+        assumption: NodeId,
+        cache: &mut FxHashMap<(NodeId, NodeId), NodeId>,
+    ) -> NodeId {
+        if assumption.is_true() || matches!(value, NodeId::TRUE | NodeId::FALSE) {
+            return value;
+        }
+        if assumption.is_false() {
+            return NodeId::FALSE;
+        }
+        if value == assumption {
+            return NodeId::TRUE;
+        }
+        if value == assumption.not() {
+            return NodeId::FALSE;
+        }
+        if let Some(&result) = cache.get(&(value, assumption)) {
+            return result;
+        }
+
+        let value_node = self.shared.node(value);
+        let assumption_node = self.shared.node(assumption);
+        let result = match value_node.var.cmp(&assumption_node.var) {
+            Ordering::Less => {
+                let children = value_node.children.map(value, |value| {
+                    self.restrict_cached(value, assumption, cache)
+                });
+                self.create_node(value_node.var.clone(), children)
+            }
+            Ordering::Greater => {
+                // The value does not depend on this variable. Existentially quantify it out of the
+                // assumption, and continue with the remaining variables.
+                let mut quantified_assumption = NodeId::FALSE;
+                for child in assumption_node.children.nodes() {
+                    quantified_assumption =
+                        self.or(quantified_assumption, child.negate(assumption));
+                }
+                self.restrict_cached(value, quantified_assumption, cache)
+            }
+            Ordering::Equal => {
+                // Split both trees into matching ranges. Replace any ranges that are unreachable
+                // under the assumption with the first reachable child, simplifying them out of the
+                // resulting marker.
+                let mut fallback = None;
+                value_node.children.apply(
+                    value,
+                    &assumption_node.children,
+                    assumption,
+                    |value, assumption| {
+                        if assumption.is_false() {
+                            NodeId::FALSE
+                        } else {
+                            let result = self.restrict_cached(value, assumption, cache);
+                            fallback.get_or_insert(result);
+                            result
+                        }
+                    },
+                );
+                let Some(fallback) = fallback else {
+                    return NodeId::FALSE;
+                };
+                let children = value_node.children.apply(
+                    value,
+                    &assumption_node.children,
+                    assumption,
+                    |value, assumption| {
+                        if assumption.is_false() {
+                            fallback
+                        } else {
+                            self.restrict_cached(value, assumption, cache)
+                        }
+                    },
+                );
+                self.create_node(value_node.var.clone(), children)
+            }
+        };
+
+        cache.insert((value, assumption), result);
+        result
+    }
+
     /// Returns a new tree where the only nodes remaining are non-`extra`
     /// nodes.
     ///

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

@@ -1268,6 +1268,20 @@ impl MarkerTree {
         )
     }
 
+    /// Restrict this marker by assuming that `assumption` is true.
+    ///
+    /// The returned marker is equivalent to this marker wherever `assumption` is true, but may
+    /// have a different value outside of that context. Before evaluating the simplified marker,
+    /// callers should conjoin `assumption` to restore its standalone meaning.
+    ///
+    /// For example, restricting
+    /// `sys_platform == 'linux' and python_version < '3.11'` under the assumption
+    /// `sys_platform == 'linux'` produces `python_version < '3.11'`.
+    #[must_use]
+    pub fn restrict(self, assumption: Self) -> Self {
+        Self(INTERNER.lock().restrict(self.0, assumption.0))
+    }
+
     /// 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 extras will be removed.
@@ -1402,14 +1416,14 @@ impl MarkerTree {
     }
 
     fn simplify_extras_with_impl(self, is_extra: &impl Fn(&ExtraName) -> bool) -> Self {
-        Self(INTERNER.lock().restrict(self.0, &|var| match var {
+        Self(INTERNER.lock().restrict_by(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) -> Self {
-        Self(INTERNER.lock().restrict(self.0, &|var| match var {
+        Self(INTERNER.lock().restrict_by(self.0, &|var| match var {
             Variable::Extra(name) => is_extra(name.extra()).then_some(false),
             _ => None,
         }))
@@ -1948,6 +1962,57 @@ mod test {
         );
     }
 
+    #[test]
+    fn restrict() {
+        let environment = m(
+            "(platform_machine == 'x86_64' and sys_platform == 'darwin') or \
+             (platform_machine == 'x86_64' and sys_platform == 'linux') or \
+             (platform_machine == 'AMD64' and sys_platform == 'win32')",
+        );
+        let marker = m(
+            "((platform_machine == 'x86_64' and sys_platform == 'darwin') or \
+             (platform_machine == 'x86_64' and sys_platform == 'linux') or \
+             (platform_machine == 'AMD64' and sys_platform == 'win32')) and \
+             python_version < '3.11'",
+        );
+
+        let simplified = marker.restrict(environment);
+        assert_eq!(simplified, m("python_version < '3.11'"));
+
+        let mut reconstructed = simplified;
+        reconstructed.and(environment);
+        assert_eq!(reconstructed, marker);
+        assert_eq!(environment.restrict(environment), MarkerTree::TRUE);
+
+        let marker = m("python_version >= '3.12'");
+        let assumption = m("sys_platform == 'linux' or python_version >= '3.12'");
+        assert_eq!(marker.restrict(assumption), marker);
+
+        for (marker, assumption) in [
+            ("python_version < '3.11'", "sys_platform == 'linux'"),
+            ("sys_platform == 'linux'", "python_version < '3.11'"),
+            (
+                "sys_platform == 'linux' or python_version < '3.11'",
+                "sys_platform == 'darwin' or python_version >= '3.10'",
+            ),
+            (
+                "extra == 'foo' and sys_platform == 'linux'",
+                "extra == 'foo' or sys_platform == 'darwin'",
+            ),
+            ("python_version < '3.11'", "python_version >= '3.12'"),
+        ] {
+            let marker = m(marker);
+            let assumption = m(assumption);
+            let simplified = marker.restrict(assumption);
+
+            let mut expected = marker;
+            expected.and(assumption);
+            let mut reconstructed = simplified;
+            reconstructed.and(assumption);
+            assert_eq!(reconstructed, expected);
+        }
+    }
+
     #[test]
     fn release_only() {
         assert!(m("python_full_version > '3.10' or python_full_version <= '3.10'").is_true());

crates/uv-resolver/src/lock/mod.rs

@@ -1195,14 +1195,7 @@ impl Lock {
 
     /// The marker describing the universe of this resolution.
     fn fork_markers_union(&self) -> MarkerTree {
-        if self.fork_markers.is_empty() {
-            return self.requires_python.to_marker_tree();
-        }
-        let mut fork_markers_union = MarkerTree::FALSE;
-        for fork_marker in &self.fork_markers {
-            fork_markers_union.or(fork_marker.pep508());
-        }
-        fork_markers_union
+        fork_markers_union(&self.fork_markers, &self.requires_python)
     }
 
     /// Checks whether the fork markers cover the entire supported marker space.
@@ -2901,10 +2894,20 @@ impl TryFrom<LockWire> for Lock {
             unambiguous_package_ids.insert(dist.id.name.clone(), dist.id.clone());
         }
 
+        let fork_markers = wire
+            .fork_markers
+            .into_iter()
+            .map(|simplified_marker| simplified_marker.into_marker(&wire.requires_python))
+            .map(UniversalMarker::from_combined)
+            .collect::<Vec<_>>();
+        let environment = SimplifiedMarkerTree::new(
+            &wire.requires_python,
+            fork_markers_union(&fork_markers, &wire.requires_python),
+        );
         let packages = wire
             .packages
             .into_iter()
-            .map(|dist| dist.unwire(&wire.requires_python, &unambiguous_package_ids))
+            .map(|dist| dist.unwire(&wire.requires_python, environment, &unambiguous_package_ids))
             .collect::<Result<Vec<_>, _>>()?;
         let supported_environments = wire
             .supported_environments
@@ -2916,12 +2919,6 @@ impl TryFrom<LockWire> for Lock {
             .into_iter()
             .map(|simplified_marker| simplified_marker.into_marker(&wire.requires_python))
             .collect();
-        let fork_markers = wire
-            .fork_markers
-            .into_iter()
-            .map(|simplified_marker| simplified_marker.into_marker(&wire.requires_python))
-            .map(UniversalMarker::from_combined)
-            .collect();
         let mut options = wire.options;
         if options.exclude_newer.exclude_newer_span.is_some() {
             options.exclude_newer.exclude_newer = None;
@@ -3999,6 +3996,7 @@ impl PackageWire {
     fn unwire(
         self,
         requires_python: &RequiresPython,
+        environment: SimplifiedMarkerTree,
         unambiguous_package_ids: &FxHashMap<PackageName, PackageId>,
     ) -> Result<Package, LockError> {
         // Consistency check
@@ -4022,7 +4020,7 @@ impl PackageWire {
 
         let unwire_deps = |deps: Vec<DependencyWire>| -> Result<Vec<Dependency>, LockError> {
             deps.into_iter()
-                .map(|dep| dep.unwire(requires_python, unambiguous_package_ids))
+                .map(|dep| dep.unwire(requires_python, environment, unambiguous_package_ids))
                 .collect()
         };
 
@@ -5730,13 +5728,12 @@ impl Dependency {
                 .collect::<Array>();
             table.insert("extra", value(extra_array));
         }
-        // Avoid writing edge markers that are always fulfilled.
-        if !self
+        // Avoid restating the resolution's environment on every dependency edge.
+        if let Some(marker) = self
             .simplified_marker
             .as_simplified_marker_tree()
-            .negate()
-            .is_disjoint(simplified_environment)
-            && let Some(marker) = self.simplified_marker.try_to_string()
+            .restrict(simplified_environment)
+            .try_to_string()
         {
             table.insert("marker", value(marker));
         }
@@ -5793,13 +5790,16 @@ impl DependencyWire {
     fn unwire(
         self,
         requires_python: &RequiresPython,
+        environment: SimplifiedMarkerTree,
         unambiguous_package_ids: &FxHashMap<PackageName, PackageId>,
     ) -> Result<Dependency, LockError> {
-        let complexified_marker = self.marker.into_marker(requires_python);
+        let mut simplified_marker = self.marker;
+        simplified_marker.and(environment);
+        let complexified_marker = simplified_marker.into_marker(requires_python);
         Ok(Dependency {
             package_id: self.package_id.unwire(unambiguous_package_ids)?,
             extra: self.extra,
-            simplified_marker: self.marker,
+            simplified_marker,
             complexified_marker: UniversalMarker::from_combined(complexified_marker),
         })
     }
@@ -6948,6 +6948,21 @@ fn each_element_on_its_line_array(elements: impl Iterator<Item = impl Into<Value
     array
 }
 
+/// Return the PEP 508 marker space covered by the resolution.
+fn fork_markers_union(
+    fork_markers: &[UniversalMarker],
+    requires_python: &RequiresPython,
+) -> MarkerTree {
+    if fork_markers.is_empty() {
+        return requires_python.to_marker_tree();
+    }
+    let mut environment = MarkerTree::FALSE;
+    for fork_marker in fork_markers {
+        environment.or(fork_marker.pep508());
+    }
+    environment
+}
+
 /// Returns the simplified string-ified version of each marker given.
 ///
 /// Note that the marker strings returned will include conflict markers if they

crates/uv/tests/lock/lock.rs

@@ -33073,7 +33073,7 @@ fn windows_arm() -> Result<()> {
         version = "0.1.0"
         source = { virtual = "." }
         dependencies = [
-            { name = "pywin32", marker = "platform_machine == 'AMD64' and sys_platform == 'win32'" },
+            { name = "pywin32", marker = "sys_platform == 'win32'" },
         ]
 
         [package.metadata]