Anyone have any suggestions on how to refactor this code? (It works but feels messy)

[wintermeyer]

I am not sure if is appropriate to ask for help for a refactoring problem in this forum. In case it isn’t I’ll delete the post right away.

I wrote a little Phoenix 1.6 LiveView demo which simulates a basic stock market.

Here’s the code:

defmodule DemoWeb.StockWatchLive do
  use DemoWeb, :live_view

  def mount(_params, _session, socket) do
    if connected?(socket), do: Process.send_after(self(), :tick, 2500)

    {:ok, update_world(socket)}
  end

  def render(assigns) do
    ~H"""
    <p>Available money: <%= @balance %><br/>
    Value of the portfolio: <%= @portfolio_value %></p>
    <table>
      <thead>
        <tr>
          <th>Name</th>
          <th>Price</th>
          <th>Portfolio</th>
          <th colspan="2"></th>
        </tr>
      </thead>
      <tbody>
      <%= for {stock_name, value} <- @stocks do %>
        <tr>
          <td><%= stock_name %></td>
          <td><%= value %></td>
          <td><%= Map.get(@portfolio, stock_name) %></td>
          <td width="20%">
            <%= if value <= @balance do %>
              <button phx-click="buy" phx-value-ref={stock_name}>Buy!</button>
            <% end %>
          </td>
          <td width="20%">
            <%= if Map.has_key?(@portfolio, stock_name) && Map.get(@portfolio, stock_name) > 0 do %>
              <button phx-click="sell" phx-value-ref={stock_name}>Sell!</button>
            <% end %>
          </td>
        </tr>
      <% end %>
      </tbody>
    </table>
    """
  end

  def handle_info(:tick, socket) do
    Process.send_after(self(), :tick, 2500)

    {:noreply, update_world(socket)}
  end

  def handle_event("buy", %{"ref" => stock_name}, socket) do
    {:noreply, execute_order(stock_name, 1, socket)}
  end

  def handle_event("sell", %{"ref" => stock_name}, socket) do
    {:noreply, execute_order(stock_name, -1, socket)}
  end

  defp stock_price(socket, stock_name) do
    Map.get(socket.assigns.stocks, stock_name)
  end

  defp portfolio_value(socket) do
    if Map.has_key?(socket.assigns, :portfolio) do
      portfolio = socket.assigns.portfolio

      stocks_value_list =
        for {stock_name, amount} <-
              portfolio,
            into: [],
            do: stock_price(socket, stock_name) * amount

      Enum.sum(stocks_value_list)
    else
      0
    end
  end

  defp execute_order(stock_name, amount, socket) do
    price = stock_price(socket, stock_name)

    portfolio =
      socket.assigns.portfolio
      |> Map.update(stock_name, 1, &(&1 + amount))

    balance = socket.assigns.balance + -1 * amount * price

    socket = assign(socket, portfolio: portfolio)

    portfolio_value = portfolio_value(socket)

    socket
    |> assign(portfolio_value: portfolio_value)
    |> assign(balance: balance)
  end

  defp update_world(socket) do
    if Map.has_key?(socket.assigns, :stocks) do
      stocks =
        for {k, v} <-
              socket.assigns.stocks,
            into: %{},
            do: {k, random_fluctuation(v)}

      socket
      |> assign(stocks: stocks)
      |> assign(portfolio_value: portfolio_value(socket))
    else
      initialize_world(socket)
    end
  end

  defp random_fluctuation(value) do
    case Enum.random(1..4) do
      4 -> value + Enum.random(-5..6)
      _ -> value
    end
  end

  defp initialize_world(socket) do
    socket
    |> assign(
      stocks: %{
        "Stock A" => 100,
        "Stock B" => 200,
        "Stock C" => 300,
        "Stock D" => 400
      }
    )
    |> assign(portfolio: %{})
    |> assign(balance: 1_000)
    |> assign(portfolio_value: 0)
  end
end

The code doesn’t feel as clean and tidy as I’d like it to be. It works but is feels messy. But I don’t have a good idea how to refactor it.

How can I improve it?

[gregvaughn]

One quick win I see is that by using the default param of Map.get/3 you can remove some explicit conditionals. For example, this should be equivalent to your function

  defp portfolio_value(socket) do
    for {stock_name, amount} <- Map.get(socket.assigns, :portfolio, []), reduce: 0 do
      running_sum -> running_sum + (stock_price(stock_name) * amount)
    end
  end

Minor note: using into: [] is redundant in a for comprehension. That is its default behavior.

[al2o3cr]

@wintermeyer

portfolio_value = portfolio_value(socket)

A general suggestion: this code entangles “domain stuff” (the concept of “portfolio value”) with “implementation stuff” (storing data on socket.assigns). Even if the functions are simple and all live in the same file, it’ll improve readability to keep the two sets of ideas apart.

