tako開発日記 #2: 仮想リストでチャットのメモリ残留を1/10に

最近のtakoはパフォーマンス改善スプリントの真っ最中で、これはその中で残っていた宿題のひとつ。Rust+gpuiで書いているターミナルアプリtakoで、チャットビューを開いて閉じてもメモリが戻らない問題が残っていた。会話本文をgpui::listで仮想化して根治したので、原因の特定から座標系の作り直しまでを書く。

残留していたのはレイアウトノードだった

gpuiのレイアウトはtaffyが持っている。要素を組むと測定用のレイアウトノードが作られるが、これがnode_context_dataに残り、TaffyTree::clear()では消えない。つまり1フレームで作った要素の数に比例してヒープが積み上がる。

厄介だったのは、効くのが「会話の長さ」ではなく「1タブにチャットペインが何枚あるか」だったこと。takoは会話をCHAT_TAIL(50件)で頭打ちにしているので、1本の会話がいくら伸びても組む量は増えない。ところが表示中のペインは全部が自分の会話を毎フレーム組むので、masterとworkerを並べて動かす実運用の構成が直撃していた。

会話本文をgpui::listへ移す

ListStateを作り、itemを返すクロージャをgpui::listに渡す。itemは「発話1件」または「末尾の付随要素(提案カード/作業中インジケータ/承認カード)」で、旧経路が本文コンテナへ縦に並べていた順と1:1になるようにした。

let list_state = self.chat_body_list_state(pane_id, count);
if snap_to_bottom {
    list_state.scroll_to_end();
}
let app = cx.entity().downgrade();
gpui::list(list_state, move |ix, _window, cx| {
    let Some(app) = app.upgrade() else {
        return div().into_any_element();
    };
    app.update(cx, |app, cx| app.render_chat_item(pane_id, ix, cx))
})
.with_sizing_behavior(gpui::ListSizingBehavior::Auto)
.py(px(CHAT_BODY_PADDING_Y))
.flex_1()
.into_any_element()

itemの中身は「呼ばれた時点の状態から引き直す」形にした。listのクロージャはrenderが返ったあとのprepaintで呼ばれるので、materialをrender側で確定させてRcで渡している。

/// 仮想リストの item が引く「このフレームの並び」
pub(crate) struct ChatRenderPlan {
    /// 表示する発話
    pub(crate) messages: Vec<ChatMessage>,
    /// 並べる item
    pub(crate) items: Vec<ChatItemKind>,
    /// 発話ごとの行範囲（索引の座標系）
    pub(crate) lines: Vec<ChatMessageLines>,
    /// 狭幅レイアウト（ペイン幅で決まるので item 側からは分からない）
    pub(crate) compact: bool,
    pub(crate) activity: Option<String>,
}

つまずき1: 選択とコピーの座標系が描画に引きずられる

いちばん時間を使ったのがここ。仮想化前は「描いた順に行を索引へ積む」でよかった。行番号はself.texts.len()、つまりそれまでに描いた行数だったからだ。

// 旧: 描画が行番号の正だった
let line = self.texts.len();

仮想化すると可視の発話しか描かないので、この行番号は「画面に見えている中での何行目か」になってしまう。⌘Aで全選択したときの範囲も、コピーした本文も、クリック位置のヒットテストも全部ずれる。

解決は座標系の出どころを描画から引き剥がすこと。行テキストは文書全体ぶんを先に作る専用の経路を用意し、そこを唯一の正にした。

