zip_iteratorとstd::iter_swap

少し頑張って解決しようとした問題が実は解決できないことに気づいてしまったお話

TL;DR: std::iter_swapをフック可能にしてほしい。

背景

配列が二つあったとして、その片方をキーとして使って、両方の配列を同時にsortしたいと言う気持ちがある。

v1 = {4, 2, 5, 3, 1};
v2 = {1, 2, 3, 4, 5};

sort_by_key(v1.begin(), v1.end(), v2.begin());
// v2 == {5, 2, 4, 1, 3}

これは、普通はstd::pair<int, int>のようなものを作ってstd::sortと適当なコンパレータを組み合わせて行う。

std::vector<std::pair<int, int>> v =
    {{4,1}, {2,2}, {5,3}, {3,4}, {1,5}};
std::sort(std::begin(v), std::end(v),
    [](const auto& lhs, const auto& rhs) -> bool {
        return lhs.first < rhs.first;
    });

だがpairでなくて配列を二つ持っていることもあろう。どうするか。

std::vector<int> v1 = {4, 2, 5, 3, 1};
std::vector<int> v2 = {1, 2, 3, 4, 5};

例えばtransformする手はある。

std::vector<std::pair<int, int>> v(v1.size());
std::transform(std::begin(v1), std::end(v1), std::begin(v2),
    std::begin(v), [](const int& l, const int& r) {
        return std::make_pair(l, r);
    });
std::sort(std::begin(v), std::end(v),
    [](const auto& lhs, const auto& rhs) -> bool {
        return lhs.first < rhs.first;
    });
std::transform(std::begin(v), std::end(v), std::begin(v1),
    [](const auto& x) -> int {return x.first;});
std::transform(std::begin(v), std::end(v), std::begin(v2),
    [](const auto& x) -> int {return x.second;});

長い。それにコピーが入ってしまう。あまり嬉しくはない。

もちろん自分でsortを実装すれば回避できるのだが、あまりやりたいことではない。 せっかく質のいい実装があるならそれを使えばいいのだ。

と言うわけで、複数のイテレータを一つにまとめるzip_iteratorを作ることにした。片方を並べ替えるときにもう片方がついてくるようにすればいい。 これをする上で必要なのは、zip_iteratorから取り出した参照を何らかの形で書き込み可能なようにすることである。これに関しては、std::tieが返すオブジェクトと同じものを用意すればいいだろう。

zip_iteratorの実装

zip_iteratorは複数のイテレータをまとめたものなので、std::tuple<Iterator...>をラップしてイテレータのように振る舞わせれば良い。

イテレータとして振る舞わせるためには、基本的にはインクリメント・デクリメントとoperator*operator->による値の取り出しがあれば良い。

インクリメントは順に全部やっていけばいいので、例えばシンプルに再帰で実装するなら以下のようになる。

template<std::size_t N, typename ... Iterators>
struct increment_impl
{
    static void invoke(std::tuple<Iterators...>& iters)
    {
        ++(std::get<N-1>(iters));
        return increment_impl<N-1, Iterators...>::invoke(iters);
    }
};

template<typename ... Iterators>
struct increment_impl<0, Iterators...>
{
    static void invoke(std::tuple<Iterators...>& iters) {return;}
};

template<typename ... Iterators>
zip_iterator<Iterators...>& zip_iterator<Iterators...>::operator++()
{
    increment_impl<sizeof...(Iterators), Iterators...
        >::invoke(this->iters_);
    return *this;
}

operator*は少し面倒だ。 これはIterator::referenceを返すので、zip_iteratorの気持ち的にはstd::tuple<Iterator::reference...>を作って返したい。

どうするか。困った時は最終形から始めて、必要なものを明らかにしよう。 std::tuple<Iterators...>からstd::tuple<Iterators::reference ...>になればいいので、最終的には以下のような関数があれば良いことになる。

template<typename ... Iterator>
std::tuple<typename std::iterator_traits<Iterator>::reference ...>
get_reference(const std::tuple<Iterator>& iters)
{
    return std::forward_as_tuple(
        *(std::get<0>(iters)), *(std::get<1>(iters)), ...
        *(std::get<sizeof...(Iterator)-1>(iters))
    );
}

