Insert element inside a iteration

I have code like this:

fn main() {
    let mut tree = BTreeMap::<usize, bool>::new();
    let required_size = 10;
    for (size, enable) in tree.range(required_size..) {
        if *enable {
            tree.remove(size);
            let rest = size - required_size;
            if rest > 0 {
                tree.insert(rest, true);
            }
            break;
        } else {
            continue;
        }
    }
}

This code do the following things: if I found a suitable size, and it is enabled, I remove it from the tree, and if there's any size left, I add the left size back to the tree.

The code will not compile because I borrow the tree as mutable when it is also borrowed as immutable, but in my case, if I found a suitable size, the iteration will end, so I think it is safe to remove and insert elements at this time, so I wrote like this:

fn main() {
    let mut tree = BTreeMap::<usize, bool>::new();
    let required_size = 10;
    for (size, enable) in tree.range(required_size..) {
        if *enable {
            // added code 
            let tree: *mut BTreeMap<usize, bool> = unsafe { std::mem::transmute(&tree as *const _ as *mut BTreeMap<usize, bool>) };
            // added code
            let tree = unsafe { &mut *tree };
            tree.remove(size);
            let rest = size - required_size;
            if rest > 0 {
                tree.insert(rest, true);
            }
            break;
        } else {
            continue;
        }
    }
}

Is it dangerous to write like this?

Yes, it's dangerous. Don't do this. In your computation of rest, size is a dangling pointer because you removed it from the tree.

The problem is that size is a reference that points into the BTreeMap. To fix it, you can make a copy of the usize before calling remove.

 fn main() {
     let mut tree = BTreeMap::<usize, bool>::new();
     let required_size = 10;
     for (size, enable) in tree.range(required_size..) {
         if *enable {
+            let size = *size;
+            tree.remove(&size);
-            tree.remove(size);
             let rest = size - required_size;
             if rest > 0 {
                 tree.insert(rest, true);
             }
             break;
         } else {
             continue;
         }
     }
 }
3 Likes

of course it is, it breaks rust's aliasing rules which means UB

1 Like

It seems that I misunderstand something, in your code, why I can remove an element in tree, isn't the tree borrowed at that time?

You seem to assume the problem is because of the iteration, but read the error:

error[E0502]: cannot borrow `tree` as mutable because it is also borrowed as immutable
 --> src/main.rs:7:13
  |
5 |     for (size, enable) in tree.range(required_size..) {
  |                           --------------------------- immutable borrow occurs here
6 |         if *enable {
7 |             tree.remove(size);
  |             ^^^^^^^^^^^^^^^^^ mutable borrow occurs here
8 |             let rest = size - required_size;
  |                        ---- immutable borrow later used here

It doesn't talk about the next iteration. If the next iteration was a problem you would get something like:

error[E0502]: cannot borrow `tree` as mutable because it is also borrowed as immutable
  --> src/main.rs:10:17
   |
5  |     for (size, enable) in tree.range(required_size..) {
   |                           ---------------------------
   |                           |
   |                           immutable borrow occurs here
   |                           immutable borrow later used here
...
10 |                 tree.insert(rest, true);
   |                 ^^^^^^^^^^^^^^^^^^^^^^^ mutable borrow occurs here

Instead, the problem is with the use of size in the line let rest = size - required_size;, as highlighted by the error. size has type &usize and points inside the BTreeMap, so you can't mutate the map while it exists. However there's no need for it to be a reference since you only care about its value. Removing that reference, as shown by @alice, fixes the problem.

No, as you already said you break out of the loop in that iteration, so there's no next iteration, the iterator is no longer used and thus there's no borrow remaining that can conflict with the remove/insert.

4 Likes

Thanks, the compiler is more clever than I think, and I thought too much and misunderstood the error message.