From 553de7477f9df447350dad7ba982d13d371ff549 Mon Sep 17 00:00:00 2001
From: doorgan <dorgandash@gmail.com>
Date: Fri, 30 Jan 2026 22:22:27 -0300
Subject: [PATCH 1/3] fix: ambiguous map update syntax
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Code like `%{a do :ok end | b c, d => e}` is ambiguous because it can be
interpreted as:

`%{(... | b(c, d)) => e}` - what Spitfires parsed
`%{... | b(c, d) => e}` - what Elixir chooses

Elixir is aware of this ambiguity and warns about it:
```
warning: missing parentheses on expression following operator "|", you must add parentheses to avoid ambiguities
```

Since this is ambiguous syntax, Spitfire it not entirely wrong here, but
we should follow what Elixir produces anyways.

The only way I got this to trigger(from examples in the property tests
in Lukasz prs) is with the lhs of `|` being a do/end block, I'm not 100%
sure why, but withut it the Elixir parser straight out returns a syntax error.

We were allowing `=>` to be consumed inside the no-parens call arguments
and the map entry parser never got a chance to see it.
The fix marks map entry contexts (map/struct literals and map updates) with
`map_context` and blocks `=>` from being consumed inside nested expressions.
That keeps `=>` at the entry level and allows the existing ambiguity handling
in `parse_pipe_op/2` to match Elixir’s choice.
---
 lib/spitfire.ex         | 145 +++++++++++++++++++++++++++++++++++++---
 lib/spitfire/context.ex |  73 ++++++++++++++++++++
 test/spitfire_test.exs  |   6 ++
 3 files changed, 214 insertions(+), 10 deletions(-)
 create mode 100644 lib/spitfire/context.ex

diff --git a/lib/spitfire.ex b/lib/spitfire.ex
index a644a80..fd5af95 100644
--- a/lib/spitfire.ex
+++ b/lib/spitfire.ex
@@ -2,6 +2,8 @@ defmodule Spitfire do
   @moduledoc """
   Spitfire parser
   """
+
+  import Spitfire.Context
   import Spitfire.Tracer
   import Spitfire.While
   import Spitfire.While2
@@ -248,6 +250,7 @@ defmodule Spitfire do
 
   @terminals MapSet.new([:eol, :eof, :"}", :")", :"]", :">>"])
   @terminals_with_comma MapSet.put(@terminals, :",")
+
   defp parse_expression(parser, assoc \\ @lowest, is_list \\ false, is_map \\ false, is_top \\ false, is_stab \\ false)
 
   defp parse_expression(parser, {associativity, precedence}, is_list, is_map, is_top, is_stab) do
@@ -314,9 +317,13 @@ defmodule Spitfire do
           {parser, is_valid} = validate_peek(parser, current_token_type(parser))
 
           if is_valid do
-            while (is_nil(Map.get(parser, :stab_state)) and not MapSet.member?(terminals, peek_token(parser))) &&
-                    (current_token(parser) != :do and peek_token(parser) != :eol) &&
-                    calc_prec(parser, associativity, precedence) <- {left, parser} do
+            while is_nil(Map.get(parser, :stab_state)) and
+                    not MapSet.member?(terminals, peek_token(parser)) and
+                    current_token(parser) != :do and
+                    peek_token(parser) != :eol and
+                    calc_prec(parser, associativity, precedence) and
+                    not (block_assoc_op?(parser, is_map) and peek_token_type(parser) == :assoc_op) <-
+                    {left, parser} do
               parser = consume_fuel(parser)
               peek_token_type = peek_token_type(parser)
 
@@ -728,6 +735,7 @@ defmodule Spitfire do
       assoc_meta = current_meta(parser)
       parser = parser |> next_token() |> eat_eoe()
       {value, parser} = parse_expression(parser, @assoc_op, false, false, false)
+      parser = Map.put(parser, :last_assoc_meta, assoc_meta)
 
       key =
         case key do
@@ -742,7 +750,7 @@ defmodule Spitfire do
     end
   end
 
