The Erlangelist

(not only) Erlang related musings
Hi, I'm Saša Jurić, a software developer with many years of professional experience in programming of web and desktop applications using various languages, such as Elixir, Erlang, Ruby, JavaScript, C# and C++. I'm also the author of the Elixir in Action book. In this blog you can read about Elixir, Erlang, and other programming related topics. You can subscribe to the feed, follow me on Twitter or fork me on GitHub.

Understanding Elixir Macros, Part 5 - Reshaping the AST

2014-06-29

Last time I presented a basic version of deftraceable macro that allows us to write traceable functions. The final version of the macro has some remaining issues, and today we’ll tackle one of those - arguments pattern matching.

Today’s exercise should demonstrate that we have to carefully consider our assumptions about possible inputs to our macros can receive.

The problem

As I hinted the last time, the current version of deftraceable doesn’t work with pattern matched arguments. Let’s demonstrate the problem:

iex(1)> defmodule Tracer do ... end

iex(2)> defmodule Test do
          import Tracer

          deftraceable div(_, 0), do: :error
        end
** (CompileError) iex:5: unbound variable _

So what happened? The deftraceable macro blindly assumes that input arguments are plain variables or constants. Hence, when you call deftracable div(a, b), do: … the generated code will contain:

passed_args = [a, b] |> Enum.map(&inspect/1) |> Enum.join(",")

This will work as expected, but if one argument is an anonymous variable (_), then we generate the following code:

passed_args = [_, 0] |> Enum.map(&inspect/1) |> Enum.join(",")

This is obviously not correct, and therefore we get the unbound variable error.

So what’s the solution? We shouldn’t assume anything about input arguments. Instead, we should take each argument into a dedicated variable generated by the macro. Or to say it with code, if our macro is called with:

deftraceable fun(pattern1, pattern2, ...)

We should generate the function head:

def fun(pattern1 = arg1, pattern2 = arg2, ...)

This allows us to take argument values into our internal temp variables, and print the contents of those variables.

The solution

So let’s implement this. First, I’m going to show you the top-level sketch of the solution:

defmacro deftraceable(head, body) do
  {fun_name, args_ast} = name_and_args(head)

  # Decorates input args by adding "= argX" to each argument.
  # Also returns a list of argument names (arg1, arg2, ...)
  {arg_names, decorated_args} = decorate_args(args_ast)

  head = ??   # Replace original args with decorated ones

  quote do
    def unquote(head) do
      ... # unchanged

      # Use temp variables to make a trace message
      passed_args = unquote(arg_names) |> Enum.map(&inspect/1) |> Enum.join(",")

      ... # unchanged
    end
  end
end

First, we extract name and args from the head (we resolved this in previous article). Then we have to inject = argX into the args_ast and take back the modified args (which we’ll put into decorated_args).

We also need pure names of generated variables (or more exactly their AST), since we’ll use these to collect argument values. The variable arg_names will essentially contain quote do [arg_1, arg_2, …] end which can be easily injected into the tree.

So let’s implement the rest. First, let’s see how we can decorate arguments:

defp decorate_args(args_ast) do
  for {arg_ast, index} <- Enum.with_index(args_ast) do
    # Dynamically generate quoted identifier
    arg_name = Macro.var(:"arg#{index}", __MODULE__)

    # Generate AST for patternX = argX
    full_arg = quote do
      unquote(arg_ast) = unquote(arg_name)
    end

    {arg_name, full_arg}
  end
  |> Enum.unzip
end

Most of the action takes place in the for comprehension. Essentially we go through input AST fragment of each variable, and compute the temp name (quoted argX) relying on the Macro.var/2 function which can transform an atom into a quoted variable that has a name of that atom. The second argument to Macro.var/2 ensures that the variable is hygienic. Although we’ll inject arg1, arg2, … variables into the caller context, the caller won’t see these variables. In fact, a user of deftraceable can freely use these names for some local variables without interfering with temps introduced by our macro.

Finally, at the end of the comprehension we return a tuple consisting of the temp’s name, and the quoted full pattern - (e.g. _ = arg1, or 0 = arg2). The little dance after the comprehension with unzip and to_tuple ensures that decorate_args returns the result in form of {arg_names, decorated_args}.

