この記事は Jij Advent Calendar 2025、Rust Advent Calendar 2025 シリーズ1、およびHaskell Advent Calendar 2025の3日目の記事です。

更新履歴

2025-12-03 00:00 - 初版自動公開
2025-12-03 08:40 - 軽微な誤字の修正

はじめに：Haskeller × Rustacean = ∞

筆者は比較的年季の入ったHaskellerですが、現在は業務ではRustを書いてくらしています。
昨年転職一ヶ月後にこんな記事を書いたりしていたように、Rustでのプログラミングはかなり関数型と命令型のいいとこどりという感じでとても快適な一方、Haskellならこんな書き方が出来るなあというような箇所もあり、刺激的な日々を送っています。

そんな中でHaskellerの知識とRustaceanの知識を組み合わせて1000倍の速度向上を達成した事例があったので、是非紹介したいと思います^[1]。

文脈：数理最適化用ドメイン特化言語JijModeling

Jijでは、数理最適化モデリングライブラリJijModelingを提供しています。
これはPythonライブラリとしてPyPIから無償で利用できるほか、弊社の提供するクラウド数理最適化プラットフォームJijZeptでも最適化問題を記述するための基盤として用いられています。

https://www.jijzept.com/ja/

https://pypi.org/project/jijmodeling/

数理最適化とは、作業シフトであったり配送計画であったりといった現実世界の社会の「計画」業務を数式を使って表し、何らかの意味で「最適」な解を求める、というものです。
特に、JijModelingは数理最適化の問題をPythonで記述しソルバに渡すためのツールです。
以下の議論では、「JijModelingはPythonで記述された数式と入力データをソルバが解けるような形に変換するコンパイラのようなもの」と理解しておけば大丈夫です。

JijModelingはPythonライブラリですが、内部的にはRustで実装されています。
2025年11月現在、Jijでは次期バージョンであるJijModeling 2の開発をフルスクラッチで進めており、ベータリリースを重ねています。

JijModeling 2では、単なるPythonライブラリから自立したプログラミング言語へと脱皮することを企図し^[2]、構文木の定式化から見直しが行われたり、自前の型システムを新たに設計するなどの刷新が図られています。
ここでも、Haskellの主要処理系であるGHCの実装で用いられているTrees that growと呼ばれる手法を適用するなど、Haskell周辺で発明された関数的な手法が応用されています。

JijModelingは数理最適化用のドメイン特化言語ですが、より具体的には、Pythonコードとして記述された数理最適化問題を、OMMXと呼ばれるソルバや言語に依存しない数理最適化問題用の中間形式にコンパイルするものです。

旧来のJijModelingでは、Pythonで記述された問題を直接OMMXへとコンパイルしていましたが、JijModeling 2ではコンテナに対する map / filter 関数や畳み込み演算など、高階関数を使ってより高級なプログラミングが可能なように設計されているため、式の抽象構文木であるExpressionとOMMX形式との間にJijModeling 2固有の値表現（Value）が中間表現として挟まる形に再設計が行われました。

この定義を簡略化したものが以下です：

