[Linter] Public mut tx context (#16878)

## Description 

This linter checks public functions to ensure they use `&mut TxContext`
instead of `&TxContext `for better upgradability. Here's a step-by-step
breakdown of how to implement this linter:

Implement the main logic :

`check_function_parameters`: Iterate through function parameters.
`is_immutable_tx_context`: Check if a parameter is an immutable
`TxContext`.
`is_tx_context_type`: Check if a type is `TxContext`.
`is_sui_tx_context`: Verify if the type is from the `sui::tx_context`
module.

## Test plan
Added more use case including false positive, false negative case

## Release notes

- [ ] Protocol: 
- [ ] Nodes (Validators and Full nodes): 
- [ ] Indexer: 
- [ ] JSON-RPC: 
- [ ] GraphQL: 
- [X] CLI: Move will now lint against using `&TxContext` instead of
`&mut TxContext` in public functions
- [ ] Rust SDK:

---------

Co-authored-by: jamedzung <[email protected]>
Co-authored-by: Todd Nowacki <[email protected]>

external-crates/move/crates/move-compiler/src/sui_mode/linters/mod.rs

@@ -20,6 +20,7 @@ pub mod custom_state_change;
 pub mod freeze_wrapped;
 pub mod freezing_capability;
 pub mod missing_key;
+pub mod public_mut_tx_context;
 pub mod public_random;
 pub mod self_transfer;
 pub mod share_owned;
@@ -72,6 +73,7 @@ pub const COLLECTION_EQUALITY_FILTER_NAME: &str = "collection_equality";
 pub const PUBLIC_RANDOM_FILTER_NAME: &str = "public_random";
 pub const MISSING_KEY_FILTER_NAME: &str = "missing_key";
 pub const FREEZING_CAPABILITY_FILTER_NAME: &str = "freezing_capability";
+pub const PREFER_MUTABLE_TX_CONTEXT_FILTER_NAME: &str = "prefer_mut_tx_context";
 
 pub const RANDOM_MOD_NAME: &str = "random";
 pub const RANDOM_STRUCT_NAME: &str = "Random";
@@ -90,6 +92,7 @@ pub enum LinterDiagnosticCode {
     PublicRandom,
     MissingKey,
     FreezingCapability,
+    PreferMutableTxContext,
 }
 
 pub fn known_filters() -> (Option<Symbol>, Vec<WarningFilter>) {
@@ -149,6 +152,12 @@ pub fn known_filters() -> (Option<Symbol>, Vec<WarningFilter>) {
             LinterDiagnosticCode::FreezingCapability as u8,
             Some(FREEZING_CAPABILITY_FILTER_NAME),
         ),
+        WarningFilter::code(
+            Some(LINT_WARNING_PREFIX),
+            LinterDiagnosticCategory::Sui as u8,
+            LinterDiagnosticCode::PreferMutableTxContext as u8,
+            Some(PREFER_MUTABLE_TX_CONTEXT_FILTER_NAME),
+        ),
     ];
 
     (Some(ALLOW_ATTR_CATEGORY.into()), filters)
@@ -169,7 +178,10 @@ pub fn linter_visitors(level: LintLevel) -> Vec<Visitor> {
         ],
         LintLevel::All => {
             let mut visitors = linter_visitors(LintLevel::Default);
-            visitors.extend([freezing_capability::WarnFreezeCapability.visitor()]);
+            visitors.extend([
+                freezing_capability::WarnFreezeCapability.visitor(),
+                public_mut_tx_context::PreferMutableTxContext.visitor(),
+            ]);
             visitors
         }
     }

external-crates/move/crates/move-compiler/src/sui_mode/linters/public_mut_tx_context.rs

