Implement Range for OpStack

triton-air/proptest-regressions/table/processor.txt

@@ -0,0 +1,10 @@
+# Seeds for failure cases proptest has generated in the past. It is
+# automatically read and these particular cases re-run before any
+# novel cases are generated.
+#
+# It is recommended to check this file in to source control so that
+# everyone who runs the test benefits from these saved cases.
+cc 4aaff1f88a9cf5299763b9d8c18c56556418aeeccbd4bed18e6c3600edb7b911 # shrinks to input = _CannotGetStackWeightForOutOfRangeIndexArgs { index: 16 }
+cc d4911feba1af17da8508124e5019c16b94f50a263acc87a90b11fd41ffcc47a0 # shrinks to input = _CannotGetHelperVariableForOutOfRangeIndexArgs { index: 6 }
+cc fbb7a017123a4358a5596b823f694b7fa1b6894bc6ac316f1964f054b2a2bd2f # shrinks to input = _CannotGetOpStackColumnForOutOfRangeIndexArgs { index: 16 }
+cc 30aae9f964fd0d8b2baaa655c86be5fe89a1a1c32704320aa3620972ded5f67d # shrinks to input = _CannotGetIndicatorPolynomialForOutOfRangeIndexArgs { index: 16 }

triton-isa/src/op_stack.rs

@@ -4,6 +4,7 @@ use std::fmt::Result as FmtResult;
 use std::num::TryFromIntError;
 use std::ops::Index;
 use std::ops::IndexMut;
+use std::ops::{Range, RangeInclusive};
 
 use arbitrary::Arbitrary;
 use get_size::GetSize;
@@ -183,13 +184,57 @@ impl Index<usize> for OpStack {
     }
 }
 
+impl Index<Range<usize>> for OpStack {
+    type Output = [BFieldElement];
+
+    fn index(&self, range: Range<usize>) -> &Self::Output {
+        if range.end <= self.stack.len() && range.start < range.end {
+            &self.stack[range]
+        } else {
+            panic!("a range is out of bounds")
+        }
+    }
+}
+
+impl Index<RangeInclusive<usize>> for OpStack {
+    type Output = [BFieldElement];
+
+    fn index(&self, range: RangeInclusive<usize>) -> &Self::Output {
+        if range.end() <= &self.stack.len() && range.start() < range.end() {
+            &self.stack[range]
+        } else {
+            panic!("a range is out of bounds")
+        }
+    }
+}
+
 impl IndexMut<usize> for OpStack {
     fn index_mut(&mut self, index: usize) -> &mut Self::Output {
         let top_of_stack = self.len() - 1;
         &mut self.stack[top_of_stack - index]
     }
 }
 
+impl IndexMut<Range<usize>> for OpStack {
+    fn index_mut(&mut self, range: Range<usize>) -> &mut Self::Output {
+        if range.end <= self.stack.len() && range.start < range.end {
+            &mut self.stack[range]
+        } else {
+            panic!("a range is out of bounds")
+        }
+    }
+}
+
+impl IndexMut<RangeInclusive<usize>> for OpStack {
+    fn index_mut(&mut self, range: RangeInclusive<usize>) -> &mut Self::Output {
+        if range.end() <= &self.stack.len() && range.start() < range.end() {
+            &mut self.stack[range]
+        } else {
+            panic!("a range is out of bounds")
+        }
+    }
+}
+
 impl Index<OpStackElement> for OpStack {
     type Output = BFieldElement;
 
@@ -198,12 +243,60 @@ impl Index<OpStackElement> for OpStack {
     }
 }
 
+impl Index<Range<OpStackElement>> for OpStack {
+    type Output = [BFieldElement];
+
+    fn index(&self, range: Range<OpStackElement>) -> &Self::Output {
+        let (start, end) = (usize::from(range.start), usize::from(range.end));
+        if end <= self.stack.len() && start < end {
+            &self.stack[start..end]
+        } else {
+            panic!("a range is out of bounds")
+        }
+    }
+}
+
+impl Index<RangeInclusive<OpStackElement>> for OpStack {
+    type Output = [BFieldElement];
+
+    fn index(&self, range: RangeInclusive<OpStackElement>) -> &Self::Output {
+        let (start, end) = (usize::from(range.start()), usize::from(range.end()));
+        if end <= self.stack.len() && start < end {
+            &self.stack[start..=end]
+        } else {
+            panic!("a range is out of bounds")
+        }
+    }
+}
+
 impl IndexMut<OpStackElement> for OpStack {
     fn index_mut(&mut self, stack_element: OpStackElement) -> &mut Self::Output {
         &mut self[usize::from(stack_element)]
     }
 }
 
