array_chunk の Generator 版

ふと array_chunk を Generator で実装するとどうなるかなと思ったので書いてみた。

あまり考えずに書いた版は次の通り。

<?php
$chunk = function ($input, $size, $preserve_keys) {
    $arr = [];
    foreach ($input as $key => $val) {
        if ($preserve_keys) {
            $arr[$key] = $val;
        } else {
            $arr[] = $val;
        }
        if (count($arr) >= $size) {
            yield $arr;
            $arr = [];
        }
    }
    if (count($arr)) {
        yield $arr;
    }
};

$range = function ($i, $n) {
    for (; $i<=$n; $i++) {
        yield $i;
    }
};

foreach ($chunk($range(100, 109), 3, true) as $key => $arr) {
    $type = gettype($arr);
    echo "$type $key ... ";
    foreach ($arr as $k => $i) {
        echo "$k => $i, ";
    }
    echo "\n";
}

output

array 0 ... 0 => 100, 1 => 101, 2 => 102, 
array 1 ... 3 => 103, 4 => 104, 5 => 105, 
array 2 ... 6 => 106, 7 => 107, 8 => 108, 
array 3 ... 9 => 109, 

が、入力として Generator を渡しているのに $arr 変数に foreach の結果を蓄えているのは微妙な気がしたので Generator を返す Generator として実装した版が次の通り。

<?php
$chunk = function ($input, $size, $preserve_keys) {
    if (is_array($input)) {
        $input = new \ArrayIterator($input);
    }
    while ($input->valid()) {
        yield call_user_func(function () use ($input, $size, $preserve_keys) {
            for (; $size>0 && $input->valid(); $size--, $input->next()) {
                if ($preserve_keys) {
                    yield $input->key() => $input->current();
                } else {
                    yield $input->current();
                }
            }
        });
    }
};

$range = function ($i, $n) {
    for (; $i<=$n; $i++) {
        yield $i;
    }
};

foreach ($chunk($range(100, 109), 3, true) as $key => $iterator) {
    $class = get_class($iterator);
    echo "$class $key ... ";
    foreach ($iterator as $k => $i) {
        echo "$k => $i, ";
    }
    echo "\n";
}

output

Generator 0 ... 0 => 100, 1 => 101, 2 => 102, 
Generator 1 ... 3 => 103, 4 => 104, 5 => 105, 
Generator 2 ... 6 => 106, 7 => 107, 8 => 108, 
Generator 3 ... 9 => 109, 

ただし、この版だと内側の Generator を回さなければ外側の Generator がどれだけ回っていいかわからないので、次のように使うと無限ループする。

foreach ($chunk($range(100, 109), 3, true) as $key => $iterator) {
    $class = get_class($iterator);
    echo "$class $key ...\n";
}

これを改善した版が次の通り。

<?php
$chunk = function ($input, $size, $preserve_keys) {
    if (is_array($input)) {
        $input = new \ArrayIterator($input);
    }
    while ($input->valid()) {
        yield $g = call_user_func(function () use ($input, $size, $preserve_keys) {
            for (; $size>0 && $input->valid(); $size--, $input->next()) {
                if ($preserve_keys) {
                    yield $input->key() => $input->current();
                } else {
                    yield $input->current();
                }
            }
        });
        while ($g->valid()) {
            $g->next();
        }
    }
};

$range = function ($i, $n) {
    for (; $i<=$n; $i++) {
        yield $i;
    }
};

foreach ($chunk($range(100, 109), 3, true) as $key => $iterator) {
    $class = get_class($iterator);
    echo "$class $key ... ";
    foreach ($iterator as $k => $i) {
        echo "$k => $i, ";
    }
    echo "\n";
}

echo "\n";

foreach ($chunk($range(100, 109), 3, true) as $key => $iterator) {
    $class = get_class($iterator);
    echo "$class $key ...\n";
}

output

Generator 0 ... 0 => 100, 1 => 101, 2 => 102, 
Generator 1 ... 3 => 103, 4 => 104, 5 => 105, 
Generator 2 ... 6 => 106, 7 => 107, 8 => 108, 
Generator 3 ... 9 => 109, 

Generator 0 ...
Generator 1 ...
Generator 2 ...
Generator 3 ...