Fix some smart cells getting cleared on notebook sync (#3161)

Author: jonatanklosko

lib/livebook/session/data_sync.ex

@@ -121,7 +121,11 @@ defmodule Livebook.Session.DataSync do
   defp key(%Cell.Code{} = cell),
     do: {:code_cell, cell.source, Map.take(cell, cell_syncable_fields(cell))}
 
-  defp key(%Cell.Smart{} = cell), do: {:smart_cell, cell.source, cell.attrs}
+  # Note that smart cell attributes in memory may be different than
+  # onces after notebook export end import (e.g. atoms are stringified).
+  # We could encode them here, but source should be enough as a key,
+  # since attrs generally determine the source.
+  defp key(%Cell.Smart{} = cell), do: {:smart_cell, cell.source}
 
   defp cell_syncable_fields(%Cell.Code{}),
     do: [:language, :reevaluate_automatically, :continue_on_error]
@@ -200,14 +204,52 @@ defmodule Livebook.Session.DataSync do
 
   defp do_annotate_moves([], _ops_map), do: []
 
-  defp annotate_updates([{:del, key1, item1}, {:ins, key2, item2} | ops])
+  defp annotate_updates(ops) do
+    # We could simply match on consecutive :del and :ins, however the
+    # diff may include a sequence of :del ops and then a sequence of
+    # :ins ops. They may not match in a zip way either, for example,
+    # the first :del may not match the insert, but the second :del
+    # does, so we should keep the first :del and make thse second one
+    # into an :upd.
+    #
+    # To do this, we go through the ops and accumulate consecutive
+    # :del ops, then once we get to an :ins, we traverse those :del
+    # ops in the same order, looking for a matching one. If a :del
+    # is matching, we flush an :upd and continue going thorugh ops.
+    # Otherwise, we flush the non-matching :del as is and try the
+    # next :del in the sequence. Finally, if there is no matching
+    # :del, we flush :ins as is and continue going through the ops.
+    do_annotate_updates(ops, [], [])
+  end
+
+  defp do_annotate_updates([{:del, key, item} | ops], pending_dels, acc) do
+    do_annotate_updates(ops, [{:del, key, item} | pending_dels], acc)
+  end
+
+  defp do_annotate_updates([{:ins, key, item} | ops], pending_dels, acc) do
+    find_matching_del(Enum.reverse(pending_dels), {:ins, key, item}, ops, acc)
+  end
+
+  defp do_annotate_updates([op | ops], pending_dels, acc) do
+    do_annotate_updates(ops, [], [op | pending_dels ++ acc])
+  end
+
+  defp do_annotate_updates([], pending_dels, acc) do
+    Enum.reverse(pending_dels ++ acc)
+  end
+
+  defp find_matching_del([{:del, key1, item1} | dels], {:ins, key2, item2}, ops, acc)
        when elem(key1, 0) == elem(key2, 0) do
-    # If the same block type, then we change it into an update.
-    [{:upd, key1, key2, item1, item2} | annotate_updates(ops)]
+    do_annotate_updates(ops, Enum.reverse(dels), [{:upd, key1, key2, item1, item2} | acc])
   end
 
-  defp annotate_updates([op | ops]), do: [op | annotate_updates(ops)]
-  defp annotate_updates([]), do: []
+  defp find_matching_del([del | dels], ins, ops, acc) do
+    find_matching_del(dels, ins, ops, [del | acc])
+  end
+
+  defp find_matching_del([], ins, ops, acc) do
+    do_annotate_updates(ops, [], [ins | acc])
+  end
 
   defp diff_to_data_ops(diff, data, client_id) do
     acc = %{

test/livebook/session/data_sync_test.exs

@@ -203,18 +203,20 @@ defmodule Livebook.Session.DataSyncTest do
       Notebook.new()
       | sections: [
           %{Section.new() | id: "s1", name: "Section 1"},
-          %{Section.new() | id: "s2", name: "Section 2", parent_id: "s1"}
+          %{Section.new() | id: "s2", name: "Section 2"},
+          %{Section.new() | id: "s3", name: "Section 3", parent_id: "s1"}
         ]
     }
 
     after_notebook = %{
       Notebook.new()
       | sections: [
-          %{Section.new() | id: "a-s2", name: "Section 2"}
+          %{Section.new() | id: "a-s2", name: "Section 2"},
+          %{Section.new() | id: "a-s3", name: "Section 3"}
         ]
     }
 
-    expected_ops = [{:unset_section_parent, @cid, "s2"}, {:delete_section, @cid, "s1", false}]
+    expected_ops = [{:unset_section_parent, @cid, "s3"}, {:delete_section, @cid, "s1", false}]
     assert_ops_and_result(before_notebook, after_notebook, expected_ops)
   end
 
@@ -361,6 +363,46 @@ defmodule Livebook.Session.DataSyncTest do
     assert_ops_and_result(before_notebook, after_notebook, expected_ops)
   end
 
+  test "consecutive code cells source updated" do
+    before_notebook = %{
+      Notebook.new()
+      | sections: [
+          %{
+            Section.new()
+            | id: "s1",
+              cells: [
+                %{Cell.new(:code) | id: "c1", source: "x = 1"},
+                %{Cell.new(:code) | id: "c2", source: "y = 2"}
+              ]
+          }
+        ]
+    }
+
+    after_notebook = %{
+      Notebook.new()
+      | sections: [
+          %{
+            Section.new()
+            | id: "a-s1",
+              cells: [
+                %{Cell.new(:code) | id: "a-c1", source: "x = 2"},
+                %{Cell.new(:code) | id: "a-c2", source: "y = 3"}
+              ]
+          }
+        ]
+    }
+
+    delta1 = Delta.diff("x = 1", "x = 2")
+    delta2 = Delta.diff("y = 2", "y = 3")
+
+    expected_ops = [
+      {:apply_cell_delta, @cid, "c1", :primary, delta1, nil, 0},
+      {:apply_cell_delta, @cid, "c2", :primary, delta2, nil, 0}
+    ]
+
+    assert_ops_and_result(before_notebook, after_notebook, expected_ops)
+  end
+
   test "code cell attrs updated" do
     before_notebook = %{
       Notebook.new()