Implement parsing of pushw and env io operations

assembly/src/v1/parsers/io_ops.rs

@@ -8,8 +8,48 @@ use super::{parse_element_param, AssemblyError, BaseElement, FieldElement, Opera
 /// In cases when the immediate value is 0, PUSH operation is replaced with PAD. Also, in cases
 /// when immediate value is 1, PUSH operation is replaced with PAD INCR because in most cases
 /// this will be more efficient than doing a PUSH.
+///
+/// # Errors
+///
+/// This function expects an assembly op with exactly one immediate value that is a valid field
+/// element in decimal or hexadecimal representation. It will return an error if the immediate
+/// value is invalid or missing.
 pub fn parse_push(span_ops: &mut Vec<Operation>, op: &Token) -> Result<(), AssemblyError> {
+    validate_op_len(op, 2, 2)?;
+
     let value = parse_element_param(op, 1)?;
+    push_value(span_ops, value);
+
+    Ok(())
+}
+
+/// Appends PUSH operations to the span block until all elements of the provided word are pushed
+/// onto the stack.
+///
+/// All immediate values are handled in the same way as for the single element "push" operation.
+///
+/// # Errors
+///
+/// This function expects an assembly op with 4 immediate values that are valid field elements
+/// in decimal or hexadecimal representation. It will return an error if the assembly instruction's
+/// immediate values are invalid.
+pub fn parse_pushw(span_ops: &mut Vec<Operation>, op: &Token) -> Result<(), AssemblyError> {
+    validate_op_len(op, 5, 5)?;
+
+    for idx in 1..=4 {
+        let value = parse_element_param(op, idx)?;
+        push_value(span_ops, value);
+    }
+
+    Ok(())
+}
+
+/// This is a helper function that appends a PUSH operation to the span block which puts the
+/// provided value parameter onto the stack.
+///
+/// When the value is 0, PUSH operation is replaced with PAD. When the value is 1, PUSH operation
+/// is replaced with PAD INCR because in most cases this will be more efficient than doing a PUSH.
+fn push_value(span_ops: &mut Vec<Operation>, value: BaseElement) {
     if value == BaseElement::ZERO {
         span_ops.push(Operation::Pad);
     } else if value == BaseElement::ONE {
@@ -18,20 +58,33 @@ pub fn parse_push(span_ops: &mut Vec<Operation>, op: &Token) -> Result<(), Assem
     } else {
         span_ops.push(Operation::Push(value));
     }
-    Ok(())
-}
-
-/// TODO: implement
-pub fn parse_pushw(_span_ops: &mut Vec<Operation>, _op: &Token) -> Result<(), AssemblyError> {
-    unimplemented!()
 }
 
 // ENVIRONMENT INPUTS
 // ================================================================================================
 
-/// TODO: implement
-pub fn parse_env(_span_ops: &mut Vec<Operation>, _op: &Token) -> Result<(), AssemblyError> {
-    unimplemented!()
+/// Appends machine operations to the current span block according to the requested environment
+/// assembly instruction.
+///
+/// "env.sdepth" pushes the current depth of the stack onto the top of the stack, which is handled
+/// directly by the SDEPTH operation.
+///
+/// # Errors
+///
+/// This function expects a valid assembly environment op that specifies the environment input to
+/// be handled. It will return an error if the assembly instruction is malformed or the environment
+/// input is unrecognized.
+pub fn parse_env(span_ops: &mut Vec<Operation>, op: &Token) -> Result<(), AssemblyError> {
+    validate_op_len(op, 2, 2)?;
+
+    match op.parts()[1] {
+        "sdepth" => {
+            span_ops.push(Operation::SDepth);
+        }
+        _ => return Err(AssemblyError::invalid_op(op)),
+    }
+
+    Ok(())
 }
 
 // NON-DETERMINISTIC INPUTS
@@ -64,14 +117,12 @@ pub fn parse_readw(_span_ops: &mut Vec<Operation>, _op: &Token) -> Result<(), As
 /// "mem.store" is a write operation that saves the top 4 elements of the stack to memory and
 /// leaves them on the stack.
 pub fn parse_mem(span_ops: &mut Vec<Operation>, op: &Token) -> Result<(), AssemblyError> {
-    match op.num_parts() {
-        0 | 1 => Err(AssemblyError::missing_param(op)),
-        2 | 3 => match op.parts()[1] {
-            "push" | "load" => parse_mem_read(span_ops, op),
-            "pop" | "store" => parse_mem_write(span_ops, op),
-            _ => Err(AssemblyError::invalid_op(op)),
-        },
-        _ => Err(AssemblyError::extra_param(op)),
+    validate_op_len(op, 2, 3)?;
+
+    match op.parts()[1] {
+        "push" | "load" => parse_mem_read(span_ops, op),
+        "pop" | "store" => parse_mem_write(span_ops, op),
+        _ => Err(AssemblyError::invalid_op(op)),
     }
 }
 
@@ -151,6 +202,26 @@ fn push_mem_addr(span_ops: &mut Vec<Operation>, op: &Token) -> Result<(), Assemb
     Ok(())
 }
 
+// HELPERS
+// ================================================================================================
+
+/// This is a helper function that validates the length of the assembly io instruction and returns
+/// an error if params are missing or there are extras.
+///
+/// The minimum and maximum number of valid instruction parts for the op are provided by the
+/// caller.
+///
+/// # Errors
+///
+/// This function will return an AssemblyError if the assembly op has too many or too few parts.
+fn validate_op_len(op: &Token, min_parts: usize, max_parts: usize) -> Result<(), AssemblyError> {
+    match op.num_parts() {
+        too_few if too_few < min_parts => Err(AssemblyError::missing_param(op)),
+        too_many if too_many > max_parts => Err(AssemblyError::extra_param(op)),
+        _ => Ok(()),
+    }
+}
+
 // TESTS
 // ================================================================================================
 
@@ -163,21 +234,114 @@ mod tests {
         let mut span_ops: Vec<Operation> = Vec::new();
         let op_0 = Token::new("push.0", 0);
         let op_1 = Token::new("push.1", 0);
-        let op_other = Token::new("push.2", 0);
+        let op_dec = Token::new("push.135", 0);
+        let op_hex = Token::new("push.0x7b", 0);
         let expected = vec![
             Operation::Pad,
             Operation::Pad,
             Operation::Incr,
-            Operation::Push(BaseElement::new(2)),
+            Operation::Push(BaseElement::new(135)),
+            Operation::Push(BaseElement::new(123)),
         ];
 
         parse_push(&mut span_ops, &op_0).expect("Failed to parse push.0");
         parse_push(&mut span_ops, &op_1).expect("Failed to parse push.1");
-        parse_push(&mut span_ops, &op_other).expect("Failed to parse push.2");
+        parse_push(&mut span_ops, &op_dec).expect("Failed to parse push of decimal element 123");
+        parse_push(&mut span_ops, &op_hex).expect("Failed to parse push of hex element 0x7b");
 
         assert_eq!(span_ops, expected);
     }
 
+    #[test]
+    fn push_invalid() {
+        // fails when immediate value is invalid or missing
+        let mut span_ops: Vec<Operation> = Vec::new();
+        let param_idx = 0;
+
+        // value missing
+        let op_no_val = Token::new("pushw", param_idx);
+        assert!(parse_push(&mut span_ops, &op_no_val).is_err());
+
+        // invalid value
+        let op_val_invalid = Token::new("push.abc", param_idx);
+        assert!(parse_push(&mut span_ops, &op_val_invalid).is_err());
+
+        // extra value
+        let op_extra_val = Token::new("pushw.0.1", param_idx);
+        assert!(parse_push(&mut span_ops, &op_extra_val).is_err());
+    }
+
+    #[test]
+    fn pushw() {
+        // pushes a word of 4 immediate values in decimal or hexadecimal onto the stack
+        let mut span_ops: Vec<Operation> = Vec::new();
+        let op = Token::new("pushw.1.23.0x1C8.0", 0);
+        let expected = vec![
+            Operation::Pad,
+            Operation::Incr,
+            Operation::Push(BaseElement::new(23)),
+            Operation::Push(BaseElement::new(456)),
+            Operation::Pad,
+        ];
+        parse_pushw(&mut span_ops, &op).expect("Failed to parse pushw");
+
+        assert_eq!(span_ops, expected);
+    }
+
+    #[test]
+    fn pushw_invalid() {
+        // fails when immediate values are invalid or missing
+        let mut span_ops: Vec<Operation> = Vec::new();
+        let param_idx = 0;
+
+        // no values
+        let op_no_vals = Token::new("pushw", param_idx);
+        assert!(parse_pushw(&mut span_ops, &op_no_vals).is_err());
+
+        // insufficient values provided
+        let op_val_missing = Token::new("pushw.0.1.2", param_idx);
+        assert!(parse_pushw(&mut span_ops, &op_val_missing).is_err());
+
+        // invalid value
+        let op_val_invalid = Token::new("push.0.1.2.abc", param_idx);
+        assert!(parse_pushw(&mut span_ops, &op_val_invalid).is_err());
+
+        // extra value
+        let op_extra_val = Token::new("pushw.0.1.2.3.4", param_idx);
+        assert!(parse_pushw(&mut span_ops, &op_extra_val).is_err());
+    }
+
+    #[test]
+    fn env_sdepth() {
+        // pushes the current depth of the stack onto the top of the stack
+        let mut span_ops = vec![Operation::Push(BaseElement::ONE); 8];
+        let op = Token::new("env.sdepth", 0);
+        let mut expected = span_ops.clone();
+        expected.push(Operation::SDepth);
+
+        parse_env(&mut span_ops, &op).expect("Failed to parse env.sdepth with empty stack");
+        assert_eq!(span_ops, expected);
+    }
+
+    #[test]
+    fn env_invalid() {
+        // fails when env op variant is invalid or missing or has too many immediate values
+        let mut span_ops: Vec<Operation> = Vec::new();
+        let param_idx = 0;
+
+        // missing env var
+        let op_no_val = Token::new("env", param_idx);
+        assert!(parse_mem(&mut span_ops, &op_no_val).is_err());
+
+        // invalid env var
+        let op_val_invalid = Token::new("env.invalid", param_idx);
+        assert!(parse_push(&mut span_ops, &op_val_invalid).is_err());
+
+        // extra value
+        let op_extra_val = Token::new("env.sdepth.0", param_idx);
+        assert!(parse_push(&mut span_ops, &op_extra_val).is_err());
+    }
+
     #[test]
     fn mem_pop() {
         // stores the top 4 elements of the stack in memory
@@ -298,4 +462,23 @@ mod tests {
 
         assert_eq!(&span_ops_addr, &expected_addr);
     }
+
+    #[test]
+    fn mem_invalid() {
+        // fails when mem op variant is invalid or missing or has too many immediate values
+        let mut span_ops: Vec<Operation> = Vec::new();
+        let param_idx = 0;
+
+        // missing variant
+        let op_no_val = Token::new("mem", param_idx);
+        assert!(parse_mem(&mut span_ops, &op_no_val).is_err());
+
+        // invalid variant
+        let op_val_invalid = Token::new("mem.abc", param_idx);
+        assert!(parse_push(&mut span_ops, &op_val_invalid).is_err());
+
+        // extra value
+        let op_extra_val = Token::new("mem.push.0.1", param_idx);
+        assert!(parse_push(&mut span_ops, &op_extra_val).is_err());
+    }
 }