Add extract_function/4 based on AST modification with Sourceror.Zipper

Author: robmckinnon

apps/language_server/lib/language_server/experimental/code_mod/extract_function.ex

@@ -0,0 +1,238 @@
+defmodule ElixirLS.LanguageServer.Experimental.CodeMod.ExtractFunction do
+  @moduledoc """
+  Elixir refactoring functions.
+  """
+
+  alias Sourceror.Zipper, as: Z
+
+  @doc """
+  Return zipper containing AST with extracted function.
+  """
+  def extract_function(zipper, start_line, end_line, function_name) do
+    {quoted, acc} = extract_lines(zipper, start_line, end_line, function_name)
+    zipper = Z.zip(quoted)
+
+    declares = vars_declared(function_name, [], acc.lines) |> Enum.uniq()
+    used = vars_used(function_name, [], acc.lines) |> Enum.uniq()
+    args = Enum.map(used -- declares, fn var -> {var, [], nil} end)
+    returns = declares |> Enum.filter(&(&1 in acc.vars))
+    {zipper, extracted} = return_declared(zipper, returns, function_name, args, acc.lines)
+
+    enclosing = acc.def
+
+    zipper
+    |> top_find(fn
+      {:def, _meta, [{^enclosing, _, _}, _]} -> true
+      _ -> false
+    end)
+    |> Z.insert_right(extracted)
+    |> fix_block()
+    |> Z.root()
+  end
+
+  @doc """
+  Return zipper containing AST for lines in the range from-to.
+  """
+  def extract_lines(zipper, start_line, end_line, replace_with \\ nil) do
+    remove_range(zipper, start_line, end_line, %{
+      lines: [],
+      def: nil,
+      def_end: nil,
+      vars: [],
+      replace_with: replace_with
+    })
+  end
+
+  defp next_remove_range(zipper, from, to, acc) do
+    if next = Z.next(zipper) do
+      remove_range(next, from, to, acc)
+    else
+      # return zipper with lines removed
+      {
+        elem(Z.top(zipper), 0),
+        %{acc | lines: Enum.reverse(acc.lines), vars: Enum.reverse(acc.vars)}
+      }
+    end
+  end
+
+  defp remove_range({{:def, meta, [{marker, _, _}, _]}, _list} = zipper, from, to, acc) do
+    acc =
+      if meta[:line] < from do
+        x = put_in(acc.def, marker)
+        put_in(x.def_end, meta[:end][:line])
+      else
+        acc
+      end
+
+    next_remove_range(zipper, from, to, acc)
+  end
+
+  defp remove_range({{marker, meta, children}, _list} = zipper, from, to, acc) do
+    if meta[:line] < from || meta[:line] > to || marker == :__block__ do
+      next_remove_range(
+        zipper,
+        from,
+        to,
+        if meta[:line] > to && meta[:line] < acc.def_end && is_atom(marker) && is_nil(children) do
+          put_in(acc.vars, [marker | acc.vars] |> Enum.uniq())
+        else
+          acc
+        end
+      )
+    else
+      acc = put_in(acc.lines, [Z.node(zipper) | acc.lines])
+
+      if is_nil(acc.replace_with) do
+        zipper
+        |> Z.remove()
+        |> next_remove_range(from, to, acc)
+      else
+        function_name = acc.replace_with
+        acc = put_in(acc.replace_with, nil)
+
+        zipper
+        |> Z.replace({function_name, [], []})
+        |> next_remove_range(from, to, acc)
+      end
+    end
+  end
+
+  defp remove_range(zipper, from, to, acc) do
+    next_remove_range(zipper, from, to, acc)
+  end
+
+  defp vars_declared(function_name, args, lines) do
+    function_name
+    |> new_function(args, lines)
+    |> Z.zip()
+    |> vars_declared(%{vars: []})
+  end
+
+  defp vars_declared(nil, acc) do
+    Enum.reverse(acc.vars)
+  end
+
+  defp vars_declared({{:=, _, [{var, _, nil}, _]}, _rest} = zipper, acc) when is_atom(var) do
+    zipper
+    |> Z.next()
+    |> vars_declared(put_in(acc.vars, [var | acc.vars]))
+  end
+
+  defp vars_declared(zipper, acc) do
+    zipper
+    |> Z.next()
+    |> vars_declared(acc)
+  end
+
+  defp vars_used(function_name, args, lines) do
+    function_name
+    |> new_function(args, lines)
+    |> Z.zip()
+    |> vars_used(%{vars: []})
+  end
+
+  defp vars_used(nil, acc) do
+    Enum.reverse(acc.vars)
+  end
+
+  defp vars_used({{marker, _meta, nil}, _rest} = zipper, acc) when is_atom(marker) do
+    zipper
+    |> Z.next()
+    |> vars_used(put_in(acc.vars, [marker | acc.vars]))
+  end
+
+  defp vars_used(zipper, acc) do
+    zipper
+    |> Z.next()
+    |> vars_used(acc)
+  end
+
+  defp return_declared(zipper, nil = _declares, function_name, args, lines) do
+    {zipper, new_function(function_name, args, lines)}
+  end
+
+  defp return_declared(zipper, [var], function_name, args, lines) when is_atom(var) do
+    zipper =
+      zipper
+      |> top_find(fn
+        {^function_name, [], []} -> true
+        _ -> false
+      end)
+      |> Z.replace({:=, [], [{var, [], nil}, {function_name, [], args}]})
+
+    {zipper, new_function(function_name, args, Enum.concat(lines, [{var, [], nil}]))}
+  end
+
+  defp return_declared(zipper, declares, function_name, args, lines) when is_list(declares) do
+    declares = Enum.reduce(declares, {}, fn var, acc -> Tuple.append(acc, {var, [], nil}) end)
+
+    zipper =
+      zipper
+      |> top_find(fn
+        {^function_name, [], []} -> true
+        _ -> false
+      end)
+      |> Z.replace(
+        {:=, [],
+         [
+           {:__block__, [],
+            [
+              declares
+            ]},
+           {function_name, [], args}
+         ]}
+      )
+
+    {zipper,
+     new_function(
+       function_name,
+       args,
+       Enum.concat(lines, [
+         {:__block__, [],
+          [
+            declares
+          ]}
+       ])
+     )}
+  end
+
+  defp new_function(function_name, args, lines) do
+    {:def, [do: [], end: []],
+     [
+       {function_name, [], args},
+       [
+         {
+           {:__block__, [], [:do]},
+           {:__block__, [], lines}
+         }
+       ]
+     ]}
+  end
+
+  defp fix_block(zipper) do
+    zipper
+    |> top_find(fn
+      {:{}, [], _children} -> true
+      _ -> false
+    end)
+    |> case do
+      nil ->
+        zipper
+
+      {{:{}, [], [block | defs]}, meta} ->
+        {
+          {
+            block,
+            {:__block__, [], defs}
+          },
+          meta
+        }
+    end
+  end
+
+  defp top_find(zipper, function) do
+    zipper
+    |> Z.top()
+    |> Z.find(function)
+  end
+end

