From 6f9b0a2b96f2ca49d194139ebdfead2d95ada78d Mon Sep 17 00:00:00 2001
From: Andrew Dirksen <andrew@dirksen.com>
Date: Mon, 19 Nov 2018 19:44:13 -0800
Subject: [PATCH] optimize push and pop for StackWithLimit

---
 src/common/stack.rs | 82 +++++++++++++++++++++++++++++++++++----------
 src/runner.rs       |  2 +-
 2 files changed, 66 insertions(+), 18 deletions(-)

diff --git a/src/common/stack.rs b/src/common/stack.rs
index 1c553e5..85e0901 100644
--- a/src/common/stack.rs
+++ b/src/common/stack.rs
@@ -1,4 +1,5 @@
 use core::marker::PhantomData;
+use core::ptr;
 use core::usize;
 
 #[allow(unused_imports)]
@@ -77,24 +78,49 @@ impl<T> StackWithLimit<T> {
 	/// # Errors
 	///
 	/// Returns Err(StackOverflow) if stack is already full.
+	#[inline]
 	pub(crate) fn push(&mut self, value: T) -> Result<(), StackOverflow> {
-		let ret = if self.stack.len() < self.limit {
-			self.stack.push(value);
-			Ok(())
-		} else {
-			Err(StackOverflow)
-		};
-		debug_assert!(
-			self.stack.len() <= self.limit,
-			"Stack length should never be larger than stack limit."
-		);
-		ret
+		if self.stack.len() == self.stack.capacity() {
+			if self.stack.len() == self.limit {
+				return Err(StackOverflow);
+			}
+			// grows exponentially, just like Vec normally does
+			// https://doc.rust-lang.org/1.26.0/src/alloc/raw_vec.rs.html#462
+			let desired_len = self
+				.stack
+				.len()
+				.checked_mul(2)
+				.unwrap_or(usize::MAX)
+				.min(self.limit)
+				.max(1);
+			let additional_len = desired_len - self.stack.len();
+			self.stack.reserve_exact(additional_len);
+		}
+		debug_assert!(self.stack.len() < self.limit);
+		debug_assert!(self.stack.len() < self.stack.capacity());
+		let len = self.stack.len();
+		unsafe {
+			ptr::write(self.stack.get_unchecked_mut(len), value);
+			self.stack.set_len(len + 1);
+		}
+		Ok(())
 	}
 
 	pub(crate) fn pop(&mut self) -> Option<T> {
 		self.stack.pop()
 	}
 
+	/// Remove and Return top element. Does not check for emptyness.
+	/// If this is called on a zero length stack, bad things will happen.
+	/// Do not call this method unless you can prove the stack has length.
+	#[inline]
+	pub(crate) unsafe fn pop_unchecked(&mut self) -> T {
+		debug_assert!(self.stack.len() > 0);
+		let len = self.stack.len();
+		self.stack.set_len(len - 1);
+		ptr::read(self.stack.get_unchecked(self.stack.len()))
+	}
+
 	/// Return optional reference to item `depth` distance away from top
 	///
 	/// `bstack.get_relative_to_top(0)` gets the top of the stack
@@ -360,12 +386,34 @@ mod test {
 			StackSizeInitial(StackSize::from_element_count(0)),
 			StackSizeLimit(StackSize::from_element_count(usize::MAX)),
 		);
-		println!(
-			"{}",
-			StackSizeLimit(StackSize::<i32>::from_element_count(usize::MAX))
-				.0
-				.element_count()
-		);
+
 		exersize(bstack);
 	}
+
+	// Make sure the stack resizes properly.
+	#[test]
+	fn must_resize() {
+		let mut bstack = StackWithLimit::<i32>::new(
+			StackSizeInitial(StackSize::from_element_count(2)),
+			StackSizeLimit(StackSize::from_element_count(4)),
+		);
+		bstack.push(2).unwrap();
+		bstack.push(4).unwrap();
+		bstack.push(8).unwrap();
+		bstack.push(16).unwrap();
+		bstack.push(16).unwrap_err();
+		assert_eq!(bstack.pop(), Some(16));
+		assert_eq!(bstack.pop(), Some(8));
+		assert_eq!(bstack.pop(), Some(4));
+		assert_eq!(bstack.pop(), Some(2));
+	}
+
+	#[test]
+	fn pop_unchecked() {
+		let mut bstack = StackWithLimit::<i32>::with_size(StackSize::from_element_count(20));
+		bstack.push(8).unwrap();
+		bstack.push(0).unwrap();
+		assert_eq!(unsafe { bstack.pop_unchecked() }, 0);
+		assert_eq!(unsafe { bstack.pop_unchecked() }, 8);
+	}
 }
diff --git a/src/runner.rs b/src/runner.rs
index c578821..6f0bb05 100644
--- a/src/runner.rs
+++ b/src/runner.rs
@@ -1321,7 +1321,7 @@ impl ValueStack {
 
 	#[inline]
 	fn pop(&mut self) -> RuntimeValueInternal {
-		self.0.pop().expect("pre-validated")
+		unsafe { self.0.pop_unchecked() }
 	}
 
 	#[inline]