Use prewalk to inject value into AST

lib/ast_mangler.ex

@@ -1,24 +1,8 @@
 defmodule ASTMangler do
-  def expand([], _), do: []
-  def expand({fun, context, args}, replacement) do
-    new_args = replace(args, replacement)
-    {fun, context, new_args}
+  def expand(ast, replacement) do
+    Macro.prewalk(ast, fn(node) -> update(node, replacement) end)
   end
-  def expand([do: thing], replacement) do
-    [do: expand(thing, replacement)]
-  end
-  def expand(n, _), do: n
 
-  def replace(nil, _), do: nil
-  def replace(args, b) do
-    args
-    |> Enum.find_index(fn(x) -> x == :__ end)
-    |> replace(args, b)
-  end
-  def replace(nil, ast, b) when is_list(ast) do
-    Enum.map(ast, fn(element) -> expand(element, b) end)
-  end
-  def replace(index, list, b) when is_integer(index) do
-    List.update_at(list, index, fn(_)-> b end)
-  end
+  def update(:__, replacement), do: replacement
+  def update(node, _), do: node
 end