From 7fb63d9bf4c59f111f3336226d56641a6c92e00d Mon Sep 17 00:00:00 2001
From: Arno Dirlam <arnodirlam@googlemail.com>
Date: Tue, 10 Dec 2024 23:44:16 +0100
Subject: [PATCH] Defd_: derive fun info for omitted default arguments (lower
 arities)

---
 lib/dx/defd_.ex          |  2 ++
 lib/dx/defd_/compiler.ex | 25 +++++++++++++++++++++++--
 test/dx/defd__test.exs   | 14 ++++++--------
 3 files changed, 31 insertions(+), 10 deletions(-)

diff --git a/lib/dx/defd_.ex b/lib/dx/defd_.ex
index 9cb6f87..82ec8e4 100644
--- a/lib/dx/defd_.ex
+++ b/lib/dx/defd_.ex
@@ -76,6 +76,8 @@ defmodule Dx.Defd_ do
   end
   ```
 
+  `args` options will also be derived for functions with omitted default arguments.
+
   2. Using the `@moduledx_` module attribute for module-wide defaults (can only be set once per module):
 
   ```elixir
diff --git a/lib/dx/defd_/compiler.ex b/lib/dx/defd_/compiler.ex
index fe9d249..a635421 100644
--- a/lib/dx/defd_/compiler.ex
+++ b/lib/dx/defd_/compiler.ex
@@ -3,7 +3,7 @@ defmodule Dx.Defd_.Compiler do
 
   alias Dx.Defd.Ast
 
-  def __compile__(%Macro.Env{module: module}, moduledx_, fun_infos) do
+  def __compile__(%Macro.Env{module: module}, moduledx_, defd_s) do
     existing_clauses =
       case Module.get_definition(module, {:__dx_fun_info, 2}) do
         {:v1, :def, _meta, clauses} ->
@@ -29,9 +29,30 @@ defmodule Dx.Defd_.Compiler do
           []
       end
 
+    # derive fun info for omitting default arguments
+    fun_infos =
+      Enum.reduce(defd_s, defd_s, fn
+        {{name, _arity} = key, %{defaults: defaults, fun_info: fun_info}}, acc ->
+          defaults
+          |> Map.keys()
+          |> Enum.sort(:desc)
+          |> Enum.reduce({fun_info, acc}, fn arg_index, {fun_info, acc} ->
+            fun_info = %{
+              fun_info
+              | args: List.delete_at(fun_info.args, arg_index),
+                arity: fun_info.arity - 1
+            }
+
+            acc = Map.put_new(acc, {name, fun_info.arity}, fun_info)
+            {fun_info, acc}
+          end)
+          |> elem(1)
+          |> Map.put(key, fun_info)
+      end)
+
     annotated_clauses =
       Enum.flat_map(fun_infos, fn
-        {{name, arity}, %{fun_info: fun_info}} ->
+        {{name, arity}, fun_info} ->
           quote do
             def __dx_fun_info(unquote(name), unquote(arity)) do
               unquote(Macro.escape(fun_info))
diff --git a/test/dx/defd__test.exs b/test/dx/defd__test.exs
index 2d9d491..a131a49 100644
--- a/test/dx/defd__test.exs
+++ b/test/dx/defd__test.exs
@@ -108,12 +108,10 @@ defmodule Dx.Defd_Test do
                fun_info(Moduledx_Test, :run, 1)
     end
 
-    test "with default args" do
+    test "infers lower arity fun info for default args" do
       defmodule DefaultArgsTest do
         use Dx.Defd_
 
-        @moduledx_ args: %{all: :preload_scope}
-
         @dx_ args: %{0 => :preload_scope, 1 => :fn}
         defd_ run(enum, mapper \\ nil) do
           if mapper do
@@ -126,16 +124,16 @@ defmodule Dx.Defd_Test do
 
       assert %FunInfo{
                args: [
-                 %ArgInfo{preload_scope: true}
+                 %ArgInfo{preload_scope: true},
+                 %ArgInfo{fn: %FunInfo{}, preload_scope: false}
                ]
-             } = fun_info(DefaultArgsTest, :run, 1)
+             } = fun_info(DefaultArgsTest, :run, 2)
 
       assert %FunInfo{
                args: [
-                 %ArgInfo{preload_scope: true},
-                 %ArgInfo{fn: %FunInfo{}, preload_scope: false}
+                 %ArgInfo{preload_scope: true}
                ]
-             } = fun_info(DefaultArgsTest, :run, 2)
+             } = fun_info(DefaultArgsTest, :run, 1)
     end
   end
 end