From 7f1d7968368734e02b152e2e907dc7af9e1604c8 Mon Sep 17 00:00:00 2001
From: guipublic <47281315+guipublic@users.noreply.github.com>
Date: Tue, 5 Dec 2023 19:18:22 +0100
Subject: [PATCH] fix: pub is required on return for entry points (#3616)

# Description

## Problem\*

Resolves #3344 and #3632

## Summary\*



## 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.
---
 .../noirc_frontend/src/hir/resolution/resolver.rs     | 11 ++++-------
 .../simple_contract/src/main.nr                       |  5 +++++
 2 files changed, 9 insertions(+), 7 deletions(-)

diff --git a/compiler/noirc_frontend/src/hir/resolution/resolver.rs b/compiler/noirc_frontend/src/hir/resolution/resolver.rs
index 4e7ed7e2ea9..0967bc097fb 100644
--- a/compiler/noirc_frontend/src/hir/resolution/resolver.rs
+++ b/compiler/noirc_frontend/src/hir/resolution/resolver.rs
@@ -791,8 +791,8 @@ impl<'a> Resolver<'a> {
             });
         }
 
-        // 'pub_allowed' also implies 'pub' is required on return types
-        if self.pub_allowed(func)
+        // 'pub' is required on return types for entry point functions
+        if self.is_entry_point_function(func)
             && return_type.as_ref() != &Type::Unit
             && func.def.return_visibility == Visibility::Private
         {
@@ -847,12 +847,9 @@ impl<'a> Resolver<'a> {
     }
 
     /// True if the 'pub' keyword is allowed on parameters in this function
+    /// 'pub' on function parameters is only allowed for entry point functions
     fn pub_allowed(&self, func: &NoirFunction) -> bool {
-        if self.in_contract {
-            !func.def.is_unconstrained
-        } else {
-            func.name() == MAIN_FUNCTION
-        }
+        self.is_entry_point_function(func)
     }
 
     fn is_entry_point_function(&self, func: &NoirFunction) -> bool {
diff --git a/test_programs/compile_success_contract/simple_contract/src/main.nr b/test_programs/compile_success_contract/simple_contract/src/main.nr
index 88edd4ac2c3..fea0ae08c22 100644
--- a/test_programs/compile_success_contract/simple_contract/src/main.nr
+++ b/test_programs/compile_success_contract/simple_contract/src/main.nr
@@ -11,4 +11,9 @@ contract Foo {
     open internal fn skibbidy(x: Field) -> pub Field {
         x * 5
     }
+    // Regression for issue #3344
+    #[contract_library_method]
+    fn foo(x : u8) -> u8 {
+        x
+    }
 }