Building a runtime reflection system for Rust 🦀️ (Part 3)

Part 3: dyn Method

Welcome to the third and final installment of our series on how we implemented a runtime reflection system in Rust.

Introduction

So far, we've shown how we came up with a fairly simple Class and Instance model for thinking about runtime Rust classes. In Part 1, we used these for type checking, and in Part 2 we added support for reading attributes off of a struct.

In this post, we pick up where we left off with attribute getters, and expand into method calls. In some ways, the same techniques we used for attributes work just as well here. We can store a map from method name to functions implementing them. However, there's a curveball: the Rust Fn* traits. We'll talk through the wrong turns we took, and the tidbits of Rust knowledge we picked up along the way.

Method Calls

Now that we have classes, instances, and attributes, the next obvious step is to add methods.

In Oso policies, it is possible to call both class and instance methods, with and without arguments.

So given the struct:

struct Cat;

impl Cat {
    /// A class method (note lack of `self`).
    fn meow() -> String {
       "meowww".to_string()
    }

    /// An instance method.
    fn feed(&self, food: &str) -> String {
        if food == "tuna" { "purr".to_string() } else { Self::meow() }
    }
}

We should be able to write policy logic:

favourite_food(cat: Cat, food) if cat.feed(food) != Cat.meow();

Which says that the input food is the cat's favourite food if the result of feeding the cat is not the same as the result of the cat meowing.

Step 1: Zero arguments

Let's start with a simple implementation for methods that take zero arguments. The approach for implementing zero-argument methods is extremely similar to how we'd implement attribute getters, and we've actually done this already in Part 2:

/// Class definitions
struct Class {
    ...

    /// Map from attribute name to the attribute lookup
    methods: HashMap<&'static str, InstanceMethod>
}

struct InstanceMethod(Arc<dyn Fn(&Instance) -> PolarValue);

Similarly, we need to add a method onto our ClassBuilder struct to allow us to register new methods:

