From 89ca95be10b4ec1b65af8694085a0bd891b3e71b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jonatan=20K=C5=82osko?= <jonatanklosko@gmail.com>
Date: Tue, 18 Jan 2022 15:13:50 +0100
Subject: [PATCH] Reuse DOM elements when replacing frame outputs (#881)

* Reuse DOM elements when replacing frame outputs

* Keep less persistent indices
---
 lib/livebook_web/live/output.ex               |  2 +-
 .../live/output/frame_component.ex            | 45 ++++++++++++++++---
 .../live/session_live/cell_component.ex       |  1 +
 3 files changed, 40 insertions(+), 8 deletions(-)

diff --git a/lib/livebook_web/live/output.ex b/lib/livebook_web/live/output.ex
index 76d788add..371bf1e74 100644
--- a/lib/livebook_web/live/output.ex
+++ b/lib/livebook_web/live/output.ex
@@ -9,7 +9,7 @@ defmodule LivebookWeb.Output do
   def outputs(assigns) do
     ~H"""
     <%= for {idx, output} <- Enum.reverse(@outputs) do %>
-      <div class="max-w-full" id={"output-wrapper-#{idx}"}
+      <div class="max-w-full" id={"output-wrapper-#{@dom_id_map[idx] || idx}"}
         data-element="output"
         data-border={border?(output)}
         data-wrapper={wrapper?(output)}>
diff --git a/lib/livebook_web/live/output/frame_component.ex b/lib/livebook_web/live/output/frame_component.ex
index f639bfb53..06887db1c 100644
--- a/lib/livebook_web/live/output/frame_component.ex
+++ b/lib/livebook_web/live/output/frame_component.ex
@@ -3,7 +3,7 @@ defmodule LivebookWeb.Output.FrameComponent do
 
   @impl true
   def mount(socket) do
-    {:ok, assign(socket, counter: 0, empty: true)}
+    {:ok, assign(socket, counter: 0, output_count: 0, persistent_idx_map: %{})}
   end
 
   @impl true
@@ -15,7 +15,11 @@ defmodule LivebookWeb.Output.FrameComponent do
 
     socket =
       if socket.assigns.counter == 0 do
-        assign(socket, empty: outputs == [], counter: 1)
+        assign(socket,
+          counter: 1,
+          output_count: length(outputs),
+          persistent_idx_map: idx_map(outputs)
+        )
       else
         socket
       end
@@ -26,24 +30,50 @@ defmodule LivebookWeb.Output.FrameComponent do
           assign(socket, outputs: outputs)
 
         :replace ->
-          socket
-          |> assign(outputs: outputs, empty: outputs == [])
-          |> update(:counter, &(&1 + 1))
+          prev_output_count = socket.assigns.output_count
+          prev_persistent_idx_map = socket.assigns.persistent_idx_map
+
+          output_count = length(outputs)
+          persistent_idx_map = idx_map(outputs)
+
+          socket =
+            assign(socket,
+              outputs: outputs,
+              output_count: output_count,
+              persistent_idx_map: persistent_idx_map
+            )
+
+          less_outputs? = prev_output_count > output_count
+          appended_outputs? = prev_output_count > map_size(prev_persistent_idx_map)
+
+          # If there are outputs that we need to remove, increase the counter.
+          # Otherwise we reuse DOM element ids via persistent_idx_map
+          if less_outputs? or appended_outputs? do
+            update(socket, :counter, &(&1 + 1))
+          else
+            socket
+          end
 
         :append ->
           socket
-          |> assign(empty: false)
           |> update(:outputs, &(outputs ++ &1))
+          |> update(:output_count, &(length(outputs) + &1))
       end
 
     {:ok, socket}
   end
 
+  defp idx_map(outputs) do
+    outputs
+    |> Enum.with_index()
+    |> Map.new(fn {{output_idx, _}, idx} -> {output_idx, idx} end)
+  end
+
   @impl true
   def render(assigns) do
     ~H"""
     <div id={"frame-output-#{@id}"}>
-      <%= if @empty do %>
+      <%= if @output_count == 0 do %>
         <div class="text-gray-300 p-4 rounded-lg border border-gray-200">
           Empty output frame
         </div>
@@ -51,6 +81,7 @@ defmodule LivebookWeb.Output.FrameComponent do
         <div id={"frame-outputs-#{@id}-#{@counter}"} phx-update="append">
           <LivebookWeb.Output.outputs
             outputs={@outputs}
+            dom_id_map={@persistent_idx_map}
             socket={@socket}
             session_id={@session_id}
             input_values={@input_values}
diff --git a/lib/livebook_web/live/session_live/cell_component.ex b/lib/livebook_web/live/session_live/cell_component.ex
index d4fc55ba3..9194027d0 100644
--- a/lib/livebook_web/live/session_live/cell_component.ex
+++ b/lib/livebook_web/live/session_live/cell_component.ex
@@ -142,6 +142,7 @@ defmodule LivebookWeb.SessionLive.CellComponent do
         phx-update="append">
         <LivebookWeb.Output.outputs
           outputs={@cell_view.outputs}
+          dom_id_map={%{}}
           socket={@socket}
           session_id={@session_id}
           runtime={@runtime}