Build your own lazy operation evaluator in Elixir

May 11, 2023

I’ve always been fascinated by how Elixir’s Stream module handles lazy operations by storing the operation data in a struct. 🤯

It’s so simple but so brilliant!

You can see that if you dig into the source code:

defmodule Stream do
 # ...
 defstruct enum: nil, funs: [], accs: [], done: nil
 # ...
end

It defines a struct with funs (among other things) to store the functions to apply lazily.

And I know Stream is not the only module to do that. Ecto.Multi also stores operations in a struct, and I’m sure there are others.

Lazy Math

So, I wanted to see what it’d be like to write a simple implementation of a lazy math evaluator.

Unsurprisingly, Elixir makes it easy to do this.

Let’s take a look:

defmodule LazyMath do
  defstruct initial: 0, ops: []

  def new(initial), do: %LazyMath{initial: initial}
end

We first define a LazyMath module with struct definition that has the initial value set to 0 and an empty list of ops (operations).

We also add a new/1 function that will be a helper to initialize our struct.

Now let’s add some operations: add/2, subtract/2, multiply/2 and divide/2:

  def add(math = %{ops: ops}, number) do
    %LazyMath{math | ops: [{:add, number} | ops]}
  end

  def subtract(math = %{ops: ops}, number) do
    %LazyMath{math | ops: [{:sub, number} | ops]}
  end

  def multiply(math = %{ops: ops}, number) do
    %LazyMath{math | ops: [{:mult, number} | ops]}
  end

  def divide(math = %{ops: ops}, number) do
    %LazyMath{math | ops: [{:div, number} | ops]}
  end

As you can see, they all look very similar (and we could probably refactor a common private function). But what’s interesting is that we’re storing a representation of the operation instead of performing the operation.

So, when we want to add a number, we prepend an {:add, number} tuple to the list of existing ops, and return the updated %LazyMath{} struct.

  def add(math = %{ops: ops}, number) do
    # store `{:add, number}` instead of adding `number` to existing total
    %LazyMath{math | ops: [{:add, number} | ops]}
  end

The rest of the operations work the exact same way (though we store different tuples).

Now, let’s see how we can evaluate all of the operations:

  def evaluate(%LazyMath{initial: init, ops: ops}) do
    ops
    |> Enum.reverse()
    |> Enum.reduce(init, fn
      {:add, number}, acc_total -> acc_total + number
      {:sub, number}, acc_total -> acc_total - number
      {:mult, number}, acc_total -> acc_total * number
      {:div, number}, acc_total -> div(acc_total, number)
    end)
  end

Our evaluate/1 function takes an existing %LazyMath{} struct, pattern matching the initial value and the operations.

We then reverse the list of operations, so we can apply them in the correct order — remember we were prepending new operations before.

Finally, we Enum.reduce/3 over the list of operations (now in order), passing the initial value, and then we pattern match on the operation tuple to perform the actual operation on the accumulated total.

Here’s the full module:

defmodule LazyMath do
  defstruct initial: 0, ops: []

  def new(initial), do: %LazyMath{initial: initial}

  def add(math = %{ops: ops}, number) do
    %LazyMath{math | ops: [{:add, number} | ops]}
  end

  def subtract(math = %{ops: ops}, number) do
    %LazyMath{math | ops: [{:sub, number} | ops]}
  end

  def multiply(math = %{ops: ops}, number) do
    %LazyMath{math | ops: [{:mult, number} | ops]}
  end

  def divide(math = %{ops: ops}, number) do
    %LazyMath{math | ops: [{:div, number} | ops]}
  end

  def evaluate(%LazyMath{initial: init, ops: ops}) do
    ops
    |> Enum.reverse()
    |> Enum.reduce(init, fn
      {:add, number}, acc_total -> acc_total + number
      {:sub, number}, acc_total -> acc_total - number
      {:mult, number}, acc_total -> acc_total * number
      {:div, number}, acc_total -> div(acc_total, number)
    end)
  end
end

Let’s test how lazy we are:

result =
  LazyMath.new(0)
  |> LazyMath.add(5)
  |> LazyMath.subtract(2)
  |> LazyMath.multiply(2)
  |> LazyMath.divide(3)
# => %LazyMath{initial: 0, ops: [div: 3, mult: 2, sub: 2, add: 5]}

As you can see, our result hasn’t evaluated any math operations yet. Instead, it stored the initial value along with the list of operations. 🥳

Finally, we can evaluate the result:

result |> LazyMath.evaluate()
# => 2

Pretty cool, right?