enum Value {
  Literal(f64),
  Polynomial(Polynomial),
  Bool(bool),
  String(String),
  Array(Array<Value>),
  JaggedArray(JaggedArray<Value>),
  Closure(Box<dyn Fn(Value) -> Value + Send + Sync + Clone + 'static>),
}

数値リテラルや多項式の他、高次元配列（Array）やシェイプが一定でないものも許容するJagged配列（JaggedArray）、そして関数クロージャに相当するClosureなどがあることがわかると思います。
特に、クロージャの表現として Rust のクロージャをそのまま使うHigher-Order Abstract Syntax（HOAS）と呼ばれる形式を採用しています。
Rustでは「関数クロージャっぽく振る舞う型」を抽象化したトレイトが複数あります：

Fn(A1, ..., An) -> B: 型 A1, ..., Anの引数をとり、B型の値を返す、何回でも呼べる関数。
- 特に、Fn(...) -> ... は（抜け穴を使わない限り）可変状態を含んでいないことが静的に保証された関数です。
FnMut(A1, ..., An) -> B: Fn(A1, ...) -> Bと似ているが、内部に可変状態を保持しており呼び出すたびに状態が変更されうる関数。
FnOnce(A1, ..., An) -> B: 一回しか呼び出すことができない関数。

ここでは、値は何回でも複製されたり呼び出されたりし得、また可変状態が含まれていないほうが望ましいため、Fn(..) を使ってクロージャを表現しています。
残りの制約は、Valueがクローン可能であるためにCloneが、スレッド間通信のために Send, Syncが、所有権の兼ね合いで 'static というライフタイム制約がそれぞれ課されています。

このように、所有権に関わる制約に気を付けさえすれば、関数型プログラミング言語のように第一級の値として関数クロージャを手軽に扱えるのはRustの大きな強みです。

!

実際には、Cloneはそのままではトレイトオブジェクト（Box<dyn ...> みたいなやつ）にできません。
こういった場合、dyn-cloneクレートを使って次のようにこれらのトレイトをまとめたdyn対応なトレイトを定義してそのトレイトオブジェクトを扱います。

use dyn_clone::{clone_trait_object, DynClone};

pub trait ClosureBodyLike: Fn(Value) -> Value + DynClone + Send + Sync {}

clone_trait_object!(ClosureBodyLike);

impl<F> ClosureBodyLike for F
where
    F: Fn(Value) -> Value + DynClone + Send + Sync,
{
}

pub enum Value {
  // ...
  Closure(Box<dyn ClosureBodyLike + 'static>)
}

これはわかりづらいので、以下の説明ではBox<dyn ...>の中にCloneを展開した形で書いていますが、実際にはこのような対処が適宜必要となることに留意してください。

問題：大規模最適化問題のコンパイルが遅すぎる

今回の問題は、この「高級な中間表現」の導入に伴うパフォーマンスの悪化でした。
十分な機能が揃ったところでJijModeling 1で用いられていたベンチマークスイートを移植したところ、あるケースでは旧来数百ミリ秒程度で終わっていたころ267秒もかかるようになっていました。
具体的にはこれはMIPLIBというベンチマーク集のsupportcase18というもので、数万程度の変数・パラメータを含むものです。
流石にこれは時間がかかりすぎているということで、まずは同僚にお願いしてベンチマークの結果を詳しく分析してもらいました。

flamegraphによるプロファイル結果の可視化

ベンチマークの常道として、まずflamegraphによるコストの可視化がまず行われました。
以下がその当初の flamegraph（に書き込みを加えたもの）です。

flamegraphの例

flamegraphの各アイテムは個別のコールスタックの地点に対応し、各横軸が実行時間に占める割合になっており、呼び出しごとに横に分割されコールスタックが縦に積み上がっていく形です。
幅の太いスタックを見付けたら、そこから屠っていけばいいという訳です。
テモート^[3]でRustプログラムのflamegraphを生成するには、cargo-flamegraph を使います。

flamegraphはRust固有のものではなく、たとえばHaskellでもghc-prof-flamegraphなどを使えば、プロファイルビルド済のバイナリの出力からflamegraphを生成することができます。

最初の課題：ValueをCloneしすぎだったので参照を使おう

同僚のIagoさんが上のflamegraphを分析したところ、一番大きな割合を占めているのは巨大配列の添え字を取る関数（overloaded_subscript）で、特にCloneに膨大な時間が費されていることがわかりました。
この関数では、エラー時に参考のために評価前の値や添え字を含めるようにしており、ここでmoveが起きないように値はすべて参照の形で受け取っていました。
そのため、一部の配列操作でValueをCloneする必要が生じていたのです。
そもそも、エラーメッセージに巨大な配列の値を入れてもユーザにとって嬉しいことは何もなく、また評価前の式の情報は同じ範囲で既にスタックに詰まれる仕組みになっています。
なので、ここの部分ではコールスタックを保存しないようにし、値も即値でとるようにしたところ、267秒→167秒まで短縮することができました。

ここまでのまとめ：

無駄なCloneを排除：267秒→167秒

本丸：遅延配列をつかって配列の変形に伴うアロケーションを節約

しかし、これでもまだ旧来の数百ミリ秒に比べて遅すぎます。
そこで改善後のflamegraphを再び眺めながら同僚と検討してみました。

今度は大半がJaggedArrayの操作に

添付画像だと見えなくなってしまっていますが、実際のflamegraphでは個別のアイテムをクリックすると拡大表示できます。
左上と真ん中右側の山がとにかく大きく、真ん中左側と一番右も似たような形をしていますね。
これらが大部分を占めているので確認してみると、多い順に以下のようなものが占めていることがわかりました：

Jagged Arrayの要素を関数でmapを取る操作
部分的な添え字に沿ってJagged Arrayの像を取る操作
Jagged ArrayそのもののClone

……いずれも Jagged Array に関連するものです。
先にも軽く触れましたが、Jagged Array というのは、次元は定まっているがシェイプが一様でないような多重配列です。
たとえば、次は3次元の Jagged Array の例です：

[ 
    [[1,2,3],[4,5]],
    [],
    [[6]]
]

空な成分を無視すれば次元は3の配列ですが、各軸の長さは一定ではないので、numpyなどのシェイプが定まった（regularな）多重配列（テンソル）と区別し、Jagged Array と呼ばれます。
Jagged Array というのは一つの呼び名にすぎず、Ragged Array や Irregular Array と呼ばれることもあります。

数理最適化を実務適用する上では、一般に数万以上の要素からなる多重配列やJagged Arrayが現れますが、今回のsupportcase18もまさにそうしたケースに該当します。
そして、その配列の中の一つの要素が必要であってもValueを即値で取るために配列が丸ごとcloneされてしまったり、配列に対するmapが呼び出されるたびに新たに配列をアロケートして引き回したり、といったことが発生し、結果的に大幅な速度低下に繋がっていたのです。

添え字の際にcloneを避けるだけであれば、配列を Arc や Rc などの参照カウンタで包んで、インデックスを参照で取った後に要素だけcloneすれば十分でしょう。
Arcはアトミックな参照カウンタの機能を提供しており、スレッドセーフなRcの亜種です。
また、旧来のJijModeling 1系統では、極力配列は参照で持つ形にして大規模配列のcloneやアロケーションを避けていました。
しかし、JijModeling 2ではmapやzip、部分的な添え字に沿った像を取る操作、軸に沿った和（[[1,2,3],[4,5,6]].sum(axis = 1)とすると [6,15] になるような）など、配列を変換する高レベルな関数が提供されるようになります。
こうした操作をサポートしようと思うと、配列を参照で持ったり、参照カウンタで包んだりしても操作の度にアロケーションが走ることになってしまいます。
つまり、JijModeling 2ではそれ以前とは根本的に異なる解決策が必要になったのです。

……というところで、筆者のHaskellerとしての知識の出番が訪れました。
Haskellの世界では、こうした「配列やストリーム上の関数」に伴うアロケーションを極力減らすための手法として、融合変換（Fusion）という手段が実用的に用いられています。
その名の通り、配列やストリームなどに対する複数回の演算を一つの演算に融合し、中間データのアロケーションを節約するという手法です。
ひとくちに融合変換といっても色々ありますが、Haskellでは一次元配列に対するストリーム融合（stream fusion）が盛んに用いられており、産業的にも日々（もはや意識されないほどに）応用されています。
これは、配列に対する演算を一旦ストリームを表すより一般的なデータ型を介する形で実装しておきつつ、コンパイラにストリームと配列との相互変換を融合して消滅させるような等式的な書き換え規則を与えて最適化させる、という手法です。

今回最初に考えたのも、このStream Fusionの手法をRustでもできないか？ということでした。
しかし、HaskellでStream Fusionが手軽に可能なのは、Haskellが純粋な言語であり、副作用を気にせずに等式的な書き換えによる最適化を行うことができる、という事情があるためです。
また、ストリームを表す単一の型を実装するには、存在型やGADTsと呼ばれる進んだ型システムの機構が必要になります。

こうした高級な型レベルの機能がRustにはまだ存在しないため、Rustのストリーム型にあたるIteratorは単一の型ではなく、Iteratorトレイトを実装するデータ型を個別に定義し、それらを合成する形で実現されています。
JijModeling2では融合変換の対象になるコンテナや演算は限られているので、ここまでオーバーキルをせずとも、「配列とその上の演算」で定義されるようなこんな型を定義すれば、似たようなことはできそうです：

pub enum ValueArray {
    RawArray(ArrayD<Value>),
    Map{func: Box<dyn Fn(Value) -> Value + Clone>, arg: Box<Self>},
    Zip{func: ..., lhs: Box<Self>, rhs: Box<Self>},
    Sum{axis: Vec<usize>, arg: Box<Self>}
}

しかし、これでは新しい演算を実装するたびにValueArrayのバリアントが新しく生えることになってしまいますし、効率的な評価を実現するには、各演算に対応するバリアント（MapやZip、Sumなど）の組み合わせごとに、いちいち人間が「融合」する処理を書いてやる必要がでてきてしまい、あまりスケールしなさそうです。

どうしようかな……と考えていたところで思い当たったのが、遅延配列による融合変換の手法です。
遅延配列は、Haskellの並行多重配列ライブラリである Repaやmassivで用いられている手法であり、配列を直接メモリで表すのではなく、シェイプの情報と、「添え字から要素への関数」の組で表現するものです。
たとえば、通常の（regularな）多重配列の遅延版は以下のように定式化できます：

#[derive(Clone)]
struct DelayedArray<A> {
    shape: Vec<usize>,
    gen_fun: Arc<dyn Fn(&[usize]) -> A + Sync + Send>,
}

このような表現と、実際の配列の間には互いに同型が成り立つことがわかるでしょう。
たとえば、Rustの多重配列ライブラリである ndarray には、シェイプと添え字から要素への関数を引数として多重配列を生成する shape_shape_fn 関数が存在します。
逆に、多重配列が与えられれば、shapeの情報と添え字を取る操作によって DelayedArray<A> を復元することができます。
そして、多重配列の「等しさ」とは何かと考えてみると、それは「形状が同じで、同じ添え字には同じ値が紐付けられていること」（難しくいえば多重配列の等値性は外延性によって定められること）であると言えそうなので、この定式化によって情報は完全に保たれていそうです。

gen_fun が Arc に包まれているのは、gen_funのクロージャに元となる配列からの情報が move されて保存されるためです。
生のクロージャを Clone すると内側の Array が Cloneされてしまうため、Arcで包むことでクロージャが直接Cloneされるのを回避しているのです^[4]。

impl<A: Clone + Send + Sync + 'static> From<ArrayD<A>> for DelayedArray<A> {
    fn from(value: ArcArray<A, IxDyn>) -> Self {
        Self {
            shape: value.raw_dim(),
            gen_fun: {
                let value = value.clone();
                Arc::new(move |ix| value.get(ix).unwrap().clone())
            },
        }
    }
}

残るタスクは、この表現に対してmapやfilterなどを実装することだけです。特に大事なのは、DelayedArrayの表現だけを使って諸々の処理を実装することです。
というのも、今回の用途はアロケーションを極力避けることですので、演算の間に中間的な配列が生成されては元も子もないからです。
幸い、表現の変化によって情報が落ちることはないわけですから、この DelayedArray の表現だけをつかって、map や軸に沿った sum なども取れるはずです。
これは、添え字に対する算数が途中で必要になる、という意味で算数程度には難しい^[5]ですが、逆に言えば算数をするくらいで実装できる、ということです。たとえば、例えば、DelayedArray に対する map や、添え字に対応する値を取得する get の実装は次のようになります：

impl<A> DelayedArray<A> {
    pub fn map<B>(self, f: impl Fn(A) -> B + Send + Sync + Clone + 'static) -> DelayedArray<B> {
        DelayedArray {
            shape: self.shape,
            gen_fun: Arc::new(move |ix| f((self.gen_fun)(ix))),
        }
    }

