From 598230d9427cf988fc6da8fe9e1eb2b7c00a2fa6 Mon Sep 17 00:00:00 2001
From: Ary Borenszweig <asterite@gmail.com>
Date: Thu, 24 Oct 2024 13:24:43 -0300
Subject: [PATCH] fix: formatter didn't format `>>=` well (#6337)

# Description

## Problem

Resolves
https://github.com/noir-lang/noir_json_parser/pull/16#issuecomment-2435374508

## Summary

I thought for "op assign" like `x += 1` there will always come one or
two tokens, then `=`, but that's not the case for `>>=` (it's `> >=`).

## Additional Context

## Documentation

Check one:
- [x] No documentation needed.
- [ ] Documentation included in this PR.
- [ ] **[For Experimental Features]** Documentation to be submitted in a
separate PR.

# PR Checklist

- [x] I have tested the changes locally.
- [x] I have formatted the changes with [Prettier](https://prettier.io/)
and/or `cargo fmt` on default settings.
---
 tooling/nargo_fmt/src/formatter/statement.rs | 25 +++++++++++++++-----
 1 file changed, 19 insertions(+), 6 deletions(-)

diff --git a/tooling/nargo_fmt/src/formatter/statement.rs b/tooling/nargo_fmt/src/formatter/statement.rs
index 036066e346f..28107294909 100644
--- a/tooling/nargo_fmt/src/formatter/statement.rs
+++ b/tooling/nargo_fmt/src/formatter/statement.rs
@@ -182,12 +182,15 @@ impl<'a, 'b> ChunkFormatter<'a, 'b> {
             if formatter.is_at(Token::Assign) {
                 formatter.write_token(Token::Assign);
             } else {
-                while formatter.token != Token::Assign {
-                    formatter.write_current_token();
-                    formatter.bump();
-                    formatter.skip_comments_and_whitespace();
-                }
-                formatter.write_token(Token::Assign);
+                // This is something like `x += 1`, which is parsed as an
+                // Assign with an InfixExpression as its right-hand side: `x = x + 1`.
+                // There will always be two tokens here, like `+ =` or `> >=`.
+                formatter.write_current_token();
+                formatter.bump();
+                formatter.skip_comments_and_whitespace();
+                formatter.write_current_token();
+                formatter.bump();
+
                 is_op_assign = true;
             }
             formatter.write_space();
@@ -435,6 +438,16 @@ mod tests {
         assert_format(src, expected);
     }
 
+    #[test]
+    fn format_shift_right_assign() {
+        let src = " fn foo() { x  >>=  2 ; } ";
+        let expected = "fn foo() {
+    x >>= 2;
+}
+";
+        assert_format(src, expected);
+    }
+
     #[test]
     fn format_comptime_let_statement() {
         let src = " fn foo() { comptime  let  x  =  1 ; } ";