Repository: tikv/tikv
File: components/raftstore/src/store/txn_ext.rs (Rust)

Code before:
        // Assert regions are sorted by key in ascending order.
        if cfg!(debug_assertions) {
            for (r1, r2) in regions.iter().zip(regions.iter().skip(1)) {
                assert!(r1.get_start_key() < r2.get_start_key());
            }
        }

        let mut res: Vec<PeerPessimisticLocks> = regions
            .iter()
            .map(|_| PeerPessimisticLocks::default())
            .collect();
        // Locks that are marked deleted still need to be moved to the new regions,
        // and the deleted mark should also be cleared.
        // Refer to the comment in `PeerPessimisticLocks` for details.
        let removed_locks = self.map.drain_filter(|key, _| {
            let key = &**key.as_encoded();
            let (start_key, end_key) = (derived.get_start_key(), derived.get_end_key());
            key < start_key || (!end_key.is_empty() && key >= end_key)
        });
        for (key, (lock, _)) in removed_locks {
            let idx = match regions
                .binary_search_by_key(&&**key.as_encoded(), |region| region.get_start_key())
            {
                Ok(idx) => idx,
                Err(idx) => idx - 1,
            };
            let size = key.len() + lock.memory_size();
            self.memory_size -= size;
            res[idx].map.insert(key, (lock, false));
            res[idx].memory_size += size;
        }
        res
    }

    #[cfg(test)]
    fn from_locks(locks: impl IntoIterator<Item = (Key, (PessimisticLock, bool))>) -> Self {
        let mut res = PeerPessimisticLocks::default();
        for (key, (locks, is_deleted)) in locks {
            res.memory_size += key.len() + locks.memory_size();
            res.map.insert(key, (locks, is_deleted));
        }
        res
    }
}

impl<'a> IntoIterator for &'a PeerPessimisticLocks {
    type Item = (&'a Key, &'a (PessimisticLock, bool));
    type IntoIter = std::collections::hash_map::Iter<'a, Key, (PessimisticLock, bool)>;

    fn into_iter(self) -> Self::IntoIter {
        self.map.iter()

Review comment (performance):
The implementation clones both key and value for removed locks, which is inefficient. Consider using a two-pass approach: first collect keys to remove, then remove and collect them to avoid unnecessary clones.

Code after:
        // Assert regions are sorted by key in ascending order.
        if cfg!(debug_assertions) {
            for (r1, r2) in regions.iter().zip(regions.iter().skip(1)) {
                assert!(r1.get_start_key() < r2.get_start_key());
            }
        }

        let mut res: Vec<PeerPessimisticLocks> = regions
            .iter()
            .map(|_| PeerPessimisticLocks::default())
            .collect();
        // Locks that are marked deleted still need to be moved to the new regions,
        // and the deleted mark should also be cleared.
        // Refer to the comment in `PeerPessimisticLocks` for details.
        // There is no drain_filter for BtreeMap, so extra clone are needed.
        let mut removed_locks = Vec::new();
        self.map.retain(|key, value| {
            let key_ref = key.as_encoded().as_slice();
            let (start_key, end_key) = (derived.get_start_key(), derived.get_end_key());
            if key_ref < start_key || (!end_key.is_empty() && key_ref >= end_key) {
                removed_locks.push((key.clone(), value.clone()));
                false
            } else {
                true
            }
        });

        for (key, (lock, _)) in removed_locks.into_iter() {
            let idx = match regions
                .binary_search_by_key(&&**key.as_encoded(), |region| region.get_start_key())
            {
                Ok(idx) => idx,
                Err(idx) => idx - 1,
            };
            let size = key.len() + lock.memory_size();
            self.memory_size -= size;
            res[idx].map.insert(key, (lock, false));
            res[idx].memory_size += size;
        }
        res
    }

    #[cfg(test)]
    fn from_locks(locks: impl IntoIterator<Item = (Key, (PessimisticLock, bool))>) -> Self {
        let mut res = PeerPessimisticLocks::default();
        for (key, (locks, is_deleted)) in locks {
            res.memory_size += key.len() + locks.memory_size();
            res.map.insert(key, (locks, is_deleted));
        }
        res
    }