From d70acc33926c6321ac6f909a2b3be4afebd59c50 Mon Sep 17 00:00:00 2001
From: Zach Allaun <zach.allaun@gmail.com>
Date: Thu, 11 Jul 2024 15:14:44 -0400
Subject: [PATCH] Add `Zipper.at/2`

---
 lib/sourceror/zipper.ex        | 68 ++++++++++++++++++++++++++++++++++
 test/support/cursor_support.ex | 21 +++++++++++
 test/zipper_test.exs           | 67 ++++++++++++++++++++++++++++++++-
 3 files changed, 155 insertions(+), 1 deletion(-)
 create mode 100644 test/support/cursor_support.ex

diff --git a/lib/sourceror/zipper.ex b/lib/sourceror/zipper.ex
index 49cb634..1373b7b 100644
--- a/lib/sourceror/zipper.ex
+++ b/lib/sourceror/zipper.ex
@@ -83,6 +83,74 @@ defmodule Sourceror.Zipper do
   @spec zip(tree) :: t
   def zip(node), do: new(node)
 
+  @doc """
+  Creates a `zipper` from a tree `node` focused at the innermost descendant containing `position`.
+
+  Returns `{:ok, zipper}` if `position` is within `node`, else `:error`.
+
+  Modifying `node` prior to using `at/2` is not recommended as added or
+  changed descendants may not contain accurate position metadata used to
+  find the focus.
+  """
+  @spec at(Macro.t(), Sourceror.position()) :: {:ok, t} | :error
+  def at(node, position) when is_list(position) do
+    with {:ok, path} <- fetch_path_to(node, position) do
+      case path do
+        [{node, [], []}, {{:__block__, _, [node]}, _, _} = block_wrapper | ancestors] ->
+          {:ok, new_from_path([block_wrapper | ancestors])}
+
+        _ ->
+          {:ok, new_from_path(path)}
+      end
+    end
+  end
+
+  defp new_from_path([{node, [], []}]) do
+    new(node)
+  end
+
+  defp new_from_path([{node, left, right} | ancestors]) do
+    path = %{left: left, right: right, parent: new_from_path(ancestors)}
+    new(node, path, nil)
+  end
+
+  defp fetch_path_to(node, position) do
+    if node_contains?(node, position) do
+      {:ok, path_to(position, [{node, [], []}])}
+    else
+      :error
+    end
+  end
+
+  defp path_to(position, [{parent, _parent_left, _parent_right} | _] = path) do
+    {left, node_and_right} =
+      parent
+      |> children()
+      |> Enum.split_while(fn child ->
+        not node_contains?(child, position)
+      end)
+
+    case node_and_right do
+      [] ->
+        path
+
+      [node | right] ->
+        reversed_left = Enum.reverse(left)
+        path_to(position, [{node, reversed_left, right} | path])
+    end
+  end
+
+  defp node_contains?(node, position) do
+    case Sourceror.get_range(node) do
+      %Sourceror.Range{} = range ->
+        Sourceror.compare_positions(position, range.start) in [:gt, :eq] and
+          Sourceror.compare_positions(position, range.end) == :lt
+
+      nil ->
+        false
+    end
+  end
+
   @doc """
   Walks the `zipper` to the top of the current subtree and returns that `zipper`.
   """
diff --git a/test/support/cursor_support.ex b/test/support/cursor_support.ex
new file mode 100644
index 0000000..80cdfb5
--- /dev/null
+++ b/test/support/cursor_support.ex
@@ -0,0 +1,21 @@
+defmodule SourcerorTest.CursorSupport do
+  @moduledoc false
+
+  @doc """
+  Extracts the `[:line, :column]` position of the first `|` character found in the text.
+
+  Returns `{position, text}`.
+  """
+  def pop_cursor(text) do
+    case String.split(text, "|", parts: 2) do
+      [prefix, suffix] ->
+        lines = String.split(prefix, "\n")
+        cursor_line_length = lines |> List.last() |> String.length()
+        position = [line: length(lines), column: cursor_line_length + 1]
+        {position, prefix <> suffix}
+
+      _ ->
+        raise ArgumentError, "Could not find cursor in:\n\n#{text}"
+    end
+  end
+end
diff --git a/test/zipper_test.exs b/test/zipper_test.exs
index 76ba00b..d7265b9 100644
--- a/test/zipper_test.exs
+++ b/test/zipper_test.exs
@@ -1,8 +1,9 @@
 defmodule SourcerorTest.ZipperTest do
   use ExUnit.Case, async: true
-
   doctest Sourceror.Zipper, import: true, except: [:moduledoc]
 
+  import SourcerorTest.CursorSupport, only: [pop_cursor: 1]
+
   alias Sourceror.Zipper, as: Z
 
   describe "zip/1" do
@@ -905,4 +906,68 @@ defmodule SourcerorTest.ZipperTest do
       assert "20" == new_zipper |> Z.node() |> Sourceror.to_string()
     end
   end
+
+  describe "at/2" do
+    defp zipper_at_cursor(code_with_cursor) do
+      {position, code} = pop_cursor(code_with_cursor)
+      code |> Sourceror.parse_string!() |> Z.at(position)
+    end
+
+    test "creates a zipper focused on an inner literal at the given position" do
+      assert {:ok, zipper} =
+               zipper_at_cursor("""
+               def foo do
+                 [1, 2, |3, 4, 5]
+               end
+               """)
+
+      assert {:__block__, _, [3]} = zipper |> Z.node()
+
+      assert [
+               {:__block__, _, [1]},
+               {:__block__, _, [2]},
+               {:__block__, _, [3]},
+               {:__block__, _, [4]},
+               {:__block__, _, [5]}
+             ] = zipper |> Z.up() |> Z.node()
+
+      assert {:def, _, _} = zipper |> Z.root()
+    end
+
+    test "creates a zipper focused on a container if position isn't in any children" do
+      assert {:ok, zipper} =
+               zipper_at_cursor("""
+               def foo do
+                 [1|, 2, 3, 4, 5]
+               end
+               """)
+
+      assert {:__block__, _, [[_, _, _, _, _]]} = zipper |> Z.node()
+      assert {:def, _, _} = zipper |> Z.root()
+    end
+
+    test "creates a zipper focused on an alias segment" do
+      assert {:ok, zipper} = zipper_at_cursor("alias Foo.{Bar, |Baz}")
+
+      assert {:__aliases__, _, [:Baz]} = zipper |> Z.node()
+      assert {:__aliases__, _, [:Bar]} = zipper |> Z.left() |> Z.node()
+      assert {:alias, _, _} = zipper |> Z.root()
+    end
+
+    test "creates a zipper focused on a qualified call" do
+      assert {:ok, zipper} = zipper_at_cursor("Foo.|bar(1, 2, 3)")
+
+      assert {:., _, [{:__aliases__, _, _}, :bar]} = zipper |> Z.node()
+      assert {:__block__, _, [1]} = zipper |> Z.right() |> Z.node()
+    end
+
+    test "returns :error if there is no node containing the given position" do
+      assert :error =
+               zipper_at_cursor("""
+               def foo do
+                 [1, 2, 3, 4, 5]
+               end|
+               """)
+    end
+  end
 end