Pythonにrep文を追加してみた

Pythonにrep文を追加してみた

はじめに

この記事はEEIC(東京大学工学部 電気電子工学科/電子情報工学科)3年の後期実験「大規模ソフトウェアを手探る」のレポートとして書かれています。

環境

• OS: Ubuntu 20.04 on WSL2
• Python: Python3.7.9

動機

repマクロとは, 指定した回数の繰り返し処理を簡潔に書くためのマクロである. repマクロが用いられるのは, 競技プログラミングでC++を使用する場合などである. n回の繰り返し処理を行う場合, for文だと

for (int i = 0; i < n; i++) {
    処理
}

と記述するところを, repマクロだと

rep(i , n) {
    処理
}

と短く記述することができる. このように, 指定した回数の繰り返し処理を簡潔に記述するための機能がPythonにもあれば便利だと考え, Pythonにrepマクロと似た機能をもつrep文を追加しようと試みた.

仕様

当初の目標としたのは, n回の繰り返し処理を

REP i, n:
    処理

と記述する仕様である. だが, 実際にできたのは

REP i, range(n):
    処理

と記述する仕様であった.

ビルド

ビルドについてはを参照

コード変更

変更箇所

変更を加えたファイルは以下の5つである.

• Grammar/Grammar
• Parser/Python.asdl
• Python/ast.c
• Python/symtable.c
• Python/compile.c

変更を加えるべきファイルを探すうえでは, https://www.python.org/dev/peps/pep-0306/ のChecklistが参考になるだろう.

Grammar/Grammar

ファイル名からして, 文法を規定するファイルだと思われる. https://docs.python.org/3.7/reference/grammar.html によると,

This is the full Python grammar, as it is read by the parser generator and used to parse Python source files

とのことである.

今回追加したいrep文は, for文を似た機能を有する. そこで, for_stmtの記述を参考に, Grammar/Grammarの70行目にrep_stmtを追加する.

# 変更前
compound_stmt: if_stmt | while_stmt | for_stmt | try_stmt | with_stmt | funcdef | classdef | decorated | async_stmt

# 変更後
compound_stmt: if_stmt | while_stmt | for_stmt | try_stmt | with_stmt | funcdef | classdef | decorated | async_stmt | rep_stmt

さらに, その下に, 以下の1行を追加する.

rep_stmt: 'REP' expr ',' testlist ':' suite ['else' ':' suite]

Parser/Python.asdl

Paser/Python.asdlは, https://www.python.org/dev/peps/pep-0306/によると,

Parser/Python.asdl may need changes to match the Grammar. Run make to regenerate Include/Python-ast.h and Python/Python-ast.c.

とのことである. 抽象構文木(AST)に関係するファイルの再生成に関係がありそうだ.

35行目のfor文に関する記述を参考に, 以下の1行を追加する.

| Rep(expr target, expr iter, stmt* body, stmt* orelse)

Python/ast.c

ファイル名からして, 抽象構文木(AST)の生成に関係がありそうだ. https://www.python.org/dev/peps/pep-0306/によると,

Python/ast.c will need changes to create the AST objects involved with the Grammar change.

とのことである.

ast_for_for_stmt()関数を参考に, ast_for_rep_stmt()関数を追加する.