    pub fn get(&self, ix: &[usize]) -> Option<A> {
        let ok = ix
            .iter()
            .zip(&self.shape)
            .all(|(i, dim)| i < dim);
        if ok {
            Some((&self.gen_fun)(ix))
        } else {
            None
        }
    }
}

mapするには単純にもともとの遅延配列のアクセス関数と f を合成すればよく、また単純に添え字を取るにはgen_funを呼んで要素を「計算」させればよい、という訳ですね。
非自明な例として、軸に沿った総和関数sum_axesの例を見てみましょう：

pub fn sum_axes(self, axes: &[usize]) -> DelayedArray<A> {
   let ndim = self.ndim();
   // 軸をソートして重複を除き、範囲外の軸は除外
    let axes = axes
        .iter()
        .copied()
        .sorted()
        .take_while(|i| *i < ndim)
        .unique()
        .collect::<Vec<_>>();
    if axes.is_empty() {
        return self;
    }
    let mut new_shape = self.shape.as_array_view().to_vec();
    let mut removed_shape = Vec::with_capacity(axes.len());
    for ax in axes.iter().rev() {
        removed_shape.push(new_shape.remove(*ax))
    }
    removed_shape.reverse();
    DelayedArray {
        shape: IxDyn(&new_shape),
        gen_fun: {
            Arc::new(move |outer_idcs| {
                let val = (std::iter::Sum::sum)(
                    ndarray::indices(IxDyn(&removed_shape))
                        .into_iter()
                        .map(|in_ix| {
                            let mut ix: Vec<usize> = Vec::with_capacity(ndim);
                            ix.extend(outer_idcs);
                            for (ax, pos) in axes.iter().zip(in_ix.as_array_view().to_vec())
                            {
                                ix.insert(*ax, pos);
                            }
                            (self.gen_fun)(&ix)
                        }),
                );
                val
            })
        },
    }
}

中々複雑ですが、これもがんばって算数をすれば導くことができます。

さて、通常のregularな多重配列についてはこの定式化でよいですが、supportcase18で出て来るのはregularではないJagged Arrayです^[6]。
勿論、大規模なregular配列が出て来る例は幾らでもあるので、これはこれとして必要な最適化ですが、supportcase18の高速化のためには、Jagged Array 版の遅延配列も考える必要があるわけです。

Jagged Array はシェイプが定まっていないので、shapeの表現を考える必要があります。
逆に言えば、shapeの情報を過不足なく表現できれば、gen_funの部分はそのまま使えます。
遅延配列が議論されるのはregularな配列の場合がほとんどで、Jagged Array に関する議論を見付けることはできなかったので、今回は自前で考えることになりました^[7]。
結論から言えば解決策は shape も関数で次のように表現してしまうことです。

#[derive(Clone)]
pub struct DelayedJaggedArray<A> {
    ndim: usize,
    shape_fun: Arc<dyn Fn(&[usize]) -> Presence + Sync + Send>,
    gen_fun: Arc<dyn Fn(&[usize]) -> A + Sync + Send>,
}

#[derive(Clone, Copy, PartialEq, Eq, Hash, Serialize, Debug)]
pub enum Presence {
    Leaf,
    HasChildren(usize),
    Absent,
}

shape_fun は gen_fun と似ていますが、違うのはPresenceを返すことです。
Presenceは、与えられた添え字が Jagged Array の中でどういうノードなのかを表すもので、Presence::Leaf はそこに配列の要素が住んでいること、HasChildren(n) は n 個の子ノード（値かもしれないし別の配列かもしれない）を持っていること、Absent は添え字が定義域の外であることを表します。

これを使って境界検査を書き直しつつ、map などはそのまま実装すればよいわけです（sum_axesは jagged array に対してはあんまり嬉しくないので定義はありません）。
像を取る操作などもよく考えれば書くことができました（皆さんも考えてみて下さい）。

こうして定義した DelayedJaggedArray を使うようにしたことで、167秒かかっていたものが64秒と倍速で終わるようになりました。

ここまでのまとめ：

Haskell由来の遅延JaggedArrayでアロケーション節約：167秒→64秒

最後のピース：もっと細かなArcの利用、ループで書けるなら書く

当初から4倍程度の高速化が出来たとはいえ、これでもまだベースラインの数百ミリ秒には遠く及びません。
この時点でのflamegraphが次です。

ArcのdropとJaggedArrayからの変換が重い

詳しく見てみると、Arcの解放時に呼ばれる内部APIのArc::drop_slowが全体の6割、そして Jagged Array から DelayedJaggedArray に変換するコードが全体の4割弱を占めていることがわかります。

まずは Arc の解放から攻めることにしました。
ここで同僚のIagoさん・てらモスさんに相談したところ、「Arcはタダではないので避けられるなら避けたほうがよい」という結論に至りました。
もちろんアトミックな処理やカウンタの増減には非自明なコストが発生することは知識として知っていましたが、ここまで如実に時間を占める場合がある、というのは体験してみないと中々わからないところでした。
今回のような例では、map や sum_axis、要素を取る操作などを重ねるたびにクロージャが作っては壊されることになるため、大量のアトミックな操作が発生し、実行時間を圧迫していたわけです。
そこで、gen_funのクロージャ自体を Arc に包むのはやめ、遅延配列に変換する際に必ず多重配列や Jagged Array を Arc に包む方向で修正をしました。
具体的にはこんな感じです：

#[derive(Clone)]
struct DelayedArray<A> {
    shape: Vec<usize>,
    gen_fun: Box<dyn Fn(&[usize]) -> A + Clone + Sync + Send>,
}

impl<A: Clone + Send + Sync + 'static> From<ArrayD<A>> for DelayedArray<A> {
    fn from(value: ArrayD<A>) -> Self {
        let value: ArcArray<A, IxDyn> = value.into_shared();
        Self {
            shape: value.raw_dim(),
            gen_fun: {
                let value = value.clone();
                Box::new(move |ix| value.get(ix).unwrap().clone())
            },
        }
    }
}

