Support extract function call in refactor

Author: robmckinnon

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

@@ -9,14 +9,17 @@ defmodule ElixirLS.LanguageServer.Experimental.CodeMod.ExtractFunction do
   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)
+    {quoted_after_extract, acc} = extract_lines(zipper, start_line, end_line, function_name)
 
-    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)
+    new_function_zipper = new_function(function_name, [], acc.lines) |> Z.zip()
+    declared_vars = vars_declared(new_function_zipper) |> Enum.uniq()
+    used_vars = vars_used(new_function_zipper) |> Enum.uniq()
+
+    args = used_vars -- declared_vars
+    returns = declared_vars |> Enum.filter(&(&1 in acc.vars))
+
+    {zipper, extracted} =
+      add_returned_vars(Z.zip(quoted_after_extract), returns, function_name, args, acc.lines)
 
     enclosing = acc.def
 
@@ -101,11 +104,8 @@ defmodule ElixirLS.LanguageServer.Experimental.CodeMod.ExtractFunction 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: []})
+  defp vars_declared(new_function_zipper) do
+    vars_declared(new_function_zipper, %{vars: []})
   end
 
   defp vars_declared(nil, acc) do
@@ -124,11 +124,8 @@ defmodule ElixirLS.LanguageServer.Experimental.CodeMod.ExtractFunction do
     |> vars_declared(acc)
   end
 
-  defp vars_used(function_name, args, lines) do
-    function_name
-    |> new_function(args, lines)
-    |> Z.zip()
-    |> vars_used(%{vars: []})
+  defp vars_used(new_function_zipper) do
+    vars_used(new_function_zipper, %{vars: []})
   end
 
   defp vars_used(nil, acc) do
@@ -147,53 +144,53 @@ defmodule ElixirLS.LanguageServer.Experimental.CodeMod.ExtractFunction do
     |> 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}
-         ]}
-      )
+  defp add_returned_vars(zipper, _returns = [], function_name, args, lines) do
+    args = var_ast(args)
+
+    {
+      replace_function_call(zipper, function_name, {function_name, [], args}),
+      new_function(function_name, args, lines)
+    }
+  end
+
+  defp add_returned_vars(zipper, returns, function_name, args, lines) when is_list(returns) do
+    args = var_ast(args)
+    returned_vars = returned(returns)
+
+    {
+      replace_function_call(
+        zipper,
+        function_name,
+        {:=, [], [returned_vars, {function_name, [], args}]}
+      ),
+      new_function(function_name, args, Enum.concat(lines, [returned_vars]))
+    }
+  end
+
+  defp var_ast(vars) when is_list(vars) do
+    Enum.map(vars, &var_ast/1)
+  end
 
-    {zipper,
-     new_function(
-       function_name,
-       args,
-       Enum.concat(lines, [
-         {:__block__, [],
-          [
-            declares
-          ]}
-       ])
-     )}
+  defp var_ast(var) when is_atom(var) do
+    {var, [], nil}
+  end
+
+  defp returned([var]) when is_atom(var) do
+    var_ast(var)
+  end
+
+  defp returned(vars) when is_list(vars) do
+    returned = vars |> var_ast() |> List.to_tuple()
+    {:__block__, [], [returned]}
+  end
+
+  defp replace_function_call(zipper, function_name, replace_with) do
+    zipper
+    |> top_find(fn
+      {^function_name, [], []} -> true
+      _ -> false
+    end)
+    |> Z.replace(replace_with)
   end
 
   defp new_function(function_name, args, lines) do

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

@@ -21,7 +21,11 @@ defmodule ElixirLS.LanguageServer.Experimental.CodeMod.ExtractFunctionTest do
              IO.inspect(three)
              four = 4
              IO.inspect(three)
-             IO.inspect(four)
+
+             IO.inspect(four: four,
+               force_format_on_new_line_with_really_long_atom: true
+             )
+
              # comment
            end
          end
@@ -45,7 +49,11 @@ defmodule ElixirLS.LanguageServer.Experimental.CodeMod.ExtractFunctionTest do
                "    IO.inspect(three)",
                "    four = 4",
                "    IO.inspect(three)",
-               "    IO.inspect(four)",
+               "",
+               "    IO.inspect(",
+               "      four: four,",
+               "      force_format_on_new_line_with_really_long_atom: true",
+               "    )",
                "",
                "    # comment",
                "  end",
@@ -74,7 +82,11 @@ defmodule ElixirLS.LanguageServer.Experimental.CodeMod.ExtractFunctionTest do
                "    IO.inspect(three)",
                "    four = 4",
                "    IO.inspect(three)",
-               "    IO.inspect(four)",
+               "",
+               "    IO.inspect(",
+               "      four: four,",
+               "      force_format_on_new_line_with_really_long_atom: true",
+               "    )",
                "",
                "    # comment",
                "  end",
