From cd796dea4937dd1a261f154e5f2e599bbc649165 Mon Sep 17 00:00:00 2001
From: Tom French <15848336+TomAFrench@users.noreply.github.com>
Date: Mon, 26 Feb 2024 15:22:01 +0000
Subject: [PATCH] fix: correct formatting for databus visibility types (#4423)

# Description

## Problem\*

Resolves <!-- Link to GitHub Issue -->

## Summary\*

This PR fixes an issue uncovered by #4422 where we're not properly
formatting databus visibility modifiers. I've fixed this and added a new
test case for regressions.

## Additional Context



## Documentation\*

Check one:
- [x] No documentation needed.
- [ ] Documentation included in this PR.
- [ ] **[Exceptional Case]** 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/utils.rs              | 14 +++++++++++++-
 tooling/nargo_fmt/src/visitor/item.rs       | 14 ++++++++++----
 tooling/nargo_fmt/tests/expected/databus.nr |  2 ++
 tooling/nargo_fmt/tests/input/databus.nr    |  2 ++
 4 files changed, 27 insertions(+), 5 deletions(-)
 create mode 100644 tooling/nargo_fmt/tests/expected/databus.nr
 create mode 100644 tooling/nargo_fmt/tests/input/databus.nr

diff --git a/tooling/nargo_fmt/src/utils.rs b/tooling/nargo_fmt/src/utils.rs
index 5874ebdebb..94969d45e8 100644
--- a/tooling/nargo_fmt/src/utils.rs
+++ b/tooling/nargo_fmt/src/utils.rs
@@ -1,3 +1,5 @@
+use std::borrow::Cow;
+
 use crate::items::HasItem;
 use crate::rewrite;
 use crate::visitor::{FmtVisitor, Shape};
@@ -143,7 +145,7 @@ impl HasItem for Param {
     fn format(self, visitor: &FmtVisitor, shape: Shape) -> String {
         let pattern = visitor.slice(self.pattern.span());
         let visibility = match self.visibility {
-            Visibility::Public => "pub ",
+            Visibility::Public => "pub",
             Visibility::Private => "",
             Visibility::DataBus => "call_data",
         };
@@ -152,6 +154,7 @@ impl HasItem for Param {
             pattern.to_string()
         } else {
             let ty = rewrite::typ(visitor, shape, self.typ);
+            let visibility = append_space_if_nonempty(visibility.into());
             format!("{pattern}: {visibility}{ty}")
         }
     }
@@ -183,6 +186,15 @@ pub(crate) fn last_line_contains_single_line_comment(s: &str) -> bool {
     s.lines().last().map_or(false, |line| line.contains("//"))
 }
 
+pub(crate) fn append_space_if_nonempty(mut string: Cow<str>) -> Cow<str> {
+    if !string.is_empty() {
+        let inner = string.to_mut();
+        inner.push(' ');
+    }
+
+    string
+}
+
 pub(crate) fn last_line_used_width(s: &str, offset: usize) -> usize {
     if s.contains('\n') {
         last_line_width(s)
diff --git a/tooling/nargo_fmt/src/visitor/item.rs b/tooling/nargo_fmt/src/visitor/item.rs
index 1825a6e05b..28aad3c551 100644
--- a/tooling/nargo_fmt/src/visitor/item.rs
+++ b/tooling/nargo_fmt/src/visitor/item.rs
@@ -7,7 +7,10 @@ use noirc_frontend::{
 
 use crate::{
     rewrite::{self, UseTree},
-    utils::{last_line_contains_single_line_comment, last_line_used_width, FindToken},
+    utils::{
+        append_space_if_nonempty, last_line_contains_single_line_comment, last_line_used_width,
+        FindToken,
+    },
     visitor::expr::{format_seq, NewlineMode},
 };
 
@@ -119,9 +122,12 @@ impl super::FmtVisitor<'_> {
                 result.push_str("distinct ");
             }
 
-            if let Visibility::Public = func.def.return_visibility {
-                result.push_str("pub ");
-            }
+            let visibility = match func.def.return_visibility {
+                Visibility::Public => "pub",
+                Visibility::DataBus => "return_data",
+                Visibility::Private => "",
+            };
+            result.push_str(&append_space_if_nonempty(visibility.into()));
 
             let typ = rewrite::typ(self, self.shape(), func.return_type());
             result.push_str(&typ);
diff --git a/tooling/nargo_fmt/tests/expected/databus.nr b/tooling/nargo_fmt/tests/expected/databus.nr
new file mode 100644
index 0000000000..60934b60b2
--- /dev/null
+++ b/tooling/nargo_fmt/tests/expected/databus.nr
@@ -0,0 +1,2 @@
+fn main(x: pub u8, y: call_data u8) -> return_data u32 {}
+
diff --git a/tooling/nargo_fmt/tests/input/databus.nr b/tooling/nargo_fmt/tests/input/databus.nr
new file mode 100644
index 0000000000..60934b60b2
--- /dev/null
+++ b/tooling/nargo_fmt/tests/input/databus.nr
@@ -0,0 +1,2 @@
+fn main(x: pub u8, y: call_data u8) -> return_data u32 {}
+