gen_fun を Arc で包むかわりに Box で包み、そのかわりトレイト境界で Clone をクロージャに対して陽に課すことにしたわけです。
ArrayDはndarrayで提供される多重配列の型ですが、これに対して into_sharedを呼ぶと、 ArcArray<A> という配列に特化したアトミック参照カウンタつきの配列型に変換されます。ArcArrayは安価かつスレッド安全なCloneをサポートしており、Box 内のクロージャへ安全に Move できるようになります。
また、実はコンパイラの実装側では配列値は既に ArcArray の形で持つようにしていたため、into_sharedなしの ArcArra<A>からの直接変換を利用しています。

これと同様のJaggedArrayも同様に適用するようにし、Jagged Array の Arc版はないため以下のように Arc<JaggedArray<Value>> から変換させるようにし、コンパイラも Arc<JaggedArray<Value>> を保持するようにしています。

impl<A: Clone + Send + Sync + 'static> From<Arc<JaggedArray<A>>> for DelayedJaggedArray<A> {
    fn from(value: Arc<JaggedArray<A>>) -> Self {
        let shape_fun = to_shape_fn(value);
        Self {
            ndim: value.dim(),
            shape_fun,
            gen_fun: Box::new(move |ix| value.get(ix).unwrap().clone()),
        }
    }
}

