Fix put_in_map/set_map_key not setting keys properly

lib/igniter/code/map.ex

@@ -30,57 +30,59 @@ defmodule Igniter.Code.Map do
   end
 
   defp do_put_in_map(zipper, [key | rest], value, updater) do
-    cond do
-      Common.node_matches_pattern?(zipper, {:%{}, _, []}) ->
-        value = Common.use_aliases(value, zipper)
-
-        {:ok,
-         Zipper.append_child(
-           zipper,
-           mappify([key | rest], value)
-         )}
-
-      Common.node_matches_pattern?(zipper, {:%{}, _, _}) ->
-        zipper
-        |> Zipper.down()
-        |> Igniter.Code.List.move_to_list_item(fn item ->
-          if Igniter.Code.Tuple.tuple?(item) do
-            case Igniter.Code.Tuple.tuple_elem(item, 0) do
-              {:ok, first_elem} ->
-                Common.nodes_equal?(first_elem, key)
-
-              :error ->
-                false
+    Common.within(zipper, fn zipper ->
+      cond do
+        Common.node_matches_pattern?(zipper, {:%{}, _, []}) ->
+          value = Common.use_aliases(value, zipper)
+
+          {:ok,
+           Common.replace_code(
+             zipper,
+             mappify([key | rest], value)
+           )}
+
+        Common.node_matches_pattern?(zipper, {:%{}, _, _}) ->
+          zipper
+          |> Zipper.down()
+          |> Common.move_right(fn item ->
+            if Igniter.Code.Tuple.tuple?(item) do
+              case Igniter.Code.Tuple.tuple_elem(item, 0) do
+                {:ok, first_elem} ->
+                  Common.nodes_equal?(first_elem, key)
+
+                :error ->
+                  false
+              end
             end
+          end)
+          |> case do
+            :error ->
+              value = Common.use_aliases(value, zipper)
+              format = map_keys_format(zipper)
+              value = mappify(rest, value)
+
+              {:ok,
+               Zipper.append_child(
+                 zipper,
+                 {{:__block__, [format: format], [key]}, {:__block__, [], [value]}}
+               )}
+
+            {:ok, zipper} ->
+              zipper
+              |> Igniter.Code.Tuple.tuple_elem(1)
+              |> case do
+                {:ok, zipper} ->
+                  do_put_in_map(zipper, rest, value, updater)
+
+                :error ->
+                  :error
+              end
           end
-        end)
-        |> case do
-          :error ->
-            value = Common.use_aliases(value, zipper)
-            format = map_keys_format(zipper)
-            value = mappify(rest, value)
-
-            {:ok,
-             Zipper.append_child(
-               zipper,
-               {{:__block__, [format: format], [key]}, {:__block__, [], [value]}}
-             )}
-
-          {:ok, zipper} ->
-            zipper
-            |> Igniter.Code.Tuple.tuple_elem(1)
-            |> case do
-              {:ok, zipper} ->
-                do_put_in_map(zipper, rest, value, updater)
-
-              :error ->
-                :error
-            end
-        end
 
-      true ->
-        :error
-    end
+        true ->
+          :error
+      end
+    end)
   end
 
   @doc "Puts a key into a map, calling `updater` on the zipper at the value if the key is already present"
