Allow multiple loop invariants for the same function, prefer same-module

[andreistefanescu]

When multiple modules provide a loop invariant for the same target function and label, prefer the one from the same module as the target function instead of erroring. Still errors on genuinely ambiguous cases.

Moves loop invariant resolution from PackageTargets to FunctionTargetsHolder, collecting candidates during attribute scanning and resolving them at holder construction time.

---

Note

Medium Risk
Changes how loop invariant functions are selected when multiple candidates exist, which can alter verification behavior and diagnostics for specs with external invariants.

Overview
Loop invariant resolution is refactored and made more permissive. PackageTargets now collects candidate loop invariants per target function/label instead of rejecting duplicates during attribute scanning.

Final selection is deferred to FunctionTargetsHolder. A new resolve_loop_invariants pass picks a single invariant per label, preferring the invariant defined in the same module as the chosen spec (and emitting clearer errors for ambiguous cases or invariant functions targeting multiple labels). Call sites in the Boogie generator, escape analysis, stackless tests, and sui-prover model building now invoke this resolution step.

Tests updated/added. Adds an integration case where two modules provide invariants for the same spec and the same-module one is selected, and updates failure snapshots to expect the new “ambiguous loop invariants” diagnostic.

^{Written by Cursor Bugbot for commit 656c410. This will update automatically on new commits. Configure here.}

[cursor]

Cursor Bugbot has reviewed your changes and found 1 potential issue.

Bugbot Free Tier Details

Your team is on the Bugbot Free tier. On this plan, Bugbot will review limited PRs each billing cycle for each member of your team.

To receive Bugbot reviews on all of your PRs, visit the Cursor dashboard to activate Pro and start your 14-day free trial.

Bugbot Autofix prepared a fix for the issue found in the latest run.

• ✅ Fixed: Unrelated test snapshot regressed from pass to fail
  • I confirmed the regression was an accidental snapshot change in the loop-invariant PR and restored vec_map/remove.move.snap to the prior passing Verification successful output.

Or push these changes by commenting:

@cursor push 8ecc640dbc

Preview (8ecc640dbc)

diff --git a/crates/move-stackless-bytecode/src/function_target_pipeline.rs b/crates/move-stackless-bytecode/src/function_target_pipeline.rs
--- a/crates/move-stackless-bytecode/src/function_target_pipeline.rs
+++ b/crates/move-stackless-bytecode/src/function_target_pipeline.rs
@@ -44,6 +44,7 @@
     package_targets: PackageTargets,
     function_specs: BiBTreeMap<QualifiedId<FunId>, QualifiedId<FunId>>,
     datatype_invs: BiBTreeMap<QualifiedId<DatatypeId>, QualifiedId<FunId>>,
+    loop_invariants: BTreeMap<QualifiedId<FunId>, BiBTreeMap<QualifiedId<FunId>, usize>>,
     target: FunctionHolderTarget,
     prover_options: ProverOptions,
 }
@@ -186,6 +187,7 @@
             targets: BTreeMap::new(),
             function_specs: BiBTreeMap::new(),
             datatype_invs: BiBTreeMap::new(),
+            loop_invariants: BTreeMap::new(),
             prover_options,
             target,
             package_targets: package_targets.clone(),
@@ -366,11 +368,99 @@
         self.package_targets.spec_timeouts().get(id)
     }
 