/// 発話 1 件が占める「選択行」のテキストを `out` へ積む（**行の座標系の正**）。
///
/// 描画が `ChatTextIndex::push` / `push_spacer` を呼ぶ順・本数と
/// **1:1 で一致していなければならない**。仮想化後は可視ぶんの発話しか
/// 描かないので、コピー・⌘A・ヒットテストの座標系はここが唯一の出どころになる。
///
/// 要素は 1 個も作らない（文字列の連結だけなので taffy ノードは増えない）
fn push_chat_message_lines(
    &mut self,
    pane_id: PaneId,
    message: &ChatMessage,
    out: &mut Vec<String>,
) {

そのうえで、可視の発話は自分専用の局所受け皿へレイアウトを積み、描き終わってから文書全体の枠へ書き戻す。索引側は「この受け皿が担当する文書上の先頭行」を1つ持つだけでよくなった。

pub(crate) struct ChatTextIndex {
    pub(crate) texts: Vec<String>,
    /// 行番号 → 実描画レイアウト（**ヒットテストの正**。文字の無い行は None）
    pub(crate) layouts: Vec<Option<gpui::TextLayout>>,
    /// この受け皿が担当する**文書上の先頭行**
    base_line: usize,
}

// 新: 局所に積んでも行番号は文書側の座標
let line = self.base_line + self.texts.len();

ここが噛み合っているかは自動では気づけないので、「plan側が割り当てた行数」と「実描画が積んだ行数」が食い違った回数を数えるカウンタを足し、セルフテストで0であることを見張るようにした。

/// 行の割り当て（plan）と実描画の本数が食い違った回数。
/// 正常なら 0 のままで、セルフテストがそれを見張る
chat_index_mismatch: usize,

つまずき2: ListAlignment::Bottomは絵を変える

チャットは下端追従なので、素直に考えるとListAlignment::Bottomを選びたくなる。実際に試したら会話が短いときにペインの下端へ貼り付いた。旧経路は上詰めのスクロールdivだったので、絵がはっきり変わる。

下端追従は器の既定ではなく、呼び出し側がscroll_to_end()で明示すればいい。器の既定は上詰めにした。

// **`Bottom` ではなく `Top`**。`ListAlignment::Bottom` は内容がビューポートより
// 短いときに会話をペインの**下端へ貼り付ける**ので、旧経路（上詰めのスクロール div）
// と絵が変わる。下端追従は呼び出し側の `scroll_to_end` が担う。
// overdraw は画面外にも組んでおく高さで、ドラッグ選択がペインの端をまたぐときに
// 「行がまだ無い」状態を避けるため広めに取る
let state = gpui::ListState::new(count, gpui::ListAlignment::Top, px(600.0));

つまずき3: item数が変わると見ている位置を失う

新着や折りたたみの開閉でitem数は動く。ListStateは数が変わったら作り直しになるので、素朴に作り直すとスクロール位置が飛ぶ。論理スクロール位置(item番号+item内オフセット)を作り直しの前後で持ち越すようにした。

fn chat_body_list_state(&mut self, pane_id: PaneId, count: usize) -> gpui::ListState {
    if let Some((state, built_for)) = self.chat_body_lists.get(&pane_id) {
        if *built_for == count {
            return state.clone();
        }
        let state = state.clone();
        let keep = state.logical_scroll_top();
        state.reset(count);
        if keep.item_ix < count {
            state.scroll_to(keep);
        }
        self.chat_body_lists.insert(pane_id, (state.clone(), count));
        return state;
    }
    // ...
}

つまずき4: 器が変わると外側の問い合わせが全部ずれる

チャット本文に対しては、外側からいろいろな問い合わせが飛んでいた。「本文の矩形はどこか」「いま何要素並んでいるか」「どれだけ下へ送ったか」といったものだ。旧経路はScrollHandle、新経路はListStateで、同じ問いに対する答えの単位が違う。

たとえば矩形。ListState::viewport_bounds()が返すのは左右余白の内側なので、旧経路のScrollHandle::bounds()(borderボックス)と比べるには余白を足し戻す必要がある。

pub(crate) fn chat_body_bounds(&self, pane_id: PaneId) -> Option<gpui::Bounds<gpui::Pixels>> {
    if let Some((list, _)) = self.chat_body_lists.get(&pane_id) {
        let inner = list.viewport_bounds();
        if f32::from(inner.size.height) > 0.0 {
            let compact = self.chat_render.get(&pane_id).is_some_and(|p| p.compact);
            let pad = px(chat_body_padding_x(compact));
            return Some(gpui::Bounds {
                origin: gpui::point(inner.origin.x - pad, inner.origin.y),
                size: gpui::size(inner.size.width + pad * 2.0, inner.size.height),
            });
        }
    }
    self.chat_scroll_handles
        .get(&pane_id)
        .map(gpui::ScrollHandle::bounds)
}

スクロール量に至っては符号すら違う。スクロールハンドルは「下へ送ると負のオフセット」、仮想リストは「下へ送ると先頭item番号が増える」。両方を先頭表示なら0・下へ送ると負へ正規化する関数を1本置いて、比較する側からは器の違いが見えないようにした。

/// 「どれだけ下へ送ったか」の指標（器が違うと単位が違うので符号だけ揃える）
pub(crate) fn chat_scroll_mark(&self, pane_id: PaneId) -> f32 {
    if let Some((list, _)) = self.chat_body_lists.get(&pane_id) {
        let top = list.logical_scroll_top();
        return -(top.item_ix as f32 + f32::from(top.offset_in_item) / 1000.0);
    }
    self.chat_scroll_handles
        .get(&pane_id)
        .map(|h| f32::from(h.offset().y))
        .unwrap_or(0.0)
}

器の違いはこの手の関数5本(chat_body_bounds / chat_scroll_to_top / chat_scroll_to_bottom / chat_item_count / chat_scroll_mark)に閉じ込めた。ここを通す限り、呼ぶ側は仮想化されているかを知らなくていい。

なおホイールにも罠があった。本文の器をlistへ移したあと、上位のスクロール委譲は従来どおり外側のdivに残っていたが、listは自前のon_mouse_eventでホイールを処理する。合成したホイールイベントを実際に配送しないと届かない。

実測

計測は隔離インスタンスで描画を自分で回すハーネス(TAKO_VISUAL_ONLY=chat-leak)を書いた。GPUIは遮蔽されると描画を止めるので、裏で起動しただけでは残留が再現しない。入力は実transcriptのtail 50(=534個のmdブロック)。

チャットペイン数	旧経路	仮想化後
1枚	11.32 MB	2.53 MB
4枚	43.78 MB	4.10 MB
8枚	86.24 MB	8.18 MB

1フレームで整形した行数は580 / 2,320 / 4,640行から26 / 104 / 208行になった。一方で索引が持つ行数は変わっていない。これが「選択・コピーは描画に依存しない」の実測での確認になっている。

A/B比較のために旧経路は環境変数TAKO_830_NO_CHAT_VIRTUAL_LIST=1で残した。1発話の作り方は仮想化と同じ関数を通すので見た目は完全に一致する。セルフテストの新設項目は、旧経路だと「整形した行=全行」でFAILEDになることまで実測してある。

唯一の挙動差

チャットを開いた最初のフレームで、いちばん新しい発話が表示されるようになった。旧経路はScrollHandle::scroll_to_bottom()が前フレームの実測に依存するため初回は動かず、tail 50のいちばん古い発話から表示されていた(承認カードが出ていても画面外だった)。「既定は追従」の設計意図どおりなので、直った側を採った。

学び

仮想化は「描画を減らす作業」だと思って着手したが、実際にやっていたのは座標系の正をどこに置くかを決める作業だった。描いた結果から行番号を数えている限り、可視だけ描く実装とは両立しない。テキストの座標系を「要素を1個も作らない経路」として切り出せたのが効いた。

もう1つは、器を差し替えるときに同じ問いに同じ意味で答える関数を先に用意すること。矩形・要素数・スクロール量のように単位も符号も違うものを、呼ぶ側に条件分岐させると差分が全体に散る。5本の関数へ閉じ込めておくと、旧経路を環境変数で残したままA/B計測できるという副産物もついてきた。