diff --git a/src/parser/tests.rs b/src/parser/tests.rs index 7ebf5ac3..50c84efa 100644 --- a/src/parser/tests.rs +++ b/src/parser/tests.rs @@ -563,7 +563,7 @@ pub fn test_parse_invalid_scripts() { pub fn test_whamm_with_asserts() { setup_logger(); let script = r#" -my_func() -> i32{ +my_func() -> i32 { return 5; } wasm::call:alt / diff --git a/src/parser/whamm_parser.rs b/src/parser/whamm_parser.rs index 8f3645e1..61baccaf 100644 --- a/src/parser/whamm_parser.rs +++ b/src/parser/whamm_parser.rs @@ -220,38 +220,23 @@ pub fn handle_fn_def(whamm: &mut Whamm, script_count: usize, pair: Pair, e }; // Get the parameters - let mut next = pairs.next().unwrap(); - let params = if matches!(next.as_rule(), Rule::param) { - let params = handle_params(next, err); - next = pairs.next().unwrap(); - params + let mut next = pairs.next(); + let params = if let Some(n) = &next { + if matches!(n.as_rule(), Rule::param) { + let params = handle_params(n.clone(), err); + next = pairs.next(); + params + } else { + vec![] + } } else { vec![] }; - // Get the function body - let body = if matches!(next.as_rule(), Rule::block) { - let loc = LineColLocation::from(next.as_span()); - let mut pairs = next.into_inner(); - - pairs.next(); // skip over block start - handle_body(&mut pairs, loc, err) - } else { - // If didn't match, create empty body - Block { - stmts: vec![], - return_ty: None, - loc: Some(Location { - line_col, - path: None - }), - } - }; - // Get the return type - let next = pairs.next(); - let return_ty = match next { + let return_ty = match next.clone() { Some(pair) => { + next = pairs.next(); type_from_rule(pair, err) }, None => { @@ -259,6 +244,22 @@ pub fn handle_fn_def(whamm: &mut Whamm, script_count: usize, pair: Pair, e } }; + // Get the function body + let body = if let Some(n) = next { + if matches!(n.as_rule(), Rule::block) { + let loc = LineColLocation::from(n.as_span()); + let mut pairs = n.into_inner(); + handle_body(&mut pairs, loc, err) + } else { + // If didn't match, create empty body + Block::default() + } + } else { + // If didn't match, create empty body + Block::default() + }; + + // Add the new function to the current script let script: &mut Script = whamm.scripts.get_mut(script_count).unwrap(); script.fns.push(types::Fn {