-  defp(parse_comma_list(parser, precedence \\ @list_comma, is_list \\ false, is_map \\ false))
+  defp parse_comma_list(parser, precedence \\ @list_comma, is_list \\ false, is_map \\ false)
 
   defp parse_comma_list(parser, precedence, is_list, is_map) do
     trace "parse_comma_list", trace_meta(parser) do
@@ -1054,14 +1062,116 @@ defmodule Spitfire do
 
       parser = eat_eoe(parser)
 
-      {pairs, parser} = parse_comma_list(parser, @list_comma, false, true)
+      {pairs, assoc_meta, parser} = parse_map_update_pairs(parser)
+      base_meta = newlines ++ meta
 
-      ast = {token, newlines ++ meta, [lhs, pairs]}
+      case map_update_rewrite(pairs, assoc_meta, newlines, meta) do
+        {:ok, pipe_meta, key, value} ->
+          pipe_ast = {token, pipe_meta, [lhs, key]}
+          {{pipe_ast, value}, parser}
 
-      {ast, parser}
+        :error ->
+          ast = {token, base_meta, [lhs, pairs]}
+          {ast, parser}
+      end
+    end
+  end
+
+  defp parse_map_update_pairs(parser) do
+    {{pairs, capture}, parser} =
+      with_state(parser, %{map_context: true, last_assoc_meta: nil}, capture: [:last_assoc_meta]) do
+        parse_comma_list(parser, @list_comma, false, true)
+      end
+
+    assoc_meta = Map.get(capture, :last_assoc_meta)
+    {pairs, assoc_meta, parser}
+  end
+
+  defp map_update_rewrite(pairs, assoc_meta, newlines, meta) do
+    # Code like %{a do :ok end | b c, d => e} is ambiguous:
+    # It can be interpreted as either
+    # %{(... | b(c, d)) => e}
+    # or
+    # %{... | b(c, d) => e}
+    #
+    # Elixir recognizes this and decides to interpret it as the latter
+    # and emit a warning.
+    #
+    # This section rewrites the parse attempt to match Elixir's behavior.
+    case pairs do
+      [{{call, call_meta, [_, _ | _] = args}, value}] ->
+        if map_update_allowed?(call_meta) do
+          case Keyword.pop(call_meta, :assoc) do
+            {nil, _call_meta} ->
+              :error
+
+            {assoc_meta, call_meta} ->
+              key = {call, call_meta, args}
+              pipe_meta = map_update_meta(newlines, meta, assoc_meta)
+              {:ok, pipe_meta, key, value}
+          end
+        else
+          :error
+        end
+
+      [{call, call_meta, [_, _ | _] = args}] ->
+        if map_update_allowed?(call_meta) do
+          case extract_map_update_arg(args, assoc_meta) do
+            {:ok, assoc_meta, assoc_key, assoc_value, rest_args} ->
+              key = {call, call_meta, rest_args ++ [assoc_key]}
+              pipe_meta = map_update_meta(newlines, meta, assoc_meta)
+              {:ok, pipe_meta, key, assoc_value}
+
+            :error ->
+              :error
+          end
+        else
+          :error
+        end
+
+      _ ->
+        :error
+    end
+  end
+
+  defp map_update_allowed?(call_meta) do
+    not Keyword.has_key?(call_meta, :closing) and not Keyword.has_key?(call_meta, :parens)
+  end
+
+  defp map_update_meta(newlines, meta, assoc_meta) do
+    case assoc_meta do
+      nil -> newlines ++ meta
+      _ -> newlines ++ [{:assoc, assoc_meta} | meta]
+    end
+  end
+
+  defp extract_map_update_arg(args, assoc_meta) do
+    {last, rest} = List.pop_at(args, -1)
+
+    case last do
+      {assoc_key, assoc_value} ->
+        {assoc_meta, assoc_key} = extract_assoc_meta(assoc_key, assoc_meta)
+        {:ok, assoc_meta, assoc_key, assoc_value, rest}
+
+      _ ->
+        :error
+    end
+  end
+
+  defp extract_assoc_meta({name, meta, args}, fallback_meta) when is_list(meta) do
+    case Keyword.pop(meta, :assoc) do
+      {nil, _meta} ->
+        {fallback_meta, {name, meta, args}}
+
+      {assoc_meta, meta} ->
+        {assoc_meta, {name, meta, args}}
     end
   end
 