ここで、to_shape_fn は、JaggedArrayのshape関数を生成する関数です。

ここまでの脱 Arc により、64秒から一気に9秒まで高速化できました。
その時のflamegraphが以下です：

あとちょっと

よく検討すると、最初の軸の長さを取得するためだけにわざわざ全体のshapeを計算しようとしている箇所などがあり、そこを削って4秒になりました。

このあたりで目処がついたような気がしたので、いっかい遠出しておいしいハンバーガーを食べにいきました。その時に食べたおいしいハンバーガーが次です。

おいしかったハンバーガー

この時点で実行時間の大半が Jagged Array から Delayed Jagged Array への変換になっていたので、風呂に入りながら考えたところ、「to_shape_fnの実装で再帰してるのがわるいじゃん」ということに気付きました。どういうことかというと、to_shape_fnは最初に書き飛ばしたときに以下のように再帰関数として実装していました：

type ShapeFn = Box<dyn Fn(&[usize]) -> Presence + Clone + Sync + Send>;
fn to_shape_fun<A: 'static + Clone + Sync + Send>(vec: &JaggedArray<A>) -> ShapeFn {
    match vec {
        JaggedArray::Leaf(_) => Box::new({
            move |ix: &[usize]| {
                if ix.is_empty() {
                    Presence::Leaf
                } else {
                    Presence::Absent
                }
            }
        }),
        JaggedArray::Node(children) => Box::new({
            // ここで再帰してる！！！
            let children = children.iter().map(to_shape_fun).collect::<Vec<_>>();
            move |ix: &[usize]| {
                if ix.is_empty() {
                    Presence::HasChildren(children.len())
                } else {
                    let i = ix[0];
                    if i >= children.len() {
                        Presence::Absent
                    } else {
                        children[i](&ix[1..])
                    }
                }
            }
        }),
    }
}

