Improve the documentation of the compute shader example.

Also make it include `1` in the output sequence, because that is the
first number according to https://oeis.org/A006877.

Author: nnethercote

examples/runners/wgpu/src/compute.rs

@@ -1,3 +1,13 @@
+//! This crate contains code to print the integer sequence at <https://oeis.org/A006877>:
+//!
+//! > In the '3x+1' problem, these values for the starting value set new records for number of
+//! > steps to reach 1.
+//!
+//! each starting value is paired with the number of steps to reach 1 (a.k.a the "Collatz
+//! sequence length").
+//!
+//! The individual Collatz sequence lengths are computed via a computer shader on the GPU.
+
 use crate::{CompiledShaderModules, Options, maybe_watch};
 
 use std::time::Duration;
@@ -252,14 +262,16 @@ async fn start_internal(options: &Options, compiled_shader_modules: CompiledShad
     drop(data);
     readback_buffer.unmap();
 
+    // Our `max` computation excludes `1`, so print that manually as the first number in the
+    // sequence.
+    println!("1: 0");
     let mut max = 0;
     for (src, out) in src_range.zip(result.iter().copied()) {
         if out == u32::MAX {
             println!("{src}: overflowed");
             break;
         } else if out > max {
             max = out;
-            // Should produce <https://oeis.org/A006877>
             println!("{src}: {out}");
         }
     }

examples/shaders/compute-shader/src/lib.rs

@@ -7,6 +7,17 @@ use spirv_std::{glam, spirv};
 
 // Adapted from the wgpu hello-compute example
 
+/// Returns the length of the Collatz sequence (excluding the starting number) for `n`. Returns
+/// `None` if (a) `n` is zero, or (b) a number in the sequence overflows a `u32`.
+///
+/// # Examples
+///
+/// The sequence for 3 (excluding the starting number) is `[10, 5, 16, 8, 4, 2, 1]`, which has
+/// length 7.
+/// ```
+/// # use compute_shader::collatz;
+/// assert_eq!(collatz(3), Some(7));
+/// ```
 pub fn collatz(mut n: u32) -> Option<u32> {
     let mut i = 0;
     if n == 0 {
@@ -20,7 +31,8 @@ pub fn collatz(mut n: u32) -> Option<u32> {
             if n >= 0x5555_5555 {
                 return None;
             }
-            // TODO: Use this instead when/if checked add/mul can work: n.checked_mul(3)?.checked_add(1)?
+            // TODO: Use this instead when/if checked add/mul can work:
+            // n.checked_mul(3)?.checked_add(1)?
             3 * n + 1
         };
         i += 1;

examples/shaders/compute-shader/src/main.rs

@@ -1,3 +1,8 @@
+//! Code for running the Collatz problem on the CPU. See the wgpu compute runner
+//! for the equivalent GPU code.
+//!
+//! Currently it's not actually built or tested, and is just here for reference.
+
 use std::time::Instant;
 
 use compute_shader::collatz;
@@ -13,15 +18,15 @@ fn main() {
         .map(collatz)
         .collect::<Vec<_>>();
     let took = start.elapsed();
+    println!("1: 0");
     let mut max = 0;
     for (src, out) in src_range.zip(result.iter().copied()) {
         match out {
             Some(out) if out > max => {
                 max = out;
-                // Should produce <https://oeis.org/A006877>
                 println!("{src}: {out}");
             }
-            Some(_) => (),
+            Some(_) => {}
             None => {
                 println!("{src}: overflowed");
                 break;