+  defp extract_assoc_meta(key, fallback_meta) do
+    {fallback_meta, key}
+  end
+
   defp parse_access_expression(parser, lhs) do
     trace "parse_access_expression", trace_meta(parser) do
       meta = current_meta(parser)
@@ -1761,7 +1871,10 @@ defmodule Spitfire do
           {{:%{}, meta, []}, parser}
 
         true ->
-          {pairs, parser} = parse_comma_list(parser, @list_comma, false, true)
+          {pairs, parser} =
+            with_state(parser, %{map_context: true}) do
+              parse_comma_list(parser, @list_comma, false, true)
+            end
 
           parser = eat_eol_at(parser, 1)
 
@@ -1899,7 +2012,11 @@ defmodule Spitfire do
             parser = Map.put(parser, :nesting, old_nesting)
             {ast, parser}
           else
-            {pairs, parser} = parse_comma_list(parser, @list_comma, false, true)
+            {pairs, parser} =
+              with_state(parser, %{map_context: true}) do
+                parse_comma_list(parser, @list_comma, false, true)
+              end
+
             parser = eat_eol_at(parser, 1)
 
             parser =
@@ -1922,7 +2039,11 @@ defmodule Spitfire do
           old_nesting = parser.nesting
           parser = Map.put(parser, :nesting, 0)
 
-          {pairs, parser} = parse_comma_list(parser, @list_comma, false, true)
+          {pairs, parser} =
+            with_state(parser, %{map_context: true}) do
+              parse_comma_list(parser, @list_comma, false, true)
+            end
+
           parser = eat_eol_at(parser, 1)
 
           {parser, closing_meta} =
@@ -2890,6 +3011,10 @@ defmodule Spitfire do
     Map.get(@precedences, peek_token_type(parser), @lowest)
   end
 
+  defp block_assoc_op?(parser, is_map) do
+    not is_map and parser.nesting > 0 and Map.get(parser, :map_context, false)
+  end
+
   defp pop_nesting(%{nesting: nesting} = parser) do
     %{parser | nesting: nesting - 1}
   end
diff --git a/lib/spitfire/context.ex b/lib/spitfire/context.ex
new file mode 100644
index 0000000..9f93e52
--- /dev/null
+++ b/lib/spitfire/context.ex
@@ -0,0 +1,73 @@
+defmodule Spitfire.Context do
+  @moduledoc false
+
+  @doc """
+  Temporarily applies parser state updates and restores them after the block.
+
+  Use the `:capture` option to return a map of selected keys as they were
+  at the end of the block, before state is restored.
+  """
+  defmacro with_state(parser, updates, opts \\ [], do: block) do
+    quote location: :keep, generated: true do
+      parser = unquote(parser)
+      updates = unquote(updates)
+      opts = unquote(opts)
+
+      updates =
+        if is_map(updates) do
+          updates
+        else
+          Map.new(updates)
+        end
+
+      old =
+        Enum.reduce(updates, %{}, fn {key, _value}, acc ->
+          Map.put(acc, key, {Map.has_key?(parser, key), Map.get(parser, key)})
+        end)
+
+      parser = Map.merge(parser, updates)
+      result = unquote(block)
+      capture_keys = Keyword.get(opts, :capture, [])
+
+      restore = fn parser ->
+        Enum.reduce(old, parser, fn {key, {had_key, value}}, acc ->
+          if had_key do
+            Map.put(acc, key, value)
+          else
+            Map.delete(acc, key)
+          end
+        end)
+      end
+
+      case result do
+        {value, parser} ->
+          captured =
+            Enum.reduce(capture_keys, %{}, fn key, acc ->
+              Map.put(acc, key, Map.get(parser, key))
+            end)
+
+          parser = restore.(parser)
+
+          if capture_keys == [] do
+            {value, parser}
+          else
+            {{value, captured}, parser}
+          end
+
+        parser when is_map(parser) ->
+          captured =
+            Enum.reduce(capture_keys, %{}, fn key, acc ->
+              Map.put(acc, key, Map.get(parser, key))
+            end)
+
+          parser = restore.(parser)
+
+          if capture_keys == [] do
+            parser
+          else
+            {{parser, captured}, parser}
+          end
+      end
+    end
+  end
+end
diff --git a/test/spitfire_test.exs b/test/spitfire_test.exs
index 81250c4..2231c69 100644
--- a/test/spitfire_test.exs
+++ b/test/spitfire_test.exs
@@ -728,6 +728,12 @@ defmodule SpitfireTest do
       end
     end
 
