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

Code before:
    /// The locks belonging to the derived region will be kept in the given
    /// `locks` map, and the corresponding position in the returned `Vec`
    /// will be an empty map.
    pub fn group_by_regions(
        &mut self,
        regions: &[metapb::Region],
        derived: &metapb::Region,
    ) -> Vec<PeerPessimisticLocks> {
        // 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
    }

Review comment (suggestion):
The call to `as_slice()` is unnecessary here since `as_encoded()` already returns a byte slice that can be compared directly with start_key and end_key.
```suggestion
            let key_ref = key.as_encoded();
```

Code after:
    /// The locks belonging to the derived region will be kept in the given
    /// `locks` map, and the corresponding position in the returned `Vec`
    /// will be an empty map.
    pub fn group_by_regions(
        &mut self,
        regions: &[metapb::Region],
        derived: &metapb::Region,
    ) -> Vec<PeerPessimisticLocks> {
        // 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)]