RubyVM::AbstractSyntaxTree を使って AASM の定義を抽出する

## はじめに
本記事は [STORES.jp Advent Calendar 2019](https://adventar.org/calendars/4549) 21日目の記事です。

Ruby2.6 から `RubyVM::AbstractSyntaxTree` という、 Ruby コードを抽象構文木に変換できる API を提供する module が使えるようになり、今年の RubyKaigi 2019 でも関連したセッションが行われていました。
これを使って何かできないものかと考えて思いついたのが、最近業務で利用している AASM という gem を使った、状態遷移の定義を取り出す...というものでした。

そこでこの記事では、 RubyVM::AbstractSyntaxTree を利用して AASM を使った状態遷移定義を取得するコードについて書きます。[^1]

## そもそも AASM とは？
[AASM](https://github.com/aasm/aasm) はオブジェクトの状態遷移を定義・管理する機能を提供する gem です。
class にステートの一覧や状態遷移のためのイベントを DSL によって記述することで、各ステートに対応した定数や状態遷移を行うためのメソッドを生成してくれます。また、定義と異なる遷移を行おうとすると状態遷移に失敗するので、意図しない状態遷移を防ぐことができます。
具体的な使い方を簡単に解説すると、 `aasm` メソッドに渡すブロック内で

- `state` メソッドにステート名を渡すとその名前のステートを管理できるようになる
    - ステート名に対応した定数が生成される
    - `initial_state: true` オプションを付与することで、初期ステートが定義できる
- `event` メソッドにイベント名を渡すことで、ステートを遷移させるメソッドが生成される
    - どのステートからどのステートへ遷移できる、というフローはブロック内の `transitions` によって定義する

複雑になりがちな状態管理を宣言的に書くことで見通しを良くしつつ、秩序を保ちやすく実装できるので便利です。

```task.rb
class Task < ApplicationRecord
  include AASM

  aasm column: :status do
    state :waiting, initial: true
    state :in_progress, :pending, :finished

    event :start do
      transitions from: :waiting, to: :in_progress
      transitions from: :pending, to: :in_progress
    end

    event :stop do
      transitions from: :in_progress, to: :pending
    end

    event :finish do
      transitions from: :in_progress, to: :finished
    end
  end
end

task = Task.create
task.status
# => "waiting"
task.start!
task.status
# => "in_progress"
task.start!
# AASM::InvalidTransition: Event 'start' cannot transition from 'in_progress'.
```

## 状態遷移の定義を抽出する
本題です。
先ほどの例の `Task` クラスの状態遷移の定義を抽象構文木にしてごにょごにょし、初期ステート、ステートのリスト、イベントを取得できるようなオブジェクトを返すことができるようにしました。

[GitHub - ta-chibana/aasm_parser: Parsing AASM definition with RubyVM::AbstractSyntaxTree](https://github.com/ta-chibana/aasm_parser)

以下のように利用します。

```rb
aasm = AasmParser.parse_file('./task.rb')
aasm.initial_state
# => :waiting
aasm.state_names
# => [:waiting, :in_progress, :pending, :finished]
aasm.events.map(&:name)
# => [:start, :stop, :finish]
aasm.events.flat_map(&:transitions).map { |e| "#{e.from} => #{e.to}" }
# => ["waiting => in_progress", "pending => in_progress", "in_progress => pending", "in_progress => finished"]
```

内部では RubyVM::AbstractSyntaxTree を利用して `task.rb` を抽象構文木に変換し、目的の値を取得しています。
本記事では今回書いたコードのうち、初期ステートの取得を行う処理についてだけ解説します。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/88200/6dd703f9-ca60-9d80-3f7c-ea4d08fc4a86.png)



初期ステートの取得は以下の手順で段階的に行ってみました。

1. `task.rb` を抽象構文木に変換
2. `aasm` メソッドの呼び出し時に渡されている block を抽出
3. block から `state` メソッドの呼び出しを抽出
4. `state` メソッドの呼び出しから `initial_state: true` オプションが付いているものを抽出
5. `state` メソッドの呼び出しの第一引数を取得（= 初期ステート）

### 手順1: `task.rb` を抽象構文木に変換
状態遷移の定義は `aasm` メソッドに渡された block 内にあり、これを RubyVM::AbstractSyntaxTree を使って取り出せる状態にするために `task.rb` を抽象構文木に変換します。変換は以下の1行で行うことができます。

```rb
node = RubyVM::AbstractSyntaxTree.parse_file('./task.rb')
```

結果として `RubyVM::AbstractSyntaxTree::Node` のインスタンスが返されます。 `RubyVM::AbstractSyntaxTree::Node` は `#type` と `#children` を持つので、それらを利用しつつ目的のノードを探していきます。

### 手順2: `aasm` メソッドの呼び出し時に渡されている block を抽出
`RubyVM::AbstractSyntaxTree::Node` インスタンスに対して `pp node` と実行すると以下のように出力され、ノードの type やノード間の親子関係を把握するのに便利です。開発時はこの構造を見ながら目的の値を抽出する道筋を立てました。
以下は手順1で変換した結果を `pp` に渡した結果です。

```
(SCOPE@1:0-21:3
 tbl: []
 args: nil
 body:
   (CLASS@1:0-21:3 (COLON2@1:6-1:10 nil :Task)
      (CONST@1:13-1:30 :ApplicationRecord)
      (SCOPE@1:0-21:3
       tbl: []
       args: nil
       body:
         (BLOCK@2:2-20:5
            (FCALL@2:2-2:14 :include
               (ARRAY@2:10-2:14 (CONST@2:10-2:14 :AASM) nil))
            (ITER@4:2-20:5
               (FCALL@4:2-4:22 :aasm
                  (ARRAY@4:7-4:22
                     (HASH@4:7-4:22
                        (ARRAY@4:7-4:22 (LIT@4:7-4:14 :column)
                           (LIT@4:15-4:22 :status) nil)) nil))
               (SCOPE@4:23-20:5
                tbl: []
                args: nil
                body:
                  (BLOCK@5:4-19:7
                     (FCALL@5:4-5:33 :state
                        (ARRAY@5:10-5:33 (LIT@5:10-5:18 :waiting)
                           (HASH@5:20-5:33
                              (ARRAY@5:20-5:33 (LIT@5:20-5:28 :initial)
                                 (TRUE@5:29-5:33) nil)) nil))
                     (FCALL@6:4-6:43 :state
                        (ARRAY@6:10-6:43 (LIT@6:10-6:22 :in_progress)
                           (LIT@6:24-6:32 :pending) (LIT@6:34-6:43 :finished)
                           nil))
                     (ITER@8:4-11:7
                        (FCALL@8:4-8:16 :event
                           (ARRAY@8:10-8:16 (LIT@8:10-8:16 :start) nil))
                        (SCOPE@8:17-11:7
                         tbl: []
                         args: nil
                         body:
                           (BLOCK@9:6-10:50
                              (FCALL@9:6-9:50 :transitions
                                 (ARRAY@9:18-9:50
                                    (HASH@9:18-9:50
                                       (ARRAY@9:18-9:50 (LIT@9:18-9:23 :from)
                                          (LIT@9:24-9:32 :waiting)
                                          (LIT@9:34-9:37 :to)
                                          (LIT@9:38-9:50 :in_progress) nil))
                                    nil))
                              (FCALL@10:6-10:50 :transitions
                                 (ARRAY@10:18-10:50
                                    (HASH@10:18-10:50
                                       (ARRAY@10:18-10:50
                                          (LIT@10:18-10:23 :from)
                                          (LIT@10:24-10:32 :pending)
                                          (LIT@10:34-10:37 :to)
                                          (LIT@10:38-10:50 :in_progress) nil))
                                    nil)))))
                     (ITER@13:4-15:7
                        (FCALL@13:4-13:15 :event
                           (ARRAY@13:10-13:15 (LIT@13:10-13:15 :stop) nil))
                        (SCOPE@13:16-15:7
                         tbl: []
                         args: nil
                         body:
                           (FCALL@14:6-14:50 :transitions
                              (ARRAY@14:18-14:50
                                 (HASH@14:18-14:50
                                    (ARRAY@14:18-14:50 (LIT@14:18-14:23 :from)
                                       (LIT@14:24-14:36 :in_progress)
                                       (LIT@14:38-14:41 :to)
                                       (LIT@14:42-14:50 :pending) nil)) nil))))
                     (ITER@17:4-19:7
                        (FCALL@17:4-17:17 :event
                           (ARRAY@17:10-17:17 (LIT@17:10-17:17 :finish) nil))
                        (SCOPE@17:18-19:7
                         tbl: []
                         args: nil
                         body:
                           (FCALL@18:6-18:51 :transitions
                              (ARRAY@18:18-18:51
                                 (HASH@18:18-18:51
                                    (ARRAY@18:18-18:51 (LIT@18:18-18:23 :from)
                                       (LIT@18:24-18:36 :in_progress)
                                       (LIT@18:38-18:41 :to)
                                       (LIT@18:42-18:51 :finished) nil))
                                 nil)))))))))))
```

雑な見方としては、以下の部分を例にすると、1行目のノードは type が `FCALL` のノードで、後ろに続く `:include (ARRAY@2:10-2:14 (CONST@2:10-2:14 :AASM) nil)` が子のノードとなります。それぞれ `RubyVM::AbstractSyntaxTree::Node#type`, `RubyVM::AbstractSyntaxTree::Node#children` の戻り値に対応しています。

```
            (FCALL@2:2-2:14 :include
               (ARRAY@2:10-2:14 (CONST@2:10-2:14 :AASM) nil))
```

これを図にすると以下のようになり、 `include AASM` の行を表すノードであることがわかります。
（開発当初は、 type の意味については解説されたドキュメントが見つからなかったのであくまで推測だったのですが、社内で共有したところ `--dump=parsetree_with_comment` オプション付きで実行すれば各ノードの説明が表示されるとの情報をいただけました。圧倒的感謝です :bow: ）

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/88200/d7380143-380b-41aa-3e10-b3349ca47cbf.png)

```sh:--dump=parsetree_with_commentで実行
$ ruby --dump=parsetree_with_comment -e 'include AASM'
###########################################################
## Do NOT use this node dump for any purpose other than  ##
## debug and research.  Compatibility is not guaranteed. ##
###########################################################

# @ NODE_SCOPE (line: 1, location: (1,0)-(1,12))
# | # new scope
# | # format: [nd_tbl]: local table, [nd_args]: arguments, [nd_body]: body
# +- nd_tbl (local table): (empty)
# +- nd_args (arguments):
# |   (null node)
# +- nd_body (body):
#     @ NODE_FCALL (line: 1, location: (1,0)-(1,12))*
#     | # function call
#     | # format: [nd_mid]([nd_args])
#     | # example: foo(1)
#     +- nd_mid (method id): :include
#     +- nd_args (arguments):
#         @ NODE_ARRAY (line: 1, location: (1,8)-(1,12))
#         | # array constructor
#         | # format: [ [nd_head], [nd_next].. ] (length: [nd_alen])
#         | # example: [1, 2, 3]
#         +- nd_alen (length): 1
#         +- nd_head (element):
#         |   @ NODE_CONST (line: 1, location: (1,8)-(1,12))
#         |   | # constant reference
#         |   | # format: [nd_vid](constant)
#         |   | # example: X
#         |   +- nd_vid (constant): :AASM
#         +- nd_next (next element):
#             (null node)
```

これを踏まえると、 `aasm` メソッドのブロックに対応している部分は以下の部分が怪しそうな気配が感じられます。

```
            (ITER@4:2-20:5
               (FCALL@4:2-4:22 :aasm
                  (ARRAY@4:7-4:22
                     (HASH@4:7-4:22
                        (ARRAY@4:7-4:22 (LIT@4:7-4:14 :column)
                           (LIT@4:15-4:22 :status) nil)) nil))
               (SCOPE@4:23-20:5
                tbl: []
                args: nil
                body:
                  (BLOCK@5:4-19:7
                     (FCALL@5:4-5:33 :state
                        (ARRAY@5:10-5:33 (LIT@5:10-5:18 :waiting)
                           (HASH@5:20-5:33
                              (ARRAY@5:20-5:33 (LIT@5:20-5:28 :initial)
                                 (TRUE@5:29-5:33) nil)) nil))
```

このノードを抜き出すために、まず以下の条件でノードを抽出し、その後 `BLOCK` 要素を取得する流れにしました。

- `type` が `ITER` （ブロック付きメソッド呼び出し？）
- 子要素に `type` が `FCALL`（関数呼び出しを表す？）の要素を持つ
- ↑の `FCALL` 要素の最初の子要素が `:aasm`

そのために書いたコードが以下で、 `find_from_node` でノードを掘り下げつつ `aasm_node?` で `true` を返すノードが見つかる、もしくはすべてのノードを走査し終えるまで処理を行います。

```lib/aasm_parser/aasm_node_finder.rb
def find_from_node(node)
  return nil unless node.respond_to?(:type)
  return node if aasm_node?(node)

  find_from_children(node.children)
end

def find_from_children(children)
  return nil if children.blank?

  head, *tail = children
  result = find_from_node(head)
  return result if result

  find_from_children(tail)
end

def aasm_node?(node)
  return false unless node&.type == :ITER

  first_child = node.children.first
  return false unless first_child.type == :FCALL

  first_child.children.first == :aasm
end
```

条件に該当するノードが見つかれば、取得できた `ITER` ノードの2つ目の子要素である `SCOPE` ノードから `BLOCK`ノードを取得すればOKです。

### 手順3: `aasm` の block から `state` メソッドの呼び出しを抽出
手順2で取得した `BLOCK` ノードの構造を見てみます。

```
                  (BLOCK@5:4-19:7
                     (FCALL@5:4-5:33 :state
                        (ARRAY@5:10-5:33 (LIT@5:10-5:18 :waiting)
                           (HASH@5:20-5:33
                              (ARRAY@5:20-5:33 (LIT@5:20-5:28 :initial)
                                 (TRUE@5:29-5:33) nil)) nil))
                     (FCALL@6:4-6:43 :state
                        (ARRAY@6:10-6:43 (LIT@6:10-6:22 :in_progress)
                           (LIT@6:24-6:32 :pending) (LIT@6:34-6:43 :finished)
                           nil))
```

 `BLOCK` ノード直下に `state` メソッドの呼び出しを表していそうなノードが見えるので、 `BLOCK` ノードの子要素から以下の条件で抽出すれば良さそうです。

- `type` が `FCALL`
- 最初の子要素が `:state`

そのための実装が以下です（↑の条件で `select` しているだけ）。

```lib/aasm_parser/aasm/block.rb
def state_nodes
  ast_block.children.select do |node|
    node.type == :FCALL && node.children.first == :state
  end
end
```


### 手順4: `state` メソッドの呼び出しのリストから `initial_state: true` オプションが付いているものを抽出
手順3で取得できた `state` メソッドの呼び出しは、元のコードの `task.rb` でいうと以下の部分です。

```rb
    state :waiting, initial: true
    state :in_progress, :pending, :finished
```

また、取得できた `state` メソッド呼び出しのノードは以下ような構造になっています。

```
                     (FCALL@5:4-5:33 :state
                        (ARRAY@5:10-5:33 (LIT@5:10-5:18 :waiting)
                           (HASH@5:20-5:33
                              (ARRAY@5:20-5:33 (LIT@5:20-5:28 :initial)
                                 (TRUE@5:29-5:33) nil)) nil))
                     (FCALL@6:4-6:43 :state
                        (ARRAY@6:10-6:43 (LIT@6:10-6:22 :in_progress)
                           (LIT@6:24-6:32 :pending) (LIT@6:34-6:43 :finished)
                           nil))
```

どちらの `FCALL` ノードも子要素として `ARRAY` ノードを持っており、その中には `:waiting` や `:in_progress` などの Symbol が含まれていることから、引数を表していそうです。
また、1つ目の `state` メソッドの呼び出しではオプション引数として `initial: true` を渡していますが、ノードの構造では `(LIT@5:10-5:18 :waiting)` の他に `HASH` ノードが含まれていたり、 `HASH` ノードの孫要素に `(LIT@5:20-5:28 :initial)` や `(TRUE@5:29-5:33)` があり、これらが含まれていることを条件に `initial_state: true` であるかどうかが判断できそうです。

以上をまとめると `initial_state` オプションが付いている要素かどうかを調べるためには、

- `FCALL` ノードの子要素である `ARRAY` ノードの子要素に `HASH` ノードが存在する（= オプション引数が存在する）
- ↑の `HASH` ノードの子要素である `ARRAY` 要素に、子要素として `:initial` を持つ `LIT` ノードが存在する
- ↑の `LIT` ノードの次の要素が `TRUE` ノードである

という条件で判断します。これらの条件を確認するコードが以下の `initial_state?` です。

```lib/aasm_parser/aasm/state.rb
def initial_state?
  return false if options.blank?

  lit_node_index = options.find_index do |e|
    e.type == :LIT && e.children[0] == :initial
  end

  return false if lit_node_index.nil?

  options[lit_node_index + 1].type == :TRUE
end

# state メソッドに渡された引数のノードを抽出する
def state_arguments
  # @state_node に、 FCALL ノードの
  # RubyVM::AbstractSyntaxTree::Node インスタンスが格納されている
  @state_node
    .children
    .last # state メソッドに渡された引数を表す ARRAY ノード
    .children
    .compact
end

def options
  @options ||= begin
    hash_node = state_arguments.find { |e| e.type == :HASH }
    return [] if hash_node.nil?

    hash_node
      .children
      .first # オプション引数の内容を表す ARRAY ノード
      .children
  end
end

# state の名前を取得する
# 手順5で名前を取得するときに利用します
def names
  state_arguments
    .select { |e| e.type == :LIT }
    .flat_map(&:children)
end
```

`#names` を呼び出すと state の名前を抽出することができるようにしました（手順5で使います）。

### 手順5: `initial_state: true` オプション付きの `state` メソッド呼び出しの第一引数を取得
手順3で抽出した `state` メソッドの呼び出しを表すノードの中から、 `initial_state: true` オプションが付いているかを判定するには手順4の `initial_state?` を実行すれば良いです。
また、その名前を得るためには手順4のコードにチラッと出てきた `names` を利用します。 `state` メソッドには複数のステート名が渡せるので、その都合上複数の名前が返るようになっていますが、 `initial_state` オプションが付いている場合はステートは一つしか渡されていないはず？なので最初の要素を取得すれば初期ステート名を得ることができます。


```lib/aasm_parser/aasm/block.rb
def initial_state
  state = states.find(&:initial_state?)
  return nil if state.nil?

  state.names.first
end

def states
  state_nodes.map { |node| State.new(node) }
end
```

## おわりに
無事に目的の値を取得することができた一方、以下のようなあらゆる問題が存在します。

- AASM のあらゆるオプションに対応できていない（`:if`, `:unless`やコールバックが設定できます）
- AASM に破壊的な変更があるとほぼ作り直しのような状態になる
- どうしても複雑な処理になってしまうので修正が容易でない
- AST生成時に対象のコードは実行されないので、実行しないと目的の値が取れない（ステートが Symbol でなく定数・変数で記述されているなど）ケースだと詰む
- などなど...

これらの問題を解決するにはとても良い方法があって、 **RubyVM::AbstractSyntaxTree を使わずに AASM を使う** という方法が有効です。つまりこうです。

```rb
Task.aasm.states.map(&:name)
# => [:waiting, :in_progress, :pending, :finished]
Task.aasm.events.map(&:name)
# => [:start, :stop, :finish]
Task.aasm.initial_state
# => :waiting
Task.aasm.events.flat_map(&:transitions).map { |e| "#{e.from} => #{e.to}" }
# => ["waiting => in_progress", "pending => in_progress", "in_progress => pending", "in_progress => finished"]
```

これですべての問題が解決しそうです。
結果的に作らなくても良いものを生んでしまったわけですが、 RubyVM::AbstractSyntaxTree を使ってみるという目的は達せられたのでヨシとします。

感想としては、取りたいノードの構造がわかっているなら段階的に抽出するのではなく一発で抽出するようにできたのではないかなあとか、この知識を使える現実的な場面（なさそう）に遭遇したらまた挑戦してみたいなあと思うなどしました。ただ、普段書かないようなコードなので楽しく取り組めました。

以上、 RubyVM::AbstractSyntaxTree を使ってみたくてコードを書いた記録でした。

## 参考

- [プロと読み解く Ruby 2.6 NEWS ファイル](https://techlife.cookpad.com/entry/2018/12/25/110240)
- [RubyVM::AbstractSyntaxTree::Node 詳細](https://qiita.com/siman/items/3017ed399ded43229189)

[^1]: おや？と思った方、お察しください。釘に見えてしまったのです。