static stmt_ty
ast_for_rep_stmt(struct compiling *c, const node *n0, bool is_async)
{
    const node * const n = is_async ? CHILD(n0, 1) : n0;
    asdl_seq *_target, *seq = NULL, *suite_seq;
    expr_ty expression;
    expr_ty target, first;
    const node *node_target;
    /* rep_stmt: 'REP' expr ',' test ':' suite ['else' ':' suite] */
    REQ(n, rep_stmt);

    if (NCH(n) == 9) {
        seq = ast_for_suite(c, CHILD(n, 8));
        if (!seq)
            return NULL;
    }

    node_target = CHILD(n, 1);
    _target = ast_for_exprlist(c, node_target, Store);
    if (!_target)
        return NULL;
    /* Check the # of children rather than the length of _target, since
       for x, in ... has 1 element in _target, but still requires a Tuple. */
    first = (expr_ty)asdl_seq_GET(_target, 0);
    if (NCH(node_target) == 1)
        target = first;
    else
        target = Tuple(_target, Store, first->lineno, first->col_offset, c->c_arena);

    expression = ast_for_testlist(c, CHILD(n, 3));
    if (!expression)
        return NULL;
    suite_seq = ast_for_suite(c, CHILD(n, 5));
    if (!suite_seq)
        return NULL;

    if (is_async)
        return AsyncFor(target, expression, suite_seq, seq,
                        LINENO(n0), n0->n_col_offset,
                        c->c_arena);
    else
        return For(target, expression, suite_seq, seq,
                   LINENO(n), n->n_col_offset,
                   c->c_arena);
}

ast_for_stmt()関数内のswitch文に, case for_stmtを参考にcase rep_stmtを追加する.

     case rep_stmt:
            return ast_for_rep_stmt(c, ch, 0);

Python/symtable.c

https://www.python.org/dev/peps/pep-0306/によると,

Python/symbtable.c: This handles the symbol collection pass that happens immediately before the compilation pass.

とのことである. コンパイルに関係がありそうだ.

symtable_visit_stmt()関数内のswitch文に, case For_kindの記述を参考に, case Rep_kindを追加する.

    case Rep_kind:
        VISIT(st, expr, s->v.Rep.target);
        VISIT(st, expr, s->v.Rep.iter);
        VISIT_SEQ(st, stmt, s->v.Rep.body);
        if (s->v.Rep.orelse)
            VISIT_SEQ(st, stmt, s->v.Rep.orelse);
        break;

Python/compile.c

ファイル名からして, コンパイルに関係がありそうだ. https://www.python.org/dev/peps/pep-0306/によると,

Python/compile.c: You will need to create or modify the compiler_* functions to generate opcodes for your productions.

とのことである.

find_ann()関数内のswitch文に, case For_kindを参考に, case Rep_kindを追加する.

case Rep_kind:
              res = find_ann(st->v.Rep.body) ||
                  find_ann(st->v.Rep.orelse);
              break;

compiler_for()関数を参考に, compiler_rep()関数を追加する.

static int
compiler_rep(struct compiler *c, stmt_ty s)
{
    basicblock *start, *cleanup, *end;

    start = compiler_new_block(c);
    cleanup = compiler_new_block(c);
    end = compiler_new_block(c);
    if (start == NULL || end == NULL || cleanup == NULL)
        return 0;
    ADDOP_JREL(c, SETUP_LOOP, end);
    if (!compiler_push_fblock(c, LOOP, start))
        return 0;
    VISIT(c, expr, s->v.Rep.iter);
    ADDOP(c, GET_ITER);
    compiler_use_next_block(c, start);
    ADDOP_JREL(c, FOR_ITER, cleanup);
    VISIT(c, expr, s->v.Rep.target);
    VISIT_SEQ(c, stmt, s->v.Rep.body);
    ADDOP_JABS(c, JUMP_ABSOLUTE, start);
    compiler_use_next_block(c, cleanup);
    ADDOP(c, POP_BLOCK);
    compiler_pop_fblock(c, LOOP, start);
    VISIT_SEQ(c, stmt, s->v.Rep.orelse);
    compiler_use_next_block(c, end);
    return 1;
}

compiler_visit_stmt()関数内のswitch文に, case For_kindを参考に, case Rep_kindを追加する.

    case Rep_kind:
        return compiler_rep(c, s); 

再ビルド

Python-3.7.9ディレクトリに移動し, ビルドする. このとき, makeの代わりにmake regen-allを実行する. https://docs.python.org/ja/3.6/whatsnew/3.6.html#new-make-regen-all-build-targetによると, make regen-allはPython 3.6.2から追加されたコマンドであり, 必要なファイルを自動で生成してくれるらしい.