当然ながらこれはそのままでは動かない。問題は、template parameterを0, 1, ..., sizeof...(Iterator)のように展開して渡すところにあるのだが、標準ライブラリにちょうどこのための機能がある。少し前の記事でも出てきたstd::index_sequenceだ。

template<typename ... Iterator, std::size_t ... Idx>
std::tuple<typename std::iterator_traits<Iterator>::reference ...>
get_reference(const std::tuple<Iterator>& iters,
    std::index_sequence<Idx...>)
{
    return std::forward_as_tuple(*(std::get<Idx>(iters)) ...);
}

index_sequenceは、この型の値そのものは特に使い道がないが、型そのものに出てくる値が重要な珍しいタイプのクラスだ。 variadic templateならコンパイル時に展開することができるので、必要なものをまとめてindex_sequenceとして渡してやれれば、上記でしているように「コンパイル時にその場で全部展開して渡す」というような操作が可能になる。

で、これを作るためのヘルパークラスがある。

make_index_sequence - cpprefjp C++日本語リファレンス

また、pointerを返す場合でも全く同じ操作が可能だ。むしろポインタのシンタックスが値のそれと異なるためにreferenceの場合よりも簡単と言っていい。 というわけでこれらは解決した。

問題点

さて、ちょっとしたテストを書いてzip_iteratorが思った通りに動くこと、関数の戻り値などが思った通りの型になっていること、zip_iteratorを通して書き込みができることなどを確かめたのち、意気揚々とstd::sortzip_iteratorを渡した時だった。

動かない。

何が問題になっているかと言うと、案の定zip_iterator::operator*()だ。定義を見て欲しい。

template<typename ... Iterator>
class zip_iterator
{
    std::tuple<typename std::iterator_traits<Iterator>::reference...>
    operator*();
};

そう、これはIterator::referenceが入ったstd::tupleをその場で作って返すのである。 その場で作って返す。関数から返った瞬間はできたてほやほやでまだどんな値にも束縛されていない。つまりこれはrvalueだ。

ところで<algorithm>が提供するシーケンスを変更する操作の多くは、swap(*a, *b)のようなことをする。 ここでstd::swapを見ていただきたい。

template<typename T>
void std::swap(T& lhs, T& rhs);

これは左辺値を取るのである。当然だ。入れ替えるのだから、代入もするだろう。 だが、zip_iterator a, bを使ってswap(*a, *b)すると、引数は両方右辺値になってしまう!

というわけで無理になってしまったわけだ。

ご存知の通りswapはフック可能だ。 Argument Dependent Lookupなる機能があり、std::swapはユーザー定義型に対してはより効率的な実装が与えられる可能性を考慮して実装されているので、標準ライブラリ内部で使われる時は大抵以下のようになっている。

using std::swap;
swap(a, b);

こうすると、ADLによってabが所属する名前空間にあるswap関数がまず検索され、なかった場合にusingされているstd::swapが適用される。abが所属する名前空間により効率的なswap実装が提供されていれば、そちらが適用されるわけだ。

今回の問題は、返すものがstd::tupleであるということだ。ユーザー定義名前空間にあるものではない。なのでこのような解決策を取ることはできない。 もちろんユーザー定義名前空間std::tupleの薄いラッパーを用意すれば解決可能ではあるのだが、zip_iteratorが返すものが全て独自tuple型になってしまうことを考えるとあまり気持ちのいいものではない。

より細かいことをいうと、実際に<algorithm>で使われているものの多くはstd::iter_swapだ。 これはそれぞれのイテレータが指す先にあるものをswapするというものなので、これをフックできれば一番丸く収まったのだが(zip_iteratorはユーザー定義クラスだ)、何故か全ての関数がexplicitにstd::iter_swapを呼んでいた。これではADLの力は発揮できない。

operator*std::tuple<Iterator::reference...>を返すzip_iteratorstd::sortと共に使おうとするなら、 非常に薄いstd::tupleのラッパーを用意してADLがswapを解決したのち最適化パスのどこかでその薄いラッパーが消え失せることを期待する(これはまともなコンパイラを使っているなら十分抱いていい期待だと思う)しかない。一応std名前空間の禁じられた扉を開くことでも実行可能だが、そういうのは実行可能とは言わない。

おとなしくペアの配列に変換するか、thrustBoost.Computeにあるsort_by_keyを呼ぶのがよいだろう。