I recently found myself wishing for the simple convenience of a while loop in Elixir. I had a test where I wanted to loop and test a condition, and either break out of that loop, or time out after some period. Writing this in the form of a recursive function is relatively easy, but is not really fully general:

def while(predicate, timeout) 
  when is_function(predicate, 0) and is_integer(timeout) do
  while(predicate, timeout, System.monotonic_time(:milliseconds))
end
defp while(predicate, timeout, started) when timeout > 0 do
  if predicate.() do
    now = System.monotonic_time(:milliseconds)
    elapsed = now - started
    while(predicate, timeout - elapsed, now)
  else
    :ok
  end
end
defp while(_predicate, _timeout, _started) do
  exit(:timeout)
end

You could then use this in a test like so:

test "wait until node is available", %{peer: some_node} do
  assert :ok = while(fn -> Node.ping(some_node) != :ping end, 30_000)
end

This would loop until Node.ping/1 returns :ping, or 30 seconds passes, at which point it would exit and fail the test. That’s…acceptable, but it reads a bit weird, and actually has a rather awkward issue in that we can’t slow the loop down to only check once a second, rather than as fast as the runtime will execute ping. To do that, we’d need to extend the definition of while to support another function argument, which would be executed if the predicate returns true, and would allow us to insert a call to :timer.sleep/1 or something on each iteration.

Let The Incantations Begin..

As soon as I saw this, I knew I was going to end up writing a macro for it. First, while is a familiar concept coming from imperative languages, and though technically not needed in Elixir, it does still have a place as a high-level abstraction of a common form of recursive function, where you need to loop until some condition is true, and reimplementing that from scratch each time is a bit annoying. Secondly, with a macro, we could make this look just like while in an imperative language, and once I saw it in my head, I knew that at the very least it was going to make for a fun experiment. I also figured that while I was building the macro, I could push the limits a bit and see if we can add mutation to the implementation as well (not real mutability, but at least the semblance of it).

We want this macro to embrace the functional aspects of the language, while providing some of the conveniences of the imperative approach. As a bonus, building in the concept of timeouts to the implementation makes our version of while far more powerful than your average implementation! Let’s take a look at a few examples of what I had in mind:

# Simple infinite loop
while true do
  :ok
end

# Loop with initializer + accumulator
5 = while n = 1, n < 5 do 
  n -> n + 1
end

# Loop with timeout condition
while true do
  :ok
after
  30_000 ->
    exit(:timeout)
end

# Loop with timeout and initializer + accumulator
while n = 1, n < 5 do
  if n == 4 do
    break
  end
  :timer.sleep(1_000)
  n + 1
after
  5_000 ->
    exit(:timeout)
end

# Loop with timeout condition but no body
while true do after 30_000 ->
  exit(:timeout)
end

Much more familiar! There are a few syntactic differences from your standard while in something C-like, but for the most part there are really no surprises here. I also wanted this to compile down to a form that was actually optimal, i.e., I didn’t want anonymous functions for the predicate, body, and after body, where each iteration would result in at least two function calls. Instead, it should compile down to a recursive function, with the minimal amount of branching based on which parts of the construct where used, e.g., no timeout means that we don’t need to branch on the timeout value and calculate new times, etc.

Recursive Anonymous Functions in Elixir

The first problem here is that in order to use this in any arbitrary location, we can’t define a named function at the point while is invoked, as def isn’t valid within another def; instead, we have to use an anonymous function, and Elixir does not yet support recursively invoking an anonymous function by the name it is bound to (even though Erlang does). Instead, we have to transform our definition of the anonymous function to combinator form, i.e., it has no free variables, and only refers to its parameters. Here is what I mean:

# Doesn't work, refers to the free variable `while`
# which is bound _after_ the closure for the fun is created
while = fn predicate -> 
  if predicate.() do
    while(predicate)
  else
    :ok
  end
end
while.(fn -> true end)

# This does, no free variables, we pass in the fun
# to itself, so that it can call itself recursively
while = fn pred, while ->
  if predicate.() do
    while.(predicate, while)
  else
    :ok
  end
end
while.(fn -> true end, while)

So we know what we’ll have to structure the macro implementation to generate an anonymous function in combinator form. We have a solution for our first problem!

Our Initial Implementation

The next problem can be seen in the example above, the predicate is another anonymous function, but what we really want is for the predicate to be a part of the definition of the while function itself, so that we avoid the function call overhead. This is pretty straightforward though, we just will just inject it into the definition of while. Let’s stop here and implement the macro up until this point:

defmacro while(predicate, do: body) do
  quote location: :keep, generated: true do
    whilef = fn whilef ->
      if unquote(predicate) do
        unquote(body)
        whilef.(whilef)
      else
        :ok
      end
    end
    whilef.(whilef)
  end
end