pub fn add_method<F, R>(mut self, name: &'static str, f: F) -> Self
where
    F: Fn(&T) -> R,
    R: crate::ToPolar,
{
    self.class.methods.insert(name, InstanceMethod::new(f));
    self
}

Super easy! End of blog post. See you next time 😎

But wait, what about methods with multiple arguments?

Step 2: Multiple Arguments and the Fn* traits

Let's take what we have above and add in support for multiple arguments.

struct InstanceMethod(Arc<dyn Fn(&Instance, Vec<PolarValue>) -> PolarValue>);

This mostly works for what we need! Polar doesn't care about method arities (how many arguments the method accepts) – it will send over however many arguments it has as a vector. Polar supports both variable number of arguments (varargs) and keyword arguments (kwargs). The latter is only supported for host languages that also have that concept, and Rust does not.

But this isn't the end of the story. The crucial part of the AttributeGetter interface was that you could pass in any method or closure for the attribute getter, and the AttributeGetter::new method handled all the type-conversions transparently. This hid the messy, error-prone details from the user and kept the interface clean.

So let's do the same for InstanceMethod!

impl InstanceMethod {
    pub fn new<T, F, ???>(f: F) -> Self
    where
        F: Fn(&T, ???)
        F::Result: ToPolarResult,
        T: 'static,
    {
        Self(Arc::new(
            move |receiver: &Instance, args: Vec<PolarValue>| {
                let receiver = receiver
                    .downcast()
                    .map_err(|e| e.invariant().into());
                // ermm.... what next?
            },
        ))
    }
}

We've hit our first problem.

The input to an attribute getter never accepted any arguments, and we only needed it to work for all Fn(&T). In order to support multiple arguments, we now need to cover Fn(&T), Fn(&T, A), Fn(&T, A, B), and so on.

The first problem is how to convert A, B, etc. into PolarValues.

The second problem is that these are all completely distinct traits. There is no trait capturing "functions of arity 2, 3, 4...". – at least not until the feature is available in stable Rust. In the future, this might be represented with the syntax Fn<Args, Output=T>, but right now using this syntax results in:

error[E0658]: the precise format of `Fn`-family traits' type parameters is subject to change
 --> src/main.rs:2:14
  |
2 |     where F: Fn<(u32,), Output=u32> {
  |              ^^^^^^^^^^^^^^^^^^^^^^ help: use parenthetical notation instead: `Fn(u32) -> u32`
  |
  = note: see issue #29625 <https://github.com/rust-lang/rust/issues/29625> for more information

We could opt to use Rust nightly to get these features, but it's not a huge stretch to implement them ourselves.

Implementing our own Method trait

This is where as the writer it's tempting to unveil my newly-created trait, perfectly matching what we needed, and make it look like I just put fingers to keyboard to get to the definition.

In reality, we spent a sizeable chunk of our engineering effort for this project on this one trait. We made mistakes. We wrote code that we threw out. And a lot of that is because we didn't understand some of the nuances of Rust functions and the trait resolution system. Instead of papering over all of that, we thought it would be more interesting to show you what we tried.

Attempt #1

The first thing we tried was, in hindsight, a little greedy. Why not just skip straight to writing a trait to encapsulate precisely what we need?

pub trait Method {
   fn invoke(&self, receiver: Instance, args: Vec<PolarValue>) -> PolarValue;
}

This looks great, let's try implementing it for one of our Fn variants:

impl<F, T, R> Method for F
where
   F: Fn(&T) -> R,
   T: 'static,
   R: ToPolarValue,
{
    fn invoke(&self, instance: Instance, args: Vec<PolarValue>) -> PolarValue {
        debug_assert!(args.is_empty());
        let receiver = instance.downcast::<T>().unwrap();
        self(receiver).to_polar()
    }
}

Results in:

impl<F, T, R> Method for F
        ^ unconstrained type parameter
the type parameter `T` is not constrained by the impl trait, self type, or predicates

impl<F, T, R> Method for F
           ^ unconstrained type parameter
the type parameter `R` is not constrained by the impl trait, self type, or predicates

It's likely that most people have hit some variation of this error in their Rust adventures! What is going on here? T and R look pretty constrained to me? They are right there inside the definition of F: Fn(&T) -> R.

Our mistake was reading those trait bounds as: F is a function from &T to R, whereas in reality this is a regular old trait bound with slightly different syntax for the trait itself. And one function might implement multiple of these trait bounds.

E.g.

// do nothing
fn ident<T>(t: T) -> T { t }

let _ = &ident as &dyn Fn(u32) -> u32;
let _ = &ident as &dyn Fn(String) -> String;

So which trait should we use for the implementation of Method for ident?

Here's another way of looking at this problem. Suppose instead we decided to exhaustively implement our Method trait for all types we care about:

impl<F> Method for F
where
   F: Fn(&u32) -> u32,
{
    fn invoke(&self, instance: Instance, args: Vec<PolarValue>) -> PolarValue {
        debug_assert!(args.is_empty());
        let receiver = instance.downcast::<u32>().unwrap();
        self(receiver).to_polar()
    }
}

impl<F> Method for F
where
   F: Fn(&String) -> String,
{
    fn invoke(&self, instance: Instance, args: Vec<PolarValue>) -> PolarValue {
        debug_assert!(args.is_empty());
        let receiver = instance.downcast::<String>().unwrap();
        self(receiver).to_polar()
    }
}

Ignoring for now just how bad an idea this is, it doesn't even work! We get:

conflicting implementations of trait `Method`:

Look back to ident. That one function implements both Fn(String) -> String and Fn(u32) -> u32 traits. Similarly, in the above case, a function that implements both Fn(&String) -> String and Fn(&u32) -> u32 would have two possible implementations for Method. So we get conflicting implementations.

Actually, it goes even further than that. Implementing a blanket trait implementation over any two function trait bounds results in conflicting implementations. Even though such a function couldn't exist:

trait Test {}

impl<F: Fn()> Test for F { }
impl<F: Fn(String) -> String> Test for F { }

Results in:

error[E0119]: conflicting implementations of trait `Test`:
  |
4 | impl<F: Fn()> Test for F { }
  | ------------------------ first implementation here
5 | impl<F: Fn(String) -> String> Test for F { }
  | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ conflicting implementation

One day (or today, if you're on Rust nightly), Rust might let you implement Fn for your own types, support variadic functions (what do you mean Rust already supports variadic functions?), and do all kinds of fun things of the sort. But for now we're not going to get much farther with this approach.

Attempt #2

This Method trait looks too convenient to throw away entirely at the first sign of complication. Let's try something different. If our original mistake was thinking of Fn as a function instead of a trait, perhaps we can heed the wisdom of the Rust docs and use fn instead.

Building%20a%20runtime%20reflection%20system%20for%20Rust%20(Par%20e9d46e4c771847e7ba119df92cf2a8b9/Untitled.png

Based on the docs, we can use fn with both regular functions and closures! Great, let's do just that:

impl<T, R> Method for fn(&T) -> R
where
    T: 'static,
    R: ToPolarValue,
{
     fn invoke(&self, instance: Instance, args: Vec<PolarValue>) -> PolarValue {
        debug_assert!(args.is_empty());
        let receiver = instance.downcast().unwrap();
        self(receiver).to_polar()
    }
}

Works fine! Let's try it out:

let clone = |receiver: &String| -> String { receiver.clone() };
let clone_method: Box<dyn Method> = Box::new(clone);
let instance = Instance::new("hello, world!".to_string());
let result = clone_method.invoke(instance, vec![]);

And...

error[E0277]: the trait bound `[closure@src/main.rs:25:17: 25:67]: Method` is not satisfied
  --> src/main.rs:26:41
   |
26 |     let clone_method: Box<dyn Method> = Box::new(clone);
   |                                         ^^^^^^^^^^^^^^^ the trait `Method` is not implemented for `[closure@src/main.rs:25:17: 25:67]`
   |
   = note: required for the cast to the object type `dyn Method`

It doesn't work for closures? What if we change clone to a proper function fn clone(s: &u32) -> String { s.to_string() }

error[E0277]: the trait bound `for<'r> fn(&'r String) -> String {main::clone}: Method` is not satisfied
  --> src/main.rs:26:41
   |
26 |     let clone_method: Box<dyn Method> = Box::new(clone);
   |                                         ^^^^^^^^^^^^^^^ the trait `Method` is not implemented for `for<'r> fn(&'r u32) -> String {main::clone}`
   |
   = note: required for the cast to the object type `dyn Method`

We have the same problem.

Notice the syntax fn(&'r u32) -> String {main::clone} The additional block statement is the important part. Each function item gets a unique anonymous type, as does each closure. Both can be coerced into a function pointer. But this coercion doesn't happen automatically for resolving trait bounds.

This error is particularly tough to spot. And on first glance it might look like the trait is not implemented correctly. There is actually an open issue about the error messages here: https://github.com/rust-lang/rust/issues/62385.

The fix is to write:

let clone_method: Box<dyn Method> = Box::new(clone as fn(&String) -> String);

But we would be punting this work to the end users of our library, and that didn't seem acceptable

Attempt #128

Eventually, we landed back at something more similar to the unstable Rust Fn trait, where the trait has a generic Args which is a tuple of the input arguments.

pub trait Method<T, Args = ()> {
    type Result;

    fn invoke(&self, receiver: &T, args: Args) -> Self::Result;
}

You may still be wondering why we need the extra receiver: &T argument? This is due to limitations in how we can express tuples of types.

What we would like to have is a single trait Function<Args> – where Args is a tuple – then define Method as something like trait Method<T, Args>: Function<(&T, ...Args)> – where ...Args is some kind of tuple-spread operator.

There's an open issue for something like this: Variadic generics. But so far every attempt at an RFC has been closed or postponed.

So until then — or until we decide to implement it using something like heterogenous lists with frunk — we'll do it the long way.

Step 3: Implementing our method trait

We're almost at the finish line. We just need to implement our new trait for every possible number of arguments.

For example, the two-argument version looks like:

impl<F, R, T, A, B> Method<Receiver, (A, B)> for F
where
    F: Fn(&T, A, B) -> R,
{
    type Result = R;

    fn invoke(&self, receiver: &T, args: (A, B)) -> Self::Result
        (self)(receiver, args.0, args.1)
    }
}

This is pretty simple, but it's also exactly the kind of thing that a macro can make less painful. The final version looks like:

macro_rules! tuple_impls {
    ( $( $name:ident )* ) => {
        impl<Fun, Res, Receiver, $($name),*> Method<Receiver, ($($name,)*)> for Fun
        where
            Fun: Fn(&Receiver, $($name),*) -> Res + Send + Sync + 'static,
        {
            type Result = Res;

            fn invoke(&self, receiver: &Receiver, args: ($($name,)*)) -> Self::Result {
                #[allow(non_snake_case)]
                let ($($name,)*) = args;
                (self)(receiver, $($name,)*)
            }
        }
    };
}

tuple_impls! {}
tuple_impls! { A }
tuple_impls! { A B }
tuple_impls! { A B C }
// .. more macro invocations follow
// we support method arities up to 16
tuple_impls! { A B C D E F G H I J K L M N O P }

There are actually a few more pieces we need. For example, how do we go from Vec<PolarValue> to (A, B, ..)? More traits and more macros!

The end result is exactly the experience we were looking for: you can pass in Rust closures, or functions:

Cat::get_polar_class_builder()
   .add_method("feed", Cat::feed)
   .build()

Comparisons

It turns out that if you build a language on top of Rust, there's a good chance you'll end up implementing this trait. For example:

Extensions

Class Methods

We also want to support class methods:

happy(cat: Cat, food) if cat.feed(food) != Cat.meow()

Here, the type check cat: Cat is using the symbol Cat as a class tag, whereas anywhere else, Cat is a variable name bound to the Cat class. We borrowed this pattern from the corresponding Python implementation, where we bind the variable to the value of Cat, i.e. <class 'Cat'>. Since Python is dynamic and we don't go through the same class registration steps as we do in Rust, it all just works.

But we can apply the same general idea here, and create our own metaclass. And it works.

First, we need to register Class as the type metaclass:

impl crate::PolarClass for Class {}

fn metaclass() -> Class {
    Class::builder::<Class>().name("oso::host::Class").build()
}

To then dispatch class methods, we need to make sure that "instance methods" of the Class metaclass take the remaining input arguments (the receiver type here being Class), and call the actual class method.

For example: on invoking Cat.meow, this resolves to a method "meow" on Cat. Cat is a constant, and is an instance of Class. So "meow" is an instance method on oso::host::Class. And we have a special hook on the instance method dispatch, which checks whether the current class is Class :

fn get_method(&self, name: &str) -> Option<InstanceMethod> {
    if self.type_id == TypeId::of::<Class>() {
        // all methods on `Class` redirect by looking up the class method
        Some(InstanceMethod::from_class_method(name.to_string()))
    } else {
        self.instance_methods.get(name).cloned()
    }
}

Which means the Cat::meow class method gets returned instead of an instance method.

Future Work

Our implementation is in a pretty good spot for most use cases! You can see the latest version of the Oso Rust crate here.

But there's a lot more we can do from here. First, we can add support for other variants of methods. We didn't cover it here, but we currently support returning results, options, and iterators, the last of which requires using a special add_iter_method. We could also support async methods.

Furthermore, we currently restrict users to defining one method per name, even though that's not a restriction that Polar enforces. Because of the approach we took above, it would be challenging to support fully generic methods (like our ident<T> method), but we could make it possible to add a method for multiple types. We could even add support for variadic arguments if we wanted!

Summary

This concludes our series on building a runtime reflection in Rust!

We started out slowly in Part 1, leaning on Rust's built-in support for runtime dynamic types through the Any trait as the foundation of our class system. With that in place, we could do some simple runtime type checking. ✅

In Part 2, we looked at pulling attributes off structs at runtime. Here we had a terrible choice to make: require users to add #[repr(C)] to all their structs so we could do ~~dark incantations~~ pointer arithmetic and read fields dynamically? Or make attribute accessors an explicit opt-in? The latter seemed like more of a Rust approach, and we took it.

Which leaves us here in Part 3. We built on our attribute getter approach, turning it into an instance method and allowing inputs. We earned the final piece with a few scars to show for it, and we have our own runtime reflection system for Rust.

If you're interested in learning more about Oso and how we use Rust to build our language and open source policy engine, here are some handy links:

We are also hiring for software engineers, engineering managers, developer advocates and technical writers.

Want us to remind you?
We'll email you before the event with a friendly reminder.

Write your first policy