With decorate_args helper ready we can pass input arguments, and get decorated ones, together with the names of temp variables. Now we need to inject these decorated arguments into the head of the function, in place of the original arguments. In particular, we must perform following steps:

  1. Walk recursively through the AST of the input function head.
  2. Find the place where function name and arguments are specified.
  3. Replace original (input) arguments with the AST of decorated arguments

This task can be reasonably simplified if we rely on Macro.postwalk/2 function:

defmacro deftraceable(head, body) do
  {fun_name, args_ast} = name_and_args(head)

  {arg_names, decorated_args} = decorate_args(args_ast)

  # 1. Walk recursively through the AST
  head = Macro.postwalk(
    head,

    # This lambda is called for each element in the input AST and
    # has a chance of returning alternative AST
    fn
      # 2. Pattern match the place where function name and arguments are
      # specified
      ({fun_ast, context, old_args}) when (
        fun_ast == fun_name and old_args == args_ast
      ) ->
        # 3. Replace input arguments with the AST of decorated arguments
        {fun_ast, context, decorated_args}

      # Some other element in the head AST (probably a guard)
      #   -> we just leave it unchanged
      (other) -> other
    end
  )

  ... # unchanged
end

Macro.postwalk/2 walks the AST recursively, and calls the provided lambda for each node, after all of the node’s descendants have been visited. The lambda receives the AST of the element, and there we have a chance of returning something else instead of that node.

So what we do in this lambda is basically a pattern match where we’re looking for the {fun_name, context, args}. As explained in part 3, this is the quoted representation of the expression some_fun(arg1, arg2, …). Once we encounter the node that matches this pattern, we just replace input arguments with new (decorated) ones. In all other cases, we simply return the input AST, leaving the rest of the tree unchanged.

This is somewhat convoluted, but it solves our problem. Here’s the final version of the trace macro:

defmodule Tracer do
  defmacro deftraceable(head, body) do
    {fun_name, args_ast} = name_and_args(head)

    {arg_names, decorated_args} = decorate_args(args_ast)

    head = Macro.postwalk(head,
      fn
        ({fun_ast, context, old_args}) when (
          fun_ast == fun_name and old_args == args_ast
        ) ->
          {fun_ast, context, decorated_args}
        (other) -> other
      end)

    quote do
      def unquote(head) do
        file = __ENV__.file
        line = __ENV__.line
        module = __ENV__.module

        function_name = unquote(fun_name)
        passed_args = unquote(arg_names) |> Enum.map(&inspect/1) |> Enum.join(",")

        result = unquote(body[:do])

        loc = "#{file}(line #{line})"
        call = "#{module}.#{function_name}(#{passed_args}) = #{inspect result}"
        IO.puts "#{loc} #{call}"

        result
      end
    end
  end

  defp name_and_args({:when, _, [short_head | _]}) do
    name_and_args(short_head)
  end

  defp name_and_args(short_head) do
    Macro.decompose_call(short_head)
  end

  defp decorate_args([]), do: {[],[]}
  defp decorate_args(args_ast) do
    for {arg_ast, index} <- Enum.with_index(args_ast) do
      # dynamically generate quoted identifier
      arg_name = Macro.var(:"arg#{index}", __MODULE__)

      # generate AST for patternX = argX
      full_arg = quote do
        unquote(arg_ast) = unquote(arg_name)
      end

      {arg_name, full_arg}
    end
    |> Enum.unzip
  end
end

Let’s try it out:

iex(1)> defmodule Tracer do ... end

iex(2)> defmodule Test do
          import Tracer

          deftraceable div(_, 0), do: :error
          deftraceable div(a, b), do: a/b
        end

iex(3)> Test.div(5, 2)
iex(line 6) Elixir.Test.div(5,2) = 2.5

iex(4)> Test.div(5, 0)
iex(line 5) Elixir.Test.div(5,0) = :error

As you can see, it’s possible, and not extremely complicated, to get into the AST, tear it apart, and sprinkle it with some custom injected code. On the downside, the code of the resulting macro gets increasingly complex, and it becomes harder to analyze.

This concludes today’s session. Next time I’m going to discuss some aspects of in-place code generation.

Comments

Privacy policy