1 files changed, 31 insertions, 13 deletions
diff --git a/src/forth/vm.rs b/src/forth/vm.rs
index 9c7f203..e5973dd 100644
--- a/src/forth/vm.rs
+++ b/src/forth/vm.rs
@@ -172,26 +172,33 @@ impl VM {
                 self.stack.0.push(n2 / n1);
             },
             OpCode::Dup => {
-                let n = self.stack.0.pop().ok_or(RuntimeError::StackUnderflow)?;
-                self.stack.0.push(n);
-                self.stack.0.push(n);
+                let len = self.stack.0.len();
+                if len < 1 {
+                    return Err(RuntimeError::StackUnderflow)
+                }
+                self.stack.0.push(self.stack.0[len-1]);
             },
             OpCode::Drop => {
                 self.stack.0.pop().ok_or(RuntimeError::StackUnderflow)?;
             },
             OpCode::Swap => {
-                let v1 = self.stack.0.pop().ok_or(RuntimeError::StackUnderflow)?;
-                let v2 = self.stack.0.pop().ok_or(RuntimeError::StackUnderflow)?;
-                self.stack.0.push(v1);
-                self.stack.0.push(v2);
+                let top = self.stack.0.len() - 1;
+                if top < 1 {
+                    return Err(RuntimeError::StackUnderflow)
+                }
+                let (a, b) =
+                    self.stack.0.split_at_mut(top);
+                std::mem::swap(&mut a[top-1], &mut b[0]);
             },
             OpCode::Rot => {
-                let v1 = self.stack.0.pop().ok_or(RuntimeError::StackUnderflow)?;
-                let v2 = self.stack.0.pop().ok_or(RuntimeError::StackUnderflow)?;
-                let v3 = self.stack.0.pop().ok_or(RuntimeError::StackUnderflow)?;
-                self.stack.0.push(v2);
-                self.stack.0.push(v1);
-                self.stack.0.push(v3);
+                let top = self.stack.0.len() - 1;
+                if top < 2 {
+                    return Err(RuntimeError::StackUnderflow)
+                }
+                let (a, b) =
+                    self.stack.0.split_at_mut(top);
+                std::mem::swap(&mut a[top-1], &mut b[0]);
+                std::mem::swap(&mut a[top-2], &mut b[0]);
             },
             OpCode::EQ => {
                 let n1 = self.stack.0.pop().ok_or(RuntimeError::StackUnderflow)?;
@@ -578,6 +585,17 @@ mod tests {
     }
 
     #[test]
+    fn opcode_dup() -> Result<(), RuntimeError> {
+        let wordlist = wl_for(vec![
+            vec![OpCode::Num(1), OpCode::Dup],
+        ]);
+        let mut vm = VM::new(wordlist);
+        vm.run()?;
+        assert_eq!(vm.stack.0, vec![1, 1]);
+        Ok(())
+    }
+
+    #[test]
     fn opcode_swap() -> Result<(), RuntimeError> {
         let wordlist = wl_for(vec![
             vec![OpCode::Num(1), OpCode::Num(2), OpCode::Swap],