えっちな grep をつくった

H(uman-friendly) な grep コマンド hgrep をつくりました.

github.com

f:id:rhysd:20211123211157p:plain
'\w+ で検索した時の出力

ファイルを特定のパターンで検索し,マッチした箇所を構文ハイライトしたコード片で表示します.超ざっくり言うと,ripgrep で検索して bat でマッチ箇所付近を表示するような感じです.

grep -C によるコンテキスト表示に似ていますが,マッチ行が近い時は1つのコード片にまとめる,周囲何行を表示するかをヒューリスティックに少し賢く決めているなど,ちょっと出力は工夫しています.

動機

手元のリポジトリでコードを検索する時は

  • 単純に grep で検索してマッチ結果を眺める
  • grep | fzf のように検索結果を fzf で絞り込んだりプレビューする
  • vim $(grep -l ...) のように検索結果をエディタで開く

あたりを使い分けているのですが,マッチ結果全体を個々のマッチの周囲を見ながら眺めるというのがやりにくいと常々思っていました. 要は GitHub のコード検索結果のような出力を手元でも見たいなというのが,hgrep を使った動機です.

インストール方法

  • リリースページから zip をダウンロードして解凍し,なかに入っている実行ファイルを /usr/local/bin などにコピーしてください.
  • cargo を使って cargo install hgrep でインストールできます.ソースからビルドする場合は feature flags を使って自分のほしい機能だけを有効にしてインストールすることもできます(バイナリサイズや依存パッケージを減らせる).
  • Homebrew (macOS もしくは Linux),MacPortsmacOS),pkgin(NetBSD)などのパッケージマネージャもサポートしています.

詳しくはドキュメントを参照してください.

使い方

ripgrep と同じように使えます.普段 rg コマンドを使ってるなら,それを hgrep に置き換えるだけでほぼ問題ないと思います.

hgrep [options] pattern [path...]

例えば ./src 以下の なんとかPrinter を検索したければ

hgrep '\w+Printer' ./src

のようにするとマッチ箇所のコードスニペット一覧が出力されます.

また,grep -nH の出力を標準入力から受け取ることもできます.hgrep ではサポートできていない検索方法を使いたい時や,一旦ファイルに保存したり加工した結果を表示する時に便利です.

grep -nH pattern -R [path...] | hgrep

例えば ./src 以下の なんとかPrinter を検索したければ

grep -nH '\w+Printer' -R ./src | hgrep
rg -nH '\w+Printer' ./src | hgrep

カスタマイズ

プリンタを指定する

hgrep はコード片をハイライトして表示するのに使うプリンタ実装が syntectbat の2つあり,--printer オプションで指定できます.

  • syntect: sytnect ライブラリを使って自前で実装したもの(デフォルト)
  • bat: bat をライブラリとして使って実装したもの

初めは bat を表示部分に使おうと思っていたのですが,パフォーマンスや出力結果に色々不満が出てきてしまったので,結局 syntect で自前実装したという経緯があり,特に理由が無ければ syntect がおすすめです.

  • bat プリンタの2倍〜4倍ほど高速
  • マッチ範囲がハイライトされるなど,出力やレイアウトが hgrep 向けに最適化されている
  • より多くのカラーテーマに対応
  • --background オプションによる背景色の描画(bat はカラーテーマの背景色を反映しない)
  • ターミナル互換性周りの改善(├ や ┬ などの unicode 文字がうまく表示できないターミナル向けのオプションや16色しか使えないターミナルへの対応)

syntect プリンタにはこれらの利点があります.

カラーテーマを選ぶ

--list-themes オプションを使うと各カラーテーマのプレビューを表示できるので,その中から気に入ったテーマを見つけることができます.

# 背景色なしでプレビュー
hgrep --list-themes
# 背景色ありでプレビュー
hgrep --list-themes --background

デフォルトの Monokai の他にも30種類のカラーテーマが用意されています.

また,あまりカラフル過ぎると見づらいという人向けに,Cyanide や Carbonight などの色数を抑えたカラーテーマもあります

ayu-dark ayu-mirage ayu-light
ayu-dark ayu-mirage ayu-light
Carbonight predawn Material
carbonight predawn cyanide

デフォルトで使うオプションを指定する

シェルの alias コマンドを使って,hgrep コマンドを上書きしてください.

# --hidden: 隠しファイルを検索対象にする
# --theme: ayu-dark をカラーテーマに使う
# --background: 背景色を描画する
alias hgrep='hgrep --hidden --theme ayu-dark --background'

これを例えば bash なら .bash_profile あたりに設定しておくと良いと思います.

また,less などの pager を使いたい場合は,下記のように関数を使って hgrep コマンドを上書きしてください.

# less を使ってページングする
function hgrep() {
    command hgrep --term-width "$COLUMNS" "$@" | less -R
}

標準出力を別プロセスにつなぐとターミナルのウィンドウサイズが分からなくなるので,--term-width$COLUMNS で明示的に指定する必要があります.

コマンド補完

--generate-completion-script オプションでシェルの補完スクリプトを標準出力に出力します.これを補完スクリプトに設定することで,オプションなどの補完が効くようになります.

# Zsh を使っていて set comps=~/.zsh/site-functions している場合
hgrep --generate-completion-script zsh > ~/.zsh/site-functions/_hgrep

BashZsh,Fish,PowerShell,Elvish に対応しています.

実装周りの話

依存ライブラリ

