From 9036780400c7afc2e52f1f9db812eb6b1d70adf5 Mon Sep 17 00:00:00 2001
From: Zach Daniel <zach@zachdaniel.dev>
Date: Wed, 1 Jan 2025 14:54:28 -0500
Subject: [PATCH] improvement: define `generate/1` callback for maps, structs,
 keywords

---
 lib/ash/generator/generator.ex |  5 +++++
 lib/ash/type/keyword.ex        |  8 ++++++++
 lib/ash/type/map.ex            | 22 ++++++++++++++++++++++
 lib/ash/type/struct.ex         | 13 +++++++++++++
 4 files changed, 48 insertions(+)

diff --git a/lib/ash/generator/generator.ex b/lib/ash/generator/generator.ex
index 13019d103..14939ed2e 100644
--- a/lib/ash/generator/generator.ex
+++ b/lib/ash/generator/generator.ex
@@ -45,6 +45,11 @@ defmodule Ash.Generator do
              do_changeset_or_query: 5}
 
   @doc "Creates a generator map where the keys are required except the list provided"
+  def mixed_map(map, []) do
+    map = to_generators(map)
+    StreamData.fixed_map(map)
+  end
+
   def mixed_map(map, keys) do
     map = to_generators(map)
     {optional, required} = Map.split(map, keys)
diff --git a/lib/ash/type/keyword.ex b/lib/ash/type/keyword.ex
index 2a2717b3b..d97e38f98 100644
--- a/lib/ash/type/keyword.ex
+++ b/lib/ash/type/keyword.ex
@@ -146,6 +146,14 @@ defmodule Ash.Type.Keyword do
     end)
   end
 
+  @impl true
+  def generator(constraints) do
+    Ash.Type.Map.generator(constraints)
+    |> StreamData.map(fn value ->
+      Map.to_list(value)
+    end)
+  end
+
   defp check_fields(value, fields) do
     Enum.reduce(fields, {:ok, []}, fn
       {field, field_constraints}, {:ok, checked_value} ->
diff --git a/lib/ash/type/map.ex b/lib/ash/type/map.ex
index 62c8f902e..eed4c0628 100644
--- a/lib/ash/type/map.ex
+++ b/lib/ash/type/map.ex
@@ -125,6 +125,28 @@ defmodule Ash.Type.Map do
     end)
   end
 
+  @impl true
+  def generator(constraints) do
+    if constraints[:fields] do
+      optional =
+        constraints[:fields]
+        |> Enum.filter(fn {_, value} ->
+          value[:allow_nil?]
+        end)
+        |> Keyword.keys()
+
+      constraints[:fields]
+      |> Map.new(fn {key, config} ->
+        type = config[:type]
+        constraints = config[:constraints] || []
+        {key, Ash.Type.generator(type, constraints)}
+      end)
+      |> Ash.Generator.mixed_map(optional)
+    else
+      StreamData.repeatedly(%{})
+    end
+  end
+
   defp check_fields(value, fields) do
     Enum.reduce(fields, {:ok, %{}}, fn
       {field, field_constraints}, {:ok, checked_value} ->
diff --git a/lib/ash/type/struct.ex b/lib/ash/type/struct.ex
index 97d19156f..064390f0c 100644
--- a/lib/ash/type/struct.ex
+++ b/lib/ash/type/struct.ex
@@ -185,6 +185,19 @@ defmodule Ash.Type.Struct do
     end
   end
 
+  @impl true
+  def generator(constraints) do
+    if !constraints[:instance_of] do
+      raise ArgumentError,
+            "Cannot generate instances of the `:struct` type without an `:instance_of` constraint"
+    else
+      Ash.Type.Map.generator(constraints)
+      |> StreamData.map(fn value ->
+        struct(constraints[:instance_of], value)
+      end)
+    end
+  end
+
   @impl true
   def apply_constraints(value, constraints) do
     with {:ok, value} <- handle_fields(value, constraints) do