For instance, portfolio_value could be refactored to take portfolio directly:

  defp portfolio_value(portfolio, stocks) do
    portfolio
    |> Enum.map(fn {stock_name, amount} -> stock_price(stocks, stock_name) * amount; end)
    |> Enum.sum()
  end

  defp stock_price(stocks, name) do
    Map.get(stocks, name)
  end

Functions “at the boundary” like execute_order that do take socket should focus on extracting data from assigns, passing it to functions that don’t take socket, and then storing the results back into assigns.

If you squint a little, this is the “functional core, imperative shell” pattern implemented at the single-controller level.

If the logic gets complicated in the future, this division also makes it straightforward to lift the not-socket parts out into another module for readability and testability.

[wintermeyer]

Thanks Matt!

Here is the updated version:

defmodule DemoWeb.StockWatchLive do
  use DemoWeb, :live_view

  @refresh_rate 250

  def mount(_params, _session, socket) when is_map_key(socket, :stocks) do
    {:ok, update_world(socket)}
  end

  def mount(_params, _session, socket) do
    if connected?(socket), do: Process.send_after(self(), :tick, @refresh_rate)

    {:ok, initialize_world(socket)}
  end

  def render(assigns) do
    ~H"""
    <table>
      <tbody>
        <tr><td>Available money:</td><td><%= @balance %></td></tr>
        <tr><td>Value of the portfolio:</td><td><%= @portfolio_value %></td></tr>
        <tr><td>Sum:</td><td><%= @balance + @portfolio_value %></td></tr>
      </tbody>
    </table>
    <table>
      <thead>
        <tr>
          <th>Name</th>
          <th>Price</th>
          <th>Portfolio</th>
          <th colspan="2"></th>
        </tr>
      </thead>
      <tbody>
      <%= for {stock_name, value} <- @stocks do %>
        <tr>
          <td><%= stock_name %></td>
          <td><%= value %></td>
          <td><%= Map.get(@portfolio, stock_name) %></td>
          <td width="20%">
            <%= if value <= @balance do %>
              <button phx-click="buy" phx-value-ref={stock_name}>Buy!</button>
            <% end %>
          </td>
          <td width="20%">
            <%= if Map.has_key?(@portfolio, stock_name) && Map.get(@portfolio, stock_name) > 0 do %>
              <button phx-click="sell" phx-value-ref={stock_name}>Sell!</button>
            <% end %>
          </td>
        </tr>
      <% end %>
      </tbody>
    </table>
    """
  end

  def handle_info(:tick, socket) do
    Process.send_after(self(), :tick, @refresh_rate)

    {:noreply, update_world(socket)}
  end

  def handle_event("buy", %{"ref" => stock_name}, socket) do
    {:noreply, execute_order(stock_name, 1, socket)}
  end

  def handle_event("sell", %{"ref" => stock_name}, socket) do
    {:noreply, execute_order(stock_name, -1, socket)}
  end

  defp execute_order(stock_name, amount, socket) do
    %{:stocks => stocks, :portfolio => portfolio, :balance => balance} = socket.assigns

    socket
    |> assign(portfolio: add_to_portfolio(portfolio, stock_name, amount))
    |> assign(portfolio_value: portfolio_value(portfolio, stocks))
    |> assign(balance: adjust_balance(balance, stocks, stock_name, amount))
  end

  defp stock_price(stocks, stock_name) do
    Map.get(stocks, stock_name)
  end

  defp portfolio_value(portfolio, stocks) do
    portfolio
    |> Enum.map(fn {stock_name, amount} -> stock_price(stocks, stock_name) * amount end)
    |> Enum.sum()
  end

  defp add_to_portfolio(portfolio, stock_name, amount) do
    Map.update(portfolio, stock_name, 1, &(&1 + amount))
  end

  defp adjust_balance(balance, stocks, stock_name, amount) do
    balance + -1 * amount * stock_price(stocks, stock_name)
  end

  defp update_prices(stocks) do
    for {name, value} <-
          stocks,
        into: %{},
        do: {name, random_new_price(value)}
  end

  defp random_new_price(value) do
    case Enum.random(1..4) do
      4 -> value + Enum.random(-5..6)
      _ -> value
    end
  end

  defp update_world(socket) do
    %{:stocks => stocks, :portfolio => portfolio} = socket.assigns

    socket
    |> assign(stocks: update_prices(stocks))
    |> assign(portfolio_value: portfolio_value(portfolio, stocks))
  end

  defp initialize_world(socket) do
    socket
    |> assign(
      stocks: %{
        "Stock A" => 100,
        "Stock B" => 200,
        "Stock C" => 300,
        "Stock D" => 400
      }
    )
    |> assign(portfolio: %{})
    |> assign(balance: 1_000)
    |> assign(portfolio_value: 0)
  end
end