+    test "parses ambiguous map update" do
+      code = ~S'%{a do :ok end | b c, d => e}'
+
+      assert Spitfire.parse(code) == s2q(code)
+    end
+
     test "parses operators" do
       codes = [
         ~s'''

From d008bf028a693d70c575524487057f7b40d18a9e Mon Sep 17 00:00:00 2001
From: doorgan <dorgandash@gmail.com>
Date: Sat, 31 Jan 2026 14:41:22 -0300
Subject: [PATCH 2/3] chore: clarify the matching in map_update_rewrite/4

---
 lib/spitfire.ex | 35 ++++++++++++++++++++---------------
 1 file changed, 20 insertions(+), 15 deletions(-)

diff --git a/lib/spitfire.ex b/lib/spitfire.ex
index fd5af95..1d8c8f9 100644
--- a/lib/spitfire.ex
+++ b/lib/spitfire.ex
@@ -1062,16 +1062,16 @@ defmodule Spitfire do
 
       parser = eat_eoe(parser)
 
-      {pairs, assoc_meta, parser} = parse_map_update_pairs(parser)
+      {entries, assoc_meta, parser} = parse_map_update_pairs(parser)
       base_meta = newlines ++ meta
 
-      case map_update_rewrite(pairs, assoc_meta, newlines, meta) do
+      case map_update_rewrite(entries, assoc_meta, newlines, meta) do
         {:ok, pipe_meta, key, value} ->
           pipe_ast = {token, pipe_meta, [lhs, key]}
           {{pipe_ast, value}, parser}
 
         :error ->
-          ast = {token, base_meta, [lhs, pairs]}
+          ast = {token, base_meta, [lhs, entries]}
           {ast, parser}
       end
     end
@@ -1087,18 +1087,22 @@ defmodule Spitfire do
     {pairs, assoc_meta, parser}
   end
 
-  defp map_update_rewrite(pairs, assoc_meta, newlines, meta) do
-    # Code like %{a do :ok end | b c, d => e} is ambiguous:
-    # It can be interpreted as either
-    # %{(... | b(c, d)) => e}
-    # or
-    # %{... | b(c, d) => e}
-    #
-    # Elixir recognizes this and decides to interpret it as the latter
-    # and emit a warning.
-    #
-    # This section rewrites the parse attempt to match Elixir's behavior.
-    case pairs do
+  # Code like %{a do :ok end | b c, d => e} is ambiguous:
+  # It can be interpreted as either
+  # %{(... | b(c, d)) => e}
+  # or
+  # %{... | b(c, d) => e}
+  #
+  # The ambiguity only exists for calls with 2+ args (a comma that could
+  # belong to the call or to the map-update list).
+  #
+  # Elixir recognizes this and decides to interpret it as the latter
+  # and emit a warning.
+  #
+  # This rewrites the parse attempt to match Elixir's behavior.
+  defp map_update_rewrite(entries, assoc_meta, newlines, meta) do
+    case entries do
+      # %{... | (b(c, d)) => e}
       [{{call, call_meta, [_, _ | _] = args}, value}] ->
         if map_update_allowed?(call_meta) do
           case Keyword.pop(call_meta, :assoc) do
@@ -1114,6 +1118,7 @@ defmodule Spitfire do
           :error
         end
 
+      # %{... | b(c, d => e)}
       [{call, call_meta, [_, _ | _] = args}] ->
         if map_update_allowed?(call_meta) do
           case extract_map_update_arg(args, assoc_meta) do

From 725126e30e71e75eb462bac66aed1cef4c1d7786 Mon Sep 17 00:00:00 2001
From: doorgan <dorgandash@gmail.com>
Date: Sat, 31 Jan 2026 15:20:42 -0300
Subject: [PATCH 3/3] fix: parser scoping in with_context, and add tests

---
 lib/spitfire.ex                | 16 +++---
 lib/spitfire/context.ex        | 95 +++++++++++++++-------------------
 test/spitfire/context_test.exs | 59 +++++++++++++++++++++
 3 files changed, 109 insertions(+), 61 deletions(-)
 create mode 100644 test/spitfire/context_test.exs

diff --git a/lib/spitfire.ex b/lib/spitfire.ex
index 1d8c8f9..7822c79 100644
--- a/lib/spitfire.ex
+++ b/lib/spitfire.ex
@@ -1079,9 +1079,9 @@ defmodule Spitfire do
 
   defp parse_map_update_pairs(parser) do
     {{pairs, capture}, parser} =
-      with_state(parser, %{map_context: true, last_assoc_meta: nil}, capture: [:last_assoc_meta]) do
+      with_state(parser, %{map_context: true, last_assoc_meta: nil}, [capture: [:last_assoc_meta]], fn parser ->
         parse_comma_list(parser, @list_comma, false, true)
-      end
+      end)
 
     assoc_meta = Map.get(capture, :last_assoc_meta)
     {pairs, assoc_meta, parser}
@@ -1877,9 +1877,9 @@ defmodule Spitfire do
 
         true ->
           {pairs, parser} =
-            with_state(parser, %{map_context: true}) do
+            with_state(parser, %{map_context: true}, fn parser ->
               parse_comma_list(parser, @list_comma, false, true)
-            end
+            end)
 
           parser = eat_eol_at(parser, 1)
 
@@ -2018,9 +2018,9 @@ defmodule Spitfire do
             {ast, parser}
           else
             {pairs, parser} =
-              with_state(parser, %{map_context: true}) do
+              with_state(parser, %{map_context: true}, fn parser ->
                 parse_comma_list(parser, @list_comma, false, true)
-              end
+              end)
 
             parser = eat_eol_at(parser, 1)
 
@@ -2045,9 +2045,9 @@ defmodule Spitfire do
           parser = Map.put(parser, :nesting, 0)
 
           {pairs, parser} =
-            with_state(parser, %{map_context: true}) do
+            with_state(parser, %{map_context: true}, fn parser ->
               parse_comma_list(parser, @list_comma, false, true)
-            end
+            end)
 
           parser = eat_eol_at(parser, 1)
 
diff --git a/lib/spitfire/context.ex b/lib/spitfire/context.ex
index 9f93e52..d32019d 100644
--- a/lib/spitfire/context.ex
+++ b/lib/spitfire/context.ex
@@ -7,67 +7,56 @@ defmodule Spitfire.Context do
   Use the `:capture` option to return a map of selected keys as they were
   at the end of the block, before state is restored.
   """
-  defmacro with_state(parser, updates, opts \\ [], do: block) do
-    quote location: :keep, generated: true do
-      parser = unquote(parser)
-      updates = unquote(updates)
-      opts = unquote(opts)
-
-      updates =
-        if is_map(updates) do
-          updates
-        else
-          Map.new(updates)
-        end
-
-      old =
-        Enum.reduce(updates, %{}, fn {key, _value}, acc ->
-          Map.put(acc, key, {Map.has_key?(parser, key), Map.get(parser, key)})
-        end)
+  def with_state(parser, updates, fun) when is_function(fun, 1) do
+    with_state(parser, updates, [], fun)
+  end
 
-      parser = Map.merge(parser, updates)
-      result = unquote(block)
-      capture_keys = Keyword.get(opts, :capture, [])
+  def with_state(parser, updates, opts, fun) when is_function(fun, 1) do
+    old =
+      Enum.reduce(updates, %{}, fn {key, _value}, acc ->
+        Map.put(acc, key, {Map.has_key?(parser, key), Map.get(parser, key)})
+      end)
 
-      restore = fn parser ->
-        Enum.reduce(old, parser, fn {key, {had_key, value}}, acc ->
-          if had_key do
-            Map.put(acc, key, value)
-          else
-            Map.delete(acc, key)
-          end
-        end)
-      end
+    parser = Map.merge(parser, updates)
+    result = fun.(parser)
+    capture_keys = Keyword.get(opts, :capture, [])
 
-      case result do
-        {value, parser} ->
-          captured =
-            Enum.reduce(capture_keys, %{}, fn key, acc ->
-              Map.put(acc, key, Map.get(parser, key))
-            end)
+    restore = fn parser ->
+      Enum.reduce(old, parser, fn {key, {had_key, value}}, acc ->
+        if had_key do
+          Map.put(acc, key, value)
+        else
+          Map.delete(acc, key)
+        end
+      end)
+    end
 
-          parser = restore.(parser)
+    capture = fn parser ->
+      Enum.reduce(capture_keys, %{}, fn key, acc ->
+        Map.put(acc, key, Map.get(parser, key))
+      end)
+    end
 
-          if capture_keys == [] do
-            {value, parser}
-          else
-            {{value, captured}, parser}
-          end
+    case result do
+      {value, parser} ->
+        captured = capture.(parser)
+        restored = restore.(parser)
 
-        parser when is_map(parser) ->
-          captured =
-            Enum.reduce(capture_keys, %{}, fn key, acc ->
-              Map.put(acc, key, Map.get(parser, key))
-            end)
+        if capture_keys == [] do
+          {value, restored}
+        else
+          {{value, captured}, restored}
+        end
 
-          parser = restore.(parser)
+      parser when is_map(parser) ->
+        captured = capture.(parser)
+        restored = restore.(parser)
 
-          if capture_keys == [] do
-            parser
-          else
-            {{parser, captured}, parser}
-          end
-      end
+        if capture_keys == [] do
+          restored
+        else
+          {{parser, captured}, restored}
+        end
     end
   end
 end
diff --git a/test/spitfire/context_test.exs b/test/spitfire/context_test.exs
new file mode 100644
index 0000000..13d51ff
--- /dev/null
+++ b/test/spitfire/context_test.exs
@@ -0,0 +1,59 @@
+defmodule Spitfire.ContextTest do
+  use ExUnit.Case, async: true
+
+  import Spitfire.Context, only: [with_state: 3, with_state: 4]
+
+  test "restores updated keys after block" do
+    parser = %{mode: :original, keep: :ok}
+
+    restored =
+      with_state(parser, %{mode: :temp, new_key: :temp}, fn parser ->
+        assert parser[:mode] == :temp
+        assert parser[:new_key] == :temp
+        Map.merge(parser, %{mode: :inner, new_key: :inner})
+      end)
+
+    assert restored == %{mode: :original, keep: :ok}
+  end
+
+  test "captures state before restore" do
+    parser = %{mode: :original}
+
+    {{value, captured}, restored} =
+      with_state(parser, %{mode: :temp, flag: false}, [capture: [:mode, :flag]], fn parser ->
+        {:ok, Map.merge(parser, %{mode: :final, flag: true})}
+      end)
+
+    assert value == :ok
+    assert captured == %{mode: :final, flag: true}
+    assert restored == %{mode: :original}
+  end
+
+  test "nested calls restore to outer then original" do
+    parser = %{mode: :base}
+
+    {value, restored} =
+      with_state(parser, %{mode: :outer, outer: true}, fn parser ->
+        assert parser[:mode] == :outer
+        assert parser[:outer]
+
+        {inner_value, parser} =
+          with_state(parser, %{mode: :inner, inner: true}, fn parser ->
+            assert parser[:mode] == :inner
+            assert parser[:outer]
+            assert parser[:inner]
+            {:inner, Map.put(parser, :mode, :inner_set)}
+          end)
+
+        assert inner_value == :inner
+        assert parser[:mode] == :outer
+        assert parser[:outer]
+        refute Map.has_key?(parser, :inner)
+
+        {:outer, parser}
+      end)
+
+    assert value == :outer
+    assert restored == %{mode: :base}
+  end
+end