knqyf263's blog

自分のためのメモとして残しておくためのブログ。

goimportsのソースコードを読んでみた

概要

go generate用のツールを作る時にパッケージを自動でimportしたくなったので、goimportsのソースコードを読んでみました。 Goの標準ツールということで凄い技術でやっているんだろうな...と漠然と思っており、自分なんかに理解できるだろうかという不安があったのですが、読んでみたら泥臭いことを丁寧にやっている感じでした。 コードは綺麗だし参考になるところだらけなのですが、割と普通のことをやっている感じなので必要以上に恐れずに一度読んでみると良いのではないか、と思いました。

せっかく読んだので重要そうに感じたところだけまとめておきます。 大分省略したにも関わらず長くなって誰にも読まれない文章に昇華されてしまって残念です。自分用のメモということで。

バージョン

自分が読んだときのコミットハッシュは以下です。

コミットハッシュ:3c07937fe18c27668fd78bbaed3d6b8b39e202ea

GitHub - golang/tools at 3c07937fe18c27668fd78bbaed3d6b8b39e202ea

流れ

最初にgoimportsの重要な(だと自分が思った)処理の部分の流れについてまとめておきます。

  • AST取得
    • goimportsの対象ファイルのASTを取得する
  • AST解析
    • ASTを解析して、import一覧と使っているパッケージ一覧(とシンボル)を取得する
      • 前者:import "github.com/knqyf263/fooなどのimport文
      • 後者:foo.Barならパッケージがfoo、シンボルがBar(structなどは除外する)
  • import pathからパッケージ名の取得
    • github.com/knqyf263/go-foo のようなimport pathから実際のパッケージ名を取得しておく(import pathは go-foo だがパッケージ名は foo など、異なったりするため)
    • 単に github.com/knqyf263/go-foo/foo.go 等のGoのファイルを開いてAST解析し、パッケージ名を取得するだけ
  • 不要なimportの削除
    • 上で取得したimport一覧のうち、使われているパッケージの一覧に含まれないパッケージを探して削除
  • importの自動追加
    • 同一パッケージ内の探索
      • 同一パッケージ内の別ファイルをパースする
        • foo/bar.goに対してgoimportsを適用したらfoo/baz.goやfoo/qux.goをパース
      • 同一パッケージ内でもし同じパッケージをimportしていれば、それを優先して使う
        • foo/bar.goでqux.Printをしていて、foo/baz.goでimport "github.com/knqyf263/qux をしていればそれを優先して利用
    • $GOROOTと$GOPATHの探索
      • 最初に全てを探索してmapを作る
        • vendorやinternalは一切関係なしに、find $GOPATH/src -name "*.go" するようなイメージ(あくまでイメージで実際はいろいろ枝刈りしている)
        • この処理は一度のみ行われる
      • パッケージ名がimport path(の最後の2つの要素)内に含まれているものを候補として抽出
        • fooというパッケージ名を探している場合、 github.com/knqyf263/go-foo はfooという文字列を含むため候補となる
        • あくまで候補であり github.com/foo/bar などもfooを含むため抽出される
      • 対象のファイルからパス的に近いものを優先するようにする
        • foo/bar.go内で使われているquxを探す場合、 ../vendor/github.com/knqyf263/qux の方が ../../../../github.com/knqyf263/qux よりもスラッシュの数が少ないため優先される
      • symbolの確認
        • 候補のパッケージのexportされているsymbolを全て取得し、全て存在するか確認する
        • qux.Printを使っている場合に、 github.com/knqyf263/qux のパッケージでexportされているsymbolがHogeFugaだけであれば、このパッケージは対象外となる
    • import文の追加
      • 上で条件を満たすパッケージが見つかれば、import文をASTのライブラリを用いて追加する

詳細

mirrorがGitHubにあるのでそちらを見ていきます。

github.com

自分が重要だと思ったところだけ流して書いていくので、細かいところが気になった人は自分で読んでみると良いと思います。

