|
| 1 | +defmodule Table.Zipper do |
| 2 | + @moduledoc false |
| 3 | + |
| 4 | + # An enumerable that zips several enumerables. |
| 5 | + # |
| 6 | + # This enumerable proxies traversal to the underlying enumerables, |
| 7 | + # so it keeps the same properties, such as optimized slicing. |
| 8 | + |
| 9 | + defstruct [:enumerables, :fun] |
| 10 | + |
| 11 | + @doc """ |
| 12 | + Returns an enumerable that zips corresponding elements from a |
| 13 | + collection of enumerables into a tuple. |
| 14 | + """ |
| 15 | + @spec zip(list(Enumerable.t())) :: Enumerable.t() |
| 16 | + def zip(enumerables) when is_list(enumerables) do |
| 17 | + zip_with(enumerables, &List.to_tuple/1) |
| 18 | + end |
| 19 | + |
| 20 | + @doc """ |
| 21 | + Returns an enumerable that zips corresponding elements from a |
| 22 | + collection using the `zip_fun` function. |
| 23 | + """ |
| 24 | + @spec zip_with(list(Enumerable.t()), (list() -> term())) :: Enumerable.t() |
| 25 | + def zip_with(enumerables, zip_fun) when is_list(enumerables) do |
| 26 | + %__MODULE__{enumerables: enumerables, fun: zip_fun} |
| 27 | + end |
| 28 | + |
| 29 | + defimpl Enumerable do |
| 30 | + def count(%{enumerables: []}), do: {:ok, 0} |
| 31 | + |
| 32 | + def count(zipper) do |
| 33 | + zipper.enumerables |
| 34 | + |> Enum.reduce_while(:infinity, fn enumerable, min_count -> |
| 35 | + case Enumerable.count(enumerable) do |
| 36 | + {:ok, count} -> {:cont, min(count, min_count)} |
| 37 | + _ -> {:halt, nil} |
| 38 | + end |
| 39 | + end) |
| 40 | + |> case do |
| 41 | + nil -> {:error, __MODULE__} |
| 42 | + count -> {:ok, count} |
| 43 | + end |
| 44 | + end |
| 45 | + |
| 46 | + def member?(_zipper, _element), do: {:error, __MODULE__} |
| 47 | + |
| 48 | + def reduce(zipper, acc, fun) do |
| 49 | + zipper.enumerables |
| 50 | + |> stream_zip_with(zipper.fun) |
| 51 | + |> Enumerable.reduce(acc, fun) |
| 52 | + end |
| 53 | + |
| 54 | + def slice(%{enumerables: []}), do: {:ok, 0, fn _start, _length -> [] end} |
| 55 | + |
| 56 | + def slice(zipper) do |
| 57 | + zipper.enumerables |
| 58 | + |> Enum.reduce_while({[], [], []}, fn enumerable, {sizes, fun2s, fun3s} -> |
| 59 | + case Enumerable.slice(enumerable) do |
| 60 | + {:ok, size, fun} when is_function(fun, 2) -> |
| 61 | + {:cont, {[size | sizes], [fun | fun2s], nil}} |
| 62 | + |
| 63 | + {:ok, size, fun} when is_function(fun, 3) -> |
| 64 | + {:cont, {[size | sizes], [(&fun.(&1, &2, 1)) | fun2s], fun3s && [fun | fun3s]}} |
| 65 | + |
| 66 | + _ -> |
| 67 | + {:halt, nil} |
| 68 | + end |
| 69 | + end) |
| 70 | + |> case do |
| 71 | + # TODO: rely only on 3-arity functions on Elixir v1.18 |
| 72 | + {sizes, fun2s, nil} -> |
| 73 | + fun = fn start, length -> |
| 74 | + fun2s |
| 75 | + |> Enum.reduce([], fn fun, slices -> [fun.(start, length) | slices] end) |
| 76 | + |> stream_zip_with(zipper.fun) |
| 77 | + |> Enum.to_list() |
| 78 | + end |
| 79 | + |
| 80 | + {:ok, Enum.min(sizes), fun} |
| 81 | + |
| 82 | + {sizes, _fun2s, fun3s} -> |
| 83 | + fun = fn start, length, step -> |
| 84 | + fun3s |
| 85 | + |> Enum.reduce([], fn fun, slices -> [fun.(start, length, step) | slices] end) |
| 86 | + |> stream_zip_with(zipper.fun) |
| 87 | + |> Enum.to_list() |
| 88 | + end |
| 89 | + |
| 90 | + {:ok, Enum.min(sizes), fun} |
| 91 | + |
| 92 | + nil -> |
| 93 | + {:error, __MODULE__} |
| 94 | + end |
| 95 | + end |
| 96 | + |
| 97 | + # --- Backports --- |
| 98 | + |
| 99 | + # TODO: remove once we require Elixir v1.12 |
| 100 | + # Source https://github.com/elixir-lang/elixir/blob/b63f8f541e9d8951dbbcb39a8551bd74a3fe9a59/lib/elixir/lib/stream.ex#L1210-L1342 |
| 101 | + defp stream_zip_with(enumerables, zip_fun) when is_function(zip_fun, 1) do |
| 102 | + if is_list(enumerables) and :lists.all(&is_list/1, enumerables) do |
| 103 | + &zip_list(enumerables, &1, &2, zip_fun) |
| 104 | + else |
| 105 | + &zip_enum(enumerables, &1, &2, zip_fun) |
| 106 | + end |
| 107 | + end |
| 108 | + |
| 109 | + defp zip_list(_enumerables, {:halt, acc}, _fun, _zip_fun) do |
| 110 | + {:halted, acc} |
| 111 | + end |
| 112 | + |
| 113 | + defp zip_list(enumerables, {:suspend, acc}, fun, zip_fun) do |
| 114 | + {:suspended, acc, &zip_list(enumerables, &1, fun, zip_fun)} |
| 115 | + end |
| 116 | + |
| 117 | + defp zip_list(enumerables, {:cont, acc}, fun, zip_fun) do |
| 118 | + case zip_list_heads_tails(enumerables, [], []) do |
| 119 | + {heads, tails} -> zip_list(tails, fun.(zip_fun.(heads), acc), fun, zip_fun) |
| 120 | + :error -> {:done, acc} |
| 121 | + end |
| 122 | + end |
| 123 | + |
| 124 | + defp zip_list_heads_tails([[head | tail] | rest], heads, tails) do |
| 125 | + zip_list_heads_tails(rest, [head | heads], [tail | tails]) |
| 126 | + end |
| 127 | + |
| 128 | + defp zip_list_heads_tails([[] | _rest], _heads, _tails) do |
| 129 | + :error |
| 130 | + end |
| 131 | + |
| 132 | + defp zip_list_heads_tails([], heads, tails) do |
| 133 | + {:lists.reverse(heads), :lists.reverse(tails)} |
| 134 | + end |
| 135 | + |
| 136 | + defp zip_enum(enumerables, acc, fun, zip_fun) do |
| 137 | + step = fn x, acc -> |
| 138 | + {:suspend, :lists.reverse([x | acc])} |
| 139 | + end |
| 140 | + |
| 141 | + enum_funs = |
| 142 | + Enum.map(enumerables, fn enum -> |
| 143 | + {&Enumerable.reduce(enum, &1, step), [], :cont} |
| 144 | + end) |
| 145 | + |
| 146 | + do_zip_enum(enum_funs, acc, fun, zip_fun) |
| 147 | + end |
| 148 | + |
| 149 | + # This implementation of do_zip_enum/4 works for any number of streams to zip |
| 150 | + defp do_zip_enum(zips, {:halt, acc}, _fun, _zip_fun) do |
| 151 | + do_zip_close(zips) |
| 152 | + {:halted, acc} |
| 153 | + end |
| 154 | + |
| 155 | + defp do_zip_enum(zips, {:suspend, acc}, fun, zip_fun) do |
| 156 | + {:suspended, acc, &do_zip_enum(zips, &1, fun, zip_fun)} |
| 157 | + end |
| 158 | + |
| 159 | + defp do_zip_enum([], {:cont, acc}, _callback, _zip_fun) do |
| 160 | + {:done, acc} |
| 161 | + end |
| 162 | + |
| 163 | + defp do_zip_enum(zips, {:cont, acc}, callback, zip_fun) do |
| 164 | + try do |
| 165 | + do_zip_next(zips, acc, callback, [], [], zip_fun) |
| 166 | + catch |
| 167 | + kind, reason -> |
| 168 | + do_zip_close(zips) |
| 169 | + :erlang.raise(kind, reason, __STACKTRACE__) |
| 170 | + else |
| 171 | + {:next, buffer, acc} -> |
| 172 | + do_zip_enum(buffer, acc, callback, zip_fun) |
| 173 | + |
| 174 | + {:done, _acc} = other -> |
| 175 | + other |
| 176 | + end |
| 177 | + end |
| 178 | + |
| 179 | + # do_zip_next/6 computes the next tuple formed by |
| 180 | + # the next element of each zipped stream. |
| 181 | + defp do_zip_next( |
| 182 | + [{_, [], :halt} | zips], |
| 183 | + acc, |
| 184 | + _callback, |
| 185 | + _yielded_elems, |
| 186 | + buffer, |
| 187 | + _zip_fun |
| 188 | + ) do |
| 189 | + do_zip_close(:lists.reverse(buffer, zips)) |
| 190 | + {:done, acc} |
| 191 | + end |
| 192 | + |
| 193 | + defp do_zip_next([{fun, [], :cont} | zips], acc, callback, yielded_elems, buffer, zip_fun) do |
| 194 | + case fun.({:cont, []}) do |
| 195 | + {:suspended, [elem | next_acc], fun} -> |
| 196 | + next_buffer = [{fun, next_acc, :cont} | buffer] |
| 197 | + do_zip_next(zips, acc, callback, [elem | yielded_elems], next_buffer, zip_fun) |
| 198 | + |
| 199 | + {_, [elem | next_acc]} -> |
| 200 | + next_buffer = [{fun, next_acc, :halt} | buffer] |
| 201 | + do_zip_next(zips, acc, callback, [elem | yielded_elems], next_buffer, zip_fun) |
| 202 | + |
| 203 | + {_, []} -> |
| 204 | + # The current zipped stream terminated, so we close all the streams |
| 205 | + # and return {:halted, acc} (which is returned as is by do_zip/3). |
| 206 | + do_zip_close(:lists.reverse(buffer, zips)) |
| 207 | + {:done, acc} |
| 208 | + end |
| 209 | + end |
| 210 | + |
| 211 | + defp do_zip_next( |
| 212 | + [{fun, zip_acc, zip_op} | zips], |
| 213 | + acc, |
| 214 | + callback, |
| 215 | + yielded_elems, |
| 216 | + buffer, |
| 217 | + zip_fun |
| 218 | + ) do |
| 219 | + [elem | rest] = zip_acc |
| 220 | + next_buffer = [{fun, rest, zip_op} | buffer] |
| 221 | + do_zip_next(zips, acc, callback, [elem | yielded_elems], next_buffer, zip_fun) |
| 222 | + end |
| 223 | + |
| 224 | + defp do_zip_next([] = _zips, acc, callback, yielded_elems, buffer, zip_fun) do |
| 225 | + # "yielded_elems" is a reversed list of results for the current iteration of |
| 226 | + # zipping. That is to say, the nth element from each of the enums being zipped. |
| 227 | + # It needs to be reversed and passed to the zipping function so it can do it's thing. |
| 228 | + {:next, :lists.reverse(buffer), callback.(zip_fun.(:lists.reverse(yielded_elems)), acc)} |
| 229 | + end |
| 230 | + |
| 231 | + defp do_zip_close(zips) do |
| 232 | + :lists.foreach(fn {fun, _, _} -> fun.({:halt, []}) end, zips) |
| 233 | + end |
| 234 | + end |
| 235 | +end |
0 commit comments