uv-resolver: almost connect everything together

This almost connects everything together. Specifically, we propagate
activated extras through the resolution graph and simplify conflict
markers using our world knowledge about which extras aren't allowed to
be combined together.

One thing that's still missing here is the filtering logic to ensure
that conflict markers are respected when traversing the graph from the
lock file.

crates/uv-pypi-types/src/conflicts.rs

@@ -24,7 +24,7 @@ impl Conflicts {
     }
 
     /// Returns an iterator over all sets of conflicting sets.
-    pub fn iter(&self) -> impl Iterator<Item = &'_ ConflictSet> + '_ {
+    pub fn iter(&self) -> impl Iterator<Item = &'_ ConflictSet> + Clone + '_ {
         self.0.iter()
     }
 
@@ -75,7 +75,7 @@ impl ConflictSet {
     }
 
     /// Returns an iterator over all conflicting items.
-    pub fn iter(&self) -> impl Iterator<Item = &'_ ConflictItem> + '_ {
+    pub fn iter(&self) -> impl Iterator<Item = &'_ ConflictItem> + Clone + '_ {
         self.0.iter()
     }
 

crates/uv-resolver/src/graph_ops.rs

@@ -1,9 +1,12 @@
-use petgraph::graph::NodeIndex;
+use petgraph::graph::{EdgeIndex, NodeIndex};
 use petgraph::visit::EdgeRef;
 use petgraph::{Direction, Graph};
-use rustc_hash::{FxBuildHasher, FxHashMap};
+use rustc_hash::{FxBuildHasher, FxHashMap, FxHashSet};
 use std::collections::hash_map::Entry;
 
+use uv_normalize::ExtraName;
+
+use crate::resolution::ResolutionGraphNode;
 use crate::universal_marker::UniversalMarker;
 
 /// Determine the markers under which a package is reachable in the dependency tree.
@@ -79,3 +82,57 @@ pub(crate) fn marker_reachability<T>(
 
     reachability
 }
+
+/// Traverse the given dependency graph and propagate activated markers.
+///
+/// For example, given an edge like `foo[x1] -> bar`, then it is known that
+/// `x1` is activated. This in turn can be used to simplify any downstream
+/// conflict markers with `extra == "x1"` in them.
+pub(crate) fn simplify_conflict_markers(graph: &mut Graph<ResolutionGraphNode, UniversalMarker>) {
+    // The set of activated extras (and TODO, in the future, groups)
+    // for each node. The ROOT nodes don't have any extras activated.
+    let mut activated: FxHashMap<NodeIndex, FxHashSet<ExtraName>> =
+        FxHashMap::with_capacity_and_hasher(graph.node_count(), FxBuildHasher);
+
+    // Collect the root nodes.
+    //
+    // Besides the actual virtual root node, virtual dev dependencies packages are also root
+    // nodes since the edges don't cover dev dependencies.
+    let mut queue: Vec<_> = graph
+        .node_indices()
+        .filter(|node_index| {
+            graph
+                .edges_directed(*node_index, Direction::Incoming)
+                .next()
+                .is_none()
+        })
+        .collect();
+
+    let mut assume_by_edge: FxHashMap<EdgeIndex, FxHashSet<ExtraName>> = FxHashMap::default();
+    while let Some(parent_index) = queue.pop() {
+        for child_edge in graph.edges_directed(parent_index, Direction::Outgoing) {
+            // TODO: The below seems excessively clone-y.
+            // Consider tightening this up a bit.
+            let target = child_edge.target();
+            let mut extras: FxHashSet<ExtraName> =
+                activated.get(&parent_index).cloned().unwrap_or_default();
+            if let Some(extra) = graph[parent_index].extra() {
+                extras.insert(extra.clone());
+            }
+            if let Some(extra) = graph[target].extra() {
+                extras.insert(extra.clone());
+            }
+            activated.entry(target).or_default().extend(extras.clone());
+            assume_by_edge
+                .entry(child_edge.id())
+                .or_default()
+                .extend(extras);
+            queue.push(child_edge.target());
+        }
+    }
+    for (edge_id, extras) in assume_by_edge {
+        for extra in &extras {
+            graph[edge_id].assume_extra(extra);
+        }
+    }
+}

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

@@ -3580,6 +3580,9 @@ impl Dependency {
         if let Some(marker) = self.simplified_marker.try_to_string() {
             table.insert("marker", value(marker));
         }
+        if let Some(conflict_marker) = self.complexified_marker.conflict().try_to_string() {
+            table.insert("conflict-marker", value(conflict_marker));
+        }
 
         table
     }
@@ -3615,7 +3618,8 @@ struct DependencyWire {
     extra: BTreeSet<ExtraName>,
     #[serde(default)]
     marker: SimplifiedMarkerTree,
-    // FIXME: Add support for representing conflict markers.
+    #[serde(default)]
+    conflict_marker: MarkerTree,
 }
 
 impl DependencyWire {
@@ -3629,8 +3633,7 @@ impl DependencyWire {
             package_id: self.package_id.unwire(unambiguous_package_ids)?,
             extra: self.extra,
             simplified_marker: self.marker,
-            // FIXME: Support reading conflict markers.
-            complexified_marker: UniversalMarker::new(complexified_marker, MarkerTree::TRUE),
+            complexified_marker: UniversalMarker::new(complexified_marker, self.conflict_marker),
         })
     }
 }

crates/uv-resolver/src/resolution/output.rs

