From b5f77f4f6d00c81a553bed127f7c7952a8566574 Mon Sep 17 00:00:00 2001
From: Andrei Navumenka <indy256@gmail.com>
Date: Sun, 1 Sep 2024 12:24:32 -0400
Subject: [PATCH] Add proptests for mergeable_heap.rs

---
 rust/Cargo.toml                   |  1 +
 rust/structures/mergeable_heap.rs | 47 +++++++++++++++++++++----------
 2 files changed, 33 insertions(+), 15 deletions(-)

diff --git a/rust/Cargo.toml b/rust/Cargo.toml
index d8eb5478..adbbfb11 100644
--- a/rust/Cargo.toml
+++ b/rust/Cargo.toml
@@ -9,6 +9,7 @@ rand = "0.8.5"
 
 [dev-dependencies]
 rstest="0.22.0"
+proptest = "1.5.0"
 
 [lib]
 path = "lib.rs"
diff --git a/rust/structures/mergeable_heap.rs b/rust/structures/mergeable_heap.rs
index 1d3256b5..bbea4aa6 100644
--- a/rust/structures/mergeable_heap.rs
+++ b/rust/structures/mergeable_heap.rs
@@ -15,21 +15,22 @@ impl<V: PartialOrd> Heap<V> {
         }))
     }
 
-    fn merge(mut a: Option<Box<Heap<V>>>, mut b: Option<Box<Heap<V>>>) -> Option<Box<Heap<V>>> {
+    fn merge(a: Option<Box<Heap<V>>>, b: Option<Box<Heap<V>>>) -> Option<Box<Heap<V>>> {
         if a.is_none() {
             return b;
         }
         if b.is_none() {
             return a;
         }
-        if a.as_ref()?.value > b.as_ref()?.value {
-            swap(&mut a, &mut b);
+        let mut ha = a.unwrap();
+        let mut hb = b.unwrap();
+        if ha.value > hb.value {
+            swap(&mut ha, &mut hb);
         }
-        let mut ha = a?;
         if rand::random() {
             swap(&mut ha.left, &mut ha.right);
         }
-        ha.left = Self::merge(ha.left, b);
+        ha.left = Self::merge(ha.left, Some(hb));
         Some(ha)
     }
 
@@ -46,6 +47,8 @@ impl<V: PartialOrd> Heap<V> {
 #[cfg(test)]
 mod tests {
     use crate::structures::mergeable_heap::Heap;
+    use proptest::collection::vec;
+    use proptest::{prop_assert_eq, proptest};
     use rand::seq::SliceRandom;
     use rand::thread_rng;
     use rstest::rstest;
@@ -56,24 +59,39 @@ mod tests {
     #[case(&mut [1, 1])]
     #[case(&mut [3, 1, 2])]
     fn basic_test(#[case] seq: &mut [u32]) {
-        test(seq);
+        let heap_sorted_values = sort_with_heap(&seq);
+        seq.sort();
+        assert_eq!(heap_sorted_values, seq);
     }
 
     #[test]
     fn big_test1() {
-        let mut values = (0..10_000).collect::<Vec<u32>>();
-        values.shuffle(&mut thread_rng());
-        test(&mut values);
+        let mut seq = (0..10_000).collect::<Vec<u32>>();
+        seq.shuffle(&mut thread_rng());
+        let heap_sorted_values = sort_with_heap(&seq);
+        seq.sort();
+        assert_eq!(heap_sorted_values, seq);
     }
 
     #[test]
     fn big_test2() {
-        let mut values = vec![0; 10_000];
-        values.shuffle(&mut thread_rng());
-        test(&mut values);
+        let mut seq = vec![0; 10_000];
+        seq.shuffle(&mut thread_rng());
+        let heap_sorted_values = sort_with_heap(&seq);
+        seq.sort();
+        assert_eq!(heap_sorted_values, seq);
     }
 
-    fn test(seq: &mut [u32]) {
+    proptest! {
+        #[test]
+        fn test_random(mut seq in vec(0u32..3, 0..10)) {
+            let heap_sorted_values = sort_with_heap(&seq);
+            seq.sort();
+            prop_assert_eq!(heap_sorted_values, seq);
+        }
+    }
+
+    fn sort_with_heap(seq: &[u32]) -> Vec<u32> {
         let mut heap = seq.iter().fold(None, |h, v| Heap::add(h, v));
         let mut heap_sorted_values = Vec::new();
         while heap.is_some() {
@@ -81,7 +99,6 @@ mod tests {
             heap = updated_heap;
             heap_sorted_values.push(*min_value);
         }
-        seq.sort();
-        assert_eq!(heap_sorted_values, seq);
+        heap_sorted_values
     }
 }