lib/runbox/runtime/stage/unit_registry/alternative_registry.ex

defmodule Runbox.Runtime.Stage.UnitRegistry.AlternativeRegistry do
  @moduledoc """
  Module manages alternative registry state.

  Every unit is registered only once per alternative registry. More than one unit can be registered
  with the same key.
  """

  @doc "Returns new alternative registry"
  def new do
    %{}
  end

  @type registration_key :: term
  @type unit_id :: term
  @type t :: %{registration_key => [unit_id]}

  @doc "Registers unit in alternative reg with given key"
  @spec register_unit(t, registration_key, unit_id) :: t
  def register_unit(alt_reg, key, unit_id) do
    registered_units = Map.get(alt_reg, key, [])

    if unit_id in registered_units do
      alt_reg
    else
      Map.put(alt_reg, key, Enum.concat(registered_units, [unit_id]))
    end
  end

  @doc "Unregisters unit from alternative reg"
  @spec unregister_unit(t, unit_id) :: t
  def unregister_unit(alt_reg, unit_id) do
    Enum.reduce(alt_reg, %{}, fn {key, units0}, acc ->
      case List.delete(units0, unit_id) do
        [] ->
          acc

        units ->
          Map.put(acc, key, units)
      end
    end)
  end

  @doc "Get registered unit ids"
  @spec lookup(t, registration_key) :: [unit_id]
  def lookup(alt_reg, key) do
    Map.get(alt_reg, key, [])
  end
end