@@ -21,7 +21,7 @@ use uv_pypi_types::{
     Conflicts, HashDigest, ParsedUrlError, Requirement, VerbatimParsedUrl, Yanked,
 };
 
-use crate::graph_ops::marker_reachability;
+use crate::graph_ops::{marker_reachability, simplify_conflict_markers};
 use crate::pins::FilePins;
 use crate::preferences::Preferences;
 use crate::redirect::url_to_precise;
@@ -72,6 +72,13 @@ impl ResolutionGraphNode {
             ResolutionGraphNode::Dist(dist) => &dist.marker,
         }
     }
+
+    pub(crate) fn extra(&self) -> Option<&ExtraName> {
+        match *self {
+            ResolutionGraphNode::Dist(ref dist) => dist.extra.as_ref(),
+            _ => None,
+        }
+    }
 }
 
 impl Display for ResolutionGraphNode {
@@ -179,12 +186,18 @@ impl ResolverOutput {
         // Compute and apply the marker reachability.
         let mut reachability = marker_reachability(&graph, &fork_markers);
 
-        // Apply the reachability to the graph.
+        // Apply the reachability to the graph and imbibe world
+        // knowledge about conflicts.
         for index in graph.node_indices() {
             if let ResolutionGraphNode::Dist(dist) = &mut graph[index] {
                 dist.marker = reachability.remove(&index).unwrap_or_default();
+                dist.marker.imbibe(conflicts);
             }
         }
+        for weight in graph.edge_weights_mut() {
+            weight.imbibe(conflicts);
+        }
+        simplify_conflict_markers(&mut graph);
 
         // Discard any unreachable nodes.
         graph.retain_nodes(|graph, node| !graph[node].marker().is_false());

crates/uv-resolver/src/resolver/environment.rs

@@ -379,6 +379,7 @@ impl ResolverEnvironment {
                 ref exclude,
                 ..
             } => {
+                // FIXME: Account for groups here in addition to extras.
                 let mut conflict_marker = MarkerTree::TRUE;
                 for item in exclude.iter() {
                     if let Some(extra) = item.extra() {

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

@@ -2614,7 +2614,9 @@ pub(crate) struct ResolutionDependencyEdge {
 
 impl ResolutionDependencyEdge {
     pub(crate) fn universal_marker(&self) -> UniversalMarker {
-        // FIXME: Account for extras and groups here.
+        // We specifically do not account for conflict
+        // markers here. Instead, those are computed via
+        // a traversal on the resolution graph.
         UniversalMarker::new(self.marker.clone(), MarkerTree::TRUE)
     }
 }

crates/uv-resolver/src/universal_marker.rs

@@ -1,5 +1,8 @@
+use itertools::Itertools;
+
 use uv_normalize::ExtraName;
 use uv_pep508::{MarkerEnvironment, MarkerTree};
+use uv_pypi_types::Conflicts;
 
 /// A representation of a marker for use in universal resolution.
 ///
@@ -56,6 +59,61 @@ impl UniversalMarker {
         self.conflict_marker.and(other.conflict_marker);
     }
 
+    /// Imbibes the world knowledge expressed by `conflicts` into this marker.
+    ///
+    /// This will effectively simplify the conflict marker in this universal
+    /// marker. In particular, it enables simplifying based on the fact that no
+    /// two items from the same set in the given conflicts can be active at a
+    /// given time.
+    pub(crate) fn imbibe(&mut self, conflicts: &Conflicts) {
+        if conflicts.is_empty() {
+            return;
+        }
+        // TODO: This is constructing what could be a big
+        // marker (depending on how many conflicts there are),
+        // which is invariant throughout the lifetime of the
+        // program. But it's doing it every time this routine
+        // is called. We should refactor the caller to build
+        // a marker from the `conflicts` once.
+        let mut marker = MarkerTree::FALSE;
+        for set in conflicts.iter() {
+            for (item1, item2) in set.iter().tuple_combinations() {
+                // FIXME: Account for groups here. And extra/group
+                // combinations too.
+                let (Some(extra1), Some(extra2)) = (item1.extra(), item2.extra()) else {
+                    continue;
+                };
+
+                let operator = uv_pep508::ExtraOperator::Equal;
+                let name = uv_pep508::MarkerValueExtra::Extra(extra1.clone());
+                let expr = uv_pep508::MarkerExpression::Extra { operator, name };
+                let marker1 = MarkerTree::expression(expr);
+
+                let operator = uv_pep508::ExtraOperator::Equal;
+                let name = uv_pep508::MarkerValueExtra::Extra(extra2.clone());
+                let expr = uv_pep508::MarkerExpression::Extra { operator, name };
+                let marker2 = MarkerTree::expression(expr);
+
+                let mut pair = MarkerTree::TRUE;
+                pair.and(marker1);
+                pair.and(marker2);
+                marker.or(pair);
+            }
+        }
+        let mut marker = marker.negate();
+        marker.implies(std::mem::take(&mut self.conflict_marker));
+        self.conflict_marker = marker;
+    }
+
+    /// Assumes that a given extra is activated.
+    ///
+    /// This may simplify the conflicting marker component of this universal
+    /// marker.
+    pub(crate) fn assume_extra(&mut self, extra: &ExtraName) {
+        self.conflict_marker = std::mem::take(&mut self.conflict_marker)
+            .simplify_extras_with(|candidate| candidate == extra)
+    }
+
     /// Returns true if this universal marker will always evaluate to `true`.
     pub(crate) fn is_true(&self) -> bool {
         self.pep508_marker.is_true() && self.conflict_marker.is_true()