+impl IndexMut<Range<OpStackElement>> for OpStack {
+    fn index_mut(&mut self, range: Range<OpStackElement>) -> &mut Self::Output {
+        let (start, end) = (usize::from(range.start), usize::from(range.end));
+        if end <= self.stack.len() && start < end {
+            &mut self.stack[start..end]
+        } else {
+            panic!("a range is out of bounds")
+        }
+    }
+}
+
+impl IndexMut<RangeInclusive<OpStackElement>> for OpStack {
+    fn index_mut(&mut self, range: RangeInclusive<OpStackElement>) -> &mut Self::Output {
+        let (start, end) = (usize::from(range.start()), usize::from(range.end()));
+        if end <= self.stack.len() && start < end {
+            &mut self.stack[start..=end]
+        } else {
+            panic!("a range is out of bounds")
+        }
+    }
+}
+
 impl IntoIterator for OpStack {
     type Item = BFieldElement;
     type IntoIter = std::vec::IntoIter<Self::Item>;
@@ -1025,4 +1118,207 @@ mod tests {
         let expected_len = 2 * OpStackElement::COUNT - 1;
         prop_assert_eq!(expected_len, op_stack.len());
     }
+
+    fn setup_op_stack() -> OpStack {
+        OpStack {
+            stack: bfe_vec![1, 2, 3, 4, 5],
+            underflow_io_sequence: vec![],
+        }
+    }
+
+    #[test]
+    fn test_opstack_index_range() {
+        let op_stack = setup_op_stack();
+
+        let actual = op_stack.index(1..3);
+
+        let expected = bfe_vec![2, 3];
+
+        // Test typical range
+        assert_eq!(actual.to_vec(), expected);
+
+        // Test boundary range
+        assert_eq!(&op_stack[0..op_stack.stack.len()], &op_stack.stack[..]);
+
+        // Test out-of-bounds range (should panic)
+        let result = std::panic::catch_unwind(|| {
+            let _ = op_stack.index(3..6);
+        });
+        assert!(result.is_err());
+    }
+
+    #[test]
+    fn test_opstack_index_range_inclusive() {
+        let op_stack = setup_op_stack();
+
+        let actual = op_stack.index(1..=3);
+
+        let expected = bfe_vec!(2, 3, 4);
+
+        // Test typical inclusive range
+        assert_eq!(actual.to_vec(), expected);
+
+        // Test boundary inclusive range
+        assert_eq!(&op_stack[0..=op_stack.stack.len() - 1], &op_stack.stack[..]);
+
+        // Test out-of-bounds inclusive range (should panic)
+        let result = std::panic::catch_unwind(|| {
+            let _ = op_stack.index(2..=5);
+        });
+        assert!(result.is_err());
+    }
+
+    #[test]
+    fn test_opstack_index_mut_range_usize() {
+        let mut stack = setup_op_stack();
+
+        // Modify a middle range of elements
+        let range = 1..3;
+        let slice = stack.index_mut(range);
+        slice[0] = BFieldElement::new(20);
+        slice[1] = BFieldElement::new(30);
+
+        assert_eq!(
+            stack.stack,
+            vec![bfe!(1), bfe!(20), bfe!(30), bfe!(4), bfe!(5)]
+        );
+
+        // Test out-of-bounds access (should panic)
+        let result = std::panic::catch_unwind(move || {
+            let _ = stack.index_mut(3..6); // This should panic
+        });
+        assert!(result.is_err(), "Expected a panic for out-of-bounds range");
+    }
+
+    #[test]
+    fn test_opstack_index_mut_range_inclusive_usize() {
+        let mut stack = setup_op_stack();
+
+        // Modify an inclusive range of elements
+        let range = 1..=3;
+        let slice = stack.index_mut(range);
+        slice[0] = bfe!(25);
+        slice[1] = bfe!(35);
+        slice[2] = bfe!(45);
+
+        assert_eq!(
+            stack.stack,
+            vec![bfe!(1), bfe!(25), bfe!(35), bfe!(45), bfe!(5),]
+        );
+
+        // Test out-of-bounds access (should panic)
+        let result = std::panic::catch_unwind(move || {
+            let _ = stack.index_mut(3..=5); // This should panic
+        });
+        assert!(
+            result.is_err(),
+            "Expected a panic for out-of-bounds inclusive range"
+        );
+    }
+
+    #[test]
+    fn test_opstack_element_index_range() {
+        // Initialize an OpStack with some BFieldElement values
+        let op_stack = setup_op_stack();
+
+        // Define some ranges
+        let range = Range::<OpStackElement> {
+            start: OpStackElement::ST1,
+            end: OpStackElement::ST3,
+        };
+
+        // Test valid range indexing
+        let slice = op_stack.index(range);
+        assert_eq!(slice, &[bfe!(2), bfe!(3)]);
+
+        // Test out of bounds range
+        let out_of_bounds_range = Range::<OpStackElement> {
+            start: OpStackElement::ST3,
+            end: OpStackElement::ST6,
+        };
+        let result = std::panic::catch_unwind(|| op_stack.index(out_of_bounds_range));
+        assert!(result.is_err(), "Expected out of bounds panic");
+    }
+
+    #[test]
+    fn test_opstack_element_index_range_inclusive() {
+        let op_stack = setup_op_stack();
+
+        // Define some inclusive ranges
+        let inclusive_range =
+            RangeInclusive::<OpStackElement>::new(OpStackElement::ST1, OpStackElement::ST3);
+
+        // Test valid inclusive range indexing
+        let inclusive_slice = op_stack.index(inclusive_range);
+        assert_eq!(inclusive_slice, &[bfe!(2), bfe!(3), bfe!(4)]);
+
+        // Test out of bounds inclusive range
+        let out_of_bounds_inclusive = RangeInclusive::new(OpStackElement::ST2, OpStackElement::ST5);
+        let result = std::panic::catch_unwind(|| op_stack.index(out_of_bounds_inclusive));
+        assert!(result.is_err());
+    }
+
+    #[test]
+    fn test_opstack_element_index_mut_range() {
+        // Initialize an OpStack with some BFieldElement values
+        let mut op_stack = setup_op_stack();
+
+        // Define a mutable range and modify the values within it
+        let range = Range {
+            start: OpStackElement::ST1,
+            end: OpStackElement::ST4,
+        };
+
+        let slice = op_stack.index_mut(range);
+        slice[0] = bfe!(20);
+        slice[1] = bfe!(30);
+        slice[2] = bfe!(40);
+
+        // Verify modifications
+        assert_eq!(
+            op_stack.stack,
+            vec![bfe!(1), bfe!(20), bfe!(30), bfe!(40), bfe!(5),]
+        );
+    }
+
+    #[test]
+    #[should_panic(expected = "a range is out of bounds")]
+    fn test_op_stack_index_mut_range_out_of_bounds_panic() {
+        let mut op_stack = setup_op_stack();
+        // Test out of bounds range
+        let out_of_bounds_range = Range {
+            start: OpStackElement::ST3,
+            end: OpStackElement::ST6,
+        };
+        let _ = op_stack.index_mut(out_of_bounds_range);
+    }
+
+    #[test]
+    fn test_index_mut_range_inclusive() {
+        // Initialize an OpStack with some BFieldElement values
+        let mut op_stack = setup_op_stack();
+
+        // Define a mutable inclusive range and modify the values within it
+        let inclusive_range = RangeInclusive::new(OpStackElement::ST1, OpStackElement::ST3);
+        {
+            let inclusive_slice = op_stack.index_mut(inclusive_range);
+            inclusive_slice[0] = bfe!(25);
+            inclusive_slice[1] = bfe!(35);
+            inclusive_slice[2] = bfe!(45);
+        }
+
+        // Verify modifications
+        assert_eq!(
+            op_stack.stack,
+            vec![bfe!(1), bfe!(25), bfe!(35), bfe!(45), bfe!(5),]
+        );
+    }
+
+    #[test]
+    #[should_panic(expected = "range end index 6 out of range for slice of length 5")]
+    fn test_op_stack_index_mut_range_inclusive_out_of_bounds_panic() {
+        let mut op_stack = setup_op_stack();
+        let out_of_bounds_inclusive = RangeInclusive::new(OpStackElement::ST2, OpStackElement::ST5);
+        let _ = op_stack.index_mut(out_of_bounds_inclusive);
+    }
 }