From 51c4b7ffc06ff9e332c91f6b2ec1c4d8e9d81f8e Mon Sep 17 00:00:00 2001
From: Noel Han <noel@aidkr.com>
Date: Tue, 7 May 2024 19:05:26 +0900
Subject: [PATCH] Handle already expanded AST (#43)

---
 lib/typed_ecto_schema/syntax_sugar.ex |   6 +-
 test/typed_ecto_schema_test.exs       | 137 ++++++++++++++++++++++++++
 2 files changed, 140 insertions(+), 3 deletions(-)

diff --git a/lib/typed_ecto_schema/syntax_sugar.ex b/lib/typed_ecto_schema/syntax_sugar.ex
index 81c0545..8112ed8 100644
--- a/lib/typed_ecto_schema/syntax_sugar.ex
+++ b/lib/typed_ecto_schema/syntax_sugar.ex
@@ -152,13 +152,13 @@ defmodule TypedEctoSchema.SyntaxSugar do
     expanded = Macro.expand(unknown, env)
 
     case expanded do
-      ^unknown ->
-        unknown
-
       {:__block__, block_context, calls} ->
         new_calls = Enum.map(calls, &transform_expression(&1, env))
         {:__block__, block_context, new_calls}
 
+      ^unknown ->
+        unknown
+
       call ->
         transform_expression(call, env)
     end
diff --git a/test/typed_ecto_schema_test.exs b/test/typed_ecto_schema_test.exs
index 99a68f7..7b13b9a 100644
--- a/test/typed_ecto_schema_test.exs
+++ b/test/typed_ecto_schema_test.exs
@@ -732,6 +732,143 @@ defmodule TypedEctoSchemaTest do
              |> Ecto.Changeset.put_embed(:one, %{int: 123})
   end
 
+  test "SyntaxSugar.apply_to_block/2 work with expanded nested AST" do
+    block =
+      {:__block__, [],
+       [
+         {:field, [line: 2], [:name, :string]},
+         {:field, [line: 3], [:age, :integer]},
+         {:__block__, [],
+          [{:field, [], [:embed_flag, :boolean]}, {:field, [], [:embed_name, :string]}]}
+       ]}
+
+    result = TypedEctoSchema.SyntaxSugar.apply_to_block(block, :env)
+
+    assert {
+             :__block__,
+             [],
+             [
+               {:__block__, [],
+                [
+                  {:field, [], [:name, :string]},
+                  {{:., [], [TypedEctoSchema.TypeBuilder, :add_field]}, [],
+                   [{:__MODULE__, [], TypedEctoSchema.SyntaxSugar}, :field, :name, :string, []]}
+                ]},
+               {:__block__, [],
+                [
+                  {:field, [], [:age, :integer]},
+                  {{:., [], [TypedEctoSchema.TypeBuilder, :add_field]}, [],
+                   [{:__MODULE__, [], TypedEctoSchema.SyntaxSugar}, :field, :age, :integer, []]}
+                ]},
+               {:__block__, [],
+                [
+                  {:__block__, [],
+                   [
+                     {:field, [], [:embed_flag, :boolean]},
+                     {{:., [], [TypedEctoSchema.TypeBuilder, :add_field]}, [],
+                      [
+                        {:__MODULE__, [], TypedEctoSchema.SyntaxSugar},
+                        :field,
+                        :embed_flag,
+                        :boolean,
+                        []
+                      ]}
+                   ]},
+                  {:__block__, [],
+                   [
+                     {:field, [], [:embed_name, :string]},
+                     {{:., [], [TypedEctoSchema.TypeBuilder, :add_field]}, [],
+                      [
+                        {:__MODULE__, [], TypedEctoSchema.SyntaxSugar},
+                        :field,
+                        :embed_name,
+                        :string,
+                        []
+                      ]}
+                   ]}
+                ]}
+             ]
+           } == result
+  end
+
+  defmodule Custom.Schema do
+    @schema_function_names [
+      :field,
+      :embeds_one,
+      :embeds_many,
+      :belongs_to
+    ]
+
+    defmacro __using__(_opts) do
+      quote do
+        use TypedEctoSchema
+        import TypedEctoSchema, except: [typed_schema: 2]
+        import unquote(__MODULE__), only: [typed_schema: 2]
+      end
+    end
+
+    defmacro typed_schema(name, do: {function_name, ctx, args}) do
+      fields = {function_name, ctx, Enum.map(args, &put_null_false/1)}
+
+      quote do
+        TypedEctoSchema.typed_schema unquote(name) do
+          unquote(fields)
+        end
+      end
+    end
+
+    defp put_null_false({function_name, ctx, args})
+         when function_name in @schema_function_names do
+      {name, type, opts} =
+        case args do
+          [name, type] -> {name, type, []}
+          [name, type, opts] -> {name, type, opts}
+        end
+
+      {function_name, ctx, [name, type, Keyword.put_new(opts, :null, false)]}
+    end
+
+    defp put_null_false({:__block__, ctx, args}) do
+      {:__block__, ctx, Enum.map(args, &put_null_false/1)}
+    end
+
+    defp put_null_false(ast), do: ast
+  end
+
+  defmodule CustomMacroSchema do
+    use Custom.Schema
+
+    typed_schema "custom_macro_schemas" do
+      field(:name, :string)
+      field(:age, :integer)
+
+      (
+        field(:foo, :string)
+        field(:bar, :integer)
+      )
+    end
+  end
+
+  test "pre macro passed schema" do
+    require IEx.Helpers
+
+    assert CustomMacroSchema.__schema__(:fields) == [
+             :id,
+             :name,
+             :age,
+             :foo,
+             :bar
+           ]
+
+    assert CustomMacroSchema.__struct__() == %CustomMacroSchema{
+             id: nil,
+             name: nil,
+             age: nil,
+             foo: nil,
+             bar: nil
+           }
+  end
+
   ##
   ## Helpers
   ##