valand.devworksarticles

Rust: Type Concealment With Any Trait and FnMut

While working on machine-tree I found that Rust has a very powerful arsenal when it comes to dynamic programming and hiding type information.

First of all, what's machine-tree? It's a project aiming to have a React-like library on Rust. But unlike React which is focused on GUI, machine-tree is focused on the tree aspect. Think of a couple of machines structured in the form of a tree each capable of running independently. It is still experimental and suboptimal, so please don't mind the code.

In a React-like construct, the library caller dictates the form inside each node in the tree. The blueprint of the node is called a Component. While at the same time the library manages the tree's data structure in the memory and schedule when each node is executed.

As Above So Below

As Above, So Below

Now we have two parties with two different interests. One is the library and the other is the caller. These two different interests are represented as two different data types.

  • The caller sees a node, which is an instance of a component, as a program. It has an input from which the node will get the value from the outside world. It has an output from where the type is restricted to be components too, which later will be treated as the nodes' children. These children nodes can be a different type of component with a different type of input.

  • The library sees nodes as runnable blobs. All nodes are uniformly typed. A node is a program that receives an input of an unknown type and outputs other nodes.

Because of these two different interests, I make two representing types for a node. The first representation is a component trait:

pub trait Component
where
    Self: Sized + 'static,
{
    type Input: Sized + Clone + 'static;

    fn construct(input: &Self::Input) -> Self
    where
        Self: Sized + 'static;

    fn step(
        &mut self,
        control: &mut dyn Control,
        input: &Self::Input
    ) -> Vec<Node>;

    fn seed(input: Self::Input, key: String) -> Node {

        // convert the component into a thing called 
        // StepFn, which relates a lot to `step` fn above
        let step_fn: AbstractStepFnBrc =
            Box::new(RefCell::new(component_utils::generate_abstract_step_fn::<
                Self,
                Self::Input,
            >(&input)));
        let input: AbstractInputBox = Box::new(input);

        Node {
            // there are other fields, but hidden due to irrelevancies
            input,
            step_fn,
        }
    }
}

The Component is a trait with the Input associated type. A struct can implement the Component trait, which then will have to also implement construct and step functions. The construct function is used to create the struct when supplied with an input. The step is the method to be run when it is the node's turn to run. (When a node runs is scheduled by the library). The last function is called seed, which is typically used in the step function by calling OtherComponent::seed(some_input), which results in a child node.

The step function is similar to React's render function. seed function is similar to React.createElement or simply <OtherComponent /> in JSX.

The second representation is a Node struct.

pub struct Node {
    // there are other fields, but hidden due to irrelevancies
    pub(crate) input: AbstractInputBox,
    pub(crate) step_fn: AbstractStepFnBrc,
}

// where
pub type StepFn<Input> = Box<dyn FnMut(&mut dyn Control, &Input) -> Vec<Node>>;
pub(crate) type AbstractStepFn = StepFn<AbstractInputBox>;
pub(crate) type AbstractStepFnBrc = Box<RefCell<AbstractStepFn>>;
pub(crate) type InputBox<Input> = Box<Input>;
pub(crate) type AbstractInputBox = InputBox<dyn Any>;

Node is how the library sees an instance of a component. A tree simply contains many nodes. In this case, a node has an input and a step_fn. Both are boxed because they have varying sizes and thus need to be put in the heap. Both use Any which helps the library to ignore the actual type of Component::Input and Component::step.

In StepFn's case, it is a dyn FnMut which is a little powerful type representing a micro stateful micro machine which I will explain in more detail later. StepFn receives Control and Input. In AbstractStepFnBrc, the StepFn is wrapped into another box and a refcell which allows us to mutate what is inside of it. We can ignore Control now because it's irrelevant to this article.

Take notice that definitions of StepFn<Input> in the Node struct and step method in the Component trait are similar.

Also take notice of the default implementation of the seed function, which calls component_utils::generate_abstract_step_fn. This particular function converts component_utils::generate_abstract_step_fn.

This way, the library can do whatever it wants with Node without worrying about its internal type. Because this is exactly what a library needs to have and be -- an abstraction layer.

FnMut, Closure, A Stateful Micro Machine

A little detour here on how FnMut is such a powerful construct in Rust. Here's a little code you can access in this Rust Playground.

use std::{any::Any, cell::RefCell};

fn main() {

    // define the FnMut here
    let mut some_stored_number: u32 = 0;
    let fnmut: Box<RefCell<Box<dyn FnMut(u32) -> u32>>> =
        Box::new(RefCell::new(Box::new(move |operand: u32| {
            some_stored_number = some_stored_number
                .checked_add(operand)
                .unwrap_or(some_stored_number);
            some_stored_number
        })));

    // Make the FnMut Any
    let fnmut: Box<dyn Any> = fnmut;

    // Downcast to identify the FnMut as dyn FnMut(u32) -> u32
    let fnmut = fnmut
        .downcast::<RefCell<Box<dyn FnMut(u32) -> u32>>>()
        .unwrap();

    // Borrow from RefCell
    let mut fnmut_refmut = (*fnmut).borrow_mut();

    // Fire the FnMut
    println!("{:?}", fnmut_refmut(1));  // 1
    println!("{:?}", fnmut_refmut(1));  // 2
    println!("{:?}", fnmut_refmut(2));  // 4
    println!("{:?}", fnmut_refmut(3));  // 7
}