imports.Process からが重要なので、その前は気にしなくて良いです。

まず、goimportsのコマンド自体は以下の cmd/goimports/goimports.go になります。

tools/goimports.go at 3c07937fe18c27668fd78bbaed3d6b8b39e202ea · golang/tools · GitHub

main (cmd/goimports/goimports.go)

まずmainを見ます。 この中で gofmtMain() を呼んでいます。

tools/goimports.go at 3c07937fe18c27668fd78bbaed3d6b8b39e202ea · golang/tools · GitHub

gofmtMain (cmd/goimports/goimports.go)

flagの処理など色々やっています。
goimportsではファイルを渡したりディレクトリを渡したりできるので、そこら辺の処理もやっています。

そして何やかんやで processFile() が呼ばれます。

tools/goimports.go at 3c07937fe18c27668fd78bbaed3d6b8b39e202ea · golang/tools · GitHub

processFile (cmd/goimports/goimports.go)

指定されたファイルをOpenして中身を読み込んだりします。 それを imports.Process に渡します。 import周りの実際の処理はこのProcessで行われます。

tools/goimports.go at 3c07937fe18c27668fd78bbaed3d6b8b39e202ea · golang/tools · GitHub

結果が返ってきたらファイルに書き込んだり差分を表示したりしてgoimportsコマンドの処理は終わりです。

ではProcessの中を見ていきます。

Process (tools/imports/imports.go)

まずこのProcessですが、exportされているため自分のツールで使いたい場合は imports.Process で呼び出すことが可能です。

imports - GoDoc

ファイルの中身が渡ってきているので、それをparseします。 parseの中では色々やっていますが、やりたいこととしては抽象構文木(AST)を得ることなので、とりあえずはParseFileを呼んでいることだけ知っておけば良いかなと思います。

parser - The Go Programming Language

これで与えられたファイルのASTが手に入りました。

次にfixImportsを呼んで不要なimportの削除や、必要なパッケージのimportを行います。 一番知りたいのはこの中の処理になります。

tools/imports.go at 3c07937fe18c27668fd78bbaed3d6b8b39e202ea · golang/tools · GitHub

その後 sortImports でimport文を並び替えたり、format.Sourceソースコードを整形したりして終わりです。

fixImports (tools/imports/fix.go)

ここからは少し細かい話になっていきます。あとの関数などは全てfix.goのものです。

得られたASTを解析するための関数 visitFn を定義しています。 ASTの解析については色々な方が解説を書いているので省略します。 visitFn では2つのことをやっています。

  • ImportSpec(import文)の解析
    • import "github.com/knqyf263/foo" みたいなやつ
  • SelectorExprの解析
    • foo.Bar みたいなやつ(多分)

ImportSpecの解析

ImportSpecの場合、importPathToName を呼んでいます。 tools/fix.go at 3c07937fe18c27668fd78bbaed3d6b8b39e202ea · golang/tools · GitHub

importPathToName によってimportのpathからパッケージ名への変換が行われます。 パッケージ名が分かれば、それをdecls というmapに保存しておきます。

importPathToName は変数で、実体は importPathToNameGoPath になります。

importPathToNameGoPath

importPathToNameGoPathを見ていきます。 この関数ではimportのpathからパッケージ名への変換を行います。 これは、importのpathだけではパッケージ名が分からないためです。 例えば gopkg.in/yaml.v2 というimport pathで、パッケージ名はyamlだったりします。 そこでどうするかと言うと、ファイルを直接見に行ってパッケージ名を調べます。

標準パッケージについては事前に分かっているため、予め変数として保持しています。 これを使うと、例えば net/http というimport pathの時に http というパッケージ名がすぐに分かります。

tools/zstdlib.go at 3c07937fe18c27668fd78bbaed3d6b8b39e202ea · golang/tools · GitHub

標準パッケージにないimport pathについてはimportPathToNameGoPathParseで探します。

