How can I have pointer of trait object like rust ?

[liuzhishan]

Hi, I have a question abound trait. How can I have a pointer of a trait object like in rust, to achieve dynamic dispatch?

For example, I have a trait in rust

trait Animal {
    fn speak(&self);
}

then I can use Box<dyn Animal> to get a heap allocated object which implement the Animal trait. Similarly in c++, where I can use std::shared_ptr<Base> to get pointer of subclass.

I search the document of mojo, but only found doc of trait and generic of trait. And I found the discussion aboud Mojo and Dynamism(#466), seems dynamism maybe achieved through decrators like dynamic.

Is the dynamic dispatch available now, or is it in plan ?

[liuzhishan]

Hi, I just found a way to achieve dynamic dispatch in mojo. The main idea is to combine the Variant and trait in Mojo. With just a little runtime overhead, the performance should be very good.

I have put the code on my github: https://github.com/liuzhishan/mojo-dynamic-dispatch.

Implemention

A Variant is a sum type, such as enum in rust, or tagged union in c++ (or std::variant in c++17). It can hold a value that could take on several different, but fixed types.

We can define several struct implementing one same trait, such as Echoable. And define a Variant from all these types, which is constrained by one same trait. Then we implement the trait Echoable for the Variant.

For example:

struct MyVariant[*Ts: Echoable](
    Echoable,
):
    ...
    fn echo(self) -> String:
        var res: String = String("")
        var cur = self._get_state()[]

        @parameter
        fn each[i: Int]():
            if cur == i:
                alias T = Ts[i]
                res = self[T].echo()

        unroll[each, len(VariadicList(Ts))]()

        return res

As the above code shows, when we have a Variant at runtime, we known its underlying type. We can compare it with all the candidate types, and then get the actual object, call the trait function. This has exactly the same effect of dynamic dispatch in rust or c++.

The comparing between types is actually comparing using an int8, which is determined by the order of types in compile time. And loop of the types is unrolled in compile time using @parameter meta programming. At runtime, there are just some int8 comparing to determine which object to use. Only Exactly one of them will be true. I'm not sure how fast it is comparing to v-table in c++ or rust, but it should be really fast with only some int8 comparing.

Wow, awesome mojo! We just implement dynamic dispatch in a library, not in the compiler! Just as mojo's design principle said, implement the language in library, not in compiler.

Right now we cannot use trait as generic parameters. So Variant cannot be reused directly. Most of code are same as implementation in mojo stdlib/src/utils/variant.mojo, so I just copied the code. Maybe when macro are supported in the future, we can have a more generic implementation.

Usage

You can use the same idea to implement your dynamic dispatch (like subtype in c++).

For example:

alias AOrB = MyVariant[A, B]
var a = A()
var b = B()

var x = AOrB(a)
var y = AOrB(b)

# result: this is A
print(x.echo())

# result: this is B
print(y.echo())

[johnsoez4]

Hi liuzhishan!
I am looking for a way to dynamic dispatch in Mojo and found your example here, but is seems that changes under the hood of Mojo has broken your code? I'm using nightly/version mojo 2024.7.2911 from today.
What do you think?

P.S. on Discord, I am EzRyder (ezryder_ed, ID 771688661522841601)

[johnsoez4]

Hello @liuzhishan!
For example, the image below shows what I'm seeing in the latest nightly/version, build mojo 2024.8.105 (1fa11138).

These issues are due to changes that are present in the current nightly build, and therefore are coming to the next released version of Mojo. Here's a link to the message on Discord channel `nightly' announcing the latest nightly build: https://discord.com/channels/1087530497313357884/1224434323193594059/1268449570128662539

If you're using the released version of Mojo today, then the next released version will break your code. You can install the nightly build instead of the released version of Mojo, and if you're using VS Code, install the Mojo 🔥 (nightly) extension. You will need to disable or uninstall the VS Code extension Mojo 🔥 which corresponds to the released version of Mojo.
Hope this helps!

https://docs.modular.com/mojo/manual/get-started#next-steps

modular uninstall mojo
modular install nightly/mojo
modular update nightly/mojo

or

https://docs.modular.com/max/install#install-max

modular uninstall max
modular install nightly/max
modular update nightly/max

[johnsoez4]

Here a link to the changelog describing the unreleased changes. https://github.com/modularml/mojo/blob/nightly/docs/changelog.md

Screenshot of the changes causing the import errors shown in my previous comment.

[liuzhishan]

Hi, @johnsoez4, thanks for the link and the detailed step. I have check the latest stdlib of nightly/mojo, and found the reason is some change of the UnsafePointer and Variant.

1. initialize_pointee_move has become a member function named init_pointee_move of UnsafePointer.
2. move_from_pointee has become a member function named take_pointee of UnsafePointer.
3. move_pointee has become a member function named move_pointee_into of UnsafePointer.
4. _get_state of Variant bas been renamed to _get_discr, which is in the file stdlib/src/utils/variant.mojo.
5. Some implementation of Variant has changed.

I have add a working version of the nightly mojo, named main_nightly_2024_8_305.mojo. Please have a try.

[johnsoez4]

I just confirmed that your new file, main_nightly_2024_8_305.mojo, works for me. Today I am on build mojo 2024.8.405
Thanks @liuzhishan!

[liuzhishan]

You are welcome.