In essence, this piece of code does a couple of things:

  1. Create an FnMut(u32) -> u32, a mutable function, with a closure.

  2. Erase its type into a Box<dyn Any>.

  3. Make the program convert it back to its original type.

  4. Fire it.

When fired, you can observe that it stores the previous number.

  • The first call returns 1 from 0 (previously stored number) + 1 (argument)

  • The second call returns 2 from 1 (previously stored number) + 1 (argument)

  • The third call returns 4 from 2 (previously stored number) + 2 (argument)

  • The fourth call returns 7 from 4 (previously stored number) + 3 (argument)

This showcases how FnMut is a stateful micro-machine. It is a powerful arsenal that is utilized heavily in machine-tree.

The Type Abstraction

Let's get back to this function.

component_utils::generate_abstract_step_fn

The general idea of this function is to convert a Component, including both its input and its struct instance into a StepFn. So first let's define its type-level requirements.

fn generate_abstract_step_fn<Machine, Input>(input: &Input) -> AbstractStepFn
    where
        Input: Sized + Clone + 'static,
        Machine: Sized + Component<Input = Input> + 'static

First, it requires two type parameters, one is a type that implements the Component trait that I dearly call Machine, and the second is Input which is assigned as Machine::Input. The function also receives an input because we want to construct the Machine inside this function. Last, is the return type of the function, which is AbstractStepFn.

Now get to the content of the function.

{
    let mut self_state = Machine::construct(input);
    Box::new(move |control: &mut dyn Control, input: &AbstractInputBox| {
        let input_ref = input.downcast_ref::<Input>().unwrap();
        self_state.step(control, input_ref)
    })
}

It may seem simple but it might take some time to process. The first line is the construction of the Machine while the rest of the function is the creation of the closure/FnMut. The construction is executed instantly, but the code inside the closure is not.

The code inside the closure will be called when the StepFn is being fired by the library -- a subject to the library's scheduling method. Let's dissect the closure:

  • The closure alongside the box is the StepFn. It receives a Control and an AbstractInputBox, which is essentially Box<dyn Any>.

  • The next line attempts to downcast the AbstractInputBox into the proper type. If the Input is u32 so it will be downcasted into u32, if it's f32 it will be downcasted into an f32, and so on. Now you may think that theoretically you can supply any box in there and the unwrap call will cause panic throughout the whole application. Yes, that is possible! This is why the library's correctness should be verified to always supply a box with the correct type. How? The Node struct is the answer. The Node struct should always be constructed so that the input field and the step_fn field are from a matching Machine so that the type is always correct.

  • The last line in the closure is the self_state calling the step method. If we revisit the step method from the Component trait, we can see that it receives an input with an exact type, not the abstract one. This is why the downcasting in the previous line is important.

Sneak Peek Into machine-tree's Inner Working

The scheduling mechanism of the machine-tree is basically how the early version of React does it, without the fiber mechanism. It's quite ancient and I don't want to expose too much here. But the essence is, whenever a node decides that it requires to be "refired" (or in React's terminology "rerendered"), the scheduler will fire that specific node and then subsequently the nodes below it in breadth-first order. There's no memoization mechanism right now either.

Within the firing mechanism of one node. This is how it roughly looks:

local_render_queues
    .for_each(|currently_rendered_node_rcc|{
        // borrow the node
        let mut node_borrow_refmut = 
            currently_rendered_node_rcc.borrow_mut();
        let node_borrow = &mut *node_borrow_refmut;

        // borrow the step_fn
        let step_fn_borrow = &mut *(node_borrow.step_fn.borrow_mut());

        // create the control
        let control = Control {
            // ...fill fields
        };

        // Call the borrowed step_fn with control 
        // and input from node_borrow
        let children_nodes: Vec<NodeRcc> = 
            (step_fn_borrow)(&mut control, &node_borrow.input);
  })

A mutable step_fn reference is borrowed and then later called with the input. In this situation, the library can't afford to have nodes with all type information intact because then it will have to manage which input goes to which node. With all type information gone, this code can simply call (step_fn_borrow)(&mut control, &node_borrow.input). At the same time, a Component author also can't afford to write the downcasting. Because... ergonomics!

This highlight how neat the component_utils::generate_abstract_step_fn can be to transform step function to an abstract one.

The complete code can be found in my repo (GitHub - Kelerchian/machine-tree)