[move-compiler] Treat entry functions as public in test_only and …

…test contexts (MystenLabs#14416)

## Description 

- entry functions are callable outside of the module, but only from the
adapter level
- Allowing entry functions to be called from tests should improve
usability of entry functions in unit tests, without the need for public
wrappers

## Test Plan 

- Added Tests

---
If your changes are not user-facing and not a breaking change, you can
skip the following section. Otherwise, please indicate what changed, and
then add to the Release Notes section as highlighted during the release
process.

### Type of Change (Check all that apply)

- [ ] protocol change
- [X] user-visible impact
- [ ] breaking change for a client SDKs
- [ ] breaking change for FNs (FN binary must upgrade)
- [ ] breaking change for validators or node operators (must upgrade
binaries)
- [ ] breaking change for on-chain data layout
- [ ] necessitate either a data wipe or data migration

### Release notes

In this release, the Move compiler now allows `entry` functions to be
called within unit-test environments, `#[test]` and `#[test_only]`
functions, without the need of an additional `#[test_only] public`
wrapper function.

external-crates/move/crates/move-compiler/src/shared/program_info.rs

@@ -7,7 +7,7 @@ use move_ir_types::location::Loc;
 use move_symbol_pool::Symbol;
 
 use crate::{
-    expansion::ast::{AbilitySet, ModuleIdent, Visibility},
+    expansion::ast::{AbilitySet, Attributes, ModuleIdent, Visibility},
     naming::ast::{
         self as N, FunctionSignature, ResolvedUseFuns, StructDefinition, StructTypeParameter, Type,
     },
@@ -20,19 +20,23 @@ use crate::{
 
 #[derive(Debug, Clone)]
 pub struct FunctionInfo {
+    pub attributes: Attributes,
     pub defined_loc: Loc,
     pub visibility: Visibility,
+    pub entry: Option<Loc>,
     pub signature: FunctionSignature,
 }
 
 #[derive(Debug, Clone)]
 pub struct ConstantInfo {
+    pub attributes: Attributes,
     pub defined_loc: Loc,
     pub signature: Type,
 }
 
 #[derive(Debug, Clone)]
 pub struct ModuleInfo {
+    pub attributes: Attributes,
     pub package: Option<Symbol>,
     pub use_funs: ResolvedUseFuns,
     pub friends: UniqueMap<ModuleIdent, Loc>,
@@ -54,11 +58,14 @@ macro_rules! program_info {
         let mut modules = UniqueMap::maybe_from_iter(all_modules.map(|(mident, mdef)| {
             let structs = mdef.structs.clone();
             let functions = mdef.functions.ref_map(|fname, fdef| FunctionInfo {
+                attributes: fdef.attributes.clone(),
                 defined_loc: fname.loc(),
                 visibility: fdef.visibility.clone(),
+                entry: fdef.entry,
                 signature: fdef.signature.clone(),
             });
             let constants = mdef.constants.ref_map(|cname, cdef| ConstantInfo {
+                attributes: cdef.attributes.clone(),
                 defined_loc: cname.loc(),
                 signature: cdef.signature.clone(),
             });
@@ -67,6 +74,7 @@ macro_rules! program_info {
                 .map(|module_use_funs| module_use_funs.remove(&mident).unwrap())
                 .unwrap_or_default();
             let minfo = ModuleInfo {
+                attributes: mdef.attributes.clone(),
                 package: mdef.package_name,
                 use_funs,
                 friends: mdef.friends.ref_map(|_, friend| friend.loc),

external-crates/move/crates/move-compiler/src/to_bytecode/context.rs

@@ -21,23 +21,34 @@ use IR::Ability;
 /// Contains all of the dependencies actually used in the module
 pub struct Context<'a> {
     pub env: &'a mut CompilationEnv,
+    current_package: Option<Symbol>,
     current_module: Option<&'a ModuleIdent>,
     seen_structs: BTreeSet<(ModuleIdent, StructName)>,
     seen_functions: BTreeSet<(ModuleIdent, FunctionName)>,
     spec_info: BTreeMap<SpecId, (IR::NopLabel, BTreeMap<Var, H::SingleType>)>,
 }
 
 impl<'a> Context<'a> {
-    pub fn new(env: &'a mut CompilationEnv, current_module: Option<&'a ModuleIdent>) -> Self {
+    pub fn new(
+        env: &'a mut CompilationEnv,
+        current_package: Option<Symbol>,
+        current_module: Option<&'a ModuleIdent>,
+    ) -> Self {
         Self {
             env,
+            current_package,
             current_module,
             seen_structs: BTreeSet::new(),
             seen_functions: BTreeSet::new(),
             spec_info: BTreeMap::new(),
         }
     }
 
+    #[allow(unused)]
+    pub fn current_package(&self) -> Option<Symbol> {
+        self.current_package
+    }
+
     pub fn current_module(&self) -> Option<&'a ModuleIdent> {
         self.current_module
     }

external-crates/move/crates/move-compiler/src/to_bytecode/translate.rs

@@ -84,7 +84,7 @@ fn extract_decls(
             })
         })
         .collect();
-    let context = &mut Context::new(compilation_env, None);
+    let context = &mut Context::new(compilation_env, None, None);
     let fdecls = all_modules()
         .flat_map(|(m, mdef)| {
             mdef.functions
@@ -177,10 +177,9 @@ fn module(
         (BTreeSet<(ModuleIdent, StructName)>, IR::FunctionSignature),
     >,
 ) -> Option<AnnotatedCompiledUnit> {
-    let mut context = Context::new(compilation_env, Some(&ident));
     let G::ModuleDefinition {
         warning_filter: _warning_filter,
-        package_name: _package_name,
+        package_name,
         attributes: _attributes,
         is_source_module: _is_source_module,
         dependency_order: _dependency_order,
@@ -189,6 +188,7 @@ fn module(
         constants: gconstants,
         functions: gfunctions,
     } = mdef;
+    let mut context = Context::new(compilation_env, package_name, Some(&ident));
     let structs = struct_defs(&mut context, &ident, gstructs);
     let constants = constants(&mut context, Some(&ident), gconstants);
     let (collected_function_infos, functions) = functions(&mut context, Some(&ident), gfunctions);
@@ -281,7 +281,7 @@ fn script(
     >,
 ) -> Option<AnnotatedCompiledUnit> {
     let loc = name.loc();
-    let mut context = Context::new(compilation_env, None);
+    let mut context = Context::new(compilation_env, package_name, None);
 
     let constants = constants(&mut context, None, gconstants);
 
@@ -619,7 +619,7 @@ fn function(
         signature,
         body,
     } = fdef;
-    let v = visibility(v);
+    let v = visibility(context, v);
     let parameters = signature.parameters.clone();
     let signature = function_signature(context, signature);
     let body = match body.value {
@@ -657,7 +657,7 @@ fn function(
     )
 }
 
-fn visibility(v: Visibility) -> IR::FunctionVisibility {
+fn visibility(_context: &mut Context, v: Visibility) -> IR::FunctionVisibility {
     match v {
         Visibility::Public(_) => IR::FunctionVisibility::Public,
         Visibility::Friend(_) => IR::FunctionVisibility::Friend,

external-crates/move/crates/move-compiler/src/typing/core.rs

@@ -5,13 +5,20 @@
 use crate::{
     debug_display, diag,
     diagnostics::{codes::NameResolution, Diagnostic},
-    expansion::ast::{AbilitySet, ModuleIdent, ModuleIdent_, Visibility},
+    expansion::ast::{AbilitySet, AttributeName_, ModuleIdent, ModuleIdent_, Visibility},
     naming::ast::{
         self as N, BuiltinTypeName_, ResolvedUseFuns, StructDefinition, StructTypeParameter,
         TParam, TParamID, TVar, Type, TypeName, TypeName_, Type_, UseFunKind, Var,
     },
-    parser::ast::{Ability_, ConstantName, Field, FunctionName, Mutability, StructName},
-    shared::{program_info::*, unique_map::UniqueMap, *},
+    parser::ast::{
+        Ability_, ConstantName, Field, FunctionName, Mutability, StructName, ENTRY_MODIFIER,
+    },
+    shared::{
+        known_attributes::{KnownAttribute, TestingAttribute},
+        program_info::*,
+        unique_map::UniqueMap,
+        *,
+    },
     FullyCompiledProgram,
 };
 use move_ir_types::location::*;
@@ -236,6 +243,7 @@ impl<'env> Context<'env> {
     pub fn bind_script_constants(&mut self, constants: &UniqueMap<ConstantName, N::Constant>) {
         assert!(self.current_script_constants.is_none());
         self.current_script_constants = Some(constants.ref_map(|cname, cdef| ConstantInfo {
+            attributes: cdef.attributes.clone(),
             defined_loc: cname.loc(),
             signature: cdef.signature.clone(),
         }));
@@ -367,6 +375,24 @@ impl<'env> Context<'env> {
         }
     }
 
+    /// current_module.is_test_only || current_function.is_test_only || current_function.is_test
+    fn is_testing_context(&self) -> bool {
+        // TODO should we store this in the context?
+        let test_only = AttributeName_::Known(KnownAttribute::Testing(TestingAttribute::TestOnly));
+        let test = AttributeName_::Known(KnownAttribute::Testing(TestingAttribute::Test));
+
+        self.current_module.as_ref().is_some_and(|m| {
+            let minfo = self.module_info(m);
+            let is_test_only = minfo.attributes.contains_key_(&test_only);
+            is_test_only
+                || self.current_function.as_ref().is_some_and(|f| {
+                    let finfo = minfo.functions.get(f).unwrap();
+                    finfo.attributes.contains_key_(&test_only)
+                        || finfo.attributes.contains_key_(&test)
+                })
+        })
+    }
+
     fn module_info(&self, m: &ModuleIdent) -> &ModuleInfo {
         self.modules.module(m)
     }
@@ -794,6 +820,7 @@ pub fn make_constant_type(
     let in_current_module = m == &context.current_module;
     let (defined_loc, signature) = {
         let ConstantInfo {
+            attributes: _,
             defined_loc,
             signature,
         } = context.constant_info(m, c);
@@ -968,7 +995,11 @@ pub fn make_function_type(
     let return_ty = subst_tparams(tparam_subst, finfo.signature.return_type.clone());
 
     let defined_loc = finfo.defined_loc;
+    let public_for_testing =
+        public_testing_visibility(context.env, context.current_package, f, finfo.entry);
+    let is_testing_context = context.is_testing_context();
     match finfo.visibility {
+        _ if is_testing_context && public_for_testing.is_some() => (),
         Visibility::Internal if in_current_module => (),
         Visibility::Internal => {
             let internal_msg = format!(
@@ -978,11 +1009,12 @@ pub fn make_function_type(
                 Visibility::FRIEND,
                 Visibility::PACKAGE
             );
-            context.env.add_diag(diag!(
-                TypeSafety::Visibility,
+            visibility_error(
+                context,
+                public_for_testing,
                 (loc, format!("Invalid call to '{}::{}'", m, f)),
                 (defined_loc, internal_msg),
-            ));
+            );
         }
         Visibility::Package(loc)
             if in_current_module || context.current_module_shares_package_and_address(m) =>
@@ -1010,29 +1042,87 @@ pub fn make_function_type(
                     .map(|pkg_name| format!("{}", pkg_name))
                     .unwrap_or("<unknown package>".to_string())
             );
-            context.env.add_diag(diag!(
-                TypeSafety::Visibility,
+            visibility_error(
+                context,
+                public_for_testing,
                 (loc, format!("Invalid call to '{}::{}'", m, f)),
                 (vis_loc, internal_msg),
-            ));
+            );
         }
         Visibility::Friend(_) if in_current_module || context.current_module_is_a_friend_of(m) => {}
         Visibility::Friend(vis_loc) => {
             let internal_msg = format!(
                 "This function can only be called from a 'friend' of module '{}'",
                 m
             );
-            context.env.add_diag(diag!(
-                TypeSafety::Visibility,
+            visibility_error(
+                context,
+                public_for_testing,
                 (loc, format!("Invalid call to '{}::{}'", m, f)),
                 (vis_loc, internal_msg),
-            ));
+            );
         }
         Visibility::Public(_) => (),
     };
     (defined_loc, ty_args, params, return_ty)
 }
 
+pub enum PublicForTesting {
+    /// The function is entry, so it can be called in unit tests
+    Entry(Loc),
+    // TODO we should allow calling init in unit tests, but this would need Sui bytecode verifier
+    // support. Or we would need to name dodge init in unit tests
+    // SuiInit(Loc),
+}
+
+pub fn public_testing_visibility(
+    env: &CompilationEnv,
+    _package: Option<Symbol>,
+    _callee_name: &FunctionName,
+    callee_entry: Option<Loc>,
+) -> Option<PublicForTesting> {
+    // is_testing && (is_entry || is_sui_init)
+    if !env.flags().is_testing() {
+        return None;
+    }
+
+    // TODO support sui init functions
+    // let flavor = env.package_config(package).flavor;
+    // flavor == Flavor::Sui && callee_name.value() == INIT_FUNCTION_NAME
+    callee_entry.map(PublicForTesting::Entry)
+}
+
+fn visibility_error(
+    context: &mut Context,
+    public_for_testing: Option<PublicForTesting>,
+    (call_loc, call_msg): (Loc, impl ToString),
+    (vis_loc, vis_msg): (Loc, impl ToString),
+) {
+    let mut diag = diag!(
+        TypeSafety::Visibility,
+        (call_loc, call_msg),
+        (vis_loc, vis_msg),
+    );
+    if context.env.flags().is_testing() {
+        if let Some(case) = public_for_testing {
+            let (test_loc, test_msg) = match case {
+                PublicForTesting::Entry(entry_loc) => {
+                    let entry_msg = format!(
+                        "'{}' functions can be called in tests, \
+                    but only from testing contexts, e.g. '#[{}]' or '#[{}]'",
+                        ENTRY_MODIFIER,
+                        TestingAttribute::TEST,
+                        TestingAttribute::TEST_ONLY,
+                    );
+                    (entry_loc, entry_msg)
+                }
+            };
+            diag.add_secondary_label((test_loc, test_msg))
+        }
+    }
+    context.env.add_diag(diag)
+}
+
 //**************************************************************************************************
 // Constraints
 //**************************************************************************************************

external-crates/move/crates/move-compiler/src/typing/translate.rs

@@ -23,7 +23,11 @@ use crate::{
         *,
     },
     sui_mode,
-    typing::{ast as T, dependency_ordering},
+    typing::{
+        ast as T,
+        core::{public_testing_visibility, PublicForTesting},
+        dependency_ordering,
+    },
     FullyCompiledProgram,
 };
 use move_ir_types::location::*;
@@ -239,6 +243,11 @@ fn function(
     context.reset_for_module_item();
     context.current_function = Some(name);
     process_attributes(context, &attributes);
+    let visibility =
+        match public_testing_visibility(context.env, context.current_package, &name, entry) {
+            Some(PublicForTesting::Entry(loc)) => Visibility::Public(loc),
+            None => visibility,
+        };
     function_signature(context, &signature);
     if is_script {
         let mk_msg = || {

external-crates/move/crates/move-compiler/tests/move_check/unit_test/entry_is_public_in_test_contexts.move

@@ -0,0 +1,21 @@
+// entry functions are public if called from test or test_only
+
+module a::m {
+    entry fun internal(_ :u64) {}
+}
+
+#[test_only]
+module a::test_only {
+    fun example() {
+        // force a type error to make sure visibility is allowed
+        a::m::internal(0u8)
+    }
+}
+
+module a::tests {
+    #[test]
+    fun test() {
+        // force a type error to make sure visibility is allowed
+        a::m::internal(0u8)
+    }
+}