importPathToNameGoPathParse

build.FindOnly オプションを付けて build.Import 呼びます。 これは標準パッケージにないimport pathを解決するためのメソッドです。

build - The Go Programming Language

この結果importしたいライブラリのpathが分かるため、そのディレクトリ内のファイル名を全て取得します。 これらのファイルの中からパッケージ名を探します。 上記ファイルのうち、 .goサフィックスについており、かつ _test.go でないものを探します。 そしてGoのファイルが見つかれば、それをさきほど同様 parser.ParseFile でパースします。 このパッケージ名がdocumentationやmainの場合はskipし、条件を満たすパッケージ名が見つかるまで探していきます。

tools/fix.go at 3c07937fe18c27668fd78bbaed3d6b8b39e202ea · golang/tools · GitHub

見つかればパッケージ名をreturnして終わりです。

SelectorExprの解析

foo.Bar みたいな形の解析です(多分)。 このような形の中には、structなども含まれます。 最初にそういう場合を弾き、パッケージ名だけを取り出します(後述しますが、実はパッケージ以外にも同一パッケージ内で定義されたvarやconstも含んでいる)。 tools/fix.go at 3c07937fe18c27668fd78bbaed3d6b8b39e202ea · golang/tools · GitHub

このパッケージ名はとりあえず利用されていることがわかったため、 refs に入れます。 foo.Barであれば refs["foo"] = make(map[strinb]bool) になる。

tools/fix.go at 3c07937fe18c27668fd78bbaed3d6b8b39e202ea · golang/tools · GitHub

次に dirPackegeInfo を呼んでいます。 これも実際の関数は dirPackageInfoFile になるのでそちらを見ていきます。

dirPackageInfoFile

この関数は、goimportsの対象となっているファイルと同じパッケージにあるファイルについての情報を収集する関数になります。

最初に ioutil.ReadDir(srcDir) でgoimportsの対象ファイルと同じディレクトリ内のファイル情報を全て取得します。

tools/fix.go at 3c07937fe18c27668fd78bbaed3d6b8b39e202ea · golang/tools · GitHub

そしてfor文を回して .go で終わるなど、goに関連するファイルだけを取得します。 いつものように parser.ParseFile でASTを取得して解析します。

まず root.Declsから ast.ValueSpec なものを取り出して info.Globals に格納しています。 tools/fix.go at 3c07937fe18c27668fd78bbaed3d6b8b39e202ea · golang/tools · GitHub

これはどういうことかと言うと、同じパッケージ内であればvarやconstで宣言した物が見えるためです。 variable.Foo()とアクセスしているが、このvariableはパッケージ名ではなく単なるvarで宣言された変数という可能性があるため、このvariableを後々パッケージ名として処理しないように収集しています。

次にimport文も集めて info.Imports に格納しています。

tools/fix.go at 3c07937fe18c27668fd78bbaed3d6b8b39e202ea · golang/tools · GitHub

最後に collectReferences を使って、ast.SelectorExpr を集めています。 collectReferences ではSelectorExprのうち、Exportされているものを info.Refs に格納するようになっています。 foo.barであれば info.Refs には入らず、 foo.Barであれば入るということです。

tools/fix.go at 3c07937fe18c27668fd78bbaed3d6b8b39e202ea · golang/tools · GitHub

これで同一パッケージ内の情報収集は終わりです。

SelectorExprの解析(続き)

dirPackageInfoの結果がpackageInfoに格納されます。 次に、取得したpkgNameが decls に存在するか確認します。 これはimport文を収集したものなので、存在すれば既にこのpkgNameはimport済みということになります。

存在しない場合は、先程収集した packageInfo.Globals を確認します。 これはvarやconstに定義された名前だったので、このmapに存在すれば実はパッケージ名ではなく変数名や定数名だったことが分かります。