# Usage Example
while Node.ping(some_node) != :ping do
  :timer.sleep(1_000)
end

Not bad! A few notes:

• We use generated: true to prevent compiler warnings if we’re careless about how we generate code for the function and end up with unused bindings or something. Not strictly necessary, but useful.
• If you aren’t familiar, location: :keep makes sure the file/line from the quotation are what end up in traces, rather than the location the macro is invoked from. This means if we make a mistake, we’ll know where in the macro implementation things went south.
• If you aren’t familiar with quote and unquote, I would stop here and go read the quote docs, either via h quote in iex or on HexDocs here

There are some problems here though, not to mention we haven’t yet added support for timeouts or “mutable” state. This is usable, but not ideal.

Problem: Emulating Mutable State

One of the biggest problems lies with the predicate itself: what if we changed the example to this:

n = 1
while n < 5 do
  n + 1
end

Can you spot the problem? The issue is that running this will loop forever, as n will always equal 1, the predicate will always see 1 < 5, and n + 1 will always calculate 2. The state from the body of while does not affect the next evaluation of the predicate. We need to ensure that state is accumulated, by making the result of the body be passed along as another parameter to the whilef function; which also means we have to make the initializer part of while, so that we can formally introduce n into the scope of both the predicate and the body, otherwise we can’t ensure that n is bound to the updated value, as it is out of scope. Let’s see what our macro looks like now:

defmacro while({:=,_[i|_]}=init, predicate, do: [{:->,_,_}] = body) do
  quote location: :keep, generated: true do
    whilef = fn whilef, acc ->
      unquote(i) = acc
      if unquote(predicate) do
        acc2 =
          case unquote(i) do
            unquote(body)
          end
        whilef.(whilef, acc2)
      else
        acc
      end
    end
    whilef.(whilef, _ = unquote(init))
  end
end

# Usage Example
5 = while n = 1, n < 5 do n ->
  n + 1
end

In the above, we ensure the initializer is just a simple assignment by matching on the AST, and for simplicity right now, I’m also matching to ensure that the body accepts the accumulator value using case clause syntax (e.g. n -> ...). You’ll notice we also bound the name of the initializer, so that we can cheat a little with the unquote(i) = acc line, which, using our example, resets n to the current state of the accumulator on each iteration. The bit at the bottom, _ = unquote(init) will expand to _ = n = 1, which seems strange, but ensures that we don’t get a warning about an unused binding.

Problem: Supporting Timeouts

Of course the problem we have now is that our macro requires the use of the initializer and accumulating state in the body, still doesn’t support after, and isn’t flexible in the way I originally envisioned. This is fine though, since we’re about to fix that!

First, let’s add support for after:

