Mojo🔥の型引数で遊ぼう

# TL; DR

```python
# 型レベルで文字列結合！
let _foo: TStr["Hello, Mojo🔥"] = add(TStr["Hello, "](), TStr["Mojo🔥"]())
```

```python
# 一致しなければコンパイルエラー
let _foo: TStr["Hello, Mojo🔥"] = add(TStr["Hello, "](), TStr["Mojo🥺"]())
```

```
error: Expression [3]:26:44: 'TStr["Hello, Mojo\F0\9F\A5\BA"]' value cannot be converted to 'TStr["Hello, Mojo\F0\9F\94\A5"]' in 'let' initializer
```

- 型引数(type parameter)を使うと...
    - コンパイル時計算が簡単にできる
    - 型レベルプログラミングもできる？（現状かなり制約あり）

# はじめに

先日公開された[^earlyaccess]プログラミング言語Mojo、「Python互換の構文でPythonより35000倍速い」の謳い文句[^35000]や、「`.🔥`」という~~ネタ~~拡張子もあいまって連日話題になっています。

https://www.modular.com/mojo

しかし、Mojoの真の力は **型引数でコンパイル時計算ができること** です！

**Pythonの表現力の高い構文とコンパイル時計算が出遭ったらどうなってしまうのでしょうか？** 本記事では、現時点でメタプログラミングできたもの、できなかったものを紹介したいと思います。

:::note warn
以下の内容は執筆時点 (2023/5/17) での動作結果です。今後のリリースで変更される可能性があります！
:::


# Mojoの構文
メタプログラミングの前に、Mojo独自の構文について紹介します（所有権等については割愛）。

## `alias`

コンパイル時定数や型を宣言するためのキーワードです。

```python
from List import VariadicList

# 定数
alias x = 1
# 型
alias IntList = VariadicList[Int]
```

（ちなみに、実行時の変数は `let` (不変), `var` (可変)で宣言できます）

定数なので、変数をもとに生成することはできません。

```python
let runtime_var = 3
alias comptime_const = runtime_var * 2
```

```
error: Expression [7]:6:24: use of unknown declaration 'runtime_var'
alias comptime_const = runtime_var * 2
                       ^~~~~~~~~~~
```

## `fn`

`def` と互換性のある、関数/メソッド宣言キーワードです。速度が速い代わりに以下の制約があります。

- 引数は原則immutable
- シグネチャには型（構文は型ヒントと同様）が必須
- ローカル変数は `let`, `var` で明示的に宣言が必要
- 例外を発生させる場合シグネチャに `raises` が必要

```python
# Int 型については後述
fn double(n: Int) -> Int:
    return n * 2

# 例外を発生させる場合
fn crash() raises:
    raise Error("error!")
```

https://docs.modular.com/mojo/programming-manual.html#fn-definitions

さらに、関数やメソッドはオーバーロードが可能です。

```python
# NOTE: 例として2つ書いているだけで、実際には型引数でジェネリックに書いたほうがシンプル
fn add(a: Int, b: Int) -> Int:
    return a + b

fn add(a: F32, b: F32) -> F32:
    return a + b

print(add(1, 2))
print(add(2.5, 1.5))
```

## `struct`

`class`  に似ていますが、メソッドや属性が不変（コンパイル時に決定）なため実行速度が上がっています。

https://docs.modular.com/mojo/programming-manual.html#struct-types

```python
struct Person:
    # 事前に属性の宣言が必要
    var name: StringLiteral
    var age: Int
    
    # fnで引数を可変に扱う場合は inout の指定が必要
    fn __init__(inout self, name: StringLiteral, age: Int):
        self.name = name
        self.age = age

    fn greet(self):
        # f-stringは未実装
        print("I'm", self.name)

    fn can_drink(self) -> Bool:
        return self.age >= 20

# インスタンス生成方法はクラスと同じ
let taro = Person("Taro", 25)
taro.greet() # I'm Taro
print(taro.can_drink()) # True
```

また、Pythonの組み込み型に対応するstructも用意されています。 （例： `int` -> `Int`, `str` -> `StringLiteral`）
使えるメソッドが制限されている代わりに高速です。

https://docs.modular.com/mojo/MojoBuiltin/Int.html

https://docs.modular.com/mojo/MojoBuiltin/StringLiteral.html

## 拡張された型引数

