Optimise parent cells queueing (#1815)

lib/livebook/notebook.ex

@@ -441,16 +441,21 @@ defmodule Livebook.Notebook do
   before the given one.
 
   The cells are ordered starting from the most direct parent.
+
+  A list of cell ids can be be given, in which case parent cells
+  are computed for each cell and returned as a single list.
   """
-  @spec parent_cells_with_section(t(), Cell.id()) :: list({Cell.t(), Section.t()})
-  def parent_cells_with_section(notebook, cell_id) do
+  @spec parent_cells_with_section(t(), Cell.id() | list(Cell.id())) ::
+          list({Cell.t(), Section.t()})
+  def parent_cells_with_section(notebook, cell_ids) when is_list(cell_ids) do
+    graph = cell_dependency_graph(notebook)
+
     parent_cell_ids =
-      notebook
-      |> cell_dependency_graph()
-      |> Graph.find_path(cell_id, nil)
-      |> MapSet.new()
-      |> MapSet.delete(cell_id)
-      |> MapSet.delete(nil)
+      for cell_id <- cell_ids,
+          parent_cell_id <- Graph.find_path(graph, cell_id, nil),
+          parent_cell_id not in [cell_id, nil],
+          into: MapSet.new(),
+          do: parent_cell_id
 
     notebook
     |> cells_with_section()
@@ -458,6 +463,10 @@ defmodule Livebook.Notebook do
     |> Enum.reverse()
   end
 
+  def parent_cells_with_section(notebook, cell_id) do
+    parent_cells_with_section(notebook, [cell_id])
+  end
+
   @doc """
   Returns a list of cells (each with section) that go logically
   after the given one, and thus may depend on it.
@@ -494,6 +503,7 @@ defmodule Livebook.Notebook do
       cell should be included in the graph. If a cell is
       excluded, transitive parenthood still applies.
       By default all cells are included.
+
   """
   @spec cell_dependency_graph(t()) :: Graph.t(Cell.id())
   def cell_dependency_graph(notebook, opts \\ []) do

lib/livebook/session/data.ex

@@ -528,11 +528,13 @@ defmodule Livebook.Session.Data do
       end)
 
     if cell_ids != [] and length(cell_ids) == length(cells_with_section) do
-      cells_with_section
-      |> Enum.reduce(with_actions(data), fn {cell, section}, data_actions ->
-        data_actions
-        |> queue_prerequisite_cells_evaluation(cell.id)
-        |> queue_cell_evaluation(cell, section)
+      cell_ids = for {cell, _section} <- cells_with_section, do: cell.id
+
+      data
+      |> with_actions()
+      |> queue_prerequisite_cells_evaluation(cell_ids)
+      |> reduce(cells_with_section, fn data_actions, {cell, section} ->
+        queue_cell_evaluation(data_actions, cell, section)
       end)
       |> maybe_connect_runtime(data)
       |> update_validity_and_evaluation()
@@ -1295,9 +1297,7 @@ defmodule Livebook.Session.Data do
           cell = last_queued_cell(data, section),
           do: cell.id
 
-    reduce(data_actions, trailing_queued_cell_ids, fn data_actions, cell_id ->
-      queue_prerequisite_cells_evaluation(data_actions, cell_id)
-    end)
+    queue_prerequisite_cells_evaluation(data_actions, trailing_queued_cell_ids)
   end
 
   defp maybe_evaluate_queued({data, _} = data_actions) do
@@ -1483,18 +1483,17 @@ defmodule Livebook.Session.Data do
     )
   end
 
-  defp queue_prerequisite_cells_evaluation({data, _} = data_actions, cell_id) do
+  defp queue_prerequisite_cells_evaluation({data, _} = data_actions, cell_ids) do
     prerequisites_queue =
       data.notebook
-      |> Notebook.parent_cells_with_section(cell_id)
+      |> Notebook.parent_cells_with_section(cell_ids)
       |> Enum.filter(fn {cell, _section} ->
         info = data.cell_infos[cell.id]
         Cell.evaluable?(cell) and cell_outdated?(data, cell) and info.eval.status == :ready
       end)
       |> Enum.reverse()
 
-    data_actions
-    |> reduce(prerequisites_queue, fn data_actions, {cell, section} ->
+    reduce(data_actions, prerequisites_queue, fn data_actions, {cell, section} ->
       queue_cell_evaluation(data_actions, cell, section)
     end)
   end
@@ -1609,7 +1608,7 @@ defmodule Livebook.Session.Data do
 
     if evaluated? and reevaluate do
       data_actions
-      |> queue_prerequisite_cells_evaluation(cell.id)
+      |> queue_prerequisite_cells_evaluation([cell.id])
       |> queue_cell_evaluation(cell, section)
       |> maybe_evaluate_queued()
     else
@@ -2420,11 +2419,12 @@ defmodule Livebook.Session.Data do
         match?(%{status: :ready, validity: :stale, reevaluates_automatically: true}, info.eval)
       end)
 
+    cell_ids = for {cell, _section} <- cells_to_reevaluate, do: cell.id
+
     data_actions
+    |> queue_prerequisite_cells_evaluation(cell_ids)
     |> reduce(cells_to_reevaluate, fn data_actions, {cell, section} ->
-      data_actions
-      |> queue_prerequisite_cells_evaluation(cell.id)
-      |> queue_cell_evaluation(cell, section)
+      queue_cell_evaluation(data_actions, cell, section)
     end)
   end