Rollup merge of #154201 - Zalathar:try-set-color, r=nnethercote

JonathanBrouwer · web-flow · commit baffc5e9bc04 · 2026-03-23T12:15:01.000+01:00
Use enums to clarify `DepNodeColorMap` color marking

When a function's documentation has to explain the meaning of nested results and options, then it is often a good candidate for using a custom result enum instead.

This PR also renames `DepNodeColorMap::try_mark` to `try_set_color`, to make it more distinct from the similarly-named `DepGraph::try_mark_green`.

The difference is that `try_mark_green` is a higher-level operation that tries to determine whether a node _can_ be marked green, whereas `try_set_color` is a lower-level operation that actually records a colour for the node.

The updated docs for `try_set_color` also fix a typo: *atomicaly* → *atomically*.

r? nnethercote (or compiler)
diff --git a/compiler/rustc_middle/src/dep_graph/graph.rs b/compiler/rustc_middle/src/dep_graph/graph.rs
@@ -164,9 +164,10 @@ impl DepGraph {
         );
         assert_eq!(red_node_index, DepNodeIndex::FOREVER_RED_NODE);
         if prev_graph_node_count > 0 {
-            colors.insert_red(SerializedDepNodeIndex::from_u32(
-                DepNodeIndex::FOREVER_RED_NODE.as_u32(),
-            ));
+            let prev_index =
+                const { SerializedDepNodeIndex::from_u32(DepNodeIndex::FOREVER_RED_NODE.as_u32()) };
+            let result = colors.try_set_color(prev_index, DesiredColor::Red);
+            assert_matches!(result, TrySetColorResult::Success);
         }
 
         DepGraph {
@@ -1415,28 +1416,29 @@ impl DepNodeColorMap {
         if value <= DepNodeIndex::MAX_AS_U32 { Some(DepNodeIndex::from_u32(value)) } else { None }
     }
 
-    /// This tries to atomically mark a node green and assign `index` as the new
-    /// index if `green` is true, otherwise it will try to atomicaly mark it red.
+    /// Atomically sets the color of a previous-session dep node to either green
+    /// or red, if it has not already been colored.
     ///
-    /// This returns `Ok` if `index` gets assigned or the node is marked red, otherwise it returns
-    /// the already allocated index in `Err` if it is green already. If it was already
-    /// red, `Err(None)` is returned.
+    /// If the node already has a color, the new color is ignored, and the
+    /// return value indicates the existing color.
     #[inline(always)]
-    pub(super) fn try_mark(
+    pub(super) fn try_set_color(
         &self,
         prev_index: SerializedDepNodeIndex,
-        index: DepNodeIndex,
-        green: bool,
-    ) -> Result<(), Option<DepNodeIndex>> {
-        let value = &self.values[prev_index];
-        match value.compare_exchange(
+        color: DesiredColor,
+    ) -> TrySetColorResult {
+        match self.values[prev_index].compare_exchange(
             COMPRESSED_UNKNOWN,
-            if green { index.as_u32() } else { COMPRESSED_RED },
+            match color {
+                DesiredColor::Red => COMPRESSED_RED,
+                DesiredColor::Green { index } => index.as_u32(),
+            },
             Ordering::Relaxed,
             Ordering::Relaxed,
         ) {
-            Ok(_) => Ok(()),
-            Err(v) => Err(if v == COMPRESSED_RED { None } else { Some(DepNodeIndex::from_u32(v)) }),
+            Ok(_) => TrySetColorResult::Success,
+            Err(COMPRESSED_RED) => TrySetColorResult::AlreadyRed,
+            Err(index) => TrySetColorResult::AlreadyGreen { index: DepNodeIndex::from_u32(index) },
         }
     }
 
@@ -1454,13 +1456,28 @@ impl DepNodeColorMap {
             DepNodeColor::Unknown
         }
     }
+}
 
-    #[inline]
-    pub(super) fn insert_red(&self, index: SerializedDepNodeIndex) {
-        let value = self.values[index].swap(COMPRESSED_RED, Ordering::Release);
-        // Sanity check for duplicate nodes
-        assert_eq!(value, COMPRESSED_UNKNOWN, "tried to color an already colored node as red");
-    }
+/// The color that [`DepNodeColorMap::try_set_color`] should try to apply to a node.
+#[derive(Clone, Copy, Debug)]
+pub(super) enum DesiredColor {
+    /// Try to mark the node red.
+    Red,
+    /// Try to mark the node green, associating it with a current-session node index.
+    Green { index: DepNodeIndex },
+}
+
+/// Return value of [`DepNodeColorMap::try_set_color`], indicating success or failure,
+/// and (on failure) what the existing color is.
+#[derive(Clone, Copy, Debug)]
+pub(super) enum TrySetColorResult {
+    /// The [`DesiredColor`] was freshly applied to the node.
+    Success,
+    /// Coloring failed because the node was already marked red.
+    AlreadyRed,
+    /// Coloring failed because the node was already marked green,
+    /// and corresponds to node `index` in the current-session dep graph.
+    AlreadyGreen { index: DepNodeIndex },
 }
 
 #[inline(never)]
diff --git a/compiler/rustc_middle/src/dep_graph/serialized.rs b/compiler/rustc_middle/src/dep_graph/serialized.rs
@@ -58,7 +58,7 @@ use rustc_serialize::{Decodable, Decoder, Encodable, Encoder};
 use rustc_session::Session;
 use tracing::{debug, instrument};
 
-use super::graph::{CurrentDepGraph, DepNodeColorMap};
+use super::graph::{CurrentDepGraph, DepNodeColorMap, DesiredColor, TrySetColorResult};
 use super::retained::RetainedDepGraph;
 use super::{DepKind, DepNode, DepNodeIndex};
 use crate::dep_graph::edges::EdgesVec;
@@ -905,13 +905,14 @@ impl GraphEncoder {
         let mut local = self.status.local.borrow_mut();
 
         let index = self.status.next_index(&mut *local);
+        let color = if is_green { DesiredColor::Green { index } } else { DesiredColor::Red };
 
-        // Use `try_mark` to avoid racing when `send_promoted` is called concurrently
+        // Use `try_set_color` to avoid racing when `send_promoted` is called concurrently
         // on the same index.
-        match colors.try_mark(prev_index, index, is_green) {
-            Ok(()) => (),
-            Err(None) => panic!("dep node {:?} is unexpectedly red", prev_index),
-            Err(Some(dep_node_index)) => return dep_node_index,
+        match colors.try_set_color(prev_index, color) {
+            TrySetColorResult::Success => {}
+            TrySetColorResult::AlreadyRed => panic!("dep node {prev_index:?} is unexpectedly red"),
+            TrySetColorResult::AlreadyGreen { index } => return index,
         }
 
         self.status.bump_index(&mut *local);
@@ -923,7 +924,8 @@ impl GraphEncoder {
     /// from the previous dep graph and expects all edges to already have a new dep node index
     /// assigned.
     ///
-    /// This will also ensure the dep node is marked green if `Some` is returned.
+    /// Tries to mark the dep node green, and returns Some if it is now green,
+    /// or None if had already been concurrently marked red.
     #[inline]
     pub(crate) fn send_promoted(
         &self,
@@ -935,10 +937,10 @@ impl GraphEncoder {
         let mut local = self.status.local.borrow_mut();
         let index = self.status.next_index(&mut *local);
 
-        // Use `try_mark_green` to avoid racing when `send_promoted` or `send_and_color`
+        // Use `try_set_color` to avoid racing when `send_promoted` or `send_and_color`
         // is called concurrently on the same index.
-        match colors.try_mark(prev_index, index, true) {
-            Ok(()) => {
+        match colors.try_set_color(prev_index, DesiredColor::Green { index }) {
+            TrySetColorResult::Success => {
                 self.status.bump_index(&mut *local);
                 self.status.encode_promoted_node(
                     index,
@@ -949,7 +951,8 @@ impl GraphEncoder {
                 );
                 Some(index)
             }
-            Err(dep_node_index) => dep_node_index,
+            TrySetColorResult::AlreadyRed => None,
+            TrySetColorResult::AlreadyGreen { index } => Some(index),
         }
     }
 
diff --git a/typos.toml b/typos.toml
@@ -48,6 +48,7 @@ unstalled = "unstalled" # short for un-stalled
 # the non-empty form can be automatically fixed by `--bless`.
 #
 # tidy-alphabetical-start
+atomicaly = "atomically"
 definitinon = "definition"
 dependy = ""
 similarlty = "similarity"

Original file line number	Diff line number	Diff line change
`@@ -48,6 +48,7 @@ unstalled = "unstalled" # short for un-stalled`
`48`	`48`	# the non-empty form can be automatically fixed by `--bless`.
`49`	`49`	`#`
`50`	`50`	`# tidy-alphabetical-start`
	`51`	`+atomicaly = "atomically"`
`51`	`52`	`definitinon = "definition"`
`52`	`53`	`dependy = ""`
`53`	`54`	`similarlty = "similarity"`