defmacro while({:=,_[i|_]}=init, predicate, [do: [{:->,_,_}] = body, after: [{:->,_,[[timeout],after_body]}]) do
  quote location: :keep, generated: true do
    whilef = fn 
      whilef, acc, timeout, started when timeout > 0 ->
        unquote(i) = acc
        if unquote(predicate) do
          acc2 =
            case unquote(i) do
              unquote(body)
            end
          now = System.monotonic_time(:milliseconds)
          elapsed = now - started
          whilef.(whilef, acc2, timeout - elapsed, now)
        else
          acc
        end
      _, _, _, _ ->
        unquote(after_body)
    end
    now = System.monotonic_time(:milliseconds)
    whilef.(whilef, _ = unquote(init), unquote(timeout), now)
  end
end

# Usage Example
5 = while n = 1, n < 5 do n ->
  n + 1
after
  5_000 ->
    exit(:timeout)
end

Similar to the previous iteration of the macro, we’re matching against the AST to ensure the after clause uses the timeout -> body syntax, and that only one timeout clause is allowed. We’re not validating the timeout, or dealing with invalid constructs very well though, for that we need to do the pattern matching in the macro body so that we can catch common mistakes and raise more helpful errors. You’ll notice that we calculate the elapsed time on each iteration before continuing, that way we can just jump straight to the after clause if the timeout reaches zero.

A Complete Implementation

We’re almost done! We have a few items to address to make this a “complete” implementation:

• Put this all in a module and document it
• Support optional after, optional do, optional initializer/accumulator so that we can arrive at the various forms I originally sketched out
• Support multiple accumulator clauses within do so that pattern matching is supported.
• Compile to an optimal form, where no unnecessary function calls, branching, or work is performed
• Validate the constructs used so that invalid forms are not permitted and raise a useful compilation error, e.g. after can only accept a positive integer for a timeout, and only one timeout clause is allowed. We could theoretically support multiple, but I don’t think that is useful for this exercise.

Let’s take a look at the final form with all of those points addressed:

defmodule LanguageExtensions.While do
  @moduledoc """
  Adds support for `while` to Elixir.

  See the docs for `#{__MODULE__}.while/2` for more information.

  ## Example

      use #{__MODULE__}

      ...

      while n = 1, n < 5 do n ->
        if n == 4 do
          break
        end
        :timer.sleep(1_000)
        n + 1
      after
        5_000 ->
          :timeout
      end
  """

  defmacro __using__(_) do
    quote do
      import unquote(__MODULE__)
    end
  end

  @doc """
  Intended for use within `while`, this will break out of iteration,
  returning the last accumulator value as the result of the `while` expression.

      iex> while true do break end
      :ok
  end
  """
  defmacro break do
    quote do
      throw {unquote(__MODULE__), :break}
    end
  end

  @doc """
  A looping construct, looks like so:

      # Simple infinite loop (never terminates)
      while true do
        :ok
      end

      # Loop with accumulator
      iex> while n = 1, n < 5 do n -> n + 1 end
      5

      # Loop with timeout condition
      iex> while true do
      ...>   :ok
      ...> after
      ...>   1_000 ->
      ...>     :timeout
      ...> end
      :timeout

      # Loop with timeout and accumulator
      iex> while n = 1, n < 5 do
      ...>   if n == 4 do
      ...>     break
      ...>   end
      ...>   :timer.sleep(1_000)
      ...>    n + 1
      ...> after
      ...>   5_000 ->
      ...>     :timeout
      ...> end
      4

      # Loop with timeout condition but no body (exits with :timeout)
      iex> while true do after 1_000 -> :timeout end
      :timeout
  """
  # A simple while loop
  defmacro while(predicate, do: block) do
    quote location: :keep do
      while(_ = :ok, unquote(predicate), [do: unquote(block), after: nil])
    end
  end

  # A loop + timeout block
  defmacro while(predicate, [do: block, after: after_block]) do
    quote location: :keep do
      while(_ = :ok, unquote(predicate), [do: unquote(block), after: unquote(after_block)])
    end
  end

  # An accumulator loop
  defmacro while(init, predicate, do: block) do
    quote location: :keep do
      while(unquote(init), unquote(predicate), [do: unquote(block), after: nil])
    end
  end

  # An accumulator loop + timeout block
  defmacro while(init, predicate, [do: block, after: after_block]) do
    # Validate initializer
    {init_name, init_expr} =
      case init do
        {:=, env, [init_name | init_expr]} ->
          {init_name, {:__block__, env, init_expr}}
        {_, env, _} ->
          raise CompileError,
            description: "expected an initializer of the form `n = <expr>`, got: #{Macro.to_string(init)}",
            file: Keyword.get(env, :file, __ENV__.file),
            line: Keyword.get(env, :line, __ENV__.line)
      end

    # Validate and extract timeout/after body
    # Timeout must be a positive integer
    # After body only allows one timeout clause
    {timeout, after_block} =
      case after_block do
        [{:->, _, [[timeout], after_body]}] when is_integer(timeout) and timeout >= 0->
          {timeout, after_body}

        [{:->, env, [[_], _after_body]}] ->
          raise CompileError,
            description: "expected a positive integer timeout in `after`",
            file: Keyword.get(env, :file, __ENV__.file),
            line: Keyword.get(env, :line, __ENV__.line)

        [{:->, env, _}, _ | _] ->
          raise CompileError,
            description: "multiple timeouts are not supported in `after`",
            file: Keyword.get(env, :file, __ENV__.file),
            line: Keyword.get(env, :line, __ENV__.line)

        [{_, env, _} | _] ->
          raise CompileError,
            description: "multiple timeouts are not supported in `after`",
            file: Keyword.get(env, :file, __ENV__.file),
            line: Keyword.get(env, :line, __ENV__.line)

        nil ->
          {nil, :ok}
      end

    # Determine the type of while block we're building
    {block_type, block} =
      case block do
        # Empty block, i.e. `while true do after ... end`
        {:__block__, _, []} ->
          {:empty, :ok}

        # Has one or more accumulator patterns
        [{:->, _, [[_binding], _body]} | _] = blk ->
          {:acc, blk}

        # No accumulator
        _other ->
          {:noacc, block}
      end

    timeout_calc =
      if is_nil(timeout) do
        quote location: :keep, generated: true do
          now = start_time
        end
      else
        quote location: :keep, generated: true do
          now = System.monotonic_time(:milliseconds)
          elapsed = now - start_time
          after_time = after_time - elapsed
        end
      end


    # Construct the body of the function
    body =
      case block_type do
        # If there is no `do` body, then skip it entirely
        :empty ->
          quote location: :keep, generated: true do
            if unquote(predicate) do
              import unquote(__MODULE__), only: [break: 0]
              unquote(timeout_calc)
              f.(after_time, now, :ok, f)
            else
              :ok
            end
          end

        # If we're not accumulating, skip the unnecessary pattern match
        # we need when managing the accumulator
        :noacc ->
          quote location: :keep, generated: true do
            if unquote(predicate) do
              import unquote(__MODULE__), only: [break: 0]
              acc2 = unquote(block)
              unquote(timeout_calc)
              f.(after_time, now, acc2, f)
            else
              acc
            end
          end

        # If we're managing an accumulator, we need to use the case
        # statement to
        :acc ->
          quote location: :keep, generated: true do
            unquote(init_name) = acc
            if unquote(predicate) do
              import unquote(__MODULE__), only: [break: 0]
              acc2 =
                case unquote(init_name) do
                  unquote(block)
                end
              unquote(timeout_calc)
              f.(after_time, now, acc2, f)
            else
              acc
            end
          end
      end

    # Construct the actual function, tag the quoted AST with
    # generated: true to make sure unused bindings and such are
    # not warned
    quote location: :keep, generated: true do
      fun = fn
        # Since nil is always > 0 due to Erlang sort order
        # We can use this to represent both infinite timeouts, and
        # finite timeouts which haven't expired
        after_time, start_time, acc, f when after_time > 0 ->
          try do
            unquote(body)
          catch
            :throw, {unquote(__MODULE__), :break} ->
              acc
          end

        _after_time, _start_time, acc, _f ->
          try do
            unquote(after_block)
          catch
            :throw, {unquote(__MODULE__), :break} ->
              acc
          end
      end
      now = System.monotonic_time(:milliseconds)
      fun.(unquote(timeout), now, unquote(init_expr), fun)
    end
  end
end

That’s it! I’ve documented the code a bit with comments describing the various changes, but I hope for the most part you will be able to see the evolution from our last version. We did the following:

• Defined new function heads for while/2 and while/3 which match on the various forms of the while macro and ultimately delegate to a single implementation of while/3. This is how we make some of the parts optional and ease the implementation.
• We extracted the function head matches into the body so that we can match on valid uses of different constructs and raise compilation errors with useful descriptions and file/line numbers for debugging.
• We break apart the two major types of while usage, with timeouts, and without timeouts, and construct the body differently based on the type so that no extra work is done while iterating
• We break apart the type of while body, whether it is empty, accumulates state, or does not accumulate state, and construct the body optimally for each case. This is especially necessary if we want to support optional accumulation - the pat -> expr syntax will not be valid outside of the case block we wrap it in when building up the accumulator form, but we want to avoid the branch + match if we don’t care about the accumulator and just want to return the value of the body on each iteration. So we base the type of body we generate based on the syntax of the do block.
• Defined break/0 which just throws a special error for us to catch and break out of the loop early. We do this transparently to the user of while, and do so in such a way that the accumulated value is returned.

The final version of the macro is a bit long due to the optimizations we are performing. If one wanted a simpler implementation, you could support fewer optional parts, or be willing to calculate elapsed time on each iteration and always branch on the accumulator so that the accumlator/no accumulator forms are the same. In any case, I wanted to pull out all of the stops and see where we got.

Wrap Up

I’m not sure I would recommend using this macro generally, but it was a fun experiment, and I love the opportunity to stretch the macro muscle a bit, especially with one like this which really shows how you extend the syntax of the language via macros, just like how with and for are built.

I hope you enjoyed reading this! It isn’t often I write something up, but this seemed like too much fun to resist :)

