diff --git a/lib/ast_mangler.ex b/lib/ast_mangler.ex index eed4281..ac5f253 100644 --- a/lib/ast_mangler.ex +++ b/lib/ast_mangler.ex @@ -1,8 +1,25 @@ defmodule ASTMangler do + + def expand([do: remainder], replacements) when is_list(replacements) do + [do: expand(remainder, replacements)] + end + def expand(ast, replacements) when is_list(replacements) do + {node, _acc} = Macro.prewalk(ast, replacements, fn(node, acc) -> pre(node, acc) end) + node + end def expand(ast, replacement) do - Macro.prewalk(ast, fn(node) -> update(node, replacement) end) + expand(ast, [replacement]) end - def update(:__, replacement), do: replacement - def update(node, _), do: node + def pre(:__, [first | remainder]), do: {first, remainder} + def pre(node, acc), do: {node, acc} + + + def count(ast) do + {_node, acc} = Macro.prewalk(ast, 0, fn(node, acc) -> count(node, acc) end) + acc + end + + def count(:__, acc), do: {node, acc+1} + def count(node, acc), do: {node, acc} end diff --git a/test/ast_mangler_test.exs b/test/ast_mangler_test.exs index 9cb66b7..223fdec 100644 --- a/test/ast_mangler_test.exs +++ b/test/ast_mangler_test.exs @@ -28,4 +28,22 @@ defmodule ASTManglerTest do assert [do: {:assert, [line: 5], [{:==, [line: 5], [true, true]}]}] == ASTMangler.expand(ast, true) end + + test "multiple arguments" do + ast = [do: {:assert, [line: 5], [{:==, [line: 5], [:__, :__]}]}] + + assert [do: {:assert, [line: 5], [{:==, [line: 5], [true, false]}]}] == ASTMangler.expand(ast, [true, false]) + end + + test "counts simple blanks" do + ast = quote do: 1 + :__ + + assert ASTMangler.count(ast) == 1 + end + + test "counts multiple blanks" do + ast = [do: {:assert, [line: 5], [{:==, [line: 5], [:__, :__]}]}] + + assert ASTMangler.count(ast) == 2 + end end