Internal function with inner mutation

[jorgecarleitao]

Hi,

I am trying to use a FnMut on a struct whose behavior is only known at run-time, and I am struggling to make it work. I am not very familiar with the interior mutability yet, and I would kindly ask for some help.

The code below is a minimal example: at new, the variable transposed dictates what the method set should do. The method set is mutable and has a lifetime associated to the data that is passed to the new of A.

type Set<'a> = Box<dyn FnMut(usize, usize) -> () + 'a>;

struct A<'a> {
    data: Vec<Vec<usize>>,
    set: Set<'a>,
}

impl<'a> A<'a> {
    fn new(data: &'a [usize], transposed: bool) -> Self {
        let mut inner = Vec::with_capacity(data.len());
        let set: Set<'a> = if transposed {
            Box::new(move |i, len| {
                let a = &data[i..i+len];
                inner.push(a.to_vec());
            })
        } else {
            Box::new(move |i, len| {
                let a = data[i..i+len].iter().map(|a| vec![*a]);
                inner.extend(a);
            })
        };
        Self {
            data: inner,
            set,
        }
    }

    fn set(&mut self, i: usize, len: usize) {
        (self.set)(i, len)
    }

    fn get(&self, i: usize, j: usize) -> usize {
        self.data[i][j]
    }
}

fn main() {
    let data = vec![1, 2, 3];
    let mut a = A::new(&data, true);

    a.set(0, 1);
    assert_eq!(a.get(0, 0), 1);
}

This code does not compile because inner is was moved to the set. However, it seems that there is not intrinsic limitation here (apart from the borrow checker).

Note that I am not interested in fixing the problem in the example above: I am interested in the idiom to solve these problems in Rust.

The gist that there is a boxed function inside a struct that mutates the struct. The function itself is not shared to outside that struct.

[H2CO3]

What you want here is essentially a self-referential struct. That is not currently possible in safe Rust. You should probably re-structure your code so that instead of insisting on the closure capturing a reference to a field of its containing type (inner), you pass the field as an explicit parameter only when needed. Playground:

type Set<'a> = Box<dyn FnMut(&mut Vec<Vec<usize>>, usize, usize) -> () + 'a>;

struct A<'a> {
    data: Vec<Vec<usize>>,
    set: Set<'a>,
}

impl<'a> A<'a> {
    fn new(data: &'a [usize], transposed: bool) -> Self {
        let set: Set = if transposed {
            Box::new(move |inner, i, len| {
                let a = &data[i..i+len];
                inner.push(a.to_vec());
            })
        } else {
            Box::new(move |inner, i, len| {
                let a = data[i..i+len].iter().map(|a| vec![*a]);
                inner.extend(a);
            })
        };
        Self {
            data: Vec::with_capacity(data.len()),
            set,
        }
    }

    fn set(&mut self, i: usize, len: usize) {
        (self.set)(&mut self.data, i, len)
    }

    fn get(&self, i: usize, j: usize) -> usize {
        self.data[i][j]
    }
}

fn main() {
    let data = vec![1, 2, 3];
    let mut a = A::new(&data, true);

    a.set(0, 1);
    assert_eq!(a.get(0, 0), 1);
}

[jorgecarleitao]

Thanks a lot, H2CO3. That was exactly what I did before posing this question

However, in my use-case, I need to prepare some of inner state in different ways depending on the transpose value, and that preparation is expensive (when compared to what set does, I benchmarked it and the flamegraph confirms). Thus the reason to pass a mutable reference of the inner itself.

Below is an example closer to my use-case, where you can assume that the operation inner[0] is as expensive as the .data.extend(a.to_vec()) one (not in this example, but in my use-case).

In other words, the set method can be a small operation, in which case indexing is expensive, but it can also be an expensive operation, in which case indexing is inexpensive. I know some invariants of the Boxed function, as the initialization of the inner is controlled by me.

type Set<'a> = Box<dyn FnMut(usize, usize) -> () + 'a>;

