Now in REALITY Tech #91 SwiftSyntaxでSwiftLintのカスタムルールを作る

2023年11月16日 14:35

REALITYのiOSエンジニアのあおやまです。
SwiftLintにはカスタムルールという仕組みがあり、正規表現やSwiftSyntaxによる構文木解析を用いて、warningやerrorを出すオリジナルのルールを作成できます。
この仕組みを用いて、SwiftSyntaxでescaping closureにおける[weak self]忘れによるself句の循環参照を検出するルールを自作してみました。

これまでの経緯

iOSチームではこれまで、SwiftLintの正規表現のカスタムルールでescaping closureにおける[waek self]忘れによるself句の循環参照を検出していました。
よく機能していた一方で、正規表現を使うので柔軟性やメンテナンス性に難がありました。

そこで、SwiftSyntaxを使ったカスタムルールへの移植に取り組みました。
SwiftSyntaxのカスタムルールは構文木解析をされた結果でチェックできるので、closureがネストするパターンの検出ができない、コメントアウトを考慮しない、closureにさまざまな記法がある、などの既存の正規表現を用いたカスタムルールの問題点を解決することができました。

SwiftLintのSwift Custom Ruleを作る

SwiftLintのSwift Custom Ruleを導入し、とりあえず以下のNeedWeakSelfRuleというルールを追加してみます。
このままでは、全てのclosureに対してwarningを表示してしまうので、ViolationsSyntaxVisitorの中身を実装し、[weak self]忘れを検出した場合にviolationsに追加するようにします。

// 書き換えるextraRulesの例
// public func extraRules() -> [Rule.Type] { [NeedWeakSelfRule.self] }

// NeedWeakSelfRule.swift
import SwiftLintCore
import SwiftSyntax

struct NeedWeakSelfRule: ConfigurationProviderRule, SwiftSyntaxRule {
    var configuration = SeverityConfiguration<Self>(.warning)

    func makeVisitor(file: SwiftLintCore.SwiftLintFile) -> SwiftLintCore.ViolationsSyntaxVisitor {
        Visitor(viewMode: .sourceAccurate)
    }

    static let description = RuleDescription(
        identifier: "need_weak_self",
        name: "Need [weak self]",
        description: "Closure maybe need [weak self]",
        kind: .idiomatic,
        nonTriggeringExamples: [],
        triggeringExamples: []
    )

    final class Visitor: ViolationsSyntaxVisitor {
        override func visitPost(_ node: ClosureExprSyntax) {
            // violationsに追加することで、warningを表示させることができる
            // このままだと、全てのclosureにwarningが表示される
            violations.append(node.positionAfterSkippingLeadingTrivia)
        }
    }
}

SwiftLintではbazelを使う手法を紹介されていましたが、REALITYではSwiftLintをforkし、パッチファイルを当て、それを取り込むようにしました。
これによりbazelを導入せずに独自のカスタムツールを追加することができました。

SwiftSyntaxで[weak self]忘れによるselfの循環参照を検出

カスタムルールの追加はとてもシンプルにできたので、SwiftSyntaxを用いて[weak self]忘れによるself句の循環参照を検出します。
(ここからが大変💦)

方針

特定の名前の引数に使われているclosureのみをチェックの対象としました。
(例えばRxSwiftのemit, drive, Combineのsinkなど)
これは、SwiftSyntaxのレイヤーではその関数の引数がescaping closureかどうかの判別ができないためです。
新しくリークするコードを書いてしまった場合、このチェックする関数の名前のリストに追加していきます。

また、ネストしたclosureでは原則全てのclosureに[weak self]を付けるというルールにしました。
若干冗長になるのですが、SwiftSyntax上で循環参照の判定が難しいので安全に倒すようにしました。

// NG例 (言語的にはOK)
act1 { [weak self] in
  act2 {
    self?.act3()
  }
}
// OK例
act1 { [weak self] in
  act2 { [weak self] in
    self?.act3()
  }
}

実装

まず、ViolationsSyntaxVisitorでは、closureを呼び出す関数の名前を取得し、それがチェック対象かを判定します。
チェック対象のclosureの場合、selfを強参照しているかを検出するSyntaxVisitorを用いてclosureのコードブロックでselfを強参照しているか判定します。
そこで強参照している場合、warningを付けるという仕組みです。

final class Visitor: ViolationsSyntaxVisitor {
    override func visitPost(_ node: ClosureExprSyntax) {
        guard let parent = node.parent else { return }
        // selfをキャプチャしていればOK
        if let capturesSelfs = node.signature?.capture?.items.map({ $0.capturesSelf }),
           capturesSelfs.contains(true) {
            return
        }

        // fuga {} の書き方の場合
        let nonBlock = parent.as(FunctionCallExprSyntax.self)
        // fuga(label: {}) の書き方
        let blocked = parent.as(LabeledExprSyntax.self)?.parent?.as(LabeledExprListSyntax.self)?.parent?.as(FunctionCallExprSyntax.self)
        guard let functionCall = nonBlock ?? blocked else { return }

        let funcName: String
        if let memberAccess = functionCall.calledExpression.as(MemberAccessExprSyntax.self) {
            // hoge.fuga{} の場合のfugaの検出
            funcName = memberAccess.declName.baseName.text
        } else if let decl = functionCall.calledExpression.as(DeclReferenceExprSyntax.self) {
            // fuga{} の場合のfugaの検出
            funcName = decl.baseName.text
        } else {
            return
        }

        let checkFuncRegexList: [String] = [
            // RxSwift
            "emit", "drive", "subscribe", "flatMap", "flatMapFirst", "flatMapLatest", "asSignal", "asDrive",
            // Combine
            "sink",
        ]

        var isTarget: Bool = false
        for target in checkFuncRegexList {
            let regex = try! NSRegularExpression(pattern: target)
            let checkingResults = regex.matches(in: funcName, range: NSRange(location: 0, length: funcName.count))
            if checkingResults.count > 0 {
                isTarget = true
            }
        }
        if !isTarget { return }

        let v = SelfSearchVisitor(viewMode: .all)
        v.walk(node.statements)
        if v.detected {
            violations.append(node.positionAfterSkippingLeadingTrivia)
        }
    }
}

final class SelfSearchVisitor: SyntaxVisitor {
    var detected: Bool = false

    override func visitPost(_ node: DeclReferenceExprSyntax) {
        // "self" ではない
        if !node.isSelf { return }
        // selfを弱参照していればOK
        if let parent = node.parent,
           parent.is(OptionalChainingExprSyntax.self) { return }
        // Hoge.selfはOK
        if let parent = node.parent,
           let member = parent.as(MemberAccessExprSyntax.self),
           !member.isBaseSelf { return }
        // \.selfはOK
        if let parent = node.parent,
           parent.is(KeyPathPropertyComponentSyntax.self) { return }
        detected = true
    }

    override func visitPost(_ node: IdentifierPatternSyntax) {
        // guard let, if letはOK
        if node.identifier.text != "self" { return }
        if let parent = node.parent,
           parent.is(OptionalBindingConditionSyntax.self) { return }
        detected = true
    }
}

結果

このカスタムルールを導入し、REALITYのソースコードで実行した所、[weak self]忘れによるselfの循環参照は33箇所検出されました。
順番にコードを確認しながら修正していく予定です。

SwiftLintのSwift Custom Ruleを用いることで、ワーニングを付けたり、
CIで実行する仕組みをSwiftLintのインフラに乗りながら、SwiftSyntaxの独自のlintを簡単に実装できました。