$ cd ~/Python-3.7.9
$ make clean
$ make regen-all
$ make install

動作検証

REP i, range(n):
    処理

と記述することで, n回の繰り返し処理が行われることを確認した.

$ ./python3.7
>>> REP i, range(5):
...    print('1')
...
1
1
1
1
1

課題

今回追加したrep文は,

REP i, range(n):
    処理

という記述で, n回の繰り返し処理が行われるものである. しかし, この記述では, Python 3の元々のfor文

for i in range(n):
    処理

と比較したときに, ほとんど短くなっていない. n回の繰り返し処理を簡潔に記述するという動機に立ち返ると, やはり

REP i, n:
    処理

と記述できるようにしたいところだ.

そのためのヒントになるのではないかと考え, 今回追加したrep文のバイトコードを調べることにした. バイトコードについては, https://yigarashi.hatenablog.com/entry/cpythonが参考になると思う.

rep文を使うrep.pyを作成し, -m disというオプションを付けて, python3.7を実行した. このオプションにより, バイトコードが表示されるらしい.

# rep.py
REP i, range(5):
    print(1)

$ ./python3.7 -m dis rep.py
  1           0 SETUP_LOOP              24 (to 26)
              2 LOAD_NAME                0 (range)
              4 LOAD_CONST               0 (5)
              6 CALL_FUNCTION            1
              8 GET_ITER
        >>   10 FOR_ITER                12 (to 24)
             12 STORE_NAME               1 (i)

  2          14 LOAD_NAME                2 (print)
             16 LOAD_CONST               1 (1)
             18 CALL_FUNCTION            1
             20 POP_TOP
             22 JUMP_ABSOLUTE           10
        >>   24 POP_BLOCK
        >>   26 LOAD_CONST               2 (None)
             28 RETURN_VALUE

これと比較するため, 当初の目標としていた形のrep文を使うrep2.pyを作成し, 同様にバイトコードを調べてみる.

# rep2.py
REP i, 5:
    print(1)

$ ./python3.7 -m dis rep2.py
  1           0 SETUP_LOOP              20 (to 22)
              2 LOAD_CONST               0 (5)
              4 GET_ITER
        >>    6 FOR_ITER                12 (to 20)
              8 STORE_NAME               0 (i)

  2          10 LOAD_NAME                1 (print)
             12 LOAD_CONST               1 (1)
             14 CALL_FUNCTION            1
             16 POP_TOP
             18 JUMP_ABSOLUTE            6
        >>   20 POP_BLOCK
        >>   22 LOAD_CONST               2 (None)
             24 RETURN_VALUE

rep2.pyのバイトコードをみると, rep.pyのバイトコードの, 2行目の2 LOAD_NAME 0 (range), 4行目の6 CALL_FUNCTION 1が欠けていることがわかる. それならば, 不足した2行を追加できるように変更を加えれば, 当初の目標としていた

REP i, n:
    処理

という形のrep文を実装することができるのだろうか. バイトコードの生成に関係しているのは, Python/compile.cのcompiler_rep()関数である. この関数を変更すれば, バイトコードに変更を加えられそうだ. 追加したいバイトコードは,

• 2 LOAD_NAME 0 (range)
• 6 CALL_FUNCTION 1

の2行である. 後者については, compiler_rep()関数にADDOP(c, GET_ITER);の1行を加えれば, 追加できるかもしれない. だが, 前者を追加する方法がわからず, これ以上の変更は断念することとなった.

感想

変更するファイルが複数あり, 慣れないうちは大変だった. 聞くところによると, これでも以前よりは, Pythonの機能を変更するのが簡単になっているらしい. Pythonに新たな機能を追加するのはいい経験になった. 今回できたrep文は, 当初の目標とは違う形となってしまったが, いつの日か当初の目標通りのrep文を実装できたらと思う.