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

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.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: andreistefanescu

crates/move-stackless-bytecode/src/function_target_pipeline.rs

@@ -44,6 +44,7 @@ pub struct FunctionTargetsHolder {
     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 @@ impl FunctionTargetsHolder {
             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 @@ impl FunctionTargetsHolder {
         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 @@ impl FunctionTargetsHolder {
     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 @@ impl FunctionTargetPipeline {
         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>>>> =

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 @@ pub struct PackageTargets {
     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 @@ impl PackageTargets {
             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 @@ impl PackageTargets {
         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 @@ impl PackageTargets {
         &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> {

crates/sui-prover/tests/inputs/loop_invariant/external_multi_module.ok.move

@@ -0,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);
+    }
+}

crates/sui-prover/tests/snapshots/loop_invariant/external_multi_module.ok.move.snap

@@ -0,0 +1,5 @@
+---
+source: crates/sui-prover/tests/integration.rs
+expression: output
+---
+Verification successful

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 │ │ }
    │ ╰─^

crates/sui-prover/tests/snapshots/vec_map/remove.move.snap

@@ -2,4 +2,11 @@
 source: crates/sui-prover/tests/integration.rs
 expression: output
 ---
-Verification successful
+exiting with verification errors
+error: prover::ensures does not hold
+   ┌─ tests/inputs/vec_map/remove.move:15:3
+   │
+15 │   ensures(!m.contains(&10));
+   │   ^^^^^^^^^^^^^^^^^^^^^^^^^
+   │
+   =     at tests/inputs/vec_map/remove.move:12: foo_spec