@@ -101,7 +113,11 @@ defmodule ElixirLS.LanguageServer.Experimental.CodeMod.ExtractFunctionTest do
                "  def foo(one, two) do",
                "    {three, four} = bar(one, two)",
                "    IO.inspect(three)",
-               "    IO.inspect(four)",
+               "",
+               "    IO.inspect(",
+               "      four: four,",
+               "      force_format_on_new_line_with_really_long_atom: true",
+               "    )",
                "",
                "    # comment",
                "  end",
@@ -130,7 +146,11 @@ defmodule ElixirLS.LanguageServer.Experimental.CodeMod.ExtractFunctionTest do
                "defmodule Baz4 do",
                "  def foo(one, two) do",
                "    four = bar(one, two)",
-               "    IO.inspect(four)",
+               "",
+               "    IO.inspect(",
+               "      four: four,",
+               "      force_format_on_new_line_with_really_long_atom: true",
+               "    )",
                "",
                "    # comment",
                "  end",
@@ -142,6 +162,7 @@ defmodule ElixirLS.LanguageServer.Experimental.CodeMod.ExtractFunctionTest do
                "    IO.inspect(three)",
                "    four = 4",
                "    IO.inspect(three)",
+               "",
                "    four",
                "  end",
                "end"
@@ -150,37 +171,88 @@ defmodule ElixirLS.LanguageServer.Experimental.CodeMod.ExtractFunctionTest do
 
       Code.eval_string(source)
     end
+
+    @tag no: 5
+    test "extracts when extract partial function call", %{quoted: quoted} do
+      zipper = ExtractFunction.extract_function(Z.zip(quoted), 10, 10, :bar)
+      source = Sourceror.to_string(zipper)
+
+      assert [
+               "defmodule Baz5 do",
+               "  def foo(one, two) do",
+               "    three = 3",
+               "    IO.inspect(one)",
+               "    IO.inspect(two)",
+               "    IO.inspect(three)",
+               "    four = 4",
+               "    IO.inspect(three)",
+               "",
+               "    bar(four)",
+               "",
+               "    # comment",
+               "  end",
+               "",
+               "  def bar(four) do",
+               "    IO.inspect(",
+               "      four: four,",
+               "      force_format_on_new_line_with_really_long_atom: true",
+               "    )",
+               "  end",
+               "end"
+             ] ==
+               source |> String.split("\n")
+
+      Code.eval_string(source)
+    end
   end
 
   describe "extract_lines/3" do
+    @tag no: 20
     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" ==
+      assert "defmodule Baz20 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\n    IO.inspect(\n      four: four,\n      force_format_on_new_line_with_really_long_atom: true\n    )\n\n    # comment\n  end\nend" ==
                Sourceror.to_string(zipper)
 
       assert [
                "{:def, :foo}",
-               "{:def_end, 11}",
+               "{:def_end, 15}",
                "{:lines, [three = 3]}",
                _,
                "{:vars, [:one, :two, :three, :four]}"
              ] = lines |> Enum.map(&Sourceror.to_string(&1))
     end
 
+    @tag no: 21
     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" =
+      assert "defmodule Baz21 do\n  def foo(one, two) do\n    IO.inspect(two)\n    IO.inspect(three)\n    four = 4\n    IO.inspect(three)\n\n    IO.inspect(\n      four: four,\n      force_format_on_new_line_with_really_long_atom: true\n    )\n\n    # comment\n  end\nend" =
                Sourceror.to_string(zipper)
 
       assert [
                "{:def, :foo}",
-               "{:def_end, 11}",
+               "{:def_end, 15}",
                "{:lines, [three = 3, IO.inspect(one)]}",
                _,
                "{:vars, [:two, :three, :four]}"
              ] = lines |> Enum.map(&Sourceror.to_string(&1))
     end
+
+    @tag no: 22
+    test "extract multi-line function call to function", %{quoted: quoted} do
+      {zipper, lines} = ExtractFunction.extract_lines(Z.zip(quoted), 10, 10)
+
+      assert "defmodule Baz22 do\n  def foo(one, two) do\n    three = 3\n    IO.inspect(one)\n    IO.inspect(two)\n    IO.inspect(three)\n    four = 4\n    IO.inspect(three)\n\n    # comment\n  end\nend" =
+               Sourceror.to_string(zipper)
+
+      assert [
+               "{:def, :foo}",
+               "{:def_end, 15}",
+               "{:lines,\n [\n   IO.inspect(\n     four: four,\n     force_format_on_new_line_with_really_long_atom: true\n   )\n ]}",
+               "{:replace_with, nil}",
+               "{:vars, []}"
+             ] = lines |> Enum.map(&Sourceror.to_string(&1))
+    end
   end
 end