どちらにも合致しない場合はまだ未importのパッケージ名ということになるため、これを refs というmapに入れます。 この refs は2次元配列のようになっているため、呼ばれている方の名前も格納します。 つまりfoo.Barのようになっていて、このfooが未importであれば refs["foo"]["Bar"] = true のようになります。

tools/fix.go at 3c07937fe18c27668fd78bbaed3d6b8b39e202ea · golang/tools · GitHub

これは、単にfooというパッケージ名だと誤ってimportすることが多くなってしまいますが(fooというパッケージが複数あった場合に判別できない)、BarというSymbolがExportされているfooパッケージを探すことで精度高くimportが可能なためです。

ここまでで、対象のファイルの中で利用されているパッケージ名の一覧と、未importのパッケージ名+Symbol一覧が手に入りました。

不要なimport分の削除

先程の2つの解析により、以下の2つを得ています。

  • decls: importされているパッケージ一覧
  • refs: 実際に利用されいているパッケージ一覧

これらの差を見ることで不要なimportが判定できます。 つまり、 decls にあるパッケージ名のうち、 refs にないものを unusedImportにいれて削除対象とします。

tools/fix.go at 3c07937fe18c27668fd78bbaed3d6b8b39e202ea · golang/tools · GitHub

実際の削除処理はこちら。CGOの場合はskipなどの処理が入っていたりはしますが、先程のunusedImportを削除しています。

tools/fix.go at 3c07937fe18c27668fd78bbaed3d6b8b39e202ea · golang/tools · GitHub

import済みのパッケージ名の除去

未importのパッケージ名についてはSymbolをrefsに入れてあります。 なので、Symbolがないものについては既にimportされているパッケージであることが分かります。 このあとの処理は未importのものをimportする処理なので、それらは削除します。

削除した結果、refs が空になれば未importのパッケージがないということになるため処理を終わります。

tools/fix.go at 3c07937fe18c27668fd78bbaed3d6b8b39e202ea · golang/tools · GitHub

未importのパッケージをimportする

ここからが一番肝心の処理になります。 まず未importのパッケージ情報が入っている refs をfor文で回します。 それらをgoroutineで処理し、パッケージのimport pathを探します。

同一パッケージから探す

dirPackageInfoFile の中で同一パッケージの別ファイルの情報を集めました。 このimportの中に同じパッケージ名があれば同じimport pathを使うようにします。 ただし、同じsymbolが使われていることが前提になります。 foo.Barがgoimportsの対象ファイル内で使われていて、同一パッケージにfoo.Barがあれば同じfooであるとみなしますが、foo.Bazになっていればskipします。

同一パッケージの別ファイルでfoo.Barが使われており、 import "github.com/knqyf263/foo" になっていれば、同様のimport pathであるとみなしてこのパッケージ名に関する処理を終わります。

findImport(全体から探す)

同一パッケージ内になければ全体から探します。 findImport が呼ばれていますが、実体は findImportGoPath なのでその中を見ていきます。

findImportStdlib

最初に findImportStdlib で標準パッケージ内を探しています。

tools/fix.go at 3c07937fe18c27668fd78bbaed3d6b8b39e202ea · golang/tools · GitHub

探すときには同時にsymbolsも渡しており、全てのsymbolがexportされている場合はimport pathが返ってきます。 標準パッケージは以下のように事前にmapが定義されています。

        "fmt.Errorf":                                    "fmt",
        "fmt.Formatter":                                 "fmt",
        "fmt.Fprint":                                    "fmt",
        "fmt.Fprintf":                                   "fmt",

例えば fmt.Fprintfmt.Errorf がgoimportsの対象ファイル内で使われていれば fmt というimport pathが返ってきます。 しかし、fmt.Fprintfmt.Foo が利用されている場合は全てのsymbolが一致しないため fmt は返されません。

また、rand.Readの場合はmath/randではなくcrypto/randを使うような特別な対応も入っていたりします。