rayonripgrepsyntect あたりをライブラリとして使ってます.

  • rayon: 有名なお手軽データ並列処理ライブラリ.便利すぎる.
  • ripgrep: Rust から rg コマンドとほぼ同等の機能・パフォーマンスの grep 実装を直接使うことができます.ライブラリとしてもよく整備されています.
  • syntect: シンタックスハイライトするためのライブラリです.Sublime Text のハイライト定義やカラーテーマを使うことができ,多くのカラーテーマや言語ごとの構文ハイライトルールの資産を使うことができます.さらに,複数のスレッドで並列にハイライト処理を行えるようにもなっていて(後述),よくできてます.

パフォーマンス

最初は ripgrep と bat をライブラリとして使い,ripgrep の検索結果を bat の bat::PrettyPrinter で表示する小さいツールを考えていたのですが,残念ながらパフォーマンスが満足のいくものになりませんでした.

問題はコードをハイライトして出力する部分で,

  • そもそも構文ハイライトは重い処理です.大量の正規表現マッチを繰り返してコードの各部分のハイライト色を決定します
  • 正確に構文ハイライトするにはファイルの頭から順番にハイライトを計算していく必要があります.なのでファイルの後ろに行くほど時間がかかります
  • 構文ハイライトに使うハイライトルール定義を合計で数 MB ほどロードする必要があります

これらを直接解決するのは難しいので,十分なパフォーマンスを出すために,処理を並列化することにしました.そこで rayon を使い,ファイル単位でスレッドを割り当ててファイル内検索(grep)→ コード片にまとめる処理 → ハイライトの計算 を並列で行って,最後に出力処理だけは直列化して行っています.

ここで前述の syntect の並列ハイライト対応をうまく活かすことができました.残念ながら bat::PrettyPrinter は並列に使える実装になってませんでした.

syntect を使った場合 bat を使った場合

これによってハイライト処理が劇的に速くなり,2倍以上高速化しました.

マルチスレッド絡みのバグ

syntect で並列にハイライト処理をした時だけごく稀にクラッシュする問題を見つけて修正したりしました.普段使いでは一度だけ再現したクラッシュバグで,その後ベンチマークで高速で初期化→入力→描画の処理をひたすら繰り返すと10回に1回程度の頻度で再現させることができ,スレッド数を1に絞ると再現しなくなるというものでした.

該当箇所だけを抜き出すとこんな感じです.

use lazycell::AtomicLazyCell;

struct X {
    cell: AtomicLazyCell<Regex>,
}

impl X {
    fn new() -> Self {
        Self { cell: AtomicLazyCell::new() }
    }

    fn get_regex(&self) -> &T {
        if let Some(x) = self.cell.borrow() {
            x
        } else {
            let regex = ...;
            self.cell.fill(regex).ok();
            self.cell.borrow().unwrap()
        }
    }
}

lazycell は値の初期化を遅らせる機能を提供するライブラリで,AtomicLazyCell は最初「空」の状態で初期化され,fill() が呼ばれた時に初めて初期化されます(初期化より前に値を取ろうとするとクラッシュします).borrow() はすでに初期化された値があればその値への参照を Some で返し,そうでなければ None を返します

get_regex() メソッドは「cell が初期化されていればすでにある値を返し,初期化されていなければ値を生成して cell を初期化してから値を返す」というメソッドです.クラッシュを起こしたのは self.cell.borrow().unwrap() の部分で,直前で fill を呼んでいるにも関わらず borrow()None を返しているということを示しています.

このバグは2つ以上のスレッドがほぼ同時に self.cell.fill(regex) を実行した場合にのみ起こります.AtomicLazyCell は内部でロックを取って値を初期化するスレッドを1つに限定しているのですが,

  1. スレッド1が fill() の呼び出しを開始
  2. スレッド1が fill() の呼び出し内部で cell のロックを取得
  3. スレッド2が fill() の呼び出しを開始
  4. 既にスレッド1が cell のロックを取得しているのでスレッド2はロックを取れない
  5. fill() はロックが取れなかった場合は諦める実装になっているので,スレッド2は fill() から return する
  6. スレッド2はすぐに次の行の borrow() を実行する.しかしこの時,スレッド1はまだ cell に値をセットできていない.そのため borrow()None を返す(後続の unwrap() でクラッシュ)

という流れでした.

これは lazycell が,fill() で値を初期化しに逝く時に既にロックを誰かが取っていると即諦める仕様になっているのが問題なので,once_cell を使うことで解決できました.

once_cell は複数スレッドが同時に1つの cell を初期化しようとした時,処理をキューイングして後に来たほうを(busy loop で)待たせる実装になっているのでこの問題は起こりません.busy loop なので,待たされている側のスレッドが長時間待たされると計算資源を食いつぶしてまずいですが,少なくとも今回は Regex の初期化でそこまで時間がかかる処理ではなかったので問題ありませんでした.

まとめ

Human-friendly な出力の grep コマンド,hgrep をつくりました.完全に自分の使いたいユースケースに向けてつくってますが,もしよければ試してみてください.とりあえず欲しい機能はすべて揃っている状態です.

パフォーマンスは現時点でも問題ないですが,コード片の範囲を決める計算を O(n) から O(1) に落とす,ハイライトルール定義の読み込みを遅延する,構文ハイライトを計算する範囲を削ってハイライト処理をサボるなどいくつかアイデアはあるので,時間がある時にでも試しに実装してみようかなと思ってます.