fix: ambiguous map update syntax

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,121 @@ defmodule Spitfire do
 
       parser = eat_eoe(parser)
 
-      {pairs, parser} = parse_comma_list(parser, @list_comma, false, true)
+      {entries, assoc_meta, parser} = parse_map_update_pairs(parser)
+      base_meta = newlines ++ meta
 
-      ast = {token, newlines ++ meta, [lhs, pairs]}
+      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}
 
-      {ast, parser}
+        :error ->
+          ast = {token, base_meta, [lhs, entries]}
+          {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]], fn parser ->
+        parse_comma_list(parser, @list_comma, false, true)
+      end)
+
+    assoc_meta = Map.get(capture, :last_assoc_meta)
+    {pairs, assoc_meta, parser}
+  end
+
+  # 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
+            {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
+
+      # %{... | b(c, d => e)}
+      [{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 +1876,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}, fn parser ->
+              parse_comma_list(parser, @list_comma, false, true)
+            end)
 
           parser = eat_eol_at(parser, 1)
 
@@ -1899,7 +2017,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}, fn parser ->
+                parse_comma_list(parser, @list_comma, false, true)
+              end)
+
             parser = eat_eol_at(parser, 1)
 
             parser =
@@ -1922,7 +2044,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}, fn parser ->
+              parse_comma_list(parser, @list_comma, false, true)
+            end)
+
           parser = eat_eol_at(parser, 1)
 
           {parser, closing_meta} =
@@ -2890,6 +3016,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

lib/spitfire/context.ex

@@ -0,0 +1,62 @@
+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.
+  """
+  def with_state(parser, updates, fun) when is_function(fun, 1) do
+    with_state(parser, updates, [], fun)
+  end
+
+  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)
+
+    parser = Map.merge(parser, updates)
+    result = fun.(parser)
+    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
+
+    capture = fn parser ->
+      Enum.reduce(capture_keys, %{}, fn key, acc ->
+        Map.put(acc, key, Map.get(parser, key))
+      end)
+    end
+
+    case result do
+      {value, parser} ->
+        captured = capture.(parser)
+        restored = restore.(parser)
+
+        if capture_keys == [] do
+          {value, restored}
+        else
+          {{value, captured}, restored}
+        end
+
+      parser when is_map(parser) ->
+        captured = capture.(parser)
+        restored = restore.(parser)
+
+        if capture_keys == [] do
+          restored
+        else
+          {{parser, captured}, restored}
+        end
+    end
+  end
+end

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

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'''