Speedup large notebooks: cache cell dependencies in notebook_to_js

fonsp · Aug 5, 2024 · a973643 · a973643
1 parent 4f8741c
commit a973643
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 15 deletions.
diff --git a/src/analysis/DependencyCache.jl b/src/analysis/DependencyCache.jl
@@ -47,8 +47,25 @@ function update_dependency_cache!(notebook::Notebook, topology::NotebookTopology
 
     if notebook._cached_cell_dependencies_source !== topology
         notebook._cached_cell_dependencies_source = topology
+
         for cell in all_cells(topology)
             update_dependency_cache!(cell, topology)
         end
+
+        notebook._cached_cell_dependencies = Dict{UUID,Dict{String,Any}}(
+            id => Dict{String,Any}(
+                "cell_id" => cell.cell_id,
+                "downstream_cells_map" => Dict{String,Vector{UUID}}(
+                    String(s) => cell_id.(r)
+                    for (s, r) in cell.cell_dependencies.downstream_cells_map
+                ),
+                "upstream_cells_map" => Dict{String,Vector{UUID}}(
+                    String(s) => cell_id.(r)
+                    for (s, r) in cell.cell_dependencies.upstream_cells_map
+                ),
+                "precedence_heuristic" => cell.cell_dependencies.precedence_heuristic,
+            )
+            for (id, cell) in notebook.cells_dict
+        )
     end
 end
diff --git a/src/notebook/Notebook.jl b/src/notebook/Notebook.jl
@@ -32,6 +32,7 @@ Base.@kwdef mutable struct Notebook
     notebook_id::UUID=uuid1()
     topology::NotebookTopology
     _cached_topological_order::TopologicalOrder
+    _cached_cell_dependencies::Dict{UUID,Dict{String,Any}}=Dict{UUID,Dict{String,Any}}()
     _cached_cell_dependencies_source::Union{Nothing,NotebookTopology}=nothing
 
     # buffer will contain all unfetched updates - must be big enough
@@ -122,7 +123,7 @@ end
 function PlutoDependencyExplorer.topological_order(notebook::Notebook)
     cached = notebook._cached_topological_order
 	if cached === nothing || cached.input_topology !== notebook.topology
-        topological_order(notebook.topology)
+        notebook._cached_topological_order = topological_order(notebook.topology)
 	else
 		cached
 	end

diff --git a/src/webserver/Dynamic.jl b/src/webserver/Dynamic.jl
@@ -115,20 +115,6 @@ function notebook_to_js(notebook::Notebook)
                 "metadata" => cell.metadata,
             )
         for (id, cell) in notebook.cells_dict),
-        "cell_dependencies" => Dict{UUID,Dict{String,Any}}(
-            id => Dict{String,Any}(
-                "cell_id" => cell.cell_id,
-                "downstream_cells_map" => Dict{String,Vector{UUID}}(
-                    String(s) => cell_id.(r)
-                    for (s, r) in cell.cell_dependencies.downstream_cells_map
-                ),
-                "upstream_cells_map" => Dict{String,Vector{UUID}}(
-                    String(s) => cell_id.(r)
-                    for (s, r) in cell.cell_dependencies.upstream_cells_map
-                ),
-                "precedence_heuristic" => cell.cell_dependencies.precedence_heuristic,
-            )
-        for (id, cell) in notebook.cells_dict),
         "cell_results" => Dict{UUID,Dict{String,Any}}(
             id => Dict{String,Any}(
                 "cell_id" => cell.cell_id,
@@ -167,6 +153,7 @@ function notebook_to_js(notebook::Notebook)
             )
         end,
         "status_tree" => Status.tojs(notebook.status_tree),
+        "cell_dependencies" => notebook._cached_cell_dependencies,
         "cell_execution_order" => cell_id.(collect(notebook._cached_topological_order)),
     )
 end