From fc81592e42c2d7fe8c2914e4e3b650b1714bad71 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Philippe=20Laferri=C3=A8re?= <plafer@protonmail.com>
Date: Mon, 27 Nov 2023 08:56:13 -0500
Subject: [PATCH] `StackOutputs`: `get_stack_{item, word}()` methods (#1155)

* get_stack_item()

* get_stack_word

* adjust comment

* adjust comments

* fix comment

* fix comments
---
 CHANGELOG.md              |  1 +
 core/src/stack/outputs.rs | 28 ++++++++++++++++++++++++++++
 2 files changed, 29 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 5846011a1f..de6c937b6d 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -17,6 +17,7 @@
 #### VM Internals
 - Introduced the `Event` decorator and an associated `on_event` handler on the `Host` trait (#1119).
 - Updated Winterfell dependency to v0.7 (#1121).
+- Added methods `StackOutputs::get_stack_item()` and `StackOutputs::get_stack_word()` (#1155).
 
 ## 0.7.0 (2023-10-11)
 
diff --git a/core/src/stack/outputs.rs b/core/src/stack/outputs.rs
index a1f7baa028..1e391dec41 100644
--- a/core/src/stack/outputs.rs
+++ b/core/src/stack/outputs.rs
@@ -1,3 +1,7 @@
+use miden_crypto::Word;
+
+use crate::utils::range;
+
 use super::{
     ByteWriter, Felt, OutputError, Serializable, StackTopState, StarkField, ToElements, Vec,
     STACK_TOP_SIZE,
@@ -84,6 +88,30 @@ impl StackOutputs {
     // PUBLIC ACCESSORS
     // --------------------------------------------------------------------------------------------
 
+    /// Returns the element located at the specified position on the stack or `None` if out of
+    /// bounds.
+    pub fn get_stack_item(&self, idx: usize) -> Option<Felt> {
+        self.stack.get(idx).map(|&felt| felt.into())
+    }
+
+    /// Returns the word located starting at the specified Felt position on the stack or `None` if
+    /// out of bounds. For example, passing in `0` returns the word at the top of the stack, and
+    /// passing in `4` returns the word starting at element index `4`.
+    pub fn get_stack_word(&self, idx: usize) -> Option<Word> {
+        let word_elements: Word = {
+            let word_elements: Vec<Felt> = range(idx, 4)
+                .map(|idx| self.get_stack_item(idx))
+                // Elements need to be reversed, since a word `[a, b, c, d]` will be stored on the
+                // stack as `[d, c, b, a]`
+                .rev()
+                .collect::<Option<_>>()?;
+
+            word_elements.try_into().expect("a Word contains 4 elements")
+        };
+
+        Some(word_elements)
+    }
+
     /// Returns the stack outputs, which is state of the stack at the end of execution converted to
     /// integers.
     pub fn stack(&self) -> &[u64] {