Using Function-Traits with other Trait as Args

Basically, I want to be able to register a function that takes an argument that implements a specific trait (Input) to a Container and later run the registered functions.

Input values can be created from &Container and are bound to its lifetime. Container::run iterates through all functions and calls fun(&self).

I basically tried two implementations and ran into an issue in each.

Associated Type Lifetimes

Playground

This implementation uses a mess of associated types with lifetimes in Input. The code compiles, but I can't get the type inference in register<Fun, Arg>() working so that Arg is inferred from Fun:

pub trait Input {
    type Arg<'arg>;
    fn from_container<'arg>(c: &'arg Container) -> Self::Arg<'arg>;
}

impl<'c> Input for &'c Container {
    type Arg<'arg> = &'arg Container;
    fn from_container<'arg>(c: &'arg Container) -> Self::Arg<'arg> {
        c
    }
}

struct OtherInput<'a>(&'a Container);

impl<'q> Input for OtherInput<'q> {
    type Arg<'arg> = OtherInput<'arg>;
    fn from_container<'arg>(c: &'arg Container) -> Self::Arg<'arg> {
        OtherInput(c)
    }
}

struct Container {
    funs: Vec<Box<dyn Fn(&Self)>>,
}

impl Container {
    pub fn run(&self) {
        for fun in &self.funs {
            fun(&self);
        }
    }

    pub fn register<In, Fun>(&mut self, system: Fun)
    where
        Fun: Fn(In::Arg<'_>) + 'static,
        In: Input,
    {
        let wrap = move |c: &Container| system(In::from_container(c));
        self.funs.push(Box::new(wrap));
    }
}

fn test() {
    let mut c = Container { funs: vec![] };

    // Works fine with explicit annotations
    c.register::<&Container, _>(|c: &Container| todo!());
    c.register::<OtherInput, _>(|i: OtherInput| todo!());

    // Resolving fails here. Compiler can't infer ::<&Container, _> from the
    // argument
    c.register(|c: &Container| todo!());
    c.register(|c: OtherInput| todo!());
}

Input<'I>

Playground

This implementations handles lifetime of Arg in Input<'I>. I have a feeling this should work, but my understanding of HRTBs is a bit too vague:

pub trait Input<'i> {
    fn from_container(c: &'i Container) -> Self;
}

impl<'c> Input<'c> for &'c Container {
    fn from_container(c: &'c Container) -> Self {
        c
    }
}

struct OtherInput<'a>(&'a Container);

impl<'q> Input<'q> for OtherInput<'q> {
    fn from_container(c: &'q Container) -> Self {
        OtherInput(c)
    }
}

struct Container {
    funs: Vec<Box<dyn Fn(&Self)>>,
}

impl Container {
    pub fn run(&self) {
        for fun in &self.funs {
            fun(&self);
        }
    }

    pub fn register<Fun, In>(&mut self, system: Fun)
    where
        Fun: Fn(In) + 'static,
        for<'i> In: Input<'i>,
    {
        let wrap = move |c: &Container| system(In::from_container(c));
        self.funs.push(Box::new(wrap));
    }
}

fn test() {
    let mut c = Container { funs: vec![] };

    c.register(|c: &Container| todo!());
    c.register(|c: OtherInput| todo!());
    
    // rustc: implementation of `Input` is not general enough
    // `Input<'0>` would have to be implemented for the type `&Container`, for any lifetime `'0`...
    // ...but `Input<'1>` is actually implemented for the type `&'1 Container`, for some specific lifetime `'1`
}

Other Approaches

My initial implementations tried abstracting away Fn(In) in a trait trait Foo { run(&self, c: &Container); } but quickly ran into the same lifetime issue as approach #2.

The problem with your second strategy is that when you define a type variable In whose value is the lifetime-bearing type (e.g. &'a Container) you have then required the lifetime to be a specific lifetime (for each call to the function) rather than a HRTB lifetime (that can vary later), because to call register() you have to have (implicitly if not explicitly) a specific value of that type variable. So, by the time you've written fn register<Fun, In>, you've already made it not work by requiring that all inputs are the same type.

Sometimes (notably in borrowing async functions, which have a similar problem with returning a Future type), you can work around this by careful use of a helper trait that eliminates the In type variable, but I can’t think of a way to do that here.

Your first strategy can work — it's effectively using a possibly different type as a marker for the higher-rank type you actually want to refer to — but as you have observed, type inference won’t work with it by default. However, with sufficient bizarre hackery, you can convince the compiler to accept it. I took your code and added tricks from bevy_ecs's SystemParamFunction until it compiled:

pub trait InputKind {
    type Concrete<'i>: InputKind;
    fn from_container(c: &Container) -> Self::Concrete<'_>;
}

pub trait Registerable<IK> {
    fn run(&self, container: &Container);
}

impl<F, IK> Registerable<IK> for F
where
    IK: InputKind,
    // yes, BOTH of these Fn bounds are needed
    for<'a> &'a F: Fn(IK) + Fn(<IK as InputKind>::Concrete<'a>),
{
    fn run<'a>(&'a self, c: &'a Container) {
        fn inner<P>(f: impl Fn(P), p: P) {
            f(p)
        }
        inner(self, IK::from_container(c))
    }
}

impl InputKind for &Container {
    type Concrete<'i> = &'i Container;
    fn from_container(c: &Container) -> Self::Concrete<'_> {
        c
    }
}

struct OtherInput<'a>(&'a Container);

impl InputKind for OtherInput<'_> {
    type Concrete<'i> = OtherInput<'i>;
    fn from_container(c: &Container) -> Self::Concrete<'_> {
        OtherInput(c)
    }
}

pub struct Container {
    funs: Vec<Box<dyn Fn(&Self)>>,
}

impl Container {
    pub fn run(&self) {
        for fun in &self.funs {
            fun(self);
        }
    }

    pub fn register<Fun: Registerable<IK> + 'static, IK>(&mut self, system: Fun) {
        self.funs.push(Box::new(move |c: &Container| system.run(c)));
    }
}

fn main() {
    let mut c = Container { funs: vec![] };

    c.register(|c: &Container| todo!());
    c.register(|c: OtherInput| todo!());
}
2 Likes

Thank you very much! I took some inspiration from bevy's ECS earlier but missed the double-fn-bound trick. Everything works perfectly fine with it.

That trick was new to me too!

This topic was automatically closed 90 days after the last reply. We invite you to open a new topic if you have further questions or comments.