[naga] Use new NonMaxU32 type for Handle indices.

Define a new type, `NonMaxU32`, that can represent any `u32` value
except `u32::MAX`, defined such that `Option<NonMaxU32>` is still a
32-bit value.

Change `Handle` to use `NonMaxU32`. Adjust all code that works
directly with handle indices to assume zero-based indices, not
one-based indices.

naga/src/arena.rs

@@ -1,9 +1,11 @@
-use std::{cmp::Ordering, fmt, hash, marker::PhantomData, num::NonZeroU32, ops};
+use std::{cmp::Ordering, fmt, hash, marker::PhantomData, ops};
+
+use crate::non_max_u32::NonMaxU32;
 
 /// An unique index in the arena array that a handle points to.
-/// The "non-zero" part ensures that an `Option<Handle<T>>` has
+/// The "non-max" part ensures that an `Option<Handle<T>>` has
 /// the same size and representation as `Handle<T>`.
-type Index = NonZeroU32;
+type Index = NonMaxU32;
 
 use crate::{FastIndexSet, Span};
 
@@ -89,13 +91,12 @@ impl<T> Handle<T> {
 
     /// Returns the zero-based index of this handle.
     pub const fn index(self) -> usize {
-        let index = self.index.get() - 1;
+        self.index.get() as usize
-        index as usize
     }
 
     /// Convert a `usize` index into a `Handle<T>`.
     fn from_usize(index: usize) -> Self {
-        let handle_index = u32::try_from(index + 1)
+        let handle_index = u32::try_from(index)
             .ok()
             .and_then(Index::new)
             .expect("Failed to insert into arena. Handle overflows");
@@ -104,7 +105,7 @@ impl<T> Handle<T> {
 
     /// Convert a `usize` index into a `Handle<T>`, without range checks.
     const unsafe fn from_usize_unchecked(index: usize) -> Self {
-        Handle::new(Index::new_unchecked((index + 1) as u32))
+        Handle::new(Index::new_unchecked(index as u32))
     }
 
     /// Write this handle's index to `formatter`, preceded by `prefix`.
@@ -174,7 +175,7 @@ impl<T> Clone for Range<T> {
 
 impl<T> fmt::Debug for Range<T> {
     fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
-        write!(formatter, "[{}..{}]", self.inner.start + 1, self.inner.end)
+        write!(formatter, "[{}..{}]", self.inner.start, self.inner.end - 1)
     }
 }
 
@@ -182,9 +183,10 @@ impl<T> Iterator for Range<T> {
     type Item = Handle<T>;
     fn next(&mut self) -> Option<Self::Item> {
         if self.inner.start < self.inner.end {
+            let next = self.inner.start;
             self.inner.start += 1;
             Some(Handle {
-                index: NonZeroU32::new(self.inner.start).unwrap(),
+                index: NonMaxU32::new(next).unwrap(),
                 marker: self.marker,
             })
         } else {
@@ -209,11 +211,10 @@ impl<T> Range<T> {
     pub fn first_and_last(&self) -> Option<(Handle<T>, Handle<T>)> {
         if self.inner.start < self.inner.end {
             Some((
-                // `Range::new_from_bounds` expects a 1-based, start- and
+                // `Range::new_from_bounds` expects a start- and end-inclusive
-                // end-inclusive range, but `self.inner` is a zero-based,
+                // range, but `self.inner` is an end-exclusive range.
-                // end-exclusive range.
+                Handle::new(Index::new(self.inner.start).unwrap()),
-                Handle::new(Index::new(self.inner.start + 1).unwrap()),
+                Handle::new(Index::new(self.inner.end - 1).unwrap()),
-                Handle::new(Index::new(self.inner.end).unwrap()),
             ))
         } else {
             None
@@ -417,10 +418,7 @@ impl<T> Arena<T> {
             return Ok(());
         }
 
-        // `range.inner` is zero-based, but end-exclusive, so `range.inner.end`
+        let last_handle = Handle::new(Index::new(range.inner.end - 1).unwrap());
-        // is actually the right one-based index for the last handle within the
-        // range.
-        let last_handle = Handle::new(range.inner.end.try_into().unwrap());
         if self.check_contains_handle(last_handle).is_err() {
             return Err(BadRangeError::new(range.clone()));
         }
@@ -436,7 +434,7 @@ impl<T> Arena<T> {
         let mut index = 0;
         let mut retained = 0;
         self.data.retain_mut(|elt| {
-            let handle = Handle::new(Index::new(index as u32 + 1).unwrap());
+            let handle = Handle::from_usize(index);
             let keep = predicate(handle, elt);
 
             // Since `predicate` needs mutable access to each element,
@@ -602,7 +600,7 @@ impl<T> UniqueArena<T> {
         UniqueArenaDrain {
             inner_elts: self.set.drain(..),
             inner_spans: self.span_info.drain(..),
-            index: Index::new(1).unwrap(),
+            index: Index::new(0).unwrap(),
         }
     }
 }
@@ -636,8 +634,7 @@ impl<T: Eq + hash::Hash> UniqueArena<T> {
     /// the item's handle and a reference to it.
     pub fn iter(&self) -> impl DoubleEndedIterator<Item = (Handle<T>, &T)> {
         self.set.iter().enumerate().map(|(i, v)| {
-            let position = i + 1;
+            let index = unsafe { Index::new_unchecked(i as u32) };
-            let index = unsafe { Index::new_unchecked(position as u32) };
             (Handle::new(index), v)
         })
     }

naga/src/compact/handle_set_map.rs

@@ -1,14 +1,13 @@
 use crate::arena::{Arena, Handle, Range, UniqueArena};
 
-type Index = std::num::NonZeroU32;
+type Index = crate::non_max_u32::NonMaxU32;
 
 /// A set of `Handle<T>` values.
 pub struct HandleSet<T> {
-    /// Bound on zero-based indexes of handles stored in this set.
+    /// Bound on indexes of handles stored in this set.
     len: usize,
 
-    /// `members[i]` is true if the handle with zero-based index `i`
+    /// `members[i]` is true if the handle with index `i` is a member.
-    /// is a member.
     members: bit_set::BitSet,
 
     /// This type is indexed by values of type `T`.
@@ -27,8 +26,6 @@ impl<T> HandleSet<T> {
 
     /// Add `handle` to the set.
     pub fn insert(&mut self, handle: Handle<T>) {
-        // Note that, oddly, `Handle::index` does not return a 1-based
-        // `Index`, but rather a zero-based `usize`.
         self.members.insert(handle.index());
     }
 
@@ -40,8 +37,6 @@ impl<T> HandleSet<T> {
     }
 
     pub fn contains(&self, handle: Handle<T>) -> bool {
-        // Note that, oddly, `Handle::index` does not return a 1-based
-        // `Index`, but rather a zero-based `usize`.
         self.members.contains(handle.index())
     }
 }
@@ -66,10 +61,9 @@ impl<T: std::hash::Hash + Eq> ArenaType<T> for UniqueArena<T> {
 pub struct HandleMap<T> {
     /// The indices assigned to handles in the compacted module.
     ///
-    /// If `new_index[i]` is `Some(n)`, then `n` is the 1-based
+    /// If `new_index[i]` is `Some(n)`, then `n` is the `Index` of the
-    /// `Index` of the compacted `Handle` corresponding to the
+    /// compacted `Handle` corresponding to the pre-compacted `Handle`
-    /// pre-compacted `Handle` whose zero-based index is `i`. ("Clear
+    /// whose index is `i`.
-    /// as mud.")
     new_index: Vec<Option<Index>>,
 
     /// This type is indexed by values of type `T`.
@@ -78,11 +72,11 @@ pub struct HandleMap<T> {
 
 impl<T: 'static> HandleMap<T> {
     pub fn from_set(set: HandleSet<T>) -> Self {
-        let mut next_index = Index::new(1).unwrap();
+        let mut next_index = Index::new(0).unwrap();
         Self {
             new_index: (0..set.len)
-                .map(|zero_based_index| {
+                .map(|index| {
-                    if set.members.contains(zero_based_index) {
+                    if set.members.contains(index) {
                         // This handle will be retained in the compacted version,
                         // so assign it a new index.
                         let this = next_index;
@@ -111,11 +105,9 @@ impl<T: 'static> HandleMap<T> {
         log::trace!(
             "adjusting {} handle [{}] -> [{:?}]",
             std::any::type_name::<T>(),
-            old.index() + 1,
+            old.index(),
             self.new_index[old.index()]
         );
-        // Note that `Handle::index` returns a zero-based index,
-        // but `Handle::new` accepts a 1-based `Index`.
         self.new_index[old.index()].map(Handle::new)
     }
 
@@ -147,20 +139,18 @@ impl<T: 'static> HandleMap<T> {
     pub fn adjust_range(&self, range: &mut Range<T>, compacted_arena: &Arena<T>) {
         let mut index_range = range.zero_based_index_range();
         let compacted;
-        // Remember that the indices we retrieve from `new_index` are 1-based
+        if let Some(first) = index_range.find_map(|i| self.new_index[i as usize]) {
-        // compacted indices, but the index range we're computing is zero-based
-        // compacted indices.
-        if let Some(first1) = index_range.find_map(|i| self.new_index[i as usize]) {
             // The first call to `find_map` mutated `index_range` to hold the
             // remainder of original range, which is exactly the range we need
             // to search for the new last handle.
-            if let Some(last1) = index_range.rev().find_map(|i| self.new_index[i as usize]) {
+            if let Some(last) = index_range.rev().find_map(|i| self.new_index[i as usize]) {
-                // Build a zero-based end-exclusive range, given one-based handle indices.
+                // Build an end-exclusive range, given the two included indices
-                compacted = first1.get() - 1..last1.get();
+                // `first` and `last`.
+                compacted = first.get()..last.get() + 1;
             } else {
                 // The range contains only a single live handle, which
                 // we identified with the first `find_map` call.
-                compacted = first1.get() - 1..first1.get();
+                compacted = first.get()..first.get() + 1;
             }
         } else {
             compacted = 0..0;

naga/src/lib.rs

@@ -277,6 +277,7 @@ pub mod compact;
 pub mod error;
 pub mod front;
 pub mod keywords;
+mod non_max_u32;
 pub mod proc;
 mod span;
 pub mod valid;

naga/src/non_max_u32.rs

@@ -0,0 +1,111 @@
+//! [`NonMaxU32`], a 32-bit type that can represent any value except [`u32::MAX`].
+//!
+//! Naga would like `Option<Handle<T>>` to be a 32-bit value, which means we
+//! need to exclude some index value for use in representing [`None`]. We could
+//! have [`Handle`] store a [`NonZeroU32`], but zero is a very useful value for
+//! indexing. We could have a [`Handle`] store a value one greater than its index,
+//! but it turns out that it's not uncommon to want to work with [`Handle`]s'
+//! indices, so that bias of 1 becomes more visible than one would like.
+//!
+//! This module defines the type [`NonMaxU32`], for which `Option<NonMaxU32>` is
+//! still a 32-bit value, but which is directly usable as a [`Handle`] index
+//! type. It still uses a bias of 1 under the hood, but that fact is isolated
+//! within the implementation.
+//!
+//! [`Handle`]: crate::arena::Handle
+//! [`NonZeroU32`]: std::num::NonZeroU32
+#![allow(dead_code)]
+
+use std::num::NonZeroU32;
+
+/// An unsigned 32-bit value known not to be [`u32::MAX`].
+///
+/// A `NonMaxU32` value can represent any value in the range `0 .. u32::MAX -
+/// 1`, and an `Option<NonMaxU32>` is still a 32-bit value. In other words,
+/// `NonMaxU32` is just like [`NonZeroU32`], except that a different value is
+/// missing from the full `u32` range.
+///
+/// Since zero is a very useful value in indexing, `NonMaxU32` is more useful
+/// for representing indices than [`NonZeroU32`].
+///
+/// # Implementation
+///
+/// A `NonMaxU32` whose value is `n` is a newtype around a [`NonZeroU32`] whose
+/// value is `n + 1`. This way, the range of values `NonMaxU32` can represent,
+/// `0..=u32::MAX - 1`, is mapped to the range `1..=u32::MAX`, which is the
+/// range that [`NonZeroU32`] can represent. (And conversely, since [`u32`]
+/// addition wraps around, the unrepresentable value [`u32::MAX`] becomes the
+/// unrepresentable value `0`.)
+///
+/// [`NonZeroU32`]: std::num::NonZeroU32
+#[derive(Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Hash)]
+#[cfg_attr(feature = "serialize", derive(serde::Serialize))]
+#[cfg_attr(feature = "deserialize", derive(serde::Deserialize))]
+#[cfg_attr(
+    any(feature = "serialize", feature = "deserialize"),
+    serde(transparent)
+)]
+#[cfg_attr(feature = "arbitrary", derive(arbitrary::Arbitrary))]
+pub struct NonMaxU32(NonZeroU32);
+
+impl NonMaxU32 {
+    /// Construct a [`NonMaxU32`] whose value is `n`, if possible.
+    pub const fn new(n: u32) -> Option<Self> {
+        // If `n` is `u32::MAX`, then `n.wrapping_add(1)` is `0`,
+        // so `NonZeroU32::new` returns `None` in exactly the case
+        // where we must return `None`.
+        match NonZeroU32::new(n.wrapping_add(1)) {
+            Some(non_zero) => Some(NonMaxU32(non_zero)),
+            None => None,
+        }
+    }
+
+    /// Return the value of `self` as a [`u32`].
+    pub const fn get(self) -> u32 {
+        self.0.get() - 1
+    }
+
+    /// Construct a [`NonMaxU32`] whose value is `n`.
+    ///
+    /// # Safety
+    ///
+    /// The value of `n` must not be [`u32::MAX`].
+    pub const unsafe fn new_unchecked(n: u32) -> NonMaxU32 {
+        NonMaxU32(unsafe { NonZeroU32::new_unchecked(n + 1) })
+    }
+
+    /// Construct a [`NonMaxU32`] whose value is `index`.
+    ///
+    /// # Safety
+    ///
+    /// - The value of `index` must be strictly less than [`u32::MAX`].
+    pub const unsafe fn from_usize_unchecked(index: usize) -> Self {
+        NonMaxU32(unsafe { NonZeroU32::new_unchecked(index as u32 + 1) })
+    }
+
+    pub fn checked_add(self, n: u32) -> Option<Self> {
+        // Adding `n` to `self` produces `u32::MAX` if and only if
+        // adding `n` to `self.0` produces `0`. So we can simply
+        // call `NonZeroU32::checked_add` and let its check for zero
+        // determine whether our add would have produced `u32::MAX`.
+        Some(NonMaxU32(self.0.checked_add(n)?))
+    }
+}
+
+impl std::fmt::Debug for NonMaxU32 {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        self.get().fmt(f)
+    }
+}
+
+impl std::fmt::Display for NonMaxU32 {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        self.get().fmt(f)
+    }
+}
+
+#[test]
+fn size() {
+    use core::mem::size_of;
+    assert_eq!(size_of::<Option<NonMaxU32>>(), size_of::<u32>());
+}

naga/src/valid/function.rs

@@ -1474,7 +1474,7 @@ impl super::Validator {
 
             if self.flags.contains(super::ValidationFlags::EXPRESSIONS) {
                 if let Some(unvisited) = self.needs_visit.iter().next() {
-                    let index = std::num::NonZeroU32::new(unvisited as u32 + 1).unwrap();
+                    let index = crate::non_max_u32::NonMaxU32::new(unvisited as u32).unwrap();
                     let handle = Handle::new(index);
                     return Err(FunctionError::UnvisitedExpression(handle)
                         .with_span_handle(handle, &fun.expressions));

naga/src/valid/handles.rs

@@ -5,11 +5,12 @@ use crate::{
     Handle,
 };
 
+use crate::non_max_u32::NonMaxU32;
 use crate::{Arena, UniqueArena};
 
 use super::ValidationError;
 
-use std::{convert::TryInto, hash::Hash, num::NonZeroU32};
+use std::{convert::TryInto, hash::Hash};
 
 impl super::Validator {
     /// Validates that all handles within `module` are:
@@ -688,7 +689,7 @@ impl<T> Handle<T> {
             Ok(self)
         } else {
             let erase_handle_type = |handle: Handle<_>| {
-                Handle::new(NonZeroU32::new((handle.index() + 1).try_into().unwrap()).unwrap())
+                Handle::new(NonMaxU32::new((handle.index()).try_into().unwrap()).unwrap())
             };
             Err(FwdDepError {
                 subject: erase_handle_type(self),

naga/tests/out/dot/quad.dot

@@ -1,23 +1,23 @@
 digraph Module {
 	subgraph cluster_globals {
 		label="Globals"
-		g0 [ shape=hexagon label="[1] Handle/'u_texture'" ]
+		g0 [ shape=hexagon label="[0] Handle/'u_texture'" ]
-		g1 [ shape=hexagon label="[2] Handle/'u_sampler'" ]
+		g1 [ shape=hexagon label="[1] Handle/'u_sampler'" ]
 	}
 	subgraph cluster_ep0 {
 		label="Vertex/'vert_main'"
 		node [ style=filled ]
-		ep0_e0 [ color="#8dd3c7" label="[1] Argument[0]" ]
+		ep0_e0 [ color="#8dd3c7" label="[0] Argument[0]" ]
-		ep0_e1 [ color="#8dd3c7" label="[2] Argument[1]" ]
+		ep0_e1 [ color="#8dd3c7" label="[1] Argument[1]" ]
-		ep0_e2 [ fillcolor="#ffffb3" label="[3] Constant" ]
+		ep0_e2 [ fillcolor="#ffffb3" label="[2] Constant" ]
-		ep0_e3 [ color="#fdb462" label="[4] Multiply" ]
+		ep0_e3 [ color="#fdb462" label="[3] Multiply" ]
 		ep0_e0 -> ep0_e3 [ label="right" ]
 		ep0_e2 -> ep0_e3 [ label="left" ]
-		ep0_e4 [ fillcolor="#ffffb3" label="[5] Literal" ]
+		ep0_e4 [ fillcolor="#ffffb3" label="[4] Literal" ]
-		ep0_e5 [ fillcolor="#ffffb3" label="[6] Literal" ]
+		ep0_e5 [ fillcolor="#ffffb3" label="[5] Literal" ]
-		ep0_e6 [ color="#bebada" label="[7] Compose" ]
+		ep0_e6 [ color="#bebada" label="[6] Compose" ]
 		{ ep0_e3 ep0_e4 ep0_e5 } -> ep0_e6
-		ep0_e7 [ color="#bebada" label="[8] Compose" ]
+		ep0_e7 [ color="#bebada" label="[7] Compose" ]
 		{ ep0_e1 ep0_e6 } -> ep0_e7
 		ep0_s0 [ shape=square label="Root" ]
 		ep0_s1 [ shape=square label="Emit" ]
@@ -34,24 +34,24 @@ digraph Module {
 	subgraph cluster_ep1 {
 		label="Fragment/'frag_main'"
 		node [ style=filled ]
-		ep1_e0 [ color="#8dd3c7" label="[1] Argument[0]" ]
+		ep1_e0 [ color="#8dd3c7" label="[0] Argument[0]" ]
-		ep1_e1 [ color="#ffffb3" label="[2] Global" ]
+		ep1_e1 [ color="#ffffb3" label="[1] Global" ]
 		g0 -> ep1_e1 [fillcolor=gray]
-		ep1_e2 [ color="#ffffb3" label="[3] Global" ]
+		ep1_e2 [ color="#ffffb3" label="[2] Global" ]
 		g1 -> ep1_e2 [fillcolor=gray]
-		ep1_e3 [ color="#80b1d3" label="[4] ImageSample" ]
+		ep1_e3 [ color="#80b1d3" label="[3] ImageSample" ]
 		ep1_e2 -> ep1_e3 [ label="sampler" ]
 		ep1_e1 -> ep1_e3 [ label="image" ]
 		ep1_e0 -> ep1_e3 [ label="coordinate" ]
-		ep1_e4 [ color="#8dd3c7" label="[5] AccessIndex[3]" ]
+		ep1_e4 [ color="#8dd3c7" label="[4] AccessIndex[3]" ]
 		ep1_e3 -> ep1_e4 [ label="base" ]
-		ep1_e5 [ fillcolor="#ffffb3" label="[6] Literal" ]
+		ep1_e5 [ fillcolor="#ffffb3" label="[5] Literal" ]
-		ep1_e6 [ color="#fdb462" label="[7] Equal" ]
+		ep1_e6 [ color="#fdb462" label="[6] Equal" ]
 		ep1_e5 -> ep1_e6 [ label="right" ]
 		ep1_e4 -> ep1_e6 [ label="left" ]
-		ep1_e7 [ color="#8dd3c7" label="[8] AccessIndex[3]" ]
+		ep1_e7 [ color="#8dd3c7" label="[7] AccessIndex[3]" ]
 		ep1_e3 -> ep1_e7 [ label="base" ]
-		ep1_e8 [ color="#fdb462" label="[9] Multiply" ]
+		ep1_e8 [ color="#fdb462" label="[8] Multiply" ]
 		ep1_e3 -> ep1_e8 [ label="right" ]
 		ep1_e7 -> ep1_e8 [ label="left" ]
 		ep1_s0 [ shape=square label="Root" ]
@@ -87,11 +87,11 @@ digraph Module {
 	subgraph cluster_ep2 {
 		label="Fragment/'fs_extra'"
 		node [ style=filled ]
-		ep2_e0 [ fillcolor="#ffffb3" label="[1] Literal" ]
+		ep2_e0 [ fillcolor="#ffffb3" label="[0] Literal" ]
-		ep2_e1 [ fillcolor="#ffffb3" label="[2] Literal" ]
+		ep2_e1 [ fillcolor="#ffffb3" label="[1] Literal" ]
-		ep2_e2 [ fillcolor="#ffffb3" label="[3] Literal" ]
+		ep2_e2 [ fillcolor="#ffffb3" label="[2] Literal" ]
-		ep2_e3 [ fillcolor="#ffffb3" label="[4] Literal" ]
+		ep2_e3 [ fillcolor="#ffffb3" label="[3] Literal" ]
-		ep2_e4 [ fillcolor="#bebada" label="[5] Compose" ]
+		ep2_e4 [ fillcolor="#bebada" label="[4] Compose" ]
 		{ ep2_e0 ep2_e1 ep2_e2 ep2_e3 } -> ep2_e4
 		ep2_s0 [ shape=square label="Root" ]
 		ep2_s1 [ shape=square label="Emit" ]