Elixir で each_cons を書く

Elixir でコードを書いていて、Ruby の Enumerable#each_cons に相当する関数が欲しかったのですが、ライブラリに見当たらなかったので自前で書いて見ました。

結論として、Elixir の Stream モジュールに展開関数 Stream.unfold/2 が用意されていたので、それを利用して実装しました。

defmodule MyStream
  def each_cons(seq, n) when is_integer(n) and n > 0 do
    Stream.unfold(seq, fn seq ->
      subseq = Enum.take(seq, n)
      case length(subseq) do
        ^n ->
          {subseq, Stream.drop(seq, 1)}
        _ ->
          nil
      end
    end)
  end
end

実行。

iex> MyStream.each_cons([:a, :b, :c, :d], 2)
#Function<64.58052446/2 in Stream.unfold/2>

Stream モジュールの関数は遅延評価なので、結果を得るには評価してやらないとなりません。

iex> MyStream.each_cons([:a, :b, :c, :d], 2) |> Enum.to_list()
[[:a, :b], [:b, :c], [:c, :d]]
iex> MyStream.each_cons([:a, :b, :c, :d], 3) |> Enum.to_list()
[[:a, :b, :c], [:b, :c, :d]]

Range を与えることもできます。

iex> MyStream.each_cons(?A..?J, 3) |> Enum.to_list()
['ABC', 'BCD', 'CDE', 'DEF', 'EFG', 'FGH', 'GHI', 'HIJ']

遅延評価ということで、具体的な値を得るために Enum.to_list/1 などで評価する必要がありますが、一方で遅延評価なので終端のない列を与えることもできます。

例として、 1 から始まり 1 ずつ増える列を作ります。これも Stream.unfold/2 で作れ ます。

iex> seq = Stream.unfold(1, &{&1, &1 + 1})
#Function<64.58052446/2 in Stream.unfold/2>
iex> seq |> Enum.take(10)
# => [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]

Stream.unfold/2 の第二引数に与える関数が nil を返すと列の生成を停止しますが、単純に nil を返さない関数を与えることで無限列が生成するできます。

これを踏まえて。

iex> MyStream.each_cons(seq, 4) |> Enum.take(5)
[[1, 2, 3, 4], [2, 3, 4, 5], [3, 4, 5, 6], [4, 5, 6, 7], [5, 6, 7, 8]]

無限列なので Enum.take/2 を使って必要な数の要素だけ取得しています。

だいたいそんな感じで。