Move the expansion code into ArrayMap

Author: freshtonic

lib/graphql/execution/executor.ex

@@ -46,33 +46,17 @@ defmodule GraphQL.Execution.Executor do
     case operation.operation do
       :query        ->
         {context, result} = execute_fields(context, type, root_value, fields)
-        {:ok, expand_array_maps(result), context.errors}
+        {:ok, ArrayMap.expand_result(result), context.errors}
       :mutation     ->
         {context, result} = execute_fields_serially(context, type, root_value, fields)
-        {:ok, expand_array_maps(result), context.errors}
+        {:ok, ArrayMap.expand_result(result), context.errors}
       :subscription ->
         {:error, "Subscriptions not currently supported"}
       _             ->
         {:error, "Can only execute queries, mutations and subscriptions"}
     end
   end
 
-  defp expand_array_maps(result) when is_list(result) do
-    Enum.map(result, &expand_array_maps/1)
-  end
-  defp expand_array_maps(%ArrayMap{} = result) do
-    Enum.reduce(Enum.sort(Map.keys(result.map)), [], fn(index, acc) ->
-      [expand_array_maps(Map.get(result.map, index))] ++ acc
-    end) |> Enum.reverse
-  end
-  defp expand_array_maps(result) when is_map(result) do
-    Enum.reduce(result, %{}, fn({k, v}, acc) ->
-      Map.put(acc, expand_array_maps(k), expand_array_maps(v))
-    end)
-  end
-  defp expand_array_maps(result), do: result
-
-
   defp collect_selections(context, runtime_type, selection_set, field_fragment_map \\ %{fields: %{}, fragments: %{}}) do
     Enum.reduce selection_set[:selections], {context, field_fragment_map}, fn(selection, {context, field_fragment_map}) ->
       collect_selection(context, runtime_type, selection, field_fragment_map)

lib/graphql/util/array_map.ex

@@ -10,6 +10,13 @@ defmodule GraphQL.Util.ArrayMap do
 
   defstruct map: %{}
 
+  def new(map) do
+    if !Enum.all?(Map.keys(map), fn(key) -> is_integer(key) end) do
+      raise "all key must be integers!"
+    end
+    %__MODULE__{map: map}
+  end
+
   def put(array_map, index, value) when is_integer(index) do
     %__MODULE__{ map: Map.put(array_map.map, index, value) }
   end
@@ -29,11 +36,29 @@ defmodule GraphQL.Util.ArrayMap do
 
   def get(array_map, key, value) do
     map = Access.get(array_map.map, key, value)
-    %__MODULE__{map: value}
+    %__MODULE__{map: map}
   end
 
   def pop(array_map, key) do
     {value, map} = Access.get(array_map.map, key)
     {value, %__MODULE__{map: map}}
   end
+
+  # Converts an intermediate executor result that contains ArrayMaps into one
+  # where the array maps are converted into lists.
+  def expand_result(result) when is_list(result) do
+    Enum.map(result, &expand_result/1)
+  end
+  def expand_result(%__MODULE__{} = result) do
+    Enum.reduce(Enum.sort(Map.keys(result.map)), [], fn(index, acc) ->
+      [expand_result(Map.get(result.map, index))] ++ acc
+    end) |> Enum.reverse
+  end
+  def expand_result(result) when is_map(result) do
+    Enum.reduce(result, %{}, fn({k, v}, acc) ->
+      Map.put(acc, expand_result(k), expand_result(v))
+    end)
+  end
+  def expand_result(result), do: result
+
 end

test/graphql/util/array_map_test.exs

@@ -5,11 +5,20 @@ defmodule GraphQL.Util.ArrayMapTest do
   alias GraphQL.Util.ArrayMap
 
   test "implements Access" do
-    version_1 = %ArrayMap{map: %{0 => :zero, 1 => :one, 2 => :two}}
-    version_2 = %ArrayMap{map: %{0 => :zero, 1 => :ONE, 2 => :two}}
+    version_1 = ArrayMap.new(%{0 => :zero, 1 => :one, 2 => :two})
+    version_2 = ArrayMap.new(%{0 => :zero, 1 => :ONE, 2 => :two})
 
     version_1_updated = put_in(version_1, [1], :ONE)
 
     assert version_2 == version_1_updated
   end
+
+  test "Access works with mix of nested Maps and ArrayMaps" do
+    nested =   %{:foo => ArrayMap.new(%{0 => %{:baz => ArrayMap.new(%{0 => %{quux: 123}})}})}
+    expected = %{:foo => ArrayMap.new(%{0 => %{:baz => ArrayMap.new(%{0 => %{quux: 456}})}})}
+
+    updated = put_in(nested, [:foo, 0, :baz, 0, :quux], 456)
+
+    assert updated == expected
+  end
 end