@@ -0,0 +1,97 @@
+// Copyright (c) The Move Contributors
+// SPDX-License-Identifier: Apache-2.0
+
+//! Enforces that public functions use `&mut TxContext` instead of `&TxContext` to ensure upgradability.
+//! Detects and reports instances where a non-mutable reference to `TxContext` is used in public function signatures.
+//! Promotes best practices for future-proofing smart contract code by allowing mutation of the transaction context.
+use super::{LinterDiagnosticCategory, LinterDiagnosticCode, LINT_WARNING_PREFIX};
+
+use crate::{
+    diag,
+    diagnostics::{
+        codes::{custom, DiagnosticInfo, Severity},
+        WarningFilters,
+    },
+    expansion::ast::{ModuleIdent, Visibility},
+    naming::ast::Type_,
+    parser::ast::FunctionName,
+    shared::CompilationEnv,
+    sui_mode::{SUI_ADDR_NAME, TX_CONTEXT_MODULE_NAME, TX_CONTEXT_TYPE_NAME},
+    typing::{
+        ast as T,
+        visitor::{TypingVisitorConstructor, TypingVisitorContext},
+    },
+};
+use move_ir_types::location::Loc;
+
+const REQUIRE_MUTABLE_TX_CONTEXT_DIAG: DiagnosticInfo = custom(
+    LINT_WARNING_PREFIX,
+    Severity::Warning,
+    LinterDiagnosticCategory::Sui as u8,
+    LinterDiagnosticCode::PreferMutableTxContext as u8,
+    "prefer '&mut TxContext' over '&TxContext'",
+);
+
+pub struct PreferMutableTxContext;
+
+pub struct Context<'a> {
+    env: &'a mut CompilationEnv,
+}
+
+impl TypingVisitorConstructor for PreferMutableTxContext {
+    type Context<'a> = Context<'a>;
+
+    fn context<'a>(env: &'a mut CompilationEnv, _program: &T::Program) -> Self::Context<'a> {
+        Context { env }
+    }
+}
+
+impl TypingVisitorContext for Context<'_> {
+    fn add_warning_filter_scope(&mut self, filter: WarningFilters) {
+        self.env.add_warning_filter_scope(filter)
+    }
+    fn pop_warning_filter_scope(&mut self) {
+        self.env.pop_warning_filter_scope()
+    }
+
+    fn visit_module_custom(&mut self, ident: ModuleIdent, _mdef: &mut T::ModuleDefinition) -> bool {
+        // skip if in 'sui::tx_context'
+        ident.value.is(SUI_ADDR_NAME, TX_CONTEXT_MODULE_NAME)
+    }
+
+    fn visit_function_custom(
+        &mut self,
+        _module: ModuleIdent,
+        _function_name: FunctionName,
+        fdef: &mut T::Function,
+    ) -> bool {
+        if !matches!(&fdef.visibility, Visibility::Public(_)) {
+            return false;
+        }
+
+        for (_, _, sp!(loc, param_ty_)) in &fdef.signature.parameters {
+            if matches!(
+                param_ty_,
+                Type_::Ref(false, t) if t.value.is(SUI_ADDR_NAME, TX_CONTEXT_MODULE_NAME, TX_CONTEXT_TYPE_NAME),
+            ) {
+                report_non_mutable_tx_context(self.env, *loc);
+            }
+        }
+
+        false
+    }
+}
+
+fn report_non_mutable_tx_context(env: &mut CompilationEnv, loc: Loc) {
+    let msg = format!(
+        "'public' functions should prefer '&mut {0}' over '&{0}' for better upgradability.",
+        TX_CONTEXT_TYPE_NAME
+    );
+    let mut diag = diag!(REQUIRE_MUTABLE_TX_CONTEXT_DIAG, (loc, msg));
+    diag.add_note(
+        "When upgrading, the public function cannot be modified to take '&mut TxContext' instead \
+         of '&TxContext'. As such, it is recommended to consider using '&mut TxContext' to \
+         future-proof the function.",
+    );
+    env.add_diag(diag);
+}

external-crates/move/crates/move-compiler/tests/sui_mode/linter/custom_state_change.move

@@ -11,7 +11,7 @@ module a::test {
         id: UID
     }
 
-    #[allow(lint(self_transfer))]
+    #[allow(lint(self_transfer, prefer_mut_tx_context))]
     public fun custom_transfer_bad(o: S1, ctx: &TxContext) {
         transfer::transfer(o, tx_context::sender(ctx))
     }

external-crates/move/crates/move-compiler/tests/sui_mode/linter/suppress_public_mut_tx_context.move

@@ -0,0 +1,20 @@
+module 0x42::suppress_cases {
+    use sui::tx_context::TxContext;
+
+    #[allow(lint(prefer_mut_tx_context))]
+    public fun suppressed_function(_ctx: &TxContext) {
+    }
+
+    #[allow(lint(prefer_mut_tx_context))]
+    public fun multi_suppressed_function(_ctx: &TxContext) {
+    }
+
+    #[allow(lint(prefer_mut_tx_context))]
+    public fun suppressed_multi_param(_a: u64, _ctx: &TxContext, _b: &mut TxContext) {
+    }
+}
+
+// Mocking the sui::tx_context module
+module sui::tx_context {
+    struct TxContext has drop {}
+}

external-crates/move/crates/move-compiler/tests/sui_mode/linter/true_negative_public_mut_tx_context.move