ここで、JaggedArray は木構造で表現されており、とくにJaggedArray::Leaf が配列の要素、JaggedArray::Node が子を持つ中間配列ノードに当たります。
まさにこの JaggedArray::Node のところで to_shape_fn の再帰が走り、これが分岐のたびに掛け算になって、スタックがどんどん積まれて遅くなってるやんけ！という、まあ当たり前といえば当たり前のことに気付いたわけです。

実際には、添え字が一個与えられたときにそのパスを辿ればいいだけなので、全ての子要素に to_shape_fn して分岐する必要なんてなく、単純にループを回せばいいだけなはずです。
Haskellで同様なコードを書いた場合は、おそらく遅延評価が効いて必要なパス以外の評価は発生しない筈ですが、Rustは正格言語ですのでそのあたりがシビアに効いてくるわけです。
ということで、風呂から上がるやこれをループを使う形で次のように書き直してみました：

// JaggedArray の参照を shape の presence に変換するヘルパ関数
fn view_nested<A>(v: &JaggedArray<A>) -> Presence {
    match v {
        JaggedArray::Leaf(_) => Presence::Leaf,
        JaggedArray::Node(chs) => Presence::HasChildren(chs.len()),
    }
}

fn to_shape_fun<A: 'static + Clone + Sync + Send>(value: Arc<JaggedArray<A>>) -> ShapeFn {
    Box::new(move |ix| {
        if ix.len() > value.dim() {
            return Presence::Absent;
        }
        let mut data = &value.as_ref().data;
        let mut presence = view_nested(data);
        for &i in ix {
            match data {
                NestedVec::Leaf(_) => {
                    presence = Presence::Absent;
                    break;
                }
                NestedVec::Node(chs) => {
                    if i >= chs.len() {
                        presence = Presence::Absent;
                        break;
                    }
                    data = &chs[i];
                    presence = view_nested(data);
                }
            }
        }
        presence
    });

