From 9c803ff72cc25cf82558724d2d8f49c4d39c8451 Mon Sep 17 00:00:00 2001
From: Felipe Sere <felipesere@gmail.com>
Date: Sat, 12 Mar 2016 15:05:26 +0000
Subject: [PATCH] Use prewalk to inject value into AST

---
 lib/ast_mangler.ex | 24 ++++--------------------
 1 file changed, 4 insertions(+), 20 deletions(-)

diff --git a/lib/ast_mangler.ex b/lib/ast_mangler.ex
index 746adf9..eed4281 100644
--- a/lib/ast_mangler.ex
+++ b/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