feat: add FunctionDef::set_return_visibility

compiler/noirc_frontend/src/hir/comptime/interpreter/builtin.rs

@@ -109,6 +109,9 @@ impl<'local, 'context> Interpreter<'local, 'context> {
             "function_def_set_return_type" => {
                 function_def_set_return_type(self, arguments, location)
             }
+            "function_def_set_return_visibility" => {
+                function_def_set_return_visibility(self, arguments, location)
+            }
             "module_functions" => module_functions(self, arguments, location),
             "module_has_named_attribute" => module_has_named_attribute(self, arguments, location),
             "module_is_contract" => module_is_contract(self, arguments, location),
@@ -1871,6 +1874,26 @@ fn function_def_set_return_type(
     Ok(Value::Unit)
 }
 
+// fn set_return_visibility(self, visibility: Quoted)
+fn function_def_set_return_visibility(
+    interpreter: &mut Interpreter,
+    arguments: Vec<(Value, Location)>,
+    location: Location,
+) -> IResult<Value> {
+    let (self_argument, visibility) = check_two_arguments(arguments, location)?;
+
+    let func_id = get_function_def(self_argument)?;
+    check_function_not_yet_resolved(interpreter, func_id, location)?;
+
+    let parser = parser::visibility();
+    let visibility = parse(visibility, parser, "a visibility")?;
+
+    let func_meta = interpreter.elaborator.interner.function_meta_mut(&func_id);
+    func_meta.return_visibility = visibility;
+
+    Ok(Value::Unit)
+}
+
 // fn functions(self) -> [FunctionDefinition]
 fn module_functions(
     interpreter: &Interpreter,

compiler/noirc_frontend/src/parser/mod.rs

@@ -26,7 +26,8 @@ use noirc_errors::Span;
 pub use parser::path::path_no_turbofish;
 pub use parser::traits::trait_bound;
 pub use parser::{
-    block, expression, fresh_statement, lvalue, parse_program, parse_type, pattern, top_level_items,
+    block, expression, fresh_statement, lvalue, parse_program, parse_type, pattern,
+    top_level_items, visibility,
 };
 
 #[derive(Debug, Clone)]

compiler/noirc_frontend/src/parser/parser.rs

@@ -28,7 +28,8 @@
 use self::types::{generic_type_args, maybe_comp_time};
 use attributes::{attributes, inner_attribute, validate_secondary_attributes};
 pub use types::parse_type;
-use visibility::visibility_modifier;
+use visibility::item_visibility;
+pub use visibility::visibility;
 
 use super::{
     foldl_with_span, labels::ParsingRuleLabel, parameter_name_recovery, parameter_recovery,
@@ -53,7 +54,7 @@
 
 use chumsky::prelude::*;
 use iter_extended::vecmap;
 use lalrpop_util::lalrpop_mod;
 use noirc_errors::{Span, Spanned};
 
 mod assertion;
@@ -68,8 +69,8 @@
 mod types;
 mod visibility;
 
 // synthesized by LALRPOP
 lalrpop_mod!(pub noir_parser);
 
 #[cfg(test)]
 mod test_helpers;
@@ -96,12 +97,12 @@
 
     if cfg!(feature = "experimental_parser") {
         for parsed_item in &parsed_module.items {
             if lalrpop_parser_supports_kind(&parsed_item.kind) {
                 match &parsed_item.kind {
                     ItemKind::Import(parsed_use_tree, _visibility) => {
                         prototype_parse_use_tree(Some(parsed_use_tree), source_program);
                     }
                     // other kinds prevented by lalrpop_parser_supports_kind
                     _ => unreachable!(),
                 }
             }
@@ -118,7 +119,7 @@
     }
 
     let mut lexer = Lexer::new(input);
     lexer = lexer.skip_whitespaces(false);
     let mut errors = Vec::new();
 
     // NOTE: this is a hack to get the references working
@@ -459,7 +460,7 @@
 }
 
 fn use_statement() -> impl NoirParser<TopLevelStatement> {
-    visibility_modifier()
+    item_visibility()
         .then_ignore(keyword(Keyword::Use))
         .then(use_tree())
         .map(|(visibility, use_tree)| TopLevelStatement::Import(use_tree, visibility))
@@ -737,15 +738,6 @@
     })
 }
 
-fn optional_visibility() -> impl NoirParser<Visibility> {
-    keyword(Keyword::Pub)
-        .map(|_| Visibility::Public)
-        .or(call_data())
-        .or(keyword(Keyword::ReturnData).map(|_| Visibility::ReturnData))
-        .or_not()
-        .map(|opt| opt.unwrap_or(Visibility::Private))
-}
-
 pub fn expression() -> impl ExprParser {
     recursive(|expr| {
         expression_with_precedence(

compiler/noirc_frontend/src/parser/parser/function.rs

@@ -1,10 +1,10 @@
 use super::{
     attributes::{attributes, validate_attributes},
-    block, fresh_statement, ident, keyword, maybe_comp_time, nothing, optional_visibility,
-    parameter_name_recovery, parameter_recovery, parenthesized, parse_type, pattern,
+    block, fresh_statement, ident, keyword, maybe_comp_time, nothing, parameter_name_recovery,
+    parameter_recovery, parenthesized, parse_type, pattern,
     primitives::token_kind,
     self_parameter,
-    visibility::visibility_modifier,
+    visibility::{item_visibility, visibility},
     where_clause, NoirParser,
 };
 use crate::token::{Keyword, Token, TokenKind};
@@ -79,13 +79,9 @@ pub(super) fn function_definition(allow_self: bool) -> impl NoirParser<NoirFunct
 ///
 /// returns (is_unconstrained, visibility) for whether each keyword was present
 fn function_modifiers() -> impl NoirParser<(bool, ItemVisibility, bool)> {
-    keyword(Keyword::Unconstrained)
-        .or_not()
-        .then(visibility_modifier())
-        .then(maybe_comp_time())
-        .map(|((unconstrained, visibility), comptime)| {
-            (unconstrained.is_some(), visibility, comptime)
-        })
+    keyword(Keyword::Unconstrained).or_not().then(item_visibility()).then(maybe_comp_time()).map(
+        |((unconstrained, visibility), comptime)| (unconstrained.is_some(), visibility, comptime),
+    )
 }
 
 pub(super) fn numeric_generic() -> impl NoirParser<UnresolvedGeneric> {
@@ -142,14 +138,12 @@ pub(super) fn generics() -> impl NoirParser<UnresolvedGenerics> {
 
 pub(super) fn function_return_type() -> impl NoirParser<(Visibility, FunctionReturnType)> {
     #[allow(deprecated)]
-    just(Token::Arrow)
-        .ignore_then(optional_visibility())
-        .then(spanned(parse_type()))
-        .or_not()
-        .map_with_span(|ret, span| match ret {
+    just(Token::Arrow).ignore_then(visibility()).then(spanned(parse_type())).or_not().map_with_span(
+        |ret, span| match ret {
             Some((visibility, (ty, _))) => (visibility, FunctionReturnType::Ty(ty)),
             None => (Visibility::Private, FunctionReturnType::Default(span)),
-        })
+        },
+    )
 }
 
 fn function_parameters<'a>(allow_self: bool) -> impl NoirParser<Vec<Param>> + 'a {
@@ -158,7 +152,7 @@ fn function_parameters<'a>(allow_self: bool) -> impl NoirParser<Vec<Param>> + 'a
     let full_parameter = pattern()
         .recover_via(parameter_name_recovery())
         .then_ignore(just(Token::Colon))
-        .then(optional_visibility())
+        .then(visibility())
         .then(typ)
         .map_with_span(|((pattern, visibility), typ), span| Param {
             visibility,

compiler/noirc_frontend/src/parser/parser/visibility.rs

@@ -4,15 +4,15 @@ use chumsky::{
 };
 
 use crate::{
-    ast::ItemVisibility,
+    ast::{ItemVisibility, Visibility},
     parser::NoirParser,
     token::{Keyword, Token},
 };
 
-use super::primitives::keyword;
+use super::{call_data, primitives::keyword};
 
 /// visibility_modifier: 'pub(crate)'? 'pub'? ''
-pub(crate) fn visibility_modifier() -> impl NoirParser<ItemVisibility> {
+pub(crate) fn item_visibility() -> impl NoirParser<ItemVisibility> {
     let is_pub_crate = (keyword(Keyword::Pub)
         .then_ignore(just(Token::LeftParen))
         .then_ignore(keyword(Keyword::Crate))
@@ -25,3 +25,12 @@ pub(crate) fn visibility_modifier() -> impl NoirParser<ItemVisibility> {
 
     choice((is_pub_crate, is_pub, is_private))
 }
+
+pub fn visibility() -> impl NoirParser<Visibility> {
+    keyword(Keyword::Pub)
+        .map(|_| Visibility::Public)
+        .or(call_data())
+        .or(keyword(Keyword::ReturnData).map(|_| Visibility::ReturnData))
+        .or_not()
+        .map(|opt| opt.unwrap_or(Visibility::Private))
+}

docs/docs/noir/standard_library/meta/function_def.md

@@ -65,3 +65,12 @@ each parameter pattern to be a syntactically valid parameter.
 Mutates the function's return type to a new type. This is only valid
 on functions in the current crate which have not yet been resolved.
 This means any functions called at compile-time are invalid targets for this method.
+
+### set_return_visibility
+
+#include_code set_return_visibility noir_stdlib/src/meta/function_def.nr rust
+
+Mutates the function's return visibility to a new one that's parsed from the given
+`Quoted` value. This is only valid on functions in the current crate which have not 
+yet been resolved. This means any functions called at compile-time are invalid targets 
+for this method.

noir_stdlib/src/meta/function_def.nr

@@ -38,4 +38,9 @@ impl FunctionDefinition {
     // docs:start:set_return_type
     fn set_return_type(self, return_type: Type) {}
     // docs:end:set_return_type
+
+    #[builtin(function_def_set_return_visibility)]
+    // docs:start:set_return_visibility
+    fn set_return_visibility(self, visibility: Quoted) {}
+    // docs:end:set_return_visibility
 }

test_programs/compile_success_empty/comptime_function_definition/src/main.nr

@@ -3,7 +3,7 @@ use std::meta::type_of;
 struct Foo { x: Field, field: Field }
 
 #[function_attr]
-fn foo(w: i32, y: Field, Foo { x, field: some_field }: Foo, mut a: bool, (b, c): (i32, i32)) -> i32 {
+pub fn foo(w: i32, y: Field, Foo { x, field: some_field }: Foo, mut a: bool, (b, c): (i32, i32)) -> i32 {
     let _ = (w, y, x, some_field, a, b, c);
     1
 }
@@ -57,3 +57,15 @@ comptime fn mutate_add_one(f: FunctionDefinition) {
 fn main() {
     assert_eq(add_one(41), 42);
 }
+
+mod some_contract {
+    // No pub on the return type is an error
+    #[super::set_pub_return_type]
+    pub fn bar() -> Field {
+        1
+    }
+}
+
+fn set_pub_return_type(f: FunctionDefinition) {
+    f.set_return_visibility(quote { what the heck });
+}