tools/fix.go at 3c07937fe18c27668fd78bbaed3d6b8b39e202ea · golang/tools · GitHub

標準パッケージになかった場合は、ついに$GOROOTと$GOPATHを探しに行きます。 まず全てのpathをscanします。 これらは sync.Once を使って一度しかscanしに行かないようになっています。

tools/fix.go at 3c07937fe18c27668fd78bbaed3d6b8b39e202ea · golang/tools · GitHub

scanGoDirs

scanGoDirsの中を見てみます。

build.Default.SrcDirs()$GOROOT/src$GOPATH/src が返ってきます。 これらの下を再帰的に全て見て行きます。 少し驚いたのですが、goimports対象のファイル付近のvendorだけ見たりとかしてるのかと思ったのですが、$GOPATH/src の下は全てなめています。 つまり、github.com/knqyf263/test1/vendorgithub.com/knqyf263/test2/vendor もこの時点では一緒に扱われます。

ファイルの場合、 src直下にあるものはskipします。$GOPATH/src/main.go みたいなもの。 あとは .go で終わらないものもskipします。 そしてこのファイルの存在するディレクトリを取得し、$GOPATH/src以降をimportPathとして取り出します。 この時、上で述べたようにvendor以下も全て取得しているためVendorlessPath を呼んでvendor以降のimport pathも取り出してimportPathShortとします。 これらを合わせてpkgというstructに保存します。 例えば以下のようになります。

imports.pkg{
  dir:             "/home/knqyf263/src/github.com/future-architect/vuls/vendor/github.com/knqyf263/go-cpe/naming",
  importPath:      "github.com/future-architect/vuls/vendor/github.com/knqyf263/go-cpe/naming",
  importPathShort: "github.com/knqyf263/go-cpe/naming",
}

出来上がったpkgをdirScanというmapに保存しておきます。

これを全てのdirについて行うため、dirScanのkey数はかなりの量になります。 当然ですが、同じdirはskipされるようになっています。

tools/fix.go at 3c07937fe18c27668fd78bbaed3d6b8b39e202ea · golang/tools · GitHub

ディレクトリの場合、"testdata"や"node_modules"というディレクトリ名の場合はそれ以降探索しないようになっています。 他にもignoreの設定が可能なため、ignoreに設定されている場合もそのディレクトリ以下は探索しません。

tools/fix.go at 3c07937fe18c27668fd78bbaed3d6b8b39e202ea · golang/tools · GitHub

シンボリックの場合の処理もありますが割愛します。

探索はfastwalkという x/tools/internal で定義されているモジュールを使っています。 internalなため、外部からは利用できません。

tools/fix.go at 3c07937fe18c27668fd78bbaed3d6b8b39e202ea · golang/tools · GitHub

上記の処理により$GOROOTと$GOPATHを全てscanしてscanDirに入れ終えました。

pkgIsCandidate(候補検索)

まず最初にscanDirのうち、今回のパッケージ名の候補となりそうなpkgの一覧を取得します。 これはpkgIsCandidate の中で行われているので見てみます。

tools/fix.go at 3c07937fe18c27668fd78bbaed3d6b8b39e202ea · golang/tools · GitHub

pkgIsCandidate では最初にcanUseでvendorやinternalなどのルールに一致しているか確かめます。 子は親のvendor等を見ることができますが、子同士やさらに子のvendor等は見ることができません。 つまり ../vendor../internal は許されるが、../foo/vendor../foo/internalbar/vendorbar/internal は見ることができません。

これも予想もつかない凄いコードでやっているんだろうと勝手に思い込んでいましたが、読んでみたら普通な感じでした(シンプルにやっていて凄いとは思いましたが)。 canUse の中を見ると、自分でも頑張れそうな気持ちになれるのでおすすめです。

tools/fix.go at 3c07937fe18c27668fd78bbaed3d6b8b39e202ea · golang/tools · GitHub

