Expressive design in Elixir with polymorphic protocols

One of the things I love from Elixir is its versatility and the possibility of solving the same problem in different ways. Sometimes with a level of expressiveness that reminds me of the way I like to design when I write code in an OO programming language.

In this article we will see how to refactor a simple Kata applying polymorphism with the help of protocols and also using the pipe operator to achieve an elegant and expressive design.

Imagine we have to code a new fighting game called "Blade Fury"[1] where two warriors fight with each other and the one with the greatest power wins. Our initial requirement is to have just two warriors: a Shaolin and a Samurai, each with a strength of 10 and 6 units respectively. In this first version, the power is the value of the warrior's strength, so a Shaolin should win a fight against a Samurai.

Let's start taking a look at the tests[2]:

defmodule BladeFuryTest do
  use ExUnit.Case
  use ExMatchers

  alias BladeFury.Warrior.Samurai
  alias BladeFury.Warrior.Shaolin

  test "when a samurai fights vs a shaolin, the shaolin wins" do
    samurai = %Samurai{}
    shaolin = %Shaolin{}

    result = BladeFury.fight_vs(samurai, shaolin)

    expect result, to: eq(:warrior_two)
  end

  test "when a shaolin fights vs a samurai, the shaolin wins" do
    shaolin = %Shaolin{}
    samurai = %Samurai{}

    result = BladeFury.fight_vs(shaolin, samurai)

    expect result, to: eq(:warrior_one)
  end

  test "when a shaolin fights vs a shaolin, it's a draw" do
    shaolin_one = %Shaolin{}
    shaolin_two = %Shaolin{}

    result = BladeFury.fight_vs(shaolin_one, shaolin_two)

    expect result, to: eq(:draw)
  end
end

Following TDD principles, a simple solution to make the tests pass might be:

defmodule BladeFury do

  def fight_vs(warrior_one, warrior_two) do
    cond do
      warrior_one.strength > warrior_two.strength -> :warrior_one
      warrior_one.strength < warrior_two.strength -> :warrior_two
      true -> :draw
    end
  end
end

defmodule BladeFury.Warrior.Shaolin do
  defstruct strength: 10
end

defmodule BladeFury.Warrior.Samurai do
  defstruct strength: 6
end

Now, let's make the game a bit more interesting! Let’s add weapons which can cause damage to the opponent by increasing the warrior's power. We can start with a Katana which doubles the Samurai's power, so a Samurai with a Katana should win vs an unarmed Shaolin. We will need to adapt our code to model a warrior with a weapon. To keep things simple, we'll just pass a keyword list to the fight_vs/2 function:

test "when a samurai with a katana fights vs a shaolin, the samurai wins" do
    samurai = %Samurai{}
    shaolin = %Shaolin{}

    katana = %Katana{}

    result =
      BladeFury.fight_vs(
        [warrior: samurai, weapon: katana],
        [warrior: shaolin]
      )

    expect result, to: eq(:warrior_one)
  end

And the new solution is:

defmodule BladeFury do

  def fight_vs(warrior_with_weapon_one, warrior_with_weapon_two) do
    warrior_one_power = calculate_power(warrior_with_weapon_one)
    warrior_two_power = calculate_power(warrior_with_weapon_two)

    cond do
      warrior_one_power > warrior_two_power -> :warrior_one
      warrior_one_power < warrior_two_power -> :warrior_two
      true -> :draw
    end
  end

  def calculate_power([warrior: warrior, weapon: weapon]) do
    warrior.strength * weapon.damage
  end

  def calculate_power([warrior: warrior]) do
    warrior.strength
  end
end

defmodule BladeFury.Weapon.Katana do
  defstruct damage: 2
end

However, what might happen if a Shaolin fights armed with a Katana? Well, given the Katana doubles the warrior's power, the Shaolin should win, right? But that doesn't sound quite right, given that the Katana is a Japanese sword, therefore the Shaolin shouldn't be as proficient as the Samurai with it. What we would like to happen is that the Shaolin's power increases a bit, but not that much. Then, when a Shaolin armed with a Katana fights vs a Samurai also armed with a Katana, the result should be a draw. In other words, the Katana doubles the power of the Samurai, but only adds 2 units to the Shaolin's strength.

Our new test would be:

  test "when a samurai with a katana fights vs a shaolin with a katana, it's a draw" do
    samurai = %Samurai{}
    shaolin = %Shaolin{}

    katana = %Katana{}

    result =
      BladeFury.fight_vs(
        [warrior: samurai, weapon: katana],
        [warrior: shaolin, weapon: katana]
      )

    expect result, to: eq(:draw)
  end

If we continue using pattern matching to solve the problem, the new code might look like this:

defmodule BladeFury do
  alias BladeFury.Warrior.Samurai
  alias BladeFury.Warrior.Shaolin
  alias BladeFury.Weapon.Katana

  def fight_vs(warrior_with_weapon_one, warrior_with_weapon_two) do
    warrior_one_power = calculate_power(warrior_with_weapon_one)
    warrior_two_power = calculate_power(warrior_with_weapon_two)

    cond do
      warrior_one_power > warrior_two_power -> :warrior_one
      warrior_one_power < warrior_two_power -> :warrior_two
      true -> :draw
    end
  end

  def calculate_power(
        warrior: %Samurai{strength: warrior_strength},
        weapon: %Katana{damage: weapon_damage}
      ) do
    warrior_strength * weapon_damage
  end

  def calculate_power(
        warrior: %Shaolin{strength: warrior_strength},
        weapon: %Katana{damage: weapon_damage}
      ) do
    warrior_strength + weapon_damage
  end

  def calculate_power(warrior: warrior) do
    warrior.strength
  end
