Optimise data view recomputation for delta operations (#286)

* Optimise data view recomputation for delta operations

* Use generic access_by_id/1 for nested updates

* Use access_by_id/1 for nested notebook updates

* Use Enum.split_while

lib/livebook/notebook.ex

@@ -14,6 +14,7 @@ defmodule Livebook.Notebook do
   defstruct [:name, :version, :sections, :metadata]
 
   alias Livebook.Notebook.{Section, Cell}
+  import Livebook.Utils, only: [access_by_id: 1]
 
   @type metadata :: %{String.t() => term()}
 
@@ -88,7 +89,6 @@ defmodule Livebook.Notebook do
   @spec insert_section(t(), integer(), Section.t()) :: t()
   def insert_section(notebook, index, section) do
     sections = List.insert_at(notebook.sections, index, section)
-
     %{notebook | sections: sections}
   end
 
@@ -97,16 +97,9 @@ defmodule Livebook.Notebook do
   """
   @spec insert_cell(t(), Section.id(), integer(), Cell.t()) :: t()
   def insert_cell(notebook, section_id, index, cell) do
-    sections =
-      Enum.map(notebook.sections, fn section ->
-        if section.id == section_id do
-          %{section | cells: List.insert_at(section.cells, index, cell)}
-        else
-          section
-        end
-      end)
-
-    %{notebook | sections: sections}
+    update_in(notebook, [Access.key(:sections), access_by_id(section_id)], fn section ->
+      %{section | cells: List.insert_at(section.cells, index, cell)}
+    end)
   end
 
   @doc """
@@ -114,9 +107,8 @@ defmodule Livebook.Notebook do
   """
   @spec delete_section(t(), Section.id()) :: t()
   def delete_section(notebook, section_id) do
-    sections = Enum.reject(notebook.sections, &(&1.id == section_id))
-
-    %{notebook | sections: sections}
+    {_, notebook} = pop_in(notebook, [Access.key(:sections), access_by_id(section_id)])
+    notebook
   end
 
   @doc """
@@ -124,12 +116,15 @@ defmodule Livebook.Notebook do
   """
   @spec delete_cell(t(), Cell.id()) :: t()
   def delete_cell(notebook, cell_id) do
-    sections =
-      Enum.map(notebook.sections, fn section ->
-        %{section | cells: Enum.reject(section.cells, &(&1.id == cell_id))}
-      end)
+    {_, notebook} =
+      pop_in(notebook, [
+        Access.key(:sections),
+        Access.all(),
+        Access.key(:cells),
+        access_by_id(cell_id)
+      ])
 
-    %{notebook | sections: sections}
+    notebook
   end
 
   @doc """
@@ -137,17 +132,11 @@ defmodule Livebook.Notebook do
   """
   @spec update_cell(t(), Cell.id(), (Cell.t() -> Cell.t())) :: t()
   def update_cell(notebook, cell_id, fun) do
-    sections =
-      Enum.map(notebook.sections, fn section ->
-        cells =
-          Enum.map(section.cells, fn cell ->
-            if cell.id == cell_id, do: fun.(cell), else: cell
-          end)
-
-        %{section | cells: cells}
-      end)
-
-    %{notebook | sections: sections}
+    update_in(
+      notebook,
+      [Access.key(:sections), Access.all(), Access.key(:cells), access_by_id(cell_id)],
+      fun
+    )
   end
 
   @doc """
@@ -155,12 +144,7 @@ defmodule Livebook.Notebook do
   """
   @spec update_section(t(), Section.id(), (Section.t() -> Section.t())) :: t()
   def update_section(notebook, section_id, fun) do
-    sections =
-      Enum.map(notebook.sections, fn section ->
-        if section.id == section_id, do: fun.(section), else: section
-      end)
-
-    %{notebook | sections: sections}
+    update_in(notebook, [Access.key(:sections), access_by_id(section_id)], fun)
   end
 
   @doc """

lib/livebook/utils.ex

@@ -59,4 +59,51 @@ defmodule Livebook.Utils do
   after
     Process.unregister(name)
   end
+
+  @doc """
+  Returns a function that accesses list items by the given id.
+
+  ## Examples
+
+      iex> list = [%{id: 1, name: "Jake"}, %{id: 2, name: "Amy"}]
+      iex> get_in(list, [Livebook.Utils.access_by_id(2), Access.key(:name)])
+      "Amy"
+
+      iex> list = [%{id: 1, name: "Jake"}, %{id: 2, name: "Amy"}]
+      iex> put_in(list, [Livebook.Utils.access_by_id(2), Access.key(:name)], "Amy Santiago")
+      [%{id: 1, name: "Jake"}, %{id: 2, name: "Amy Santiago"}]
+
+  An error is raised if the accessed structure is not a list:
+
+      iex> get_in(%{}, [Livebook.Utils.access_by_id(1)])
+      ** (RuntimeError) Livebook.Utils.access_by_id/1 expected a list, got: %{}
+  """
+  @spec access_by_id(term()) ::
+          Access.access_fun(data :: struct() | map(), current_value :: term())
+  def access_by_id(id) do
+    fn
+      :get, data, next when is_list(data) ->
+        data
+        |> Enum.find(fn item -> item.id == id end)
+        |> next.()
+
+      :get_and_update, data, next when is_list(data) ->
+        case Enum.split_while(data, fn item -> item.id != id end) do
+          {prev, [item | cons]} ->
+            case next.(item) do
+              {get, update} ->
+                {get, prev ++ [update | cons]}
+
+              :pop ->
+                {item, prev ++ cons}
+            end
+
+          _ ->
+            {nil, data}
+        end
+
+      _op, data, _next ->
+        raise "Livebook.Utils.access_by_id/1 expected a list, got: #{inspect(data)}"
+    end
+  end
 end

lib/livebook_web/live/session_live.ex

@@ -4,6 +4,7 @@ defmodule LivebookWeb.SessionLive do
   import LivebookWeb.UserHelpers
 
   alias Livebook.{SessionSupervisor, Session, Delta, Notebook, Runtime}
+  import Livebook.Utils, only: [access_by_id: 1]
 
   @impl true
   def mount(%{"id" => session_id}, %{"current_user_id" => current_user_id}, socket) do
@@ -553,7 +554,7 @@ defmodule LivebookWeb.SessionLive do
         new_socket =
           socket
           |> assign_private(data: data)
-          |> assign(data_view: data_to_view(data))
+          |> assign(data_view: update_data_view(socket.assigns.data_view, data, operation))
           |> after_operation(socket, operation)
           |> handle_actions(actions)
 
@@ -788,4 +789,31 @@ defmodule LivebookWeb.SessionLive do
       changed?: info.evaluation_digest != nil and info.digest != info.evaluation_digest
     }
   end
+
+  # Updates current data_view in response to an operation.
+  # In most cases we simply recompute data_view, but for the
+  # most common ones we only update the relevant parts.
+  defp update_data_view(data_view, data, operation) do
+    case operation do
+      {:report_cell_revision, _pid, _cell_id, _revision} ->
+        data_view
+
+      {:apply_cell_delta, _pid, cell_id, _delta, _revision} ->
+        update_cell_view(data_view, data, cell_id)
+
+      _ ->
+        data_to_view(data)
+    end
+  end
+
+  defp update_cell_view(data_view, data, cell_id) do
+    {:ok, cell, section} = Notebook.fetch_cell_and_section(data.notebook, cell_id)
+    cell_view = cell_to_view(cell, data)
+
+    put_in(
+      data_view,
+      [:section_views, access_by_id(section.id), :cell_views, access_by_id(cell.id)],
+      cell_view
+    )
+  end
 end

test/livebook/utils_test.exs

@@ -0,0 +1,7 @@
+defmodule Livebook.UtilsTest do
+  use ExUnit.Case, async: true
+
+  alias Livebook.Utils
+
+  doctest Utils
+end