こうして再帰をやめた結果、4秒から一気に345ミリ秒まで短縮できました。
こうして当初の267秒から345ミリ秒まで約1000倍の高速化が達成できたわけです！ :tada:

これによって、supportcase18もJijModelingのベンチマークスイートに復活させることができ、CodSpeedでCI上で大規模配列に対するパフォーマンスリグレッションを常にチェックできるようになりました。
ちなみにCodSpeedはCPUシミュレーションや隔離環境による実行でパフォーマンスを信頼性高く継続的に計測してくれるサービスで、速度低下をおしえてくれたり、flamegraphを自動で生成してくれたり弊社では以前から利用しています（個人的にも使いたいのでHaskellに対応してほしい……）。

https://codspeed.io

ここまでのまとめ：

クロージャのArcをやめる：64秒→9秒
重い計算の省略：9秒→4秒
全通り再帰からループへ：4秒→345ミリ秒

まとめて教訓：Cloneを避けよ、Arcはホットスポットを避け細部に、再帰は必要最小限に

というわけで、HaskellerとRustaceanが知恵をあわせた結果、プロダクトを3日で267秒から345ミリ秒という1000倍高速化した話でした。
JijのValueとして 10x Faster というものがありますが、いってしまえば 10x 10x 10x Faster を達成したという訳です。
もちろん、1000倍といっているのはsupportcase18という個別のケースだけの話で、実際には計算量が漸近的に改善されているわけです。

