1 files changed, 83 insertions, 0 deletions
diff --git a/third_party/rust/smawk/tests/random_monge/mod.rs b/third_party/rust/smawk/tests/random_monge/mod.rs
new file mode 100644
index 0000000000..50a932fbeb
--- /dev/null
+++ b/third_party/rust/smawk/tests/random_monge/mod.rs
@@ -0,0 +1,83 @@
+//! Test functionality for generating random Monge matrices.
+
+// The code is put here so we can reuse it in different integration
+// tests, without Cargo finding it when `cargo test` is run. See the
+// section on "Submodules in Integration Tests" in
+// https://doc.rust-lang.org/book/ch11-03-test-organization.html
+
+use ndarray::{s, Array2};
+use num_traits::PrimInt;
+use rand::distributions::{Distribution, Standard};
+use rand::Rng;
+
+/// A Monge matrix can be decomposed into one of these primitive
+/// building blocks.
+#[derive(Copy, Clone)]
+pub enum MongePrim {
+    ConstantRows,
+    ConstantCols,
+    UpperRightOnes,
+    LowerLeftOnes,
+}
+
+impl MongePrim {
+    /// Generate a Monge matrix from a primitive.
+    pub fn to_matrix<T: PrimInt, R: Rng>(&self, m: usize, n: usize, rng: &mut R) -> Array2<T>
+    where
+        Standard: Distribution<T>,
+    {
+        let mut matrix = Array2::from_elem((m, n), T::zero());
+        // Avoid panic in UpperRightOnes and LowerLeftOnes below.
+        if m == 0 || n == 0 {
+            return matrix;
+        }
+
+        match *self {
+            MongePrim::ConstantRows => {
+                for mut row in matrix.rows_mut() {
+                    if rng.gen::<bool>() {
+                        row.fill(T::one())
+                    }
+                }
+            }
+            MongePrim::ConstantCols => {
+                for mut col in matrix.columns_mut() {
+                    if rng.gen::<bool>() {
+                        col.fill(T::one())
+                    }
+                }
+            }
+            MongePrim::UpperRightOnes => {
+                let i = rng.gen_range(0..(m + 1) as isize);
+                let j = rng.gen_range(0..(n + 1) as isize);
+                matrix.slice_mut(s![..i, -j..]).fill(T::one());
+            }
+            MongePrim::LowerLeftOnes => {
+                let i = rng.gen_range(0..(m + 1) as isize);
+                let j = rng.gen_range(0..(n + 1) as isize);
+                matrix.slice_mut(s![-i.., ..j]).fill(T::one());
+            }
+        }
+
+        matrix
+    }
+}
+
+/// Generate a random Monge matrix.
+pub fn random_monge_matrix<R: Rng, T: PrimInt>(m: usize, n: usize, rng: &mut R) -> Array2<T>
+where
+    Standard: Distribution<T>,
+{
+    let monge_primitives = [
+        MongePrim::ConstantRows,
+        MongePrim::ConstantCols,
+        MongePrim::LowerLeftOnes,
+        MongePrim::UpperRightOnes,
+    ];
+    let mut matrix = Array2::from_elem((m, n), T::zero());
+    for _ in 0..(m + n) {
+        let monge = monge_primitives[rng.gen_range(0..monge_primitives.len())];
+        matrix = matrix + monge.to_matrix(m, n, rng);
+    }
+    matrix
+}