How to create a list of trait

Mojo
10

Here is an approach using a trait object and dynamic dispatch.

from memory import ArcPointer

trait Op(Copyable, Movable):
	fn calc(self) -> Int:
		pass

fn calc_trampoline[T: Op](ptr: ArcPointer[NoneType]) -> Int:
    var data = rebind[ArcPointer[T]](ptr)
    return data[].calc()

fn clear_trampoline[T: Op](var ptr: ArcPointer[NoneType]):
    var data = rebind[ArcPointer[T]](ptr^)
    _ = data^

struct DynOp(Op):
	var data: ArcPointer[NoneType]
	var calc_func: fn(ArcPointer[NoneType]) -> Int
	var clear_func: fn(var ArcPointer[NoneType])
	
	fn __init__[T: Op](out self, var ptr: ArcPointer[T]):
		self.data = rebind[ArcPointer[NoneType]](ptr^)
		self.calc_func = calc_trampoline[T]
		self.clear_func = clear_trampoline[T]
	
	fn calc(self) -> Int:
		return self.calc_func(self.data)
	
	fn __del__(var self):
		self.clear_func(self.data^)



@fieldwise_init 
struct Add(Op):
	var a: Int
	var b: Int

	fn calc(self) -> Int:
		return self.a + self.b

@fieldwise_init  
struct Mul(Op):
	var a: Int
	var b: Int

	fn calc(self) -> Int:
		return self.a * self.b


def main():
	var add = ArcPointer(Add(1,2))
	var mul = ArcPointer(Mul(3,4))
	var l: List[DynOp] = [DynOp(add), DynOp(mul)]

	for op in l:
		print(op.calc())

Notice that this implementation is open compared to the one using variant in the sense that you can add new types that conform to the Op trait without changing existing code.

4 Likes