diff --git a/compiler/parser/src/parser/statement.rs b/compiler/parser/src/parser/statement.rs index 3d8f78cb4e..dbfcd75542 100644 --- a/compiler/parser/src/parser/statement.rs +++ b/compiler/parser/src/parser/statement.rs @@ -19,7 +19,7 @@ use super::*; use leo_errors::{ParserError, Result}; use leo_span::sym; -const ASSIGN_TOKENS: &[Token] = &[Token::Assign]; +const ASSIGN_TOKENS: &[Token] = &[Token::Assign, Token::AddAssign]; impl ParserContext<'_> { /// Returns a [`Statement`] AST node if the next tokens represent a statement. @@ -40,13 +40,18 @@ impl ParserContext<'_> { let place = self.parse_expression()?; if self.eat_any(ASSIGN_TOKENS) { + let operation = match &self.prev_token.token { + Token::Assign => AssignOperation::Assign, + Token::AddAssign => AssignOperation::Add, + _ => unreachable!("`parse_assign_statement` shouldn't produce this"), + }; + let value = self.parse_expression()?; self.expect(&Token::Semicolon)?; Ok(Statement::Assign(Box::new(AssignStatement { span: place.span() + value.span(), place, - // Currently only `=` so this is alright. - operation: AssignOperation::Assign, + operation, value, }))) } else { diff --git a/compiler/parser/src/tokenizer/lexer.rs b/compiler/parser/src/tokenizer/lexer.rs index bda8426f59..9e3b398d48 100644 --- a/compiler/parser/src/tokenizer/lexer.rs +++ b/compiler/parser/src/tokenizer/lexer.rs @@ -185,34 +185,41 @@ impl Token { let input_str = input; let mut input = input.chars().peekable(); - // Consumes a single character token. - let single = |input: &mut Peekable<_>, token| { + // Returns one token matching one character. + let match_one = |input: &mut Peekable<_>, token| { input.next(); Ok((1, token)) }; - // Consumes a character followed by `on` with `then` if found or `els` otherwise. - let followed_by = |input: &mut Peekable<_>, on, then, els| { + + // Returns one token matching one or two characters. + // If the `second` character matches, return the `second_token` that matches both characters. + // Otherwise, return the `first_token` that matches the single character. + let match_two = |input: &mut Peekable<_>, first_token, second_char, second_token| { input.next(); - Ok(if input.next_if_eq(&on).is_some() { - (2, then) + Ok(if input.next_if_eq(&second_char).is_some() { + (2, second_token) } else { - (1, els) + (1, first_token) }) }; - // Consumes a character followed by `on_1`, `on_2` or none. Outputs case_1, case_2, or els. - let three_cases = |input: &mut Peekable<_>, on_1, case_1, on_2, case_2, els| { + + // Returns one token matching one or two characters. + // If the `second` character matches, return the `second_token` that matches both characters. + // If the `third` character matches, return the `third_token` that matches both characters. + // Otherwise, return the `first_token` that matches the single character. + let match_three = |input: &mut Peekable<_>, first_token, second_char, second_token, third_char, third_token| { input.next(); - Ok(if input.next_if_eq(&on_1).is_some() { - (2, case_1) - } else if input.next_if_eq(&on_2).is_some() { - (2, case_2) + Ok(if input.next_if_eq(&second_char).is_some() { + (2, second_token) + } else if input.next_if_eq(&third_char).is_some() { + (2, third_token) } else { - (1, els) + (1, first_token) }) }; match *input.peek().ok_or_else(ParserError::lexer_empty_input)? { - x if x.is_ascii_whitespace() => return single(&mut input, Token::WhiteSpace), + x if x.is_ascii_whitespace() => return match_one(&mut input, Token::WhiteSpace), '"' => { // Find end string quotation mark. // Instead of checking each `char` and pushing, we can avoid reallocations. @@ -228,17 +235,17 @@ impl Token { return Ok((string.len() + 2, Token::StaticString(string))); } x if x.is_ascii_digit() => return Self::eat_integer(&mut input), - '!' => return followed_by(&mut input, '=', Token::NotEq, Token::Not), - '?' => return single(&mut input, Token::Question), - '&' => return followed_by(&mut input, '&', Token::And, Token::BitwiseAnd), - '(' => return single(&mut input, Token::LeftParen), - ')' => return single(&mut input, Token::RightParen), - '_' => return single(&mut input, Token::Underscore), - '*' => return followed_by(&mut input, '*', Token::Exp, Token::Mul), - '+' => return single(&mut input, Token::Add), - ',' => return single(&mut input, Token::Comma), - '-' => return followed_by(&mut input, '>', Token::Arrow, Token::Minus), - '.' => return followed_by(&mut input, '.', Token::DotDot, Token::Dot), + '!' => return match_two(&mut input, Token::Not, '=', Token::NotEq), + '?' => return match_one(&mut input, Token::Question), + '&' => return match_two(&mut input, Token::BitwiseAnd, '&', Token::And), + '(' => return match_one(&mut input, Token::LeftParen), + ')' => return match_one(&mut input, Token::RightParen), + '_' => return match_one(&mut input, Token::Underscore), + '*' => return match_two(&mut input, Token::Mul, '*', Token::Exp), + '+' => return match_two(&mut input, Token::Add, '=', Token::AddAssign), + ',' => return match_one(&mut input, Token::Comma), + '-' => return match_two(&mut input, Token::Minus, '>', Token::Arrow), + '.' => return match_two(&mut input, Token::Dot, '.', Token::DotDot), '/' => { input.next(); if input.next_if_eq(&'/').is_some() { @@ -276,17 +283,17 @@ impl Token { } return Ok((1, Token::Div)); } - ':' => return followed_by(&mut input, ':', Token::DoubleColon, Token::Colon), - ';' => return single(&mut input, Token::Semicolon), - '<' => return three_cases(&mut input, '=', Token::LtEq, '<', Token::Shl, Token::Lt), - '>' => return three_cases(&mut input, '=', Token::GtEq, '>', Token::Shr, Token::Gt), - '=' => return followed_by(&mut input, '=', Token::Eq, Token::Assign), - '[' => return single(&mut input, Token::LeftSquare), - ']' => return single(&mut input, Token::RightSquare), - '{' => return single(&mut input, Token::LeftCurly), - '}' => return single(&mut input, Token::RightCurly), - '|' => return followed_by(&mut input, '|', Token::Or, Token::BitwiseOr), - '^' => return single(&mut input, Token::Xor), + ':' => return match_two(&mut input, Token::Colon, ':', Token::DoubleColon), + ';' => return match_one(&mut input, Token::Semicolon), + '<' => return match_three(&mut input, Token::Lt, '=', Token::LtEq, '<', Token::Shl), + '>' => return match_three(&mut input, Token::Gt, '=', Token::GtEq, '>', Token::Shr), + '=' => return match_two(&mut input, Token::Assign, '=', Token::Eq), + '[' => return match_one(&mut input, Token::LeftSquare), + ']' => return match_one(&mut input, Token::RightSquare), + '{' => return match_one(&mut input, Token::LeftCurly), + '}' => return match_one(&mut input, Token::RightCurly), + '|' => return match_two(&mut input, Token::BitwiseOr, '|', Token::Or), + '^' => return match_one(&mut input, Token::Xor), _ => (), } if let Some(ident) = eat_identifier(&mut input) { diff --git a/compiler/parser/src/tokenizer/token.rs b/compiler/parser/src/tokenizer/token.rs index d741d10b8c..24cdab7583 100644 --- a/compiler/parser/src/tokenizer/token.rs +++ b/compiler/parser/src/tokenizer/token.rs @@ -47,6 +47,7 @@ pub enum Token { Gt, GtEq, Add, + AddAssign, Minus, Mul, Div, @@ -230,6 +231,7 @@ impl fmt::Display for Token { Gt => write!(f, ">"), GtEq => write!(f, ">="), Add => write!(f, "+"), + AddAssign => write!(f, "+="), Minus => write!(f, "-"), Mul => write!(f, "*"), Div => write!(f, "/"), diff --git a/tests/compiler/statements/inputs/u8.in b/tests/compiler/statements/inputs/u8.in index 1c4d9fede9..9cf7b5b0af 100644 --- a/tests/compiler/statements/inputs/u8.in +++ b/tests/compiler/statements/inputs/u8.in @@ -1,5 +1,2 @@ [main] -a: u8 = 1u8; - -[registers] -r0: bool = true; \ No newline at end of file +a: u8 = 1u8; \ No newline at end of file diff --git a/tests/compiler/statements/operations/add_assign.leo b/tests/compiler/statements/operations/add_assign.leo new file mode 100644 index 0000000000..f4d1c98bcc --- /dev/null +++ b/tests/compiler/statements/operations/add_assign.leo @@ -0,0 +1,12 @@ +/* +namespace: Compile +expectation: Pass +input_files: ../inputs/u8.in +*/ + +function main(a: u8) -> u8 { + let b: u8 = 1u8; + b += a; + + return b; +} \ No newline at end of file