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!