struct Inner {
    data: Vec<usize>,
    offset: usize,
}

impl Inner {
    fn data(&self) -> &[usize] {
        &self.data[self.offset..]
    }
}

struct A<'a> {
    data: Vec<Inner>,
    set: Set<'a>,
}

impl<'a> A<'a> {
    fn new(data: &'a [usize], offset: usize, is_simple: bool) -> Self {
        let mut inner = if is_simple {
             vec![
                Inner {data: Vec::with_capacity(data.len()), offset},
            ]
        } else {
            vec![
                Inner {data: Vec::with_capacity(data.len()), offset},
                Inner {data: Vec::with_capacity(data.len()), offset}
            ]
        };

        let set: Set<'a> = if is_simple {
            // this operation is expensive (when compared to the one inside the closure below), and thus we do it here and move its result to the closure
            let inner_data = &mut inner[0].data;
            Box::new(move |i, len| {
                let a = &data[i..i+len];
                inner_data.extend(a);
            })
        } else {
            let inner_data0 = &mut inner[0].data;
            let inner_data1 = &mut inner[1].data;
            Box::new(move |i, len| {
                let a = &data[i..i+len];
                inner_data0.extend(a);
                inner_data1.extend(a);
            })
        };
        Self {
            data: inner,
            set,
        }
    }

    fn set(&mut self, i: usize, len: usize) {
        (self.set)(i, len)
    }

    fn get_0(&self, i: usize) -> usize {
        self.data[0].data()[i]
    }
}

fn main() {
    let data = vec![1, 2, 3];
    let mut a = A::new(&data, 0, true);

    a.set(0, 1);
    assert_eq!(a.get_0(0), 1);
}

I could try using dyn, as this seems to be a situation where run-time drives behavior. Do you think it could make sense? I.e. re-formulate the problem so that we have a trait with set and get_0, and then have different implementations define different behavior. The set: Box<Fn> is to some extend a vtable, no?

[H2CO3]

I'm sorry, I don't really follow anymore. You can still prepare inner differently in new() depending on the value of the flag. I have rewritten your code above according to my previous suggestion (playground), but I don't quite understand what difference the type of the data field makes.

Or does the type of the data field somehow change dynamically? This might be what your subsequent question suggests:

@jorgecarleitao

I could try using dyn , as this seems to be a situation where run-time drives behavior. Do you think it could make sense?

But you are already using dyn – what is the other place you would try to use it in the above program?

@jorgecarleitao

The set: Box<Fn> is to some extend a vtable, no?

Yes, it is, dyn Fn is not special with respect to any other dyn Trait. Any dyn Trait is implemented using dynamic dispatch and a vtable.

[jorgecarleitao]

Sorry, I did not explain myself correctly. I edited the code above and added a comment on the relevant bit.

The data is the same (in my use-case, everything is in bytes), but how we write to it depends on run time. Since the set sometimes only writes 4 bytes (and the capacity was already prepared), the indexing becomes expensive. Thus, I am trying to offload the indexing to outside the set closure (the let inner_data = inner[0].data statement on the revised code).

A dynamic trait in this case does not really help because I do not have an interface that is fulfilled by different structs: the underlying data structure is the same, but how we write to it depends on runtime (the branch on the new in this example).

[H2CO3]

Oh, I think I understand now, thanks. I'm thinking about a solution.

---

Is there a way you could reduce the expensive indexing operation to an inexpensive one, e.g. to the equivalent of just a literal, primitive slice indexing, but sort of "stopping" before the last step? What I mean is that you could pre-calculate indexes/ranges, which is the expensive part, but then not get an actual reference to the corresponding slice. By doing this, you wouldn't need a self-referential struct, and the overhead that you pay inside the closure would only be a bounds check.

Or, if that is not possible, you could re-arrange the code even more: there would now be two separate types, where one would contain the pre-calculated slices, the other one would hold the owned data, and you would require users of these two types to keep around both of them.