Repository: MystenLabs/sui
File: external-crates/move/crates/move-compiler/src/cfgir/borrows/state.rs (Rust)

Code before:
                return Diagnostics::new();
            }
            Value::Ref(id) => id,
        };

        self.mark_mutably_used(id);
        let diags = self.writable(loc, || "Invalid mutation of reference.".into(), id);
        self.release(id);
        diags
    }

    pub fn return_(&mut self, loc: Loc, rvalues: Values) -> Diagnostics {
        let mut released = BTreeSet::new();
        for (_, _local, stored_value) in &self.locals {
            if let Value::Ref(id) = stored_value {
                released.insert(*id);
            }
        }
        released.into_iter().for_each(|id| self.release(id));

        // Check locals are not borrowed
        let mut diags = Diagnostics::new();
        for (local, stored_value) in self.locals.key_cloned_iter() {
            if let Value::NonRef = stored_value {
                let (mut full_borrows, local_borrows) = self.local_borrowed_by(&local);
                full_borrows.extend(local_borrows);
                let local_diag = Self::borrow_error(
                    &self.borrows,
                    loc,
                    &full_borrows,
                    &BTreeMap::new(),
                    ReferenceSafety::InvalidReturn,
                    || {
                        let case = match display_var(local.value()) {
                            DisplayVar::Orig(v) => format!("Local variable '{v}'"),
                            DisplayVar::MatchTmp(_) => "Local value".to_string(),
                            DisplayVar::Tmp => "Local value".to_string(),
                        };
                        format!("Invalid return. {case} is still being borrowed.")
                    },
                );
                diags.add_opt(local_diag)
            }
        }

        // check any returned reference is not borrowed
        for rvalue in rvalues {
            match rvalue {
                Value::Ref(id) if self.borrows.is_mutable(id) => {
                    self.mark_mutably_used(id);
                    let (fulls, fields) = self.borrows.borrowed_by(id);

Review comment (suggestion):
This being repeated all of these places feels like a bit of an anti-pattern; could we add a different function alongside `local_borrowed_by` that does this extension and hands back a single value? I would hate to mess this concatenation up in a single place.

Code after:
            }
        }
    }

    pub fn mutate(&mut self, loc: Loc, rvalue: Value) -> Diagnostics {
        let id = match rvalue {
            Value::NonRef => {
                assert!(
                    self.prev_had_errors,
                    "ICE borrow checking failed {:#?}",
                    loc
                );
                return Diagnostics::new();
            }
            Value::Ref(id) => id,
        };

        self.mark_mutably_used(id);
        let diags = self.writable(loc, || "Invalid mutation of reference.".into(), id);
        self.release(id);
        diags
    }

    pub fn return_(&mut self, loc: Loc, rvalues: Values) -> Diagnostics {
        let mut released = BTreeSet::new();
        for (_, _local, stored_value) in &self.locals {
            if let Value::Ref(id) = stored_value {
                released.insert(*id);
            }
        }
        released.into_iter().for_each(|id| self.release(id));

        // Check locals are not borrowed
        let mut diags = Diagnostics::new();
        for (local, stored_value) in self.locals.key_cloned_iter() {
            if let Value::NonRef = stored_value {
                let root_and_local_borrows = self.local_borrowed_by(&local);
                let local_diag = Self::borrow_error(
                    &self.borrows,
                    loc,
                    &root_and_local_borrows,
                    &BTreeMap::new(),
                    ReferenceSafety::InvalidReturn,
                    || {
                        let case = match display_var(local.value()) {
                            DisplayVar::Orig(v) => format!("Local variable '{v}'"),
                            DisplayVar::MatchTmp(_) => "Local value".to_string(),
                            DisplayVar::Tmp => "Local value".to_string(),
                        };
                        format!("Invalid return. {case} is still being borrowed.")
                    },