From 149e408da5a1a92156f2a66af93964439b24b8e2 Mon Sep 17 00:00:00 2001
From: Caio <c410.f3r@gmail.com>
Date: Sat, 2 Nov 2024 16:52:05 -0300
Subject: [PATCH] Implement `Arbitrary` for `Range*` (#3666)

Resolves #3662

cc @zhassan-aws
---
 library/kani_core/src/arbitrary.rs |  68 +++++++++++++++++
 tests/kani/Arbitrary/ops.rs        | 114 +++++++++++++++++++++++++++++
 2 files changed, 182 insertions(+)
 create mode 100644 tests/kani/Arbitrary/ops.rs

diff --git a/library/kani_core/src/arbitrary.rs b/library/kani_core/src/arbitrary.rs
index 90a2ae5e4c82..3ef474b65364 100644
--- a/library/kani_core/src/arbitrary.rs
+++ b/library/kani_core/src/arbitrary.rs
@@ -193,6 +193,74 @@ macro_rules! generate_arbitrary {
         pub mod slice {
             kani_core::slice_generator!();
         }
+
+        mod range_structures {
+            use super::{
+                Arbitrary,
+                core_path::{
+                    mem,
+                    ops::{Bound, Range, RangeFrom, RangeInclusive, RangeTo, RangeToInclusive},
+                },
+            };
+
+            impl<T> Arbitrary for Bound<T>
+            where
+                T: Arbitrary,
+            {
+                fn any() -> Self {
+                    match u8::any() {
+                        0 => Bound::Included(T::any()),
+                        1 => Bound::Excluded(T::any()),
+                        _ => Bound::Unbounded,
+                    }
+                }
+            }
+
+            impl<T> Arbitrary for Range<T>
+            where
+                T: Arbitrary,
+            {
+                fn any() -> Self {
+                    T::any()..T::any()
+                }
+            }
+
+            impl<T> Arbitrary for RangeFrom<T>
+            where
+                T: Arbitrary,
+            {
+                fn any() -> Self {
+                    T::any()..
+                }
+            }
+
+            impl<T> Arbitrary for RangeInclusive<T>
+            where
+                T: Arbitrary,
+            {
+                fn any() -> Self {
+                    T::any()..=T::any()
+                }
+            }
+
+            impl<T> Arbitrary for RangeTo<T>
+            where
+                T: Arbitrary,
+            {
+                fn any() -> Self {
+                    ..T::any()
+                }
+            }
+
+            impl<T> Arbitrary for RangeToInclusive<T>
+            where
+                T: Arbitrary,
+            {
+                fn any() -> Self {
+                    ..=T::any()
+                }
+            }
+        }
     };
 }
 
diff --git a/tests/kani/Arbitrary/ops.rs b/tests/kani/Arbitrary/ops.rs
new file mode 100644
index 000000000000..6ae97b2335f4
--- /dev/null
+++ b/tests/kani/Arbitrary/ops.rs
@@ -0,0 +1,114 @@
+// Copyright Kani Contributors
+// SPDX-License-Identifier: Apache-2.0 OR MIT
+//
+// Check that users can generate range structures
+
+extern crate kani;
+
+use std::ops::{Bound, Range, RangeFrom, RangeInclusive, RangeTo, RangeToInclusive};
+
+#[kani::proof]
+fn bound() {
+    let elem: Wrapper<Bound<u8>> = kani::any();
+    match elem.0 {
+        Bound::Included(elem) => {
+            assert!(elem < 100);
+        }
+        Bound::Excluded(elem) => {
+            assert!(elem < 100);
+        }
+        Bound::Unbounded => {}
+    }
+}
+
+#[kani::proof]
+fn range() {
+    let elem: Wrapper<Range<u8>> = kani::any();
+    assert!(elem.0.start < 100);
+    assert!(elem.0.end < 100);
+}
+
+#[kani::proof]
+fn range_from() {
+    let elem: Wrapper<RangeFrom<u8>> = kani::any();
+    assert!(elem.0.start < 100);
+}
+
+#[kani::proof]
+fn range_inclusive() {
+    let elem: Wrapper<RangeInclusive<u8>> = kani::any();
+    assert!(*elem.0.start() < 100);
+    assert!(*elem.0.end() < 100);
+}
+
+#[kani::proof]
+fn range_to() {
+    let elem: Wrapper<RangeTo<u8>> = kani::any();
+    assert!(elem.0.end < 100);
+}
+
+#[kani::proof]
+fn range_to_inclusive() {
+    let elem: Wrapper<RangeToInclusive<u8>> = kani::any();
+    assert!(elem.0.end < 100);
+}
+
+struct Wrapper<T>(T);
+
+impl kani::Arbitrary for Wrapper<Bound<u8>> {
+    fn any() -> Self {
+        let val = kani::any();
+        match val {
+            Bound::Included(elem) => {
+                kani::assume(elem < 100);
+            }
+            Bound::Excluded(elem) => {
+                kani::assume(elem < 100);
+            }
+            Bound::Unbounded => {}
+        }
+        Self(val)
+    }
+}
+
+impl kani::Arbitrary for Wrapper<Range<u8>> {
+    fn any() -> Self {
+        let val = kani::any()..kani::any();
+        kani::assume(val.start < 100);
+        kani::assume(val.end < 100);
+        Self(val)
+    }
+}
+
+impl kani::Arbitrary for Wrapper<RangeFrom<u8>> {
+    fn any() -> Self {
+        let val = kani::any()..;
+        kani::assume(val.start < 100);
+        Self(val)
+    }
+}
+
+impl kani::Arbitrary for Wrapper<RangeInclusive<u8>> {
+    fn any() -> Self {
+        let val = kani::any()..=kani::any();
+        kani::assume(*val.start() < 100);
+        kani::assume(*val.end() < 100);
+        Self(val)
+    }
+}
+
+impl kani::Arbitrary for Wrapper<RangeTo<u8>> {
+    fn any() -> Self {
+        let val = ..kani::any();
+        kani::assume(val.end < 100);
+        Self(val)
+    }
+}
+
+impl kani::Arbitrary for Wrapper<RangeToInclusive<u8>> {
+    fn any() -> Self {
+        let val = ..=kani::any();
+        kani::assume(val.end < 100);
+        Self(val)
+    }
+}