Rust: return &mut reference from downcast_ref of HashMap mutable value

user_6688415 0
Viewed 36

As the title says, we have an application which runs wasm plugins. Each plugin can register their own State in the StateRegistry. Then whenever the plugins are executed, they will modify their respective State.

The code below illustrates this:

use std::{collections::HashMap, any::Any};

type ContractId = String;
type GenericContractState = Box<dyn Any>;

// This will be WASM code
mod foo_contract {
    use super::StateRegistry;

    pub struct State {
    }

    pub fn update(states: &mut StateRegistry) {
        let state = states.states.get_mut(&"foo_contract".to_string()).unwrap();
        let state = state.downcast_mut::<State>().unwrap();
        // I'd prefer something like:
        // let state = state.lookup::<State>().unwrap()
    }
}

pub struct StateRegistry {
    pub states: HashMap<ContractId, GenericContractState>,
}

impl StateRegistry {
    fn new() -> Self {
        Self { states: HashMap::new() }
    }

    fn register(&mut self, contract_id: ContractId, state: GenericContractState) {
        self.states.insert(contract_id, state);
    }

    /*
    fn lookup<'a, S>(&'a mut self, contract_id: &ContractId) -> Option<StateRefWrapper<'a, S>> {
        match self.states.get_mut(contract_id) {
            Some(state) => {
                let ptr = state.downcast_mut::<S>();
                match ptr {
                    Some(ptr) => Some(StateRefWrapper { _mut: state, ptr }),
                    None => None,
                }
            }
            None => None,
        }
    }
    */
}

/*
struct StateRefWrapper<'a, S> {
    _mut: &'a mut Box<dyn Any>,
    ptr: &'a mut S,
}
*/

fn main() {
    let mut states = StateRegistry::new();
    let foo_state = Box::new(foo_contract::State {});
    states.register("foo_contract".to_string(), foo_state);

    foo_contract::update(&mut states);
}

The part I want to improve is that currently the plugin developer has to use this code to lookup their State in the registry:

        let state = states.states.get_mut(&"foo_contract".to_string()).unwrap();
        let state = state.downcast_mut::<State>().unwrap();

I want to create a respective method in StateRegistry that is usable like this:

        let state = state.lookup::<State>().unwrap()

My attempt (commented above) was something like this:

impl StateRegistry {
    // ...

    fn lookup<'a, S>(&'a mut self, contract_id: &ContractId) -> Option<StateRefWrapper<'a, S>> {
        match self.states.get_mut(contract_id) {
            Some(state) => {
                let ptr = state.downcast_mut::<S>();
                match ptr {
                    Some(ptr) => Some(StateRefWrapper { _mut: state, ptr }),
                    None => None,
                }
            }
            None => None,
        }
    }
}

struct StateRefWrapper<'a, S> {
    _mut: &'a mut Box<dyn Any>,
    ptr: &'a mut S,
}

How can I get this lookup() function working?

Thanks

rust
1 Answers

You can't construct your StateRefWrapper in the way you want to, because it would contain two mutable borrows of the same value. You can, however, return just a &mut S which seems sufficient:

fn lookup<'a, S: 'static>(&'a mut self, contract_id: &ContractId) -> Option<&'a mut S> {
    self.states
        .get_mut(contract_id)
        .and_then(|state| state.downcast_mut())
}

And use it like this:

pub fn update(states: &mut StateRegistry) {
    let state = states.lookup::<State>(&"foo_contract".to_string()).unwrap();
    // state.modify();
}
user_114231040
Related