From 2c754f90a64a0792c5ccdc5ce726f9417c2001d9 Mon Sep 17 00:00:00 2001
From: pinetops <tom@u2i.com>
Date: Sun, 1 Sep 2024 17:06:51 -0400
Subject: [PATCH] improvement: Cache always selected fields and use mapsets for
 building select list (#1428)

* Add attribute_names function on Info that returns MapSet

And use that when calculating select lists.
---
 lib/ash/query/query.ex                        | 55 +++++++++----------
 lib/ash/resource/info.ex                      | 10 ++++
 .../transformers/attributes_by_name.ex        | 10 +++-
 3 files changed, 43 insertions(+), 32 deletions(-)

diff --git a/lib/ash/query/query.ex b/lib/ash/query/query.ex
index 39e4ed67..c41eca39 100644
--- a/lib/ash/query/query.ex
+++ b/lib/ash/query/query.ex
@@ -845,42 +845,37 @@ defmodule Ash.Query do
   """
   def select(query, fields, opts \\ []) do
     query = new(query)
-    fields = List.wrap(fields)
+    existing_fields = Ash.Resource.Info.attribute_names(query.resource)
+    fields = MapSet.new(List.wrap(fields))
 
-    {fields, non_existent} =
-      Enum.split_with(fields, &Ash.Resource.Info.attribute(query.resource, &1))
+    valid_fields = MapSet.intersection(fields, existing_fields)
 
     query =
-      Enum.reduce(non_existent, query, fn field, query ->
-        Ash.Query.add_error(
-          query,
-          Ash.Error.Query.NoSuchAttribute.exception(resource: query.resource, attribute: field)
-        )
-      end)
+      if MapSet.size(valid_fields) != MapSet.size(fields) do
+        MapSet.difference(fields, existing_fields)
+        |> Enum.reduce(query, fn field, query ->
+          Ash.Query.add_error(
+            query,
+            Ash.Error.Query.NoSuchAttribute.exception(resource: query.resource, attribute: field)
+          )
+        end)
+      else
+        query
+      end
 
     always_select =
-      query.resource
-      |> Ash.Resource.Info.attributes()
-      |> Enum.filter(& &1.always_select?)
-      |> Enum.map(& &1.name)
+      valid_fields
+      |> MapSet.union(Ash.Resource.Info.always_selected_attribute_names(query.resource))
+      |> MapSet.union(MapSet.new(Ash.Resource.Info.primary_key(query.resource)))
 
-    if opts[:replace?] do
-      %{
-        query
-        | select:
-            Enum.uniq(fields ++ Ash.Resource.Info.primary_key(query.resource) ++ always_select)
-      }
-    else
-      %{
-        query
-        | select:
-            Enum.uniq(
-              fields ++
-                (query.select || []) ++
-                Ash.Resource.Info.primary_key(query.resource) ++ always_select
-            )
-      }
-    end
+    new_select =
+      if opts[:replace?] do
+        always_select
+      else
+        MapSet.union(MapSet.new(query.select || []), always_select)
+      end
+
+    %{query | select: MapSet.to_list(new_select)}
   end
 
   @doc """
diff --git a/lib/ash/resource/info.ex b/lib/ash/resource/info.ex
index be9b022d..3d1c179a 100644
--- a/lib/ash/resource/info.ex
+++ b/lib/ash/resource/info.ex
@@ -649,6 +649,11 @@ defmodule Ash.Resource.Info do
     Extension.get_entities(resource, [:attributes])
   end
 
+  @spec attribute_names(Spark.Dsl.t() | Ash.Resource.t()) :: MapSet.t()
+  def attribute_names(resource) do
+    Extension.get_persisted(resource, :attribute_names)
+  end
+
   @doc "Returns all attributes of a resource with lazy non-matching-defaults"
   @spec lazy_non_matching_default_attributes(
           Spark.Dsl.t() | Ash.Resource.t(),
@@ -769,6 +774,11 @@ defmodule Ash.Resource.Info do
     end
   end
 
+  @spec always_selected_attribute_names(Spark.Dsl.t() | Ash.Resource.t()) :: MapSet.t()
+  def always_selected_attribute_names(resource) do
+    Extension.get_persisted(resource, :always_selected_attribute_names)
+  end
+
   @doc "Returns all attributes, aggregates, calculations and relationships of a resource"
   @spec fields(
           Spark.Dsl.t() | Ash.Resource.t(),
diff --git a/lib/ash/resource/transformers/attributes_by_name.ex b/lib/ash/resource/transformers/attributes_by_name.ex
index 3d384c29..f5a00a90 100644
--- a/lib/ash/resource/transformers/attributes_by_name.ex
+++ b/lib/ash/resource/transformers/attributes_by_name.ex
@@ -20,7 +20,7 @@ defmodule Ash.Resource.Transformers.AttributesByName do
         |> Map.put(to_string(name), attr)
       end)
 
-    attribute_names = Enum.map(attributes, & &1.name)
+    attribute_names = Enum.map(attributes, & &1.name) |> MapSet.new()
 
     create_attributes_with_static_defaults =
       attributes
@@ -61,6 +61,11 @@ defmodule Ash.Resource.Transformers.AttributesByName do
           (is_function(attribute.update_default) or match?({_, _, _}, attribute.update_default))
       end)
 
+    always_selected_attribute_names =
+      Enum.filter(attributes, & &1.always_select?)
+      |> Enum.map(& &1.name)
+      |> MapSet.new()
+
     {:ok,
      persist(
        dsl_state,
@@ -74,7 +79,8 @@ defmodule Ash.Resource.Transformers.AttributesByName do
          update_attributes_with_static_defaults: update_attributes_with_static_defaults,
          update_attributes_with_non_matching_lazy_defaults:
            update_attributes_with_non_matching_lazy_defaults,
-         update_attributes_with_matching_defaults: update_attributes_with_matching_defaults
+         update_attributes_with_matching_defaults: update_attributes_with_matching_defaults,
+         always_selected_attribute_names: always_selected_attribute_names
        }
      )}
   end