end

We can imagine how the BladeFury module will grow pretty quickly as soon as we start adding more warriors and weapons. The problem with our design is that the BladeFury module has all the responsibilities for solving the game's logic and the rest of the modules are almost empty!

Let's refactor our model using polymorphism and protocols to distribute the responsibilities in a better way.

We want to keep the BladeFury module as simple as possible, so we will delegate the responsibility of calculating the warrior's power to each warrior. Given we know each warrior will calculate the power slightly differently we will use the Power protocol to achieve that:

defmodule BladeFury do
  import BladeFury.Warrior.Power, only: [power: 1]

  def fight_vs(warrior_one, warrior_two) do
    warrior_one_power = warrior_one |> power()
    warrior_two_power = warrior_two |> power()

    cond do
      warrior_one_power > warrior_two_power -> :warrior_one
      warrior_one_power < warrior_two_power -> :warrior_two
      true -> :draw
    end
  end
end

The warriors should implement this new protocol and tell to their weapon for whom they are increasing the power:

defprotocol BladeFury.Warrior.Power do
  def power(warrior)
end

defmodule BladeFury.Warrior.Shaolin do
  import BladeFury.Weapon, only: [increase_shaolin_power: 2]

  defstruct [:weapon, strength: 10]

  defimpl BladeFury.Warrior.Power do
    def power(warrior) do
      warrior.weapon |> increase_shaolin_power(warrior.strength)
    end
  end
end

defmodule BladeFury.Warrior.Samurai do
  import BladeFury.Weapon, only: [increase_samurai_power: 2]

  defstruct [:weapon, strength: 6]

  defimpl BladeFury.Warrior.Power do
    def power(warrior) do
      warrior.weapon |> increase_samurai_power(warrior.strength)
    end
  end
end

And the weapon now knows how to increase the power for the different warriors:

defprotocol BladeFury.Weapon do
  def increase_samurai_power(weapon, warrior_strength)
  def increase_shaolin_power(weapon, warrior_strength)
end

defmodule BladeFury.Weapon.Katana do
  defstruct damage: 2

  defimpl BladeFury.Weapon do
    def increase_samurai_power(weapon, warrior_strength) do
      weapon.damage * warrior_strength
    end
    def increase_shaolin_power(weapon, warrior_strength) do
      weapon.damage + warrior_strength
    end
  end
end

Good designs tend to be more explicit, so let's also model the unarmed scenario by creating a Weapon that doesn't affect the warrior's power.

defmodule BladeFury.Weapon.PunchKick do
  defstruct damage: 0

  alias BladeFury.Weapon.PunchKick

  defimpl BladeFury.Weapon do
    defdelegate increase_shaolin_power(weapon, warrior_strength), to: PunchKick, as: :increase_power
    defdelegate increase_samurai_power(weapon, warrior_strength), to: PunchKick, as: :increase_power
  end

  def increase_power(_weapon, warrior_strength) do
    warrior_strength
  end
end

We can see how the responsibilities are much better delegated. The fight tells each warrior to calculate its power through the Power protocol. Then each warrior implements the protocol and tells the weapon to increase its power for a specific kind of warrior through the Weapon protocol. Finally, each Weapon has to implement the protocol and knows how to affect the power for different warriors.

Now if we want to add a new Weapon, it is as simple as adding a new module which implements the Weapon protocol and defines how the weapon behaves when it's used by each kind of Warrior. It's a bit more work if we want to add a new Warrior, but not that much. We have to add the new warrior's module which should implement the Power protocol, add the new method on the Weapon protocol to handle the power calculation for this new warrior, and finally implement this new function on all existing weapons.

Last but not least, with the help of some helper functions and factories, we can write tests that are much more expressive, almost feeling like we have built a small DSL for our game engine. Plus, with the help of a custom matcher we can reuse the code that asserts on the fight result, so if the API changes in the future we will only need to change the code in one place.

defmodule BladeFuryTest do
  use ExUnit.Case
  use ExMatchers

  import CustomMatchers
  import BladeFury.Warrior.Gymnasium
  import BladeFury.Weapon.Armory
  import BladeFury.Warrior
  import BladeFuryFour, only: [fight_vs: 2]

  test "when a samurai with a katana fights vs a shaolin, the samurai wins" do
    expect(
      a_samurai() |> armed_with(a_katana())
      |> fight_vs(
        a_shaolin() |> unarmed()
      ),
      to: warrior_one_wins()
    )
  end

  test "when a samurai with a katana fights vs a shaolin with a katana, it's a draw" do
    expect(
      a_samurai() |> armed_with(a_katana())
      |> fight_vs(
        a_shaolin() |> armed_with(a_katana())
      ),
      to: be_a_draw()
    )
  end
end

You can find the code for each different step we took over here. Take a look at the different solutions and leave a comment saying which one you prefer or if you have a different solution to propose.

[1] Thanks to ChatGPT for suggesting the name
[2] I am using the library ExMatchers to write the assertions