+    /// Resolve loop invariant candidates into the final map.
+    /// When multiple invariant functions target the same (target_func, label),
+    /// prefer the one from the same module as the spec selected by this holder.
+    fn resolve_loop_invariants(&mut self, env: &GlobalEnv) {
+        for (target_id, entries) in self.package_targets.loop_invariant_candidates() {
+            // Group by label, checking for duplicate inv_funcs
+            let mut by_label: BTreeMap<usize, Vec<QualifiedId<FunId>>> = BTreeMap::new();
+            let mut func_to_label: BTreeMap<QualifiedId<FunId>, usize> = BTreeMap::new();
+            let mut has_error = false;
+
+            for (inv_id, label) in entries {
+                if let Some(&existing_label) = func_to_label.get(inv_id) {
+                    if existing_label != *label {
+                        env.diag(
+                            Severity::Error,
+                            &env.get_function(*inv_id).get_loc(),
+                            &format!(
+                                "Loop invariant function {} targets multiple labels ({} and {}) in {}",
+                                env.get_function(*inv_id).get_full_name_str(),
+                                existing_label,
+                                label,
+                                env.get_function(*target_id).get_full_name_str()
+                            ),
+                        );
+                        has_error = true;
+                    }
+                    continue;
+                }
+                func_to_label.insert(*inv_id, *label);
+                by_label.entry(*label).or_default().push(*inv_id);
+            }
+
+            if has_error {
+                continue;
+            }
+
+            // Find the spec module for this target function via function_specs.
+            // If the target itself is a spec (scenario spec, etc.), use its own module.
+            let spec_module = self
+                .function_specs
+                .get_by_right(target_id)
+                .map(|spec_id| spec_id.module_id)
+                .unwrap_or(target_id.module_id);
+
+            let mut resolved: BiBTreeMap<QualifiedId<FunId>, usize> = BiBTreeMap::new();
+
+            for (label, inv_ids) in by_label {
+                let winner = if inv_ids.len() == 1 {
+                    inv_ids[0]
+                } else {
+                    // Prefer the invariant from the same module as the spec
+                    let in_spec_module: Vec<_> = inv_ids
+                        .iter()
+                        .filter(|id| id.module_id == spec_module)
+                        .copied()
+                        .collect();
+
+                    if in_spec_module.len() == 1 {
+                        in_spec_module[0]
+                    } else {
+                        let func_names = inv_ids
+                            .iter()
+                            .map(|id| env.get_function(*id).get_full_name_str())
+                            .collect::<Vec<_>>()
+                            .join(", ");
+                        env.diag(
+                            Severity::Error,
+                            &env.get_function(*target_id).get_loc(),
+                            &format!(
+                                "Ambiguous loop invariants for label {} in {}: [{}]",
+                                label,
+                                env.get_function(*target_id).get_full_name_str(),
+                                func_names
+                            ),
+                        );
+                        continue;
+                    }
+                };
+
+                resolved.insert(winner, label);
+            }
+
+            if !resolved.is_empty() {
+                self.loop_invariants.insert(*target_id, resolved);
+            }
+        }
+    }
+
     pub fn get_loop_invariants(
         &self,
         id: &QualifiedId<FunId>,
     ) -> Option<&BiBTreeMap<QualifiedId<FunId>, usize>> {
-        self.package_targets.loop_invariants().get(id)
+        self.loop_invariants.get(id)
     }
 
     pub fn get_uninterpreted_functions(
@@ -445,8 +535,7 @@
     pub fn get_loop_inv_with_targets(
         &self,
     ) -> BiBTreeMap<QualifiedId<FunId>, BTreeSet<QualifiedId<FunId>>> {
-        self.package_targets
-            .loop_invariants()
+        self.loop_invariants
             .iter()
             .map(|(target_fun_id, invs)| {
                 (
@@ -873,6 +962,7 @@
         H1: Fn(&FunctionTargetsHolder),
         H2: Fn(usize, &dyn FunctionTargetProcessor, &FunctionTargetsHolder),
     {
+        targets.resolve_loop_invariants(env);
         hook_before_pipeline(targets);
         for (step_count, processor) in self.processors.iter().enumerate() {
             let topological_order: Vec<Either<QualifiedId<FunId>, Vec<QualifiedId<FunId>>>> =

diff --git a/crates/move-stackless-bytecode/src/package_targets.rs b/crates/move-stackless-bytecode/src/package_targets.rs
--- a/crates/move-stackless-bytecode/src/package_targets.rs
+++ b/crates/move-stackless-bytecode/src/package_targets.rs
@@ -1,5 +1,4 @@
 use crate::target_filter::TargetFilterOptions;
-use bimap::BiBTreeMap;
 use codespan_reporting::diagnostic::Severity;
 use move_binary_format::file_format::FunctionHandleIndex;
 use move_compiler::{
@@ -43,7 +42,7 @@
     spec_uninterpreted_functions: BTreeMap<QualifiedId<FunId>, BTreeSet<QualifiedId<FunId>>>,
     spec_boogie_options: BTreeMap<QualifiedId<FunId>, String>,
     spec_timeouts: BTreeMap<QualifiedId<FunId>, u64>,
-    loop_invariants: BTreeMap<QualifiedId<FunId>, BiBTreeMap<QualifiedId<FunId>, usize>>,
+    loop_invariant_candidates: BTreeMap<QualifiedId<FunId>, Vec<(QualifiedId<FunId>, usize)>>,
     module_external_attributes: BTreeMap<ModuleId, BTreeSet<ModuleExternalSpecAttribute>>,
     function_external_attributes:
         BTreeMap<QualifiedId<FunId>, BTreeSet<ModuleExternalSpecAttribute>>,
@@ -82,7 +81,7 @@
             spec_uninterpreted_functions: BTreeMap::new(),
             spec_boogie_options: BTreeMap::new(),
             spec_timeouts: BTreeMap::new(),
-            loop_invariants: BTreeMap::new(),
+            loop_invariant_candidates: BTreeMap::new(),
             module_external_attributes: BTreeMap::new(),
             function_external_attributes: BTreeMap::new(),
             module_extra_bpl: BTreeMap::new(),
@@ -200,53 +199,10 @@
         if let Some(target_func_env) =
             module_env.find_function(func_env.symbol_pool().make(fun_name.as_str()))
         {
-            if let Some(existing) = self
-                .loop_invariants
-                .get_mut(&target_func_env.get_qualified_id())
-            {
-                match existing.insert(func_env.get_qualified_id(), label) {
-                    bimap::Overwritten::Neither => {}
-                    bimap::Overwritten::Left(..) => {
-                        env.diag(
-                            Severity::Error,
-                            &func_env.get_loc(),
-                            &format!(
-                                "Duplicated Loop Invariant Function {} in {}",
-                                func_env.get_full_name_str(),
-                                fun_name
-                            ),
-                        );
-                        return;
-                    }
-                    bimap::Overwritten::Right(..) => {
-                        env.diag(
-                            Severity::Error,
-                            &func_env.get_loc(),
-                            &format!("Duplicated Loop Invariant Label {} in {}", label, fun_name),
-                        );
-                        return;
-                    }
-                    bimap::Overwritten::Both(..) | bimap::Overwritten::Pair(..) => {
-                        env.diag(
-                            Severity::Error,
-                            &func_env.get_loc(),
-                            &format!(
-                                "Duplicated Loop Invariant Function {} and Label {} in {}",
-                                func_env.get_full_name_str(),
-                                label,
-                                fun_name
-                            ),
-                        );
-                    }
-                }
-            } else {
-                self.loop_invariants
-                    .insert(target_func_env.get_qualified_id(), {
-                        let mut map = BiBTreeMap::new();
-                        map.insert(func_env.get_qualified_id(), label);
-                        map
-                    });
-            }
+            self.loop_invariant_candidates
+                .entry(target_func_env.get_qualified_id())
+                .or_default()
+                .push((func_env.get_qualified_id(), label));
         } else {
             env.diag(
                 Severity::Error,
@@ -1011,10 +967,10 @@
         &self.spec_timeouts
     }
 
-    pub fn loop_invariants(
+    pub fn loop_invariant_candidates(
         &self,
-    ) -> &BTreeMap<QualifiedId<FunId>, BiBTreeMap<QualifiedId<FunId>, usize>> {
-        &self.loop_invariants
+    ) -> &BTreeMap<QualifiedId<FunId>, Vec<(QualifiedId<FunId>, usize)>> {
+        &self.loop_invariant_candidates
     }
 
     pub fn get_module_extra_bpl(&self, module_id: &ModuleId) -> Option<&String> {

diff --git a/crates/sui-prover/tests/inputs/loop_invariant/external_multi_module.ok.move b/crates/sui-prover/tests/inputs/loop_invariant/external_multi_module.ok.move
new file mode 100644
--- /dev/null
+++ b/crates/sui-prover/tests/inputs/loop_invariant/external_multi_module.ok.move
@@ -1,0 +1,32 @@
+// Two modules provide a loop invariant for the same spec function.
+// The invariant in the same module as the spec should be selected.
+module 0x42::other_specs {
+    #[allow(unused_variable)]
+    #[spec_only(loop_inv(target = 0x42::my_specs::test_spec))]
+    #[ext(no_abort)]
+    fun loop_inv_other(i: u64, n: u64): bool {
+        // This invariant is intentionally too weak; it would fail if selected.
+        true
+    }
+}
+
+module 0x42::my_specs {
+    use prover::prover::ensures;
+
+    #[spec_only(loop_inv(target = test_spec))]
+    #[ext(no_abort)]
+    fun loop_inv(i: u64, n: u64): bool {
+        i <= n
+    }
+
+    #[spec(prove)]
+    fun test_spec(n: u64) {
+        let mut i = 0;
+
+        while (i < n) {
+            i = i + 1;
+        };
+
+        ensures(i == n);
+    }
+}

diff --git a/crates/sui-prover/tests/snapshots/loop_invariant/external_multi_module.ok.move.snap b/crates/sui-prover/tests/snapshots/loop_invariant/external_multi_module.ok.move.snap
new file mode 100644
--- /dev/null
+++ b/crates/sui-prover/tests/snapshots/loop_invariant/external_multi_module.ok.move.snap
@@ -1,0 +1,5 @@
+---
+source: crates/sui-prover/tests/integration.rs
+expression: output
+---
+Verification successful

diff --git a/crates/sui-prover/tests/snapshots/loop_invariant/external_wrong_label_2.fail.move.snap b/crates/sui-prover/tests/snapshots/loop_invariant/external_wrong_label_2.fail.move.snap
--- a/crates/sui-prover/tests/snapshots/loop_invariant/external_wrong_label_2.fail.move.snap
+++ b/crates/sui-prover/tests/snapshots/loop_invariant/external_wrong_label_2.fail.move.snap
@@ -1,13 +1,16 @@
 ---
 source: crates/sui-prover/tests/integration.rs
-assertion_line: 262
 expression: output
 ---
 exiting with bytecode transformation errors
-error: Duplicated Loop Invariant Label 0 in test_spec
-   ┌─ tests/inputs/loop_invariant/external_wrong_label_2.fail.move:13:1
+error: Ambiguous loop invariants for label 0 in loop_invariant_external_wrong_label_2_fail::test_spec: [loop_invariant_external_wrong_label_2_fail::loop_inv_1, loop_invariant_external_wrong_label_2_fail::loop_inv_2]
+   ┌─ tests/inputs/loop_invariant/external_wrong_label_2.fail.move:18:1
    │  
-13 │ ╭ fun loop_inv_2(i: u64, n: u64, s: u128): bool {
-14 │ │     i <= n && (s == (i as u128) * ((i as u128) + 1) / 2)
-15 │ │ }
+18 │ ╭ fun test_spec(n: u64): u128 {
+19 │ │     let mut s: u128 = 0;
+20 │ │     let mut i = 0;
+21 │ │ 
+   · │
+28 │ │     s
+29 │ │ }
    │ ╰─^

_{This Bugbot Autofix run was free. To enable autofix for future PRs, go to the Cursor dashboard.}