WIP: Options for sharding strategies

Author: sourcefrog

src/main.rs

@@ -560,7 +560,7 @@ fn main() -> Result<()> {
         };
     }
     if let Some(shard) = &args.shard {
-        mutants = shard.select(mutants);
+        mutants = options.sharding().shard(*shard, mutants);
     }
     if args.list {
         print!("{}", list_mutants(&mutants, &options));

src/options.rs

@@ -28,6 +28,7 @@ use crate::annotation::ResolvedAnnotation;
 use crate::config::Config;
 use crate::glob::build_glob_set;
 use crate::mutant::Mutant;
+use crate::shard::Sharding;
 use crate::{Args, BaselineStrategy, Phase, Result, ValueEnum};
 
 /// Options for mutation testing, based on both command-line arguments and the config file.
@@ -180,6 +181,10 @@ pub struct Common {
     /// Tool used to run test suites: cargo or nextest.
     #[arg(long, help_heading = "Execution")]
     pub test_tool: Option<TestTool>,
+
+    /// Sharding method to use when running with --shards.
+    #[arg(long, help_heading = "Execution", requires = "shard")]
+    pub sharding: Option<Sharding>,
 }
 
 impl Common {
@@ -190,6 +195,7 @@ impl Common {
         Common {
             emit_diffs: self.emit_diffs.or(other.emit_diffs),
             test_tool: self.test_tool.or(other.test_tool),
+            sharding: self.sharding.or(other.sharding),
         }
     }
 }
@@ -473,6 +479,10 @@ impl Options {
     pub fn test_tool(&self) -> TestTool {
         self.common.test_tool.unwrap_or_default()
     }
+
+    pub fn sharding(&self) -> Sharding {
+        self.common.sharding.unwrap_or_default()
+    }
 }
 
 /// If the first slices is non-empty, return that, otherwise the second.

src/shard.rs

@@ -1,10 +1,11 @@
-// Copyright 2023 Martin Pool
-
 //! Sharding parameters.
 
 use std::str::FromStr;
 
 use anyhow::{anyhow, ensure, Context, Error};
+use clap::ValueEnum;
+use schemars::JsonSchema;
+use serde::Deserialize;
 
 /// Select mutants for a particular shard of the total list.
 #[derive(Debug, Clone, Copy, Eq, PartialEq)]
@@ -15,17 +16,6 @@ pub struct Shard {
     pub n: usize,
 }
 
-impl Shard {
-    /// Select the mutants that should be run for this shard.
-    pub fn select<M, I: IntoIterator<Item = M>>(&self, mutants: I) -> Vec<M> {
-        mutants
-            .into_iter()
-            .enumerate()
-            .filter_map(|(i, m)| if i % self.n == self.k { Some(m) } else { None })
-            .collect()
-    }
-}
-
 impl FromStr for Shard {
     type Err = Error;
 
@@ -38,6 +28,57 @@ impl FromStr for Shard {
     }
 }
 
+/// Method for sharding.
+#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, Deserialize, ValueEnum, JsonSchema)]
+#[serde(rename_all = "snake_case")]
+pub enum Sharding {
+    /// Run mutant `i` on shard `i % k`.
+    ///
+    /// This was the default up to cargo-mutants 25.3.1.
+    ///
+    /// This distributes mutants more evenly and will likely generate more equal completion
+    /// times, but it has less locality of reference within each shard's cache and so may
+    /// cause more build time.
+    #[default]
+    RoundRobin,
+
+    /// Run consecutive ranges of mutants on each shard: `0..(n/k)` on shard 0, etc. (default)
+    ///
+    /// This makes the incremental change between each build likely to be smaller and
+    /// so may reduce build time, but it may also lead to more unbalanced shards.
+    Slice,
+}
+
+impl Sharding {
+    /// Select the mutants that should be run for this shard.
+    pub fn shard<M>(self, shard: Shard, mut mutants: Vec<M>) -> Vec<M> {
+        match self {
+            Sharding::RoundRobin => mutants
+                .into_iter()
+                .enumerate()
+                .filter_map(|(i, m)| {
+                    if i % shard.n == shard.k {
+                        Some(m)
+                    } else {
+                        None
+                    }
+                })
+                .collect(),
+            Sharding::Slice => {
+                let total = mutants.len();
+                let chunk_size = total.div_ceil(shard.n);
+                let start = shard.k * chunk_size;
+                let end = ((shard.k + 1) * chunk_size).min(total);
+                if start >= total {
+                    Vec::new()
+                } else {
+                    mutants.drain(start..end).collect()
+                }
+            }
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -74,10 +115,21 @@ mod tests {
     }
 
     #[test]
-    fn shard_select() {
-        assert_eq!(
-            Shard::from_str("1/4").unwrap().select(0..10).as_slice(),
-            &[1, 5, 9]
-        );
+    fn shard_round_robin() {
+        // This test works on ints instead of real mutants just for ease of testing.
+        let fake_mutants: Vec<usize> = (0..10).collect();
+        for (k, expect) in [
+            (0, [0, 4, 8].as_slice()),
+            (1, &[1, 5, 9]),
+            (2, &[2, 6]),
+            (3, &[3, 7]),
+        ] {
+            assert_eq!(
+                Sharding::RoundRobin
+                    .shard(Shard { k, n: 4 }, fake_mutants.clone())
+                    .as_slice(),
+                expect
+            );
+        }
     }
 }