[Swift] SwiftSyntaxが面白い ~ Enumのcaseをアルファベット順でソートしよう ~

最近SwiftSyntaxに入門したのですが、結構おもしろいことができそうだったので紹介します！
また、この記事はこちらのイベントで発表した内容となっています。

登壇資料も併せてご覧ください。

SwiftSyntaxとは

• Apple公式のライブラリ
• ソースコードを解析、生成、変換することができるライブラリ
• swift-formatやSwiftLint, mockolo, Sourceryなどのソースコードを生成したり変更したりするライブラリに使われている。

今回はSwiftSyntaxを用いてEnumのcaseをアルファベット順にソートするプログラムを書いてみます。

コードを抽象構文木(AST)に変換してみる

enum Hoge {
    case b
    case a
}

これをASTに変換すると、、

このようになります。
また、SwiftをASTに変換する際はこちらのswift-ast-explorerというツールが便利なので、よかったら使ってみてください。

SwiftSyntaxを少し学ぶ

SyntaxRewriter

既存のコードを別のコードに書き換えることができます。

final class EnumEmptyRewriter: SyntaxRewriter {
    override func visit(_ node: EnumDeclSyntax) -> DeclSyntax {
        let formattedNode = node.with(\.memberBlock.members, MemberDeclListSyntax([]))
        return super.visit(formattedNode)
    }
}

node.with(_:_:)で子ノードの値を書き換えることができます。

赤がmemberBlockの部分で青がmembersです。membersを空にしたので、
Enumのcaseが存在しないコードが生成されるはずです。

実行方法

import SwiftSyntax
import SwiftParser

let `enum` =
"""
enum E1 {
    case b
    case a
    case j
    case h
    case i
    case e
    case d
    case g
    case f
    case c
}
"""
let syntax = Parser.parse(source: `enum`) // ソースコードをASTに変換
let formatted = EnumEmptyRewriter().visit(syntax) // ASTを書き換え
print(formatted.description) // 書き換えたものをSwiftで出力

出力は

enum E1 {
}

しっかりと想定していた出力になりました。

Enumのcaseをアルファベット順でソートする

final class EnumRewriter: SyntaxRewriter {
    override func visit(_ node: EnumDeclSyntax) -> DeclSyntax {
        var members = node.memberBlock.members
        var enumCases = members.compactMap { $0.decl.as(EnumCaseDeclSyntax.self) }
        var elementLists = enumCases.map(\.elements)
        // ソート部分
        elementLists = elementLists.sorted {
            let firstText = $0.first?.identifier.text ?? ""
            let secondText = $1.first?.identifier.text ?? ""
            return firstText.localizedStandardCompare(secondText) == .orderedAscending
        }
        // 書き換え
        enumCases = zip(enumCases, elementLists).map {
            $0.with(\.elements, $1)
        }
        // 書き換え
        members = MemberDeclListSyntax(
            zip(members, enumCases).map {
                $0.with(\.decl, $1.cast(DeclSyntax.self))
            }
        )
        // 書き換え
        return super.visit(node.with(\.memberBlock.members, members))
    }
}

そして、実行してみます。

let `enum` =
"""
enum E1 {
    case b
    case a
    case j
    case h
    case i
    case e
    case d
    case g
    case f
    case c
}
"""
let syntax = Parser.parse(source: `enum`)
let formatted = EnumSortRewriter().visit(syntax)
print(formatted.description)

すると、、

しっかりと並び変わりました。

ただ、このコードだと

enum E1 {
    case b, a, j, h, i, e, d
}

こういう形のEnumには非対応です。
そこでこのように書き換えました。

final class EnumSortRewriter: SyntaxRewriter {
    override func visit(_ node: EnumDeclSyntax) -> DeclSyntax {
        var members = node.memberBlock.members
        var enumCases = members.compactMap { $0.decl.as(EnumCaseDeclSyntax.self) }

        var elementLists = enumCases.map(\.elements)

        if let elementList = elementLists.first,
           elementLists.count == 1
        {
            // enum E1 {
            //     case c1, c2, c3
            // }
            let sortedList = sortElementList(elementList)
            elementLists = [sortedList]
        } else {
            // Enum E1 {
            //     case c1
            //     case c2
            //     case c3
            // }
            elementLists = sortElementLists(elementLists)
        }

        enumCases = zip(enumCases, elementLists).map {
            $0.with(\.elements, $1)
        }

        members = MemberDeclListSyntax(
            zip(members, enumCases).map {
                $0.with(\.decl, $1.cast(DeclSyntax.self))
            }
        )

        return super.visit(node.with(\.memberBlock.members, members))
    }

    private func sortElementLists(_ elementLists: [EnumCaseElementListSyntax]) -> [EnumCaseElementListSyntax] {
        elementLists.sorted {
            let firstText = $0.first?.identifier.text ?? ""
            let secondText = $1.first?.identifier.text ?? ""
            return firstText.localizedStandardCompare(secondText) == .orderedAscending
        }
    }

    private func sortElementList(_ elementList: EnumCaseElementListSyntax) -> EnumCaseElementListSyntax {
        let elements = elementList
            .sorted {
                $0.identifier.text.localizedStandardCompare($1.identifier.text) == .orderedAscending
            }
            .enumerated()
            .map { index, element in
                var element = element
                if index == elementList.count - 1 {
                    element.trailingComma = nil
                } else {
                    element.trailingComma = .commaToken(trailingTrivia: .space)
                }
                return element
            }

        return EnumCaseElementListSyntax(elements)
    }

このようにすることで問題に対応することができました。
また、こちらのコードをライブラリ化したのでよかったらご覧ください。