Internal function with inner mutation

help
#1

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.

#2

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);
}
#3

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

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?

#4

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:

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

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.

#5

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).

#6

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.