Improve MutableState / Machines #2130
Conversation
There was a problem hiding this comment.
Awesome! 🤩
My main question is about lifetimes:
- I don't understand why
MutableDataneeds two different lifetimes. - I'm pretty sure the 3 lifetimes explained in the
Processordocstring can be reduced to 2 now, which is great!
I feel like fixing this would likely lead to a lot of changes all over the witgen codebase, so probably this should be done in a follow-up PR.
| machines: Vec<RefCell<KnownMachine<'a, T>>>, | ||
| identity_to_machine_index: BTreeMap<u64, usize>, |
There was a problem hiding this comment.
| machines: Vec<RefCell<KnownMachine<'a, T>>>, | |
| identity_to_machine_index: BTreeMap<u64, usize>, | |
| identity_to_machine: BTreeMap<u64, RefCell<KnownMachine<'a, T>>>, |
Why not this?
There was a problem hiding this comment.
take_witness_col_values uses the given order of machines. Also If we do it like that, then multiple identities calling the same machine will lead to a copy of the machine, which we don't want to have.
| pub struct MutableState<'a, 'b, T: FieldElement, Q: QueryCallback<T>> { | ||
| machines: Vec<RefCell<KnownMachine<'a, T>>>, | ||
| identity_to_machine_index: BTreeMap<u64, usize>, | ||
| query_callback: &'b Q, |
There was a problem hiding this comment.
Couldn't this be 'a as well?
There was a problem hiding this comment.
Probably yes, but it's more annoying lifetime changes because it will lead to Processor requiring fewer lifetimes etc. Will do it in another pr.
| @@ -98,7 +99,7 @@ pub struct Processor<'a, 'b, 'c, T: FieldElement, Q: QueryCallback<T>> { | |||
| /// The values of the publics | |||
| publics: BTreeMap<&'a str, T>, | |||
| /// The mutable state | |||
| mutable_state: &'c mut MutableState<'a, 'b, T, Q>, | |||
| mutable_state: &'c MutableState<'a, 'b, T, Q>, | |||
There was a problem hiding this comment.
A few lines above, there is a comment explaining the 3 lifetimes, including:
/// - `'b`: The duration of this machine's call (e.g. the mutable references of the other machines)I think this is no longer accurate?
There was a problem hiding this comment.
I think it's the lifetime of the query callback.
MutableState is the main way to get access to sub-machines during witness generation. We used to create copies of MutableState for each lookup, extracting the "current machine" where we need mutable access and creating copies of references to the other machines. This causes an allocation for each lookup (including fixed lookups, I think) which is bad for performance.
This PR changes the approach to use RefCell instead: MutableState now owns the machines and for each call to a machine, we mutably borrow that machine. The RefCell mechanism ensures that there are no recursive calls to machines and also avoids allocations.
I also changed the query callback to use non-mut references.
The first and third commits only move code around.