Concepts

metadsl inserts a layer between calling a function and computing its result, so that we can build up a bunch of calls, transform them, and then execute them all at once.

Expression

We start with a building block Expressions:

from __future__ import annotations

import dataclasses
import metadsl

class MyObject(metadsl.Expression):
    @metadsl.expression
    def do_things(self) -> MyObject:
        ...

    @metadsl.expression
    def __add__(self, other: MyObject) -> MyObject:
        ...


@metadsl.expression
def create_object(x: int) -> MyObject:
    ...

o = create_object(123)
o

__main__.MyObject(create_object, [123], {})

o.do_things() + o

__main__.MyObject(MyObject.__add__, [__main__.MyObject(MyObject.do_things, [__main__.MyObject(create_object, [123], {})], {}), __main__.MyObject(create_object, [123], {})], {})

It is useful to keep in mind the strict typing constraints here, not all of which can be faithfully checked by MyPy:

1. The arguments in a Expression should fit the signature of the _function.

2. The return type of the _function should correspond to the type of the Expression.

Let’s check the first two of these for o. We see that the expression’s function is create_object, which takes in an int and returns a MyObject. What is metadsl.E? It is a type alias for Union[T, metadsl.LiteralExpression[T]]. This represents anything that could be a python literal, or a leaf of the expression tree. The argument here is an int, which is compatible with the argument hint. And the instances holding it is of type MyObject, which is its return type.

Rewrite Strategies

We can define a possible rewrite rule for a single expression, and then have that execute repeatedly on the graph.

We can also combine several strategies and execute them together:

import metadsl_rewrite

@metadsl_rewrite.rule
def _add(x: int, y: int):
    return create_object(x) + create_object(y), lambda: create_object(x + y)

@metadsl_rewrite.rule
def _do_things(x: int):
    return create_object(x).do_things(), lambda: create_object(x * 2)


inner_strategy = metadsl_rewrite.StrategySequence(
    _add,
    _do_things
)
strategies = metadsl_rewrite.StrategyRepeat(
    metadsl_rewrite.StrategyFold(
        inner_strategy
    )
)

The requirements for these replacements is that they take in some arguments, which can match any expression in the graph. Their first return value builds up a template expression based on the inputs, that shows what it should match again. The second is a is the resulting to replace it with. Note that both should have the same type, because if you replace an expression it should not invalidate the type of something it is a part of:

We can call these on an expression and it will return a replaced version of it:

print(o+o)
print(metadsl_rewrite.execute(o + o, strategies))

MyObject.__add__(create_object(123), create_object(123))
create_object(246)

We can also have a rule that unwraps the int from the object:

@metadsl.expression
def unwrap_object(o: MyObject) -> int:
    ...

@metadsl_rewrite.rule
def _unwrap_object(i: int):
    return unwrap_object(create_object(i)), i

inner_strategy.strategies = inner_strategy.strategies + (_unwrap_object,)

metadsl_rewrite.execute(unwrap_object(o + o), strategies)

246

This is nice, because now we can write our unboxing as a replacement, which means it’s nice and type safe.