fix: properly key nested calculations and add additional tests

lib/ash/actions/read/calculations.ex

@@ -442,6 +442,11 @@ defmodule Ash.Actions.Read.Calculations do
   def rewrite(_rewrites, nil), do: nil
   def rewrite(_rewrites, []), do: []
 
+  def rewrite(rewrites, %struct{results: results} = page)
+      when struct in [Ash.Page.Keyset, Ash.Page.Offset] do
+    %{page | results: rewrite(results, rewrites)}
+  end
+
   def rewrite(rewrites, record) when not is_list(record) do
     rewrites
     |> rewrite([record])
@@ -552,6 +557,19 @@ defmodule Ash.Actions.Read.Calculations do
     end)
   end
 
+  defp rewrite_at_path(
+         records,
+         {{[{:rel, name} | rest], data, calc_name, calc_load}, source}
+       ) do
+    new_rewrites = [
+      {{rest, data, calc_name, calc_load}, source}
+    ]
+
+    Enum.map(records, fn record ->
+      Map.update!(record, name, &rewrite(new_rewrites, &1))
+    end)
+  end
+
   defp rewrite_at_path(
          records,
          {{[{:calc, type, constraints, name, load} | rest], data, calc_name, calc_load}, source}
@@ -643,13 +661,21 @@ defmodule Ash.Actions.Read.Calculations do
         if calculation.module.strict_loads? do
           []
         else
-          relationship.destination
-          |> Ash.Query.new()
-          |> get_all_rewrites(calculation, path ++ [name])
+          query = Ash.Query.new(relationship.destination)
+
+          query
+          |> get_all_rewrites(calculation, path)
+          |> Enum.map(fn {{path, data, calc_name, calc_load}, source} ->
+            {{path ++ [{:rel, name}], data, calc_name, calc_load}, source}
+          end)
         end
 
       {name, query} ->
-        get_all_rewrites(query, calculation, path ++ [name])
+        query
+        |> get_all_rewrites(calculation, path)
+        |> Enum.map(fn {{path, data, calc_name, calc_load}, source} ->
+          {{path ++ [{:rel, name}], data, calc_name, calc_load}, source}
+        end)
     end)
   end
 
@@ -687,6 +713,10 @@ defmodule Ash.Actions.Read.Calculations do
     end)
   end
 
+  # TODO: This currently must assume that all relationship loads are different if
+  # authorize?: true, because the policies have not yet been applied.
+  #
+
   def split_and_load_calculations(
         domain,
         ash_query,
@@ -1785,7 +1815,7 @@ defmodule Ash.Actions.Read.Calculations do
       strict_loads?,
       relationship_path,
       can_expression_calculation?,
-      relationship_path,
+      [],
       initial_data,
       reuse_values?,
       authorize?

test/calculations/calculation_test.exs

@@ -135,6 +135,34 @@ defmodule Ash.Test.CalculationTest do
     end
   end
 
+  defmodule BestFriendsNameDifferent do
+    use Ash.Resource.Calculation
+
+    def load(_query, _opts, _) do
+      [best_friend: [full_name: %{separator: ":"}]]
+    end
+
+    def calculate(records, _opts, _) do
+      Enum.map(records, fn record ->
+        record.best_friend && record.best_friend.full_name
+      end)
+    end
+  end
+
+  defmodule BestFriendsActive do
+    use Ash.Resource.Calculation
+
+    def load(_query, _opts, _) do
+      [best_friend: :active]
+    end
+
+    def calculate(records, _opts, _) do
+      Enum.map(records, fn record ->
+        record.best_friend && record.best_friend.active
+      end)
+    end
+  end
+
   defmodule BestFriendsFirstNameThatFails do
     use Ash.Resource.Calculation
 
@@ -564,6 +592,14 @@ defmodule Ash.Test.CalculationTest do
         public?(true)
       end
 
+      calculate :best_friends_name_different, :string, BestFriendsNameDifferent do
+        public?(true)
+      end
+
+      calculate :best_friends_active, :boolean, BestFriendsActive do
+        public?(true)
+      end
+
       calculate :names_of_best_friends_of_me, :string, NamesOfBestFriendsOfMe do
         public?(true)
         argument(:only_special, :boolean, default: false)
@@ -1030,13 +1066,61 @@ defmodule Ash.Test.CalculationTest do
     best_friends_names =
       User
       |> Ash.Query.load([:best_friends_name])
-      |> Ash.read!()
+      |> Ash.read!(authorize?: false)
       |> Enum.map(& &1.best_friends_name)
       |> Enum.sort()
 
     assert best_friends_names == [nil, "zach daniel"]
   end
 
+  test "nested calculations are loaded differently if necessary" do
+    res =
+      User
+      |> Ash.Query.load([:best_friends_name, :best_friends_name_different])
+      |> Ash.read!(authorize?: false)
+
+    best_friends_names =
+      res
+      |> Enum.map(& &1.best_friends_name)
+      |> Enum.sort()
+
+    best_friends_names_different =
+      res
+      |> Enum.map(& &1.best_friends_name_different)
+      |> Enum.sort()
+
+    assert best_friends_names == [nil, "zach daniel"]
+    assert best_friends_names_different == [nil, "zach:daniel"]
+  end
+
+  test "nested calculations inside nested relationships are loaded differently if necessary" do
+    res =
+      User
+      |> Ash.Query.load(
+        best_friends_of_me: [
+          :id,
+          :best_friends_name,
+          :best_friends_name_different,
+          best_friend: [:id]
+        ]
+      )
+      |> Ash.read!(authorize?: false)
+      |> Enum.flat_map(& &1.best_friends_of_me)
+
+    best_friends_names =
+      res
+      |> Enum.map(& &1.best_friends_name)
+      |> Enum.sort()
+
+    best_friends_names_different =
+      res
+      |> Enum.map(& &1.best_friends_name_different)
+      |> Enum.sort()
+
+    assert best_friends_names == ["zach daniel"]
+    assert best_friends_names_different == ["zach:daniel"]
+  end
+
   test "calculations must specify required fields by default" do
     assert_raise RuntimeError,
                  ~r/Invalid return from calculation, expected a value, got \`%Ash.NotLoaded{}\`/,