feat: add type parameters

Co-authored-by: Phil Chen <06fahchen@gmail.com>
Co-authored-by: Jean-Philippe Cugnet <jean-philippe@cugnet.eu>

Author: fahchen

.formatter.exs

@@ -1,7 +1,7 @@
 # SPDX-FileCopyrightText: NONE
 # SPDX-License-Identifier: CC0-1.0
 
-locals_without_parens = [field: 2, field: 3, plugin: 1, plugin: 2]
+locals_without_parens = [field: 2, field: 3, parameter: 1, plugin: 1, plugin: 2]
 
 [
   inputs: [

README.md

@@ -175,6 +175,32 @@ defmodule MyOpaqueStruct do
 end
 ```
 
+You can add type parameters:
+
+```elixir
+defmodule User do
+  use TypedStruct
+
+  typedstruct do
+    # Define a type parameter with the `parameter` macro.
+    parameter :state
+
+    # You can then use it here as `state`.
+    field :state, state, enforce: true
+    field :name, String.t()
+  end
+end
+```
+
+And use them like this:
+
+```elixir
+@type user_state() :: :registered | :confirmed | :logged_in
+
+@spec get_user_state(User.t(user_state())) :: user_state()
+def get_user_state(%User{state, _name}), do: state
+```
+
 If you often define submodules containing only a struct, you can avoid
 boilerplate code:
 
@@ -380,10 +406,29 @@ generates the following type:
 
 ```elixir
 @opaque t() :: %__MODULE__{
-          name: String.t()
+          name: String.t() | nil
         }
 ```
 
+When you specify parameters with the `parameter/1` macro, they are used in the
+definition of the type:
+
+```elixir
+typedstruct do
+  parameter :param
+
+  field :field, param
+end
+```
+
+gives the following type:
+
+```elixir
+@type t(param) :: %__MODULE__{
+        field: param | nil
+      }
+```
+
 When passing `module: ModuleName`, the whole `typedstruct` block is wrapped in a
 module definition. This way, the following definition:
 

lib/typed_struct.ex

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: 2018 Marcin Górnik <marcin.gornik@gmail.com>
 # SPDX-FileCopyrightText: 2022 Jonathan Chukinas <chukinas@gmail.com>
 # SPDX-FileCopyrightText: 2022 Balázs Jávorszky <javorszky.balazs@estyle.hu>
+# SPDX-FileCopyrightText: 2022 Phil Chen <06fahchen@gmail.com>
 #
 # SPDX-License-Identifier: MIT
 
@@ -15,6 +16,7 @@ defmodule TypedStruct do
   @accumulating_attrs [
     :ts_plugins,
     :ts_plugin_fields,
+    :ts_parameters,
     :ts_fields,
     :ts_types,
     :ts_docs,
@@ -70,6 +72,18 @@ defmodule TypedStruct do
         end
       end
 
+  You can also add type parameters:
+
+      defmodule MyModule do
+        use TypedStruct
+
+        typedstruct do
+          parameter :type
+
+          field :field, type
+        end
+      end
+
   You can create the struct in a submodule instead:
 
       defmodule MyModule do
@@ -119,7 +133,12 @@ defmodule TypedStruct do
       defstruct @ts_fields
 
       TypedStruct.__typedoc__(@ts_docs)
-      TypedStruct.__type__(@ts_types, unquote(opts))
+
+      TypedStruct.__type__(
+        Enum.reverse(@ts_parameters),
+        @ts_types,
+        unquote(opts)
+      )
     end
   end
 
@@ -159,14 +178,18 @@ defmodule TypedStruct do
   end
 
   @doc false
-  defmacro __type__(types, opts) do
+  defmacro __type__(parameters, types, opts) do
     if Keyword.get(opts, :opaque, false) do
-      quote bind_quoted: [types: types] do
-        @opaque t() :: %__MODULE__{unquote_splicing(types)}
+      quote bind_quoted: [parameters: parameters, types: types] do
+        @opaque t(unquote_splicing(parameters)) :: %__MODULE__{
+                  unquote_splicing(types)
+                }
       end
     else
-      quote bind_quoted: [types: types] do
-        @type t() :: %__MODULE__{unquote_splicing(types)}
+      quote bind_quoted: [parameters: parameters, types: types] do
+        @type t(unquote_splicing(parameters)) :: %__MODULE__{
+                unquote_splicing(types)
+              }
       end
     end
   end
@@ -199,6 +222,35 @@ defmodule TypedStruct do
     end
   end
 
+  @doc """
+  Defines a type parameter for the currently defined struct.
+
+  ## Example
+
+      typedstruct do
+        # Defines a type parameter named `type_param`
+        parameter :type_param
+
+        # The type parameter can be used as a type in the `field` macro.
+        field :a_field, type_param
+      end
+  """
+  defmacro parameter(name) do
+    quote bind_quoted: [name: name] do
+      TypedStruct.__parameter__(name, __ENV__)
+    end
+  end
+
+  @doc false
+  def __parameter__(name, %Macro.Env{module: mod}) when is_atom(name) do
+    Module.put_attribute(mod, :ts_parameters, Macro.var(name, mod))
+  end
+
+  def __parameter__(name, _env) do
+    raise ArgumentError,
+          "the name of a type parameter must be an atom, got #{inspect(name)}"
+  end
+
   @doc """
   Defines a field in a typed struct.
 
@@ -254,7 +306,7 @@ defmodule TypedStruct do
   # Checks whether some value looks like Elixir AST.
   defp ast?({name, meta, params})
        when (is_atom(name) or is_tuple(name)) and is_list(meta) and
-              is_list(params),
+              (is_list(params) or is_nil(params)),
        do: true
 
   defp ast?(_), do: false

test/support/test_struct.ex

@@ -1,5 +1,6 @@
 # SPDX-FileCopyrightText: 2018, 2020, 2025 Jean-Philippe Cugnet <jean-philippe@cugnet.eu>
 # SPDX-FileCopyrightText: 2018 Marcin Górnik <marcin.gornik@gmail.com>
+# SPDX-FileCopyrightText: 2022 Phil Chen <06fahchen@gmail.com>
 # SPDX-FileCopyrightText: 2023 Serge Aleynikov <saleyn@gmail.com>
 #
 # SPDX-License-Identifier: MIT
@@ -82,6 +83,37 @@ defmodule TypedStruct.TestStruct do
     end
   end
 
+  defmodule WithParameter do
+    @moduledoc """
+    A struct with a parameterised type.
+    """
+    use TypedStruct
+
+    typedstruct do
+      parameter :t1
+      parameter :t2
+
+      field :field_t1, t1
+      field :field_t2, t2
+      field :enforced_field_t1, t1, enforce: true
+    end
+
+    defmodule Expected do
+      @moduledoc """
+      `WithParameter` but defined manually.
+      """
+
+      @enforce_keys [:enforced_field_t1]
+      defstruct [:field_t1, :field_t2, :enforced_field_t1]
+
+      @type t(t1, t2) :: %__MODULE__{
+              field_t1: t1 | nil,
+              field_t2: t2 | nil,
+              enforced_field_t1: t1
+            }
+    end
+  end
+
   defmodule AsSubmodule do
     @moduledoc """
     A struct defined as a submodule.

test/typed_struct_test.exs

@@ -1,5 +1,6 @@
 # SPDX-FileCopyrightText: 2018, 2020, 2022, 2025 Jean-Philippe Cugnet <jean-philippe@cugnet.eu>
 # SPDX-FileCopyrightText: 2018 Marcin Górnik <marcin.gornik@gmail.com>
+# SPDX-FileCopyrightText: 2022 Phil Chen <06fahchen@gmail.com>
 #
 # SPDX-License-Identifier: MIT
 
@@ -79,6 +80,22 @@ defmodule TypedStructTest do
     assert type1 == type2
   end
 
+  test "generates a parameterized type for the struct" do
+    # Get both types and standardise them (remove line numbers and rename
+    # the second struct with the name of the first one).
+    type1 =
+      TestStruct.WithParameter
+      |> extract_first_type()
+      |> standardise(TestStruct.WithParameter)
+
+    type2 =
+      TestStruct.WithParameter.Expected
+      |> extract_first_type()
+      |> standardise(TestStruct.WithParameter.Expected)
+
+    assert type1 == type2
+  end
+
   test "generates the struct in a submodule if `module: ModuleName` is set" do
     # credo:disable-for-next-line Credo.Check.Design.AliasUsage
     assert TestStruct.AsSubmodule.Struct.__struct__() ==
@@ -163,6 +180,21 @@ defmodule TypedStructTest do
            end) =~ "undefined function field/2"
   end
 
+  test "the name of a type parameter be an atom" do
+    assert_raise ArgumentError,
+                 "the name of a type parameter must be an atom, got 3",
+                 fn ->
+                   defmodule InvalidStruct do
+                     use TypedStruct
+
+                     typedstruct do
+                       parameter 3
+                       field :field, term()
+                     end
+                   end
+                 end
+  end
+
   test "the name of a field must be an atom" do
     assert_raise ArgumentError, "a field name must be an atom, got 3", fn ->
       defmodule InvalidStruct do
@@ -237,14 +269,20 @@ defmodule TypedStructTest do
   defp standardise({:type, _, type, params}, struct),
     do: {:type, :line, type, standardise(params, struct)}
 
+  defp standardise({:user_type, _, type, params}, struct),
+    do: {:user_type, :line, type, standardise(params, struct)}
+
   defp standardise({:remote_type, _, params}, struct),
     do: {:remote_type, :line, standardise(params, struct)}
 
   defp standardise({:atom, _, struct}, struct),
     do: {:atom, :line, TestStruct}
 
+  defp standardise({:var, _, name}, _),
+    do: {:var, :line, name}
+
   defp standardise({name, type, params}, struct) when is_tuple(type),
-    do: {name, standardise(type, struct), params}
+    do: {name, standardise(type, struct), standardise(params, struct)}
 
   defp standardise({type, _, literal}, _struct),
     do: {type, :line, literal}