Skip to content

Commit

Permalink
refactor: simplify indexing into OpStack
Browse files Browse the repository at this point in the history
  • Loading branch information
@jan-ferdinand
jan-ferdinand committed Dec 12, 2023
1 parent bed148b commit 4b31b2f
Show file tree
Hide file tree
Showing 3 changed files with 126 additions and 94 deletions.
9 changes: 4 additions & 5 deletions triton-vm/src/error.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,6 @@ use twenty_first::shared_math::bfield_codec::BFieldCodec;
use twenty_first::shared_math::digest::DIGEST_LENGTH;

use crate::instruction::Instruction;
use crate::op_stack::OpStackElement;
use crate::proof_item::ProofItem;
use crate::proof_stream::ProofStream;
use crate::stark::StarkHasher;
Expand Down Expand Up @@ -64,8 +63,8 @@ pub enum InstructionError {
#[error("assertion failed: st0 must be 1")]
AssertionFailed,

#[error("vector assertion failed: stack[{0}] != stack[{}]", usize::from(.0) + DIGEST_LENGTH)]
VectorAssertionFailed(OpStackElement),
#[error("vector assertion failed: stack[{0}] != stack[{}]", .0 + DIGEST_LENGTH)]
VectorAssertionFailed(usize),

#[error("cannot swap stack element 0 with itself")]
SwapST0,
Expand Down Expand Up @@ -313,7 +312,7 @@ mod tests {
};
let_assert!(Err(err) = program.run([].into(), [].into()));
let_assert!(InstructionError::VectorAssertionFailed(index) = err.source);
assert!(1 == usize::from(index));
assert!(1 == index);
}

#[test]
Expand Down Expand Up @@ -356,7 +355,7 @@ mod tests {

let_assert!(Err(err) = program.run([].into(), [].into()));
let_assert!(InstructionError::VectorAssertionFailed(index) = err.source);
prop_assert_eq!(disturbance_index, usize::from(index));
prop_assert_eq!(disturbance_index, index);
}

#[test]
Expand Down
131 changes: 86 additions & 45 deletions triton-vm/src/op_stack.rs
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
use std::fmt::Display;
use std::fmt::Formatter;
use std::fmt::Result as FmtResult;
use std::ops::Index;
use std::ops::IndexMut;

use arbitrary::Arbitrary;
use get_size::GetSize;
Expand Down Expand Up @@ -40,8 +42,15 @@ pub const NUM_OP_STACK_REGISTERS: usize = OpStackElement::COUNT;
/// [`OpStackElement::COUNT`] elements of the op-stack, and the op-stack underflow memory is the
/// remaining elements.
#[derive(Debug, Clone, PartialEq, Eq)]
// If the op stack is empty, things have gone horribly wrong. Suppressing this lint is preferred
// to implementing a basically useless `is_empty()` method.
#[allow(clippy::len_without_is_empty)]
pub struct OpStack {
/// The underlying, actual stack. When manually accessing, be aware of reversed indexing:
/// while `op_stack[0]` is the top of the stack, `op_stack.stack[0]` is the lowest element in
/// the stack.
pub stack: Vec<BFieldElement>,

underflow_io_sequence: Vec<UnderflowIO>,
}

Expand All @@ -58,6 +67,10 @@ impl OpStack {
}
}

pub fn len(&self) -> usize {
self.stack.len()
}

