From 3eb35d8e076eda236b305322c87fa594ec97863b Mon Sep 17 00:00:00 2001
From: doorgan <dorgandash@gmail.com>
Date: Fri, 31 May 2024 00:09:15 -0300
Subject: [PATCH] Add Sourceror.Zipper.within/2

---
 lib/sourceror/zipper.ex | 11 +++++++++++
 test/zipper_test.exs    | 32 +++++++++++++++++++++++++++++++-
 2 files changed, 42 insertions(+), 1 deletion(-)

diff --git a/lib/sourceror/zipper.ex b/lib/sourceror/zipper.ex
index 77f7e35..6e2c496 100644
--- a/lib/sourceror/zipper.ex
+++ b/lib/sourceror/zipper.ex
@@ -474,4 +474,15 @@ defmodule Sourceror.Zipper do
   @spec subtree(t) :: t
   @compile {:inline, subtree: 1}
   def subtree(%Z{} = zipper), do: %{zipper | path: nil}
+
+  @doc """
+  Runs the function `fun` on the subtree of the currently focused `node` and
+  returns the updated `zipper`.
+
+  `fun` must return a zipper, which may be positioned at the top of the subtree.
+  """
+  def within(%Z{} = zipper, fun) when is_function(fun, 1) do
+    updated = zipper |> subtree() |> fun.() |> top()
+    into(updated, zipper)
+  end
 end
diff --git a/test/zipper_test.exs b/test/zipper_test.exs
index b58a15e..941badc 100644
--- a/test/zipper_test.exs
+++ b/test/zipper_test.exs
@@ -1,6 +1,7 @@
 defmodule SourcerorTest.ZipperTest do
   use ExUnit.Case, async: true
-  doctest Sourceror.Zipper, except: [:moduledoc]
+
+  doctest Sourceror.Zipper, import: true, except: [:moduledoc]
 
   alias Sourceror.Zipper, as: Z
 
@@ -685,4 +686,33 @@ defmodule SourcerorTest.ZipperTest do
       assert :ok = Z.Inspect.default_inspect_as(:as_ast)
     end
   end
+
+  describe "within/2" do
+    test "executes a function within a zipper" do
+      code = """
+      config :target, key: :change_me
+
+      config :unrelated, key: :dont_change_me
+      """
+
+      updated =
+        code
+        |> Sourceror.parse_string!()
+        |> Z.zip()
+        |> Z.find(&match?({:config, _, [{:__block__, _, [:target]} | _]}, &1))
+        |> Z.within(fn zipper ->
+          zipper
+          |> Z.find(&match?({{:__block__, _, [:key]}, _value}, &1))
+          |> Z.update(fn {key, _value} -> {key, {:__block__, [], [:changed]}} end)
+        end)
+        |> Z.root()
+        |> Sourceror.to_string()
+
+      assert updated == """
+             config :target, key: :changed
+
+             config :unrelated, key: :dont_change_me\
+             """
+    end
+  end
 end