Simple Arrays

In this example, we will show how to start building a simple Array API.

The important things to note are:

1. Our “optimization” replacement isn’t tied to any particular type of array. It is agnostic to how we actually represent arrays (in this case strings).

2. We end up with a Python class that acts like an array, but builds up an expression graph. We use this same class to build up the replacements, so they are very succinct and readble.

3. There is no metaprogramming in this example, and no inspecting of type signatures.

4. Everthing is well typed and MyPy will warn us if we are trying to call operations on arrays which are not valid.

We only define two operations, add and getitem:

from __future__ import annotations

import metadsl
import metadsl_rewrite
import typing

class Array(metadsl.Expression):
    @metadsl.expression
    def __add__(self, other: Array) -> Array:
        ...

    @metadsl.expression
    def __getitem__(self, idx: int) -> Array:
        ...

We define one replacement, that distributes indexing over addition:

register_array = metadsl_rewrite.register['array']

@register_array
@metadsl_rewrite.rule
def _distribute_add(a: Array, b: Array, c: int):
    return (
        # expression to match agains
        (a + b)[c],
        # expression it is replaced with
        lambda: a[c] + b[c]
    )

Now we define a simple “implementation” of arrays, that are just strings representing their names.

@metadsl.expression
def array(x: str) -> Array:
    ...

a = array("a")
a

__main__.Array(array, ['a'], {})

We write how to convert adding and indexing into strings, to “compile” the expression:

register_array_string = metadsl_rewrite.register['array']


@metadsl.expression
def as_string(x: Array) ->  str:
    ...

@register_array_string
@metadsl_rewrite.rule
def _string(x: str):
    return as_string(array(x)), lambda: x


@register_array_string
@metadsl_rewrite.rule
def _add(x: str, y: str):
    return (
        array(x) + array(y),
        lambda: array(f"({x} + {y})")
    )


@register_array_string
@metadsl_rewrite.rule
def _getitem(x: str, y: int):
    return (
        array(x)[y],
        lambda: array(f"{x}[{y}]")
    )

Now we can create two arrays, add them and index, and turn them into a string, which will run the replacments:

a = array("a")
b = array("b")

res = (a + b)[10]
print(metadsl_rewrite.execute(as_string(res)))

(a + b)[10]