そして、本題の型引数(type parameter)です。[PEP 695](https://peps.python.org/pep-0695/) の型ヒントを拡張したもので、

- 単なるアノテーションではなくコンパイル時に評価、検査される
- **値を入れることもできる**

という特徴を持ちます。

例えば、配列の長さを型引数に渡すことで、要素数をコンパイル時に検証可能です。

```python
# https://docs.modular.com/mojo/notebooks/HelloMojo.html#parameterization-compile-time-meta-programming より
let a = MySIMD[4](1, 2, 3, 4)
let b = MySIMD[3](5, 6, 7)

# 長さを間違えると...
let c: MySIMD[6] = a.concat(b)
```

```
error: Expression [8]:25:32: 'MySIMD[7]' value cannot be converted to 'MySIMD[6]' in 'let' initializer
    let c: MySIMD[6] = a.concat(b)
                       ~~~~~~~~^~~
```

# コンパイル時計算

関数内で `@parameter` を使うことで、型引数を **コンパイル時計算に使うことが可能です**。実質 C++のconstexprです。

```python
struct MyInt[n: Int]:
    alias value = n

fn fact[n: Int]() -> Int:
    # 型引数を使ってコンパイル時に戻り値を計算
    # NOTE: @parameterとifを同時に使う必要がある
    @parameter
    if n <= 0:
        return 1
    else:
        return n * fact[n-1]()

print(fact[6]()) # 720
```

自由度が高く、静的に決まる値であればほぼなんでもできます。~~コンパイル時に計算完了していれば35000倍越えも夢じゃない~~

詳細についてはこちらの記事が分かりやすかったです。

https://ikesyo.hatenablog.com/entry/2023/05/07/160052

# 型レベルプログラミングの試み

~~`@parameter` 、すごく便利だけど自由度が高すぎて黒魔術感があまりない~~

というわけで、次にやるのは型レベルプログラミングです。残念ながら、現状では機能が揃っておらずあまり表現力はありません...

## 型レベル足し算

発想としては先述の「長さを型情報に持てる配列型」と同じです。値を格納するためのラッパーstructを用意し、関数で型引数どうしの関係を定義します。
（言い換えると、「関数のシグネチャ」を関数として利用します）

```python
struct TInt[n: Int]:
    # 関数の引数に渡すにはインスタンスが必要なため__init__定義
    fn __init__(inout self):
        pass


fn add[n: Int, m: Int](a: TInt[n], b: TInt[m]) -> TInt[n+m]:
    return TInt[n+m]()


# 型レベル 3 + 2 == 5
let _n: TInt[5] = add(TInt[3](), TInt[2]())
```

文字列を使えば冒頭の実装になります。

```python
struct TStr[s: StringLiteral]:
    fn __init__(inout self):
        pass


fn add[n: StringLiteral, m: StringLiteral](a: TStr[n], b: TStr[m]) -> TStr[n+m]:
    return TStr[n+m]()


let _foo: TStr["Hello, Mojo🔥"] = add(TStr["Hello, "](), TStr["Mojo🔥"]())
# let _foo: TStr["Hello, Mojo🔥"] = add(TStr["Hello, "](), TStr["Mojo🥺"]()) # error: Expression [3]:26:44: 'TStr["Hello, Mojo\F0\9F\A5\BA"]' value cannot be converted to 'TStr["Hello, Mojo\F0\9F\94\A5"]' in 'let' initializer
```

## 型レベルリスト

structのaliasを使って、型引数を型として取り出すことができます。こうすることで、型引数の一部分を取り出す計算が可能です。

```python
@register_passable
struct Cons[car: Int, cdr: AnyType]:
    # 型引数を使って型を定義
    alias Car = car
    alias Cdr = cdr

# 型引数はコンパイル時に決まりインスタンスと独立して参照できる！
# (car '(1))
let _x: TInt[1] = TInt[Cons[1, NoneType].Car]()
# (cdr '(1 2))
let _y: TType[Cons[2, NoneType]] = TType[Cons[1, Cons[2, NoneType]].Cdr]()
# (car (cdr '(1 2))
let _z: TInt[2] = TInt[(Cons[1, Cons[2, NoneType]].Cdr).Car]()
```

こちらの記事に着想を得て作成しました。他にも面白いテクニックがたくさん紹介されていて圧巻です。

https://zenn.dev/yyu/articles/1eefb8f547dc1b#%E5%9E%8B%E3%83%AC%E3%83%99%E3%83%AB%E3%83%AA%E3%82%B9%E3%83%88

本家Rust版と異なり、 cdrに `Cons` `NoneType` 以外の任意の要素を入れられてしまうのが欠点です。`Cons` と `Nil`(`NoneType`の代わり) が共通のstruct `TList` を継承し、cdrに `TList` しか入らないようにすれば解決できるのですが、現在Mojoでは継承が未実装です...

## 型レベル条件分岐
型引数には3項演算子も入れられるので、型引数の値に応じて異なる型を導出することが可能です[^difftype]。内包表記や継承が使えるようになれば応用先が広がりそうです。

```python
struct TTrue:
    fn __init__(inout self):
        pass

struct TFalse:
    fn __init__(inout self):
        pass

struct IsEven[n: Int]:
    alias Type = TTrue if n % 2 == 0 else TFalse

let _t: TFalse = IsEven[7].Type()
let _s: TTrue = IsEven[8].Type()
```

## できなかったこと
一方、型レベルプログラミングに必須の再帰が封じられていました。
（言語処理系としては、コンパイルの停止を保証できるので安全です :grin:）

### :x: aliasの再帰的定義
aliasと3項演算子を見て真っ先に思い付いたアイデアです。

```python
struct Fact[n: Int]:
    alias value = Fact[n-1].value * n if n > 0 else 1
```

```
error: Expression [27]:10:28: recursive reference to declaration
    alias value = Fact[n-1].value * n if n > 0 else 1
                           ^

Expression [27]:10:5: previously used here
    alias value = Fact[n-1].value * n if n > 0 else 1
    ^
```

対策済みでした :innocent:
(別のaliasやstructを経由して間接的に参照してもエラーになりました。賢い)

`globals` を使って動的に参照するのも検討しましたが、そもそも `globals` が実装途中のようでした...

```python:globalsが動かない
# PythonのモジュールをMojoにimport
# ビルトイン関数は "builtins" 指定
let py = Python.import_module("builtins")


py.print(py.globals())
```

```
Error: <built-in function globals> returned NULL without setting an error
```

### :x: 型引数に関数呼び出しを記述

```python
fn fact[n: Int]() -> Int:
    @parameter
    if n <= 0:
        return 1
    else:
        return n * fact[n-1]()

struct MyInt[n: Int]:
    alias value = n
    fn __init__(inout self): pass

struct Fact[n: Int]:
    alias value = fact[n]()

let m: MyInt[720] = MyInt[Fact[6].value]()
```

```
error: Expression [9]:30:45: 'MyInt[fact[6]()]' value cannot be converted to 'MyInt[720]' in 'let' initializer
    let m: MyInt[720] = MyInt[Fact[6].value]()
                        ~~~~~~~~~~~~~~~~~~~~^~
```

Mojoは型引数を完全に評価するわけではなく、インスタンス化の手前で打ち切ってしまうようです。

https://docs.modular.com/mojo/roadmap.html#the-rebind-builtin

そのため、関数呼び出しや要素参照等を行うと同値判定ができませんでした。
ちなみに、上記の `rebind` 関数は型を明示的にキャストするだけなので、型レベルプログラミングには使えません。残念。

## 今後の展望
現状できることは限られていますが、今後実装される機能でまだまだパワーアップの余地があります。最後のピースがハマった暁には、真の型レベルプログラミングができるようになるでしょう。

### 実装済みだけどまだ動かない？標準関数
#### assert_param_msg

型引数 `cond` が `False` のとき、**コンパイル時に** メッセージを出力する関数です。言い換えるとコンパイル時printです。

https://docs.modular.com/mojo/MojoStdlib/Assert.html#assert_param_msg

```python
from Assert import assert_param_msg

assert_param_msg[False, "Hello, world!"]()
```

...が、なぜかコンパイルエラーが起きても何も表示されません。Playgroundの問題なのかバグなのか、ソースコードが読めない現状では何ともいえません。

#### List.get

Mojoの `List` は型引数に長さを持っているため、範囲外アクセスをコンパイルエラーで検知可能です。

が、現状Listの型引数に触れるとランタイムがクラッシュしてしまいます...（クラッシュしたら🔃ボタンでカーネル再起動が必要）

```python
l = [1,2,3]
# 型引数を使って要素をget
print(l.get[0]())
```

```
error: The Mojo REPL has crashed and attempted recovery. If the REPL behaves inconsistently, please restart to ensure correct behavior.
```

### 未実装のPython構文
※ **以下は未実装の内容に対する予想（妄想）です！実装されたとしても型レベルプログラミングに使えるかは分かりませんのでご注意を！**

Mojoは100%のPython互換を目標に掲げているため、表現力の高い構文がMojoへ（そして型引数へ）導入されていくでしょう。

個人的に期待している構文は以下の通りです。

- 型引数にコレクションのリテラルを指定
    - `in` と3項演算子の組み合わせがなにか型の変形に使えそう
    - ~~ヒープを使うから禁止されそう~~
- structの継承
    - 型レベルプログラミングがより安全になる（上記の型レベルリスト等）
    - 親子で型引数のシグネチャを変えられる場合、型推論だけで型が導出できる？
- 内包表記
    - 言わずと知れた黒魔術の相棒
    - チューリング完全なので本当になんでもできる
- セイウチ演算子 `:=`
    - 計算途中の値の使いまわしができる
    - スコープ次第では別オブジェクトの型引数に干渉できる...？
- `lambda`
    - ~~関数呼び出しがダメなので無理そう~~

# おわりに

以上、Mojoで型レベルプログラミングに挑戦した内容の紹介でした。機能が開発途中でソースコードも未公開のため、まだ手探り感が否めません...1か月くらい経ってから再挑戦ですね。

中途半端な内容なので記事にするか迷いましたが、同じようにMojoで黒魔術をやりたい方の素材にできればと思って公開しました。

というわけで、**皆さんもModularのアーリーアクセスを申請して型引数で遊びましょう**！

https://www.modular.com/get-started

[^earlyaccess]: まだOSSとして公開されていないため、実行にはPlaygroundアーリーアクセスの申請が必要です。
[^35000]: 35000倍速くなるのはかなり限定的な条件下のようですが...
[^overload]: 戻り値だけ異なるメソッドをオーバーロードすることはできません。
[^difftype]: もちろん値の条件分岐も可能です。