@@ -0,0 +1,26 @@
+// tests the lint for preferring &mut TxContext over &TxContext in public functions
+// these cases correctly should not trigger the lint
+module 0x42::true_negative {
+    use sui::tx_context::TxContext;
+
+    public fun correct_mint(_ctx: &mut TxContext) {
+    }
+
+    public fun another_correct(_a: u64, _b: &mut TxContext, _c: u64) {
+    }
+
+    fun private_function(_ctx: &TxContext) {
+    }
+
+    public fun custom_module(_b: &mut sui::mock_tx_context::TxContext) {}
+
+
+}
+
+module sui::tx_context {
+    struct TxContext has drop {}
+}
+
+module sui::mock_tx_context {
+    struct TxContext has drop {}
+}

external-crates/move/crates/move-compiler/tests/sui_mode/linter/true_positive_public_mut_tx_context.exp

@@ -0,0 +1,36 @@
+warning[Lint W99009]: prefer '&mut TxContext' over '&TxContext'
+  ┌─ tests/sui_mode/linter/true_positive_public_mut_tx_context.move:8:37
+  │
+8 │     public fun incorrect_mint(_ctx: &TxContext) {
+  │                                     ^^^^^^^^^^ 'public' functions should prefer '&mut TxContext' over '&TxContext' for better upgradability.
+  │
+  = When upgrading, the public function cannot be modified to take '&mut TxContext' instead of '&TxContext'. As such, it is recommended to consider using '&mut TxContext' to future-proof the function.
+  = This warning can be suppressed with '#[allow(lint(prefer_mut_tx_context))]' applied to the 'module' or module member ('const', 'fun', or 'struct')
+
+warning[Lint W99009]: prefer '&mut TxContext' over '&TxContext'
+   ┌─ tests/sui_mode/linter/true_positive_public_mut_tx_context.move:11:47
+   │
+11 │     public fun another_incorrect(_a: u64, _b: &TxContext, _c: u64) {
+   │                                               ^^^^^^^^^^ 'public' functions should prefer '&mut TxContext' over '&TxContext' for better upgradability.
+   │
+   = When upgrading, the public function cannot be modified to take '&mut TxContext' instead of '&TxContext'. As such, it is recommended to consider using '&mut TxContext' to future-proof the function.
+   = This warning can be suppressed with '#[allow(lint(prefer_mut_tx_context))]' applied to the 'module' or module member ('const', 'fun', or 'struct')
+
+warning[Lint W99009]: prefer '&mut TxContext' over '&TxContext'
+   ┌─ tests/sui_mode/linter/true_positive_public_mut_tx_context.move:14:54
+   │
+14 │     public fun mixed_function(_a: &CustomStruct, _b: &TxContext, _c: &mut TxContext) {}
+   │                                                      ^^^^^^^^^^ 'public' functions should prefer '&mut TxContext' over '&TxContext' for better upgradability.
+   │
+   = When upgrading, the public function cannot be modified to take '&mut TxContext' instead of '&TxContext'. As such, it is recommended to consider using '&mut TxContext' to future-proof the function.
+   = This warning can be suppressed with '#[allow(lint(prefer_mut_tx_context))]' applied to the 'module' or module member ('const', 'fun', or 'struct')
+
+warning[Lint W99009]: prefer '&mut TxContext' over '&TxContext'
+   ┌─ tests/sui_mode/linter/true_positive_public_mut_tx_context.move:20:13
+   │
+20 │         _b: &TxContext, // Should warn
+   │             ^^^^^^^^^^ 'public' functions should prefer '&mut TxContext' over '&TxContext' for better upgradability.
+   │
+   = When upgrading, the public function cannot be modified to take '&mut TxContext' instead of '&TxContext'. As such, it is recommended to consider using '&mut TxContext' to future-proof the function.
+   = This warning can be suppressed with '#[allow(lint(prefer_mut_tx_context))]' applied to the 'module' or module member ('const', 'fun', or 'struct')
+

external-crates/move/crates/move-compiler/tests/sui_mode/linter/true_positive_public_mut_tx_context.move

@@ -0,0 +1,29 @@
+// tests the lint for preferring &mut TxContext over &TxContext in public functions
+// these cases correctly should trigger the lint
+module 0x42::true_positive {
+    use sui::tx_context::TxContext;
+
+    struct CustomStruct has drop {}
+
+    public fun incorrect_mint(_ctx: &TxContext) {
+    }
+
+    public fun another_incorrect(_a: u64, _b: &TxContext, _c: u64) {
+    }
+
+    public fun mixed_function(_a: &CustomStruct, _b: &TxContext, _c: &mut TxContext) {}
+
+    fun private_function(_ctx: &TxContext) {}
+
+    public fun complex_function<T: drop>(
+        _a: u64,
+        _b: &TxContext, // Should warn
+        _c: &mut TxContext,
+        _d: &T,
+        _e: &CustomStruct
+    ) {}
+}
+
+module sui::tx_context {
+    struct TxContext has drop {}
+}