次に lastTwoComponents でimportPathの最後の2つの要素を取り出します。 github.com/knqyf263/foo なら knqyf263/foo になります。 そして、パッケージ名がこの文字列に含まれるかを調べます。 例えばパッケージ名がscanなら以下の全てにヒットします。

  • hcl/scanner
  • vuls/scan
  • text/scanner

それどころかnqみたいなパッケージ名の場合、knqyf263/foo にもヒットします。 あくまで候補でありこのあとフィルタするので間違うことはないのですが、シンプルな感じだなーと思いました。

また、これは github.com/knqyf263/foo というimport pathでパッケージ名がbarの場合は候補に選ばれません。 つまりgoimportsが勝手に補完してくれなくなります。 パッケージ名とディレクトリ名は一致させるようにしましょう(普通にGo書いてたら一致させるとは思いますが)。

他にも大文字の場合やハイフンが入ることにより一致しない場合をなくすため、小文字にしてハイフンを削除してからマッチングしたりもしています。 github.com/json-iterator/go はパッケージ名がjsoniterですがimport path内には存在しません。ハイフンを削除して初めてマッチします。

tools/fix.go at 3c07937fe18c27668fd78bbaed3d6b8b39e202ea · golang/tools · GitHub

distance(距離計算)

上の関数で候補のパッケージを取得しましたが、これらのパッケージと今処理中のパッケージの距離を計算します。 この距離というのは、どのパッケージを優先してimportするべきか、という優先度になります。 計算自体はdistanceで行われています。

もう段々分かってきたかと思いますが、このdistanceも非常にシンプルです。 まず相対パスを出します。あとはその相対パス内のスラッシュ(正確にはセパレータ)の数を数えるだけです。

例えば、"../vendor/github.com/knqyf263/foo" であれば4になります(実際の処理では最後に1足してるので5ですが)。 ../../../../github.com/knqyf263/foo であれば6になります。 つまり、パス的に近い遠いを測っています。 すぐに難しいアルゴリズムとか持ち出してこない感じで好きになりました。

tools/fix.go at 3c07937fe18c27668fd78bbaed3d6b8b39e202ea · golang/tools · GitHub

byDistanceOrImportPathShortLength(優先度順に並び替え)

上で計算したdistanceの順に並び替えます。 もしdistanceの値が同じであればimportPathShortの長さで並び替えます。 github.com/knqyf263/foogithub.com/knqyf263/foobar であれば前者が先になります。 それも同じであればあとはimportPathShortの辞書順に並び替えます。

symbolの確認

このあとの処理は候補のパッケージそれぞれについてgoroutineで並列に行います。 まず候補のパッケージがexportしているsymbolをloadExportsで全て取得します。 このloadExports内も今までどおりASTで頑張っています。

loadExportsで全てのsymbolが得られたら、処理中のパッケージが利用しているsymbolと比較します。 全てexportされているものが存在したら該当のパッケージを見つけたということで処理を終えます。 この時、importPathの最後の要素とパッケージ名が異なる場合はneedsRenameをtrueにして返します。

長くなりましたが、これでfindImportの処理は終わりです。

未importのパッケージをimportする(続き)

findImportの結果、空文字列が返ってきた場合は条件を満たすパッケージが見つからなかったことを意味するので終了します。 もし見つかっていれば、resultsに入れます。

そのresultsそれぞれに対して astutil.AddImportastutil.AddNamedImport を呼び出してソースコードにimportを追加します。 あとは追加したimport pathをまとめてreturnします。

tools/fix.go at 3c07937fe18c27668fd78bbaed3d6b8b39e202ea · golang/tools · GitHub

これで重要なところは全て終わりです。

まとめ

せっかく読んだしまとめておこうと思ったら長くなったし誰が読むんだって感じになってしまいました。 ですが、恐れずに読むと意外と普通のことをやっているだけだったりするのでブログは読まなくてもソースコードは一度読んでみることをおすすめします。