Macro Tips

A few useful tips for writing macros:

• Have two iex windows open, one for testing, one for docs. I’m constantly experimenting with quoting expressions and reviewing the AST, compiling and importing the macro I’m working on so that I can test it, and using h <whatever> to pull up docs to double check what I’m doing
• To see what kind of AST is produced for some syntax, just use one of the iex windows and type in quote do <your expr> end and it will spit out the quoted form of the expression.
• You can get the pretty printed form of a quoted expression for easier reading with quoted |> Macro.to_string |> IO.puts. As of 1.6 you can also do quoted |> Macro.to_string |> Code.format_string! |> IO.puts to format the code using the new formatter.
• If there is a compile error coming from your macro, get a quoted usage of it with quote do <expr> end, and use Macro.expand_once(quoted, __ENV__) to see a single expansion of the macro, which will help you see how the macro is actually being converted into Elixir syntax. You can chain this call as well, or just use Macro.expand(quoted, __ENV__) to see the fully expanded AST.
• Read up on quoting, unquoting, splicing, and :bind_quoted before you get started. Understanding those are critical to not only know what you are doing, but also saving yourself a lot of pain when trying to figure out why something isn’t working the way you thought it would. Things can get quite complex when you are working with nested macros, which is almost always the case anyway, as much of Elixir syntax itself is written using macros. Anytime you find yourself confused by the behavior of your macro, it pays to go back and re-read the docs, because it is almost certainly a case of something you overlooked.