この解決はCloneやアロケーションを如何に回避するかというのが鍵だった訳ですが、その中では関数表現による遅延配列という、Haskell界のテクニックが中核的な役割を果たしました。これにより、ほとんどの操作を添え字の算数や関数の合成だけで行えるようになり、中間的な配列のアロケーションが行われないようになった訳です。
Rustの型システムはこうした表現を安全に実現できるだけの柔軟性を持っていたため、（dyn-cloneなどを無視すれば）このような関数型的な表現をほぼそのまま実装できたというのは、かなり嬉しい点です。

一方でHaskellとRustの違いとして、Rustではクロージャ（に限らず任意のヒープ上のオブジェクトもそうですが）をどのような形で保持するのか、つまり Boxに包むのか、ArcやRcを使うのか、などに任意性があります。
ここでリソースの細かな制御を嫌って最初はクロージャ自体をArcで包んだわけですが、クロージャが作っては壊されるような用法だったため、結果的にArcのコストが爆発する結果になってしまっていました。
そこで、関数クロージャそのものではなく、クロージャが保持する配列などの資源の方をArcで包むようにし、クロージャは Clone 可能な Box に包むことで、リソースそのもののcloneは避けながら関数を安価に複製できるようにしたのでした。

そして最後のピースは、シェイプの計算時に無駄な計算を含む全ケース再帰をやめ、ループで直に記述することでした。
Haskellで再帰ばかりのコードを書いていたので、最適化されることを暗黙に当て込んで今でも手癖で再帰を書いてしまいますが、当たり前のはなしとしてループで書けるものは極力ループで書いたほうが効率はよいわけです。
遅延評価なんてサポートしている言語の方が珍しい（Haskellの中心人物も「次Haskell作るなら遅延評価はデフォルトじゃなくすわ」といっています）ので、丼勘定の再帰はなおさら避けたほうがいいです（これもあたりまえ……）。

Rustは非常に懐の広い言語で、関数型言語の手法を使った高速化すらできる、しかしそこにはRustならではの落とし穴もあって……、という中々面白いケーススタディを今回は提供できたのではないかと思います。
今後も、HaskellerとRustaceanの叡智を注ぎ込んでJijModelingを言語としてより進化・深化させていきたいと思います。
Happy Haskell-flavoured Rusting!

謝辞

ともに問題解決に向けて議論してくださった同僚のてらモスさんとIagoさんに感謝します。

宣伝

Jijでは各ポジションを積極的に採用しています！
現在の募集職種は、以下リンクよりご覧いただけます。
カジュアル面談からのスタートも大歓迎ですので、お気軽にご連絡ください。

脚注

もちろん1000倍というのはあるケースでそうだっただけで、実際には漸近的な計算量の改善になっています。 ↩︎
もちろん、JijModeling 2もPythonコードとして書かれる事が前提の設計になっており、Jupyter Notebookやnumpy、pandasなどの既存のPythonエコシステムとの統合も従来同様念頭に置いて設計されています。一方で、言語のコア自体はPythonに依存しない構造になっており、将来的には他の言語や自前のフロントエンド言語を持つことも可能な設計になっています。 ↩︎
筆者が十年くらいずっと流行らせようとしている、リモートの対義語。 ↩︎
何かに気付いたひとがいるかもしれませんね。この後触れるので待ってください :+1: ↩︎
数学で博士号を取った者として断言しますが、算数は難しいです。 ↩︎
実際、この DelayedArray を Value内で使うようにしても、supportcase18のベンチ結果には全く影響がありません。 ↩︎
なんかいい先行研究知ってる人がいたら教えてください。 ↩︎

Discussion