From 8ba9eeb3b4d3108cf627610fd697fc095bcf2187 Mon Sep 17 00:00:00 2001
From: Zach Daniel <zach@zachdaniel.dev>
Date: Sat, 31 Aug 2024 15:32:54 -0400
Subject: [PATCH] fix: properly clean calculation dependency values, and don't
 unload them with `lazy?`

---
 lib/ash/actions/read/calculations.ex | 67 +++++++++++++++++++---------
 lib/ash/actions/read/read.ex         |  8 +++-
 2 files changed, 53 insertions(+), 22 deletions(-)

diff --git a/lib/ash/actions/read/calculations.ex b/lib/ash/actions/read/calculations.ex
index 3d81bb35..436111a0 100644
--- a/lib/ash/actions/read/calculations.ex
+++ b/lib/ash/actions/read/calculations.ex
@@ -508,10 +508,21 @@ defmodule Ash.Actions.Read.Calculations do
         if to_load do
           Enum.map(records, fn record ->
             case record do
-              %Ash.NotLoaded{} -> record
-              %Ash.ForbiddenField{} -> record
-              nil -> record
-              record -> Map.put(record, to_load, Map.get(record.calculations, source))
+              %Ash.NotLoaded{} ->
+                record
+
+              %Ash.ForbiddenField{} ->
+                record
+
+              nil ->
+                record
+
+              record ->
+                {new_val, calculations} = Map.pop(record.calculations, source)
+
+                record
+                |> Map.put(to_load, new_val)
+                |> Map.put(:calculations, calculations)
             end
           end)
         else
@@ -527,11 +538,11 @@ defmodule Ash.Actions.Read.Calculations do
                 record
 
               record ->
-                Map.update!(
-                  record,
-                  :calculations,
-                  &Map.put(&1, to_name, Map.get(record.calculations, source))
-                )
+                {new_val, calculations} = Map.pop(record.calculations, source)
+
+                record
+                |> Map.update!(:calculations, &Map.put(&1, to_name, new_val))
+                |> Map.put(:calculations, calculations)
             end
           end)
         end
@@ -540,10 +551,21 @@ defmodule Ash.Actions.Read.Calculations do
         if to_load do
           Enum.map(records, fn record ->
             case record do
-              %Ash.NotLoaded{} -> record
-              %Ash.ForbiddenField{} -> record
-              nil -> record
-              record -> Map.put(record, to_load, Map.get(record.aggregates, source))
+              %Ash.NotLoaded{} ->
+                record
+
+              %Ash.ForbiddenField{} ->
+                record
+
+              nil ->
+                record
+
+              record ->
+                {new_val, aggregates} = Map.pop(record.aggregates, source)
+
+                record
+                |> Map.put(to_load, new_val)
+                |> Map.put(:aggregates, aggregates)
             end
           end)
         else
@@ -559,18 +581,23 @@ defmodule Ash.Actions.Read.Calculations do
                 record
 
               record ->
-                Map.update!(
-                  record,
-                  :aggregates,
-                  &Map.put(&1, to_name, Map.get(record.calculations, source))
-                )
+                {new_val, aggregates} = Map.pop(record.aggregates, source)
+
+                record
+                |> Map.update!(:aggregates, &Map.put(&1, to_name, new_val))
+                |> Map.put(:aggregates, aggregates)
             end
           end)
         end
 
       {{path, {:rel, rel_name}, _, _}, source}, records ->
-        update_at_path(records, path, fn record ->
-          Map.put(record, rel_name, Map.get(record.calculations, source))
+        records
+        |> update_at_path(path, fn record ->
+          {new_val, calculations} = Map.pop(record.calculations, source)
+
+          record
+          |> Map.put(rel_name, new_val)
+          |> Map.put(:calculations, calculations)
         end)
     end)
   end
diff --git a/lib/ash/actions/read/read.ex b/lib/ash/actions/read/read.ex
index 30154040..727a0276 100644
--- a/lib/ash/actions/read/read.ex
+++ b/lib/ash/actions/read/read.ex
@@ -1279,8 +1279,12 @@ defmodule Ash.Actions.Read do
       |> Map.update!(:calculations, fn calculations ->
         loaded_keys =
           calculations
-          |> Enum.filter(fn {_key, calc} ->
-            Ash.Resource.loaded?(initial_data, calc)
+          |> Enum.filter(fn
+            {{:__calc_dep__, _}, _} ->
+              false
+
+            {_key, calc} ->
+              Ash.Resource.loaded?(initial_data, calc)
           end)
           |> Enum.map(&elem(&1, 0))