pub(crate) fn push(&mut self, element: BFieldElement) {
self.stack.push(element);
self.record_underflow_io(UnderflowIO::Write);
Expand Down Expand Up @@ -95,7 +108,7 @@ impl OpStack {
}

pub(crate) fn assert_is_u32(&self, stack_element: OpStackElement) -> Result<()> {
let element = self.peek_at(stack_element);
let element = self[stack_element];
match element.value() <= u32::MAX as u64 {
true => Ok(()),
false => Err(FailedU32Conversion(element)),
Expand All @@ -120,42 +133,29 @@ impl OpStack {
Ok(elements)
}

pub(crate) fn peek_at(&self, stack_element: OpStackElement) -> BFieldElement {
let stack_element_index = usize::from(stack_element);
let top_of_stack_index = self.stack.len() - 1;
self.stack[top_of_stack_index - stack_element_index]
}

pub(crate) fn peek_at_top_extension_field_element(&self) -> XFieldElement {
let coefficient_0 = self.peek_at(ST0);
let coefficient_1 = self.peek_at(ST1);
let coefficient_2 = self.peek_at(ST2);
let coefficients = [coefficient_0, coefficient_1, coefficient_2];
XFieldElement::new(coefficients)
XFieldElement::new([self[0], self[1], self[2]])
}

pub(crate) fn swap_top_with(&mut self, stack_element: OpStackElement) {
let stack_element_index = usize::from(stack_element);
let top_of_stack_index = self.stack.len() - 1;
self.stack
.swap(top_of_stack_index, top_of_stack_index - stack_element_index);
pub(crate) fn swap_top_with(&mut self, st: OpStackElement) {
(self[0], self[st]) = (self[st], self[0]);
}

pub(crate) fn would_be_too_shallow(&self, stack_delta: i32) -> bool {
self.stack.len() as i32 + stack_delta < OpStackElement::COUNT as i32
self.len() as i32 + stack_delta < OpStackElement::COUNT as i32
}

/// The address of the next free address of the op-stack. Equivalent to the current length of
/// the op-stack.
pub(crate) fn pointer(&self) -> BFieldElement {
(self.stack.len() as u64).into()
(self.len() as u64).into()
}

/// The first element of the op-stack underflow memory, or 0 if the op-stack underflow memory
/// is empty.
pub(crate) fn first_underflow_element(&self) -> BFieldElement {
let default = BFieldElement::zero();
let Some(top_of_stack_index) = self.stack.len().checked_sub(1) else {
let Some(top_of_stack_index) = self.len().checked_sub(1) else {
return default;
};
let Some(underflow_start) = top_of_stack_index.checked_sub(OpStackElement::COUNT) else {
Expand All @@ -165,6 +165,47 @@ impl OpStack {
}
}

impl Index<usize> for OpStack {
type Output = BFieldElement;

fn index(&self, index: usize) -> &Self::Output {
let top_of_stack = self.len() - 1;
&self.stack[top_of_stack - index]
}
}

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 Index<OpStackElement> for OpStack {
type Output = BFieldElement;

fn index(&self, stack_element: OpStackElement) -> &Self::Output {
&self[usize::from(stack_element)]
}
}

impl IndexMut<OpStackElement> for OpStack {
fn index_mut(&mut self, stack_element: OpStackElement) -> &mut Self::Output {
&mut self[usize::from(stack_element)]
}
}

impl IntoIterator for OpStack {
type Item = BFieldElement;
type IntoIter = std::vec::IntoIter<Self::Item>;

fn into_iter(self) -> Self::IntoIter {
let mut stack = self.stack;
stack.reverse();
stack.into_iter()
}
}

/// Indicates changes to the op-stack underflow memory.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, GetSize, Serialize, Deserialize, Arbitrary)]
#[must_use = "The change to underflow memory should be handled."]
Expand Down Expand Up @@ -620,36 +661,36 @@ mod tests {
let mut op_stack = OpStack::new(digest);

// verify height
assert!(op_stack.stack.len() == 16);
assert!(op_stack.pointer().value() as usize == op_stack.stack.len());
assert!(op_stack.len() == 16);
assert!(op_stack.pointer().value() as usize == op_stack.len());

// push elements 1 thru 17
for i in 1..=17 {
op_stack.push(BFieldElement::new(i as u64));
}

// verify height
assert!(op_stack.stack.len() == 33);
assert!(op_stack.pointer().value() as usize == op_stack.stack.len());
assert!(op_stack.len() == 33);
assert!(op_stack.pointer().value() as usize == op_stack.len());

// verify that all accessible items are different
let mut container = vec![
op_stack.peek_at(ST0),
op_stack.peek_at(ST1),
op_stack.peek_at(ST2),
op_stack.peek_at(ST3),
op_stack.peek_at(ST4),
op_stack.peek_at(ST5),
op_stack.peek_at(ST6),
op_stack.peek_at(ST7),
op_stack.peek_at(ST8),
op_stack.peek_at(ST9),
op_stack.peek_at(ST10),
op_stack.peek_at(ST11),
op_stack.peek_at(ST12),
op_stack.peek_at(ST13),
op_stack.peek_at(ST14),
op_stack.peek_at(ST15),
op_stack[ST0],
op_stack[ST1],
op_stack[ST2],
op_stack[ST3],
op_stack[ST4],
op_stack[ST5],
op_stack[ST6],
op_stack[ST7],
op_stack[ST8],
op_stack[ST9],
op_stack[ST10],
op_stack[ST11],
op_stack[ST12],
op_stack[ST13],
op_stack[ST14],
op_stack[ST15],
op_stack.first_underflow_element(),
];
let len_before = container.len();
Expand All @@ -664,16 +705,16 @@ mod tests {
}

// verify height
assert!(op_stack.stack.len() == 22);
assert!(op_stack.pointer().value() as usize == op_stack.stack.len());
assert!(op_stack.len() == 22);
assert!(op_stack.pointer().value() as usize == op_stack.len());

// pop 2 XFieldElements
let _ = op_stack.pop_extension_field_element().expect("can't pop");
let _ = op_stack.pop_extension_field_element().expect("can't pop");

// verify height
assert!(op_stack.stack.len() == 16);
assert!(op_stack.pointer().value() as usize == op_stack.stack.len());
assert!(op_stack.len() == 16);
assert!(op_stack.pointer().value() as usize == op_stack.len());

// verify underflow
assert!(op_stack.would_be_too_shallow(-1));
Expand All @@ -682,7 +723,7 @@ mod tests {
#[test]
fn trying_to_access_first_underflow_element_never_panics() {
let mut op_stack = OpStack::new(Default::default());
let way_too_long = 2 * op_stack.stack.len();
let way_too_long = 2 * op_stack.len();
for _ in 0..way_too_long {
let _ = op_stack.pop();
let _ = op_stack.first_underflow_element();
Expand Down
Loading

0 comments on commit 4b31b2f

Please sign in to comment.