apps/language_server/test/experimental/code_mod/extract_function_test.exs

@@ -0,0 +1,186 @@
+defmodule ElixirLS.LanguageServer.Experimental.CodeMod.ExtractFunctionTest do
+  alias ElixirLS.LanguageServer.Experimental.CodeMod.ExtractFunction
+
+  use ExUnit.Case
+  alias Sourceror.Zipper, as: Z
+
+  setup ctx do
+    if Map.has_key?(ctx, :no_setup) do
+      {:ok, []}
+    else
+      no = Map.get(ctx, :no, "")
+
+      {:ok,
+       quoted:
+         """
+         defmodule Baz#{no} do
+           def foo(one, two) do
+             three = 3
+             IO.inspect(one)
+             IO.inspect(two)
+             IO.inspect(three)
+             four = 4
+             IO.inspect(three)
+             IO.inspect(four)
+             # comment
+           end
+         end
+         """
+         |> Sourceror.parse_string!()}
+    end
+  end
+
+  describe "extract_function" do
+    @tag no: 1
+    test "extract one line to function", %{quoted: quoted} do
+      zipper = ExtractFunction.extract_function(Z.zip(quoted), 3, 3, :bar)
+      source = Sourceror.to_string(zipper)
+
+      assert [
+               "defmodule Baz1 do",
+               "  def foo(one, two) do",
+               "    three = bar()",
+               "    IO.inspect(one)",
+               "    IO.inspect(two)",
+               "    IO.inspect(three)",
+               "    four = 4",
+               "    IO.inspect(three)",
+               "    IO.inspect(four)",
+               "",
+               "    # comment",
+               "  end",
+               "",
+               "  def bar() do",
+               "    three = 3",
+               "    three",
+               "  end",
+               "end"
+             ] ==
+               source |> String.split("\n")
+
+      Code.eval_string(source)
+    end
+
+    @tag no: 2
+    test "extract multiple lines to function", %{quoted: quoted} do
+      zipper = ExtractFunction.extract_function(Z.zip(quoted), 3, 4, :bar)
+      source = Sourceror.to_string(zipper)
+
+      assert [
+               "defmodule Baz2 do",
+               "  def foo(one, two) do",
+               "    three = bar(one)",
+               "    IO.inspect(two)",
+               "    IO.inspect(three)",
+               "    four = 4",
+               "    IO.inspect(three)",
+               "    IO.inspect(four)",
+               "",
+               "    # comment",
+               "  end",
+               "",
+               "  def bar(one) do",
+               "    three = 3",
+               "    IO.inspect(one)",
+               "    three",
+               "  end",
+               "end"
+             ] ==
+               source |> String.split("\n")
+
+      Code.eval_string(source)
+    end
+
+    @tag no: 3
+    test "extract multiple lines with multiple returns to function", %{quoted: quoted} do
+      zipper = ExtractFunction.extract_function(Z.zip(quoted), 3, 7, :bar)
+      source = Sourceror.to_string(zipper)
+
+      assert [
+               "defmodule Baz3 do",
+               "  def foo(one, two) do",
+               "    {three, four} = bar(one, two)",
+               "    IO.inspect(three)",
+               "    IO.inspect(four)",
+               "",
+               "    # comment",
+               "  end",
+               "",
+               "  def bar(one, two) do",
+               "    three = 3",
+               "    IO.inspect(one)",
+               "    IO.inspect(two)",
+               "    IO.inspect(three)",
+               "    four = 4",
+               "    {three, four}",
+               "  end",
+               "end"
+             ] ==
+               source |> String.split("\n")
+
+      Code.eval_string(source)
+    end
+
+    @tag no: 4
+    test "extract multiple lines with single return value to function", %{quoted: quoted} do
+      zipper = ExtractFunction.extract_function(Z.zip(quoted), 3, 8, :bar)
+      source = Sourceror.to_string(zipper)
+
+      assert [
+               "defmodule Baz4 do",
+               "  def foo(one, two) do",
+               "    four = bar(one, two)",
+               "    IO.inspect(four)",
+               "",
+               "    # comment",
+               "  end",
+               "",
+               "  def bar(one, two) do",
+               "    three = 3",
+               "    IO.inspect(one)",
+               "    IO.inspect(two)",
+               "    IO.inspect(three)",
+               "    four = 4",
+               "    IO.inspect(three)",
+               "    four",
+               "  end",
+               "end"
+             ] ==
+               source |> String.split("\n")
+
+      Code.eval_string(source)
+    end
+  end
+
+  describe "extract_lines/3" do
+    test "extract one line to function", %{quoted: quoted} do
+      {zipper, lines} = ExtractFunction.extract_lines(Z.zip(quoted), 3, 3)
+
+      assert "defmodule Baz do\n  def foo(one, two) do\n    IO.inspect(one)\n    IO.inspect(two)\n    IO.inspect(three)\n    four = 4\n    IO.inspect(three)\n    IO.inspect(four)\n\n    # comment\n  end\nend" ==
+               Sourceror.to_string(zipper)
+
+      assert [
+               "{:def, :foo}",
+               "{:def_end, 11}",
+               "{:lines, [three = 3]}",
+               _,
+               "{:vars, [:one, :two, :three, :four]}"
+             ] = lines |> Enum.map(&Sourceror.to_string(&1))
+    end
+
+    test "extract multiple lines to function", %{quoted: quoted} do
+      {zipper, lines} = ExtractFunction.extract_lines(Z.zip(quoted), 3, 4)
+
+      assert "defmodule Baz do\n  def foo(one, two) do\n    IO.inspect(two)\n    IO.inspect(three)\n    four = 4\n    IO.inspect(three)\n    IO.inspect(four)\n\n    # comment\n  end\nend" =
+               Sourceror.to_string(zipper)
+
+      assert [
+               "{:def, :foo}",
+               "{:def_end, 11}",
+               "{:lines, [three = 3, IO.inspect(one)]}",
+               _,
+               "{:vars, [:two, :three, :four]}"
+             ] = lines |> Enum.map(&Sourceror.to_string(&1))
+    end
+  end
+end