@@ -92,48 +94,50 @@ defmodule Igniter.Code.Map do
         value = Common.use_aliases(value, zipper)
 
         {:ok,
-         Zipper.append_child(
+         Common.replace_code(
            zipper,
            mappify([key], value)
          )}
 
       Common.node_matches_pattern?(zipper, {:%{}, _, _}) ->
-        zipper
-        |> Zipper.down()
-        |> Common.move_right(fn item ->
-          if Igniter.Code.Tuple.tuple?(item) do
-            case Igniter.Code.Tuple.tuple_elem(item, 0) do
-              {:ok, first_elem} ->
-                Common.nodes_equal?(first_elem, key)
-
-              :error ->
-                false
+        Common.within(zipper, fn zipper ->
+          zipper
+          |> Zipper.down()
+          |> Common.move_right(fn item ->
+            if Igniter.Code.Tuple.tuple?(item) do
+              case Igniter.Code.Tuple.tuple_elem(item, 0) do
+                {:ok, first_elem} ->
+                  Common.nodes_equal?(first_elem, key)
+
+                :error ->
+                  false
+              end
             end
+          end)
+          |> case do
+            :error ->
+              value = Common.use_aliases(value, zipper)
+
+              format = map_keys_format(zipper)
+
+              {:ok,
+               Zipper.append_child(
+                 zipper,
+                 {{:__block__, [format: format], [key]}, {:__block__, [], [value]}}
+               )}
+
+            {:ok, zipper} ->
+              zipper
+              |> Igniter.Code.Tuple.tuple_elem(1)
+              |> case do
+                {:ok, zipper} ->
+                  {:ok, updater.(zipper)}
+
+                :error ->
+                  :error
+              end
           end
         end)
-        |> case do
-          :error ->
-            value = Common.use_aliases(value, zipper)
-
-            format = map_keys_format(zipper)
-
-            {:ok,
-             Zipper.append_child(
-               zipper,
-               {{:__block__, [format: format], [key]}, {:__block__, [], [value]}}
-             )}
-
-          {:ok, zipper} ->
-            zipper
-            |> Igniter.Code.Tuple.tuple_elem(1)
-            |> case do
-              {:ok, zipper} ->
-                updater.(zipper)
-
-              :error ->
-                :error
-            end
-        end
 
       true ->
         :error

test/igniter/code/map_test.exs

@@ -0,0 +1,93 @@
+# SPDX-FileCopyrightText: 2025 igniter contributors <https://github.com/ash-project/igniter/graphs.contributors>
+#
+# SPDX-License-Identifier: MIT
+
+defmodule Igniter.Code.Map.Test do
+  alias Igniter.Code
+  alias Sourceror.Zipper
+
+  use ExUnit.Case
+
+  doctest Igniter.Code.Map
+
+  describe "set_map_key/4" do
+    test "inserts value into empty map" do
+      {:ok, zipper} =
+        "%{}"
+        |> Code.Common.parse_to_zipper!()
+        |> Code.Map.set_map_key(:foo, :bar, fn _valuex -> flunk() end)
+
+      assert {:ok, %{foo: :bar}} == Code.Common.expand_literal(zipper)
+    end
+
+    test "inserts value into map if not present" do
+      {:ok, zipper} =
+        "%{hello: :world}"
+        |> Code.Common.parse_to_zipper!()
+        |> Code.Map.set_map_key(:foo, :bar, &flunk/1)
+
+      assert {:ok, %{foo: :bar, hello: :world}} == Code.Common.expand_literal(zipper)
+    end
+
+    test "replaces keys in map with the updater rather than the passed value" do
+      {:ok, zipper} =
+        "%{hello: :world}"
+        |> Code.Common.parse_to_zipper!()
+        |> Code.Map.set_map_key(:hello, :this_value_is_ignored, fn %Zipper{node: :world} = zipper ->
+          Zipper.replace(zipper, :baz)
+        end)
+
+      assert {:ok, %{hello: :baz}} == Code.Common.expand_literal(zipper)
+    end
+  end
+
+  describe "put_in_map/4" do
+    test "inserts value into empty map" do
+      {:ok, zipper} =
+        "%{}"
+        |> Code.Common.parse_to_zipper!()
+        |> Code.Map.put_in_map([:foo], :bar, fn _valuex -> flunk() end)
+
+      assert {:ok, %{foo: :bar}} == Code.Common.expand_literal(zipper)
+    end
+
+    test "inserts value into empty map at nested position" do
+      {:ok, zipper} =
+        "%{}"
+        |> Code.Common.parse_to_zipper!()
+        |> Code.Map.put_in_map([:alpha, :beta, :gamma], :abc, fn _valuex -> flunk() end)
+
+      assert {:ok, %{alpha: %{beta: %{gamma: :abc}}}} == Code.Common.expand_literal(zipper)
+    end
+
+    test "inserts value into map if not present" do
+      {:ok, zipper} =
+        "%{hello: :world}"
+        |> Code.Common.parse_to_zipper!()
+        |> Code.Map.put_in_map([:foo], :bar, fn _valuex -> flunk() end)
+
+      assert {:ok, %{hello: :world, foo: :bar}} ==
+               Code.Common.expand_literal(zipper)
+    end
+
+    test "inserts nested value into map if not present" do
+      {:ok, zipper} =
+        "%{hello: :world}"
+        |> Code.Common.parse_to_zipper!()
+        |> Code.Map.put_in_map([:foo, :bar], :baz, fn _valuex -> flunk() end)
+
+      assert {:ok, %{hello: :world, foo: %{bar: :baz}}} ==
+               Code.Common.expand_literal(zipper)
+    end
+
+    test "inserts value into map if not present at nested position" do
+      {:ok, zipper} =
+        "%{alpha: %{beta: %{delta: :abd}}}"
+        |> Code.Common.parse_to_zipper!()
+        |> Code.Map.put_in_map([:alpha, :beta, :gamma], :abc, fn _valuex -> flunk() end)
+
+      assert {:ok, %{alpha: %{beta: %{delta: :abd, gamma: :abc}}}} ==
+               Code.Common.expand_literal(zipper)
+    end
+  end
+end