fix: lazy compilation and partial bundling bug (#44)

.vscode/settings.json

@@ -1,4 +1,6 @@
 {
-    "editor.formatOnSave": true,
-    "editor.tabSize": 2
-}
+  "editor.formatOnSave": true,
+  "editor.tabSize": 2,
+  "editor.fontFamily": "Fira Code",
+  "editor.fontLigatures": true
+}

crates/compiler/src/build/mod.rs

@@ -331,19 +331,19 @@ impl Compiler {
 
     let mut module_graph = context.module_graph.write();
 
-    // check if the module already exists
-    if module_graph.has_module(&module.id) {
-      return false;
-    }
-
     // mark entry module
     if matches!(kind, ResolveKind::Entry) {
       module_graph.entries.insert(module.id.clone());
     }
 
     tracing::trace!("added module to the graph: {:?}", module.id);
 
-    module_graph.add_module(module);
+    // check if the module already exists
+    if module_graph.has_module(&module.id) {
+      module_graph.replace_module(module);
+    } else {
+      module_graph.add_module(module);
+    }
 
     true
   }
@@ -383,14 +383,18 @@ fn resolve_module(
   tracing::trace!("resolving module: {:?}", resolve_param);
 
   let resolve_module_id_result = Compiler::resolve_module_id(resolve_param, context)?;
-
-  // be care of dead lock, see https://github.com/Amanieu/parking_lot/issues/212
-  let module_graph = context.module_graph.read();
+  let mut module_graph = context.module_graph.write();
 
   let res = if module_graph.has_module(&resolve_module_id_result.module_id) {
     // the module has already been handled and it should not be handled twice
     ResolveModuleResult::Built(resolve_module_id_result.module_id)
   } else {
+    // insert a dummy module to the graph to prevent the module from being handled twice
+    module_graph.add_module(Compiler::create_module(
+      resolve_module_id_result.module_id.clone(),
+      false,
+      false,
+    ));
     ResolveModuleResult::Success(Box::new(ResolvedModuleInfo {
       module: Compiler::create_module(
         resolve_module_id_result.module_id.clone(),

crates/compiler/src/generate/partial_bundling.rs

@@ -3,7 +3,11 @@ use std::sync::Arc;
 use farmfe_core::{
   context::CompilationContext,
   error::CompilationError,
-  module::module_group::{ModuleGroup, ModuleGroupGraph},
+  hashbrown::HashSet,
+  module::{
+    module_group::{ModuleGroup, ModuleGroupGraph},
+    ModuleId,
+  },
   parking_lot::Mutex,
   plugin::PluginHookContext,
   resource::{resource_pot::ResourcePot, resource_pot_map::ResourcePotMap},
@@ -61,35 +65,54 @@ fn generate_resource_pot_map(
 ) -> farmfe_core::error::Result<ResourcePotMap> {
   tracing::trace!("Starting generate_resource_pot_map");
 
-  let resource_pot_map = Mutex::new(ResourcePotMap::new());
+  let mut modules = HashSet::new();
 
   for g in module_group_graph.module_groups_mut() {
-    let resources_pots = call_partial_bundling_hook(g, context, hook_context)?;
+    modules.extend(g.modules().clone());
+  }
+
+  let resources_pots =
+    call_partial_bundling_hook(&modules.into_iter().collect(), context, hook_context)?;
 
-    let mut resource_pot_map = resource_pot_map.lock();
+  let mut resource_pot_map = ResourcePotMap::new();
+  let module_graph = context.module_graph.read();
 
-    for resource_pot in resources_pots {
-      g.add_resource_pot(resource_pot.id.clone());
-      resource_pot_map.add_resource_pot(resource_pot);
+  for mut resource_pot in resources_pots {
+    let mut module_groups = HashSet::new();
+
+    for module_id in resource_pot.modules() {
+      let module = module_graph.module(module_id).unwrap();
+      module_groups.extend(module.module_groups.clone());
     }
+
+    resource_pot.module_groups = module_groups.clone();
+
+    for module_group_id in module_groups {
+      let module_group = module_group_graph
+        .module_group_mut(&module_group_id)
+        .unwrap();
+      module_group.add_resource_pot(resource_pot.id.clone());
+    }
+
+    resource_pot_map.add_resource_pot(resource_pot);
   }
 
   tracing::trace!("generate_resource_pot_map finished");
 
-  Ok(resource_pot_map.into_inner())
+  Ok(resource_pot_map)
 }
 
 #[tracing::instrument(skip_all)]
 pub fn call_partial_bundling_hook(
-  g: &mut ModuleGroup,
+  modules: &Vec<ModuleId>,
   context: &Arc<CompilationContext>,
   hook_context: &PluginHookContext,
 ) -> farmfe_core::error::Result<Vec<ResourcePot>> {
   tracing::trace!("Starting call_partial_bundling_hook");
 
   let res = context
     .plugin_driver
-    .partial_bundling(g, context, hook_context)?
+    .partial_bundling(modules, context, hook_context)?
     .ok_or(CompilationError::PluginHookResultCheckError {
       hook_name: "partial_bundling".to_string(),
     })?;

crates/compiler/src/lib.rs

@@ -1,4 +1,6 @@
 #![deny(clippy::all)]
+#![allow(clippy::needless_collect)]
+#![allow(clippy::ptr_arg)]
 #![feature(box_patterns)]
 
 use std::sync::Arc;

crates/compiler/src/update/diff_and_patch_module_graph.rs

@@ -108,7 +108,7 @@ pub fn patch_module_graph(
       m
     };
 
-    module_graph.replace_module(&updated, module);
+    module_graph.replace_module(module);
   }
 
   removed_modules

crates/compiler/src/update/mod.rs

@@ -7,10 +7,14 @@ use farmfe_core::{
   context::CompilationContext,
   error::CompilationError,
   hashbrown::HashSet,
-  module::{module_graph::ModuleGraphEdge, Module, ModuleId},
+  module::{
+    module_graph::ModuleGraphEdge, module_group::ModuleGroupId, Module, ModuleId, ModuleType,
+  },
   plugin::{PluginResolveHookParam, ResolveKind},
   rayon::ThreadPool,
+  resource::ResourceType,
 };
+use farmfe_plugin_html::get_dynamic_resources_map;
 use farmfe_toolkit::tracing;
 
 use crate::{
@@ -43,6 +47,7 @@ pub struct UpdateResult {
   /// This code string should be returned to the client side as MIME type `application/javascript`
   pub resources: String,
   pub boundaries: HashMap<String, Vec<Vec<String>>>,
+  pub dynamic_resources_map: Option<HashMap<ModuleId, Vec<(String, ResourceType)>>>,
 }
 
 #[derive(Debug, Clone)]
@@ -111,23 +116,23 @@ impl Compiler {
 
     self.optimize_update_module_graph(&update_context).unwrap();
 
+    let previous_module_groups = {
+      let module_group_graph = self.context.module_group_graph.read();
+      module_group_graph
+        .module_groups()
+        .into_iter()
+        .map(|m| m.id.clone())
+        .collect::<HashSet<_>>()
+    };
+
     let (affected_module_groups, updated_module_ids, diff_result) =
       self.diff_and_patch_context(paths, &update_context);
-    let cloned_updated_module_ids = updated_module_ids.clone();
-
-    let cloned_context = self.context.clone();
-    std::thread::spawn(move || {
-      // TODO: manage a task queue, and run the tasks in sequence
-      regenerate_resources_for_affected_module_groups(
-        affected_module_groups,
-        &cloned_updated_module_ids,
-        &cloned_context,
-      )
-      .unwrap();
-
-      finalize_resources(&cloned_context).unwrap();
-    });
 
+    let dynamic_resources_map = self.regenerate_resources(
+      affected_module_groups,
+      previous_module_groups,
+      &updated_module_ids,
+    );
     // TODO1: only regenerate the resources for script modules.
     // TODO2: should reload when html change
     // TODO3: cover it with tests
@@ -142,6 +147,7 @@ impl Compiler {
       removed_module_ids: diff_result.removed_modules.into_iter().collect(),
       resources,
       boundaries,
+      dynamic_resources_map,
     })
   }
 
@@ -237,10 +243,10 @@ impl Compiler {
     let mut update_module_graph = update_context.module_graph.write();
 
     if update_module_graph.has_module(&module.id) {
-      return;
+      update_module_graph.replace_module(module);
+    } else {
+      update_module_graph.add_module(module);
     }
-
-    update_module_graph.add_module(module);
   }
 
   fn add_edge_to_update_module_graph(
@@ -318,6 +324,72 @@ impl Compiler {
 
     (affected_module_groups, start_points, diff_result)
   }
+
+  fn regenerate_resources(
+    &self,
+    affected_module_groups: HashSet<ModuleGroupId>,
+    previous_module_groups: HashSet<ModuleGroupId>,
+    updated_module_ids: &Vec<ModuleId>,
+  ) -> Option<HashMap<ModuleId, Vec<(String, ResourceType)>>> {
+    let mut dynamic_resources_map = None;
+    let cloned_updated_module_ids = updated_module_ids.clone();
+
+    let cloned_context = self.context.clone();
+
+    // if there are new module groups, we should run the tasks synchronously
+    if affected_module_groups
+      .iter()
+      .any(|ag| !previous_module_groups.contains(ag))
+    {
+      regenerate_resources_for_affected_module_groups(
+        affected_module_groups,
+        &cloned_updated_module_ids,
+        &cloned_context,
+      )
+      .unwrap();
+
+      finalize_resources(&cloned_context).unwrap();
+      let module_group_graph = self.context.module_group_graph.read();
+      let resource_pot_map = self.context.resource_pot_map.read();
+      let resources_map = self.context.resources_map.lock();
+      let module_graph = self.context.module_graph.read();
+      let html_entries_ids = module_graph
+        .entries
+        .clone()
+        .into_iter()
+        .filter(|m| {
+          let module = module_graph.module(m).unwrap();
+          matches!(module.module_type, ModuleType::Html)
+        })
+        .collect::<Vec<_>>();
+      let mut dynamic_resources = HashMap::new();
+
+      for html_entry_id in html_entries_ids {
+        dynamic_resources.extend(get_dynamic_resources_map(
+          &*module_group_graph,
+          &html_entry_id,
+          &*resource_pot_map,
+          &*resources_map,
+        ));
+      }
+
+      dynamic_resources_map = Some(dynamic_resources);
+    } else {
+      std::thread::spawn(move || {
+        // TODO: manage a task queue, and run the tasks in sequence
+        regenerate_resources_for_affected_module_groups(
+          affected_module_groups,
+          &cloned_updated_module_ids,
+          &cloned_context,
+        )
+        .unwrap();
+
+        finalize_resources(&cloned_context).unwrap();
+      });
+    }
+
+    dynamic_resources_map
+  }
 }
 
 /// Similar to [crate::build::resolve_module], but the resolved module may be existed in both context and update_context
@@ -349,12 +421,18 @@ fn resolve_module(
     ));
   }
 
-  let update_module_graph = update_context.module_graph.read();
+  let mut update_module_graph = update_context.module_graph.write();
   if update_module_graph.has_module(&resolve_module_id_result.module_id) {
     return Ok(ResolveModuleResult::ExistingWhenUpdate(
       resolve_module_id_result.module_id,
     ));
   }
+  // just a placeholder module, it will be replaced by the real module later
+  update_module_graph.add_module(Compiler::create_module(
+    resolve_module_id_result.module_id.clone(),
+    false,
+    false,
+  ));
 
   Ok(ResolveModuleResult::Success(Box::new(ResolvedModuleInfo {
     module: Compiler::create_module(

crates/compiler/src/update/patch_module_group_graph.rs

@@ -198,6 +198,7 @@ pub fn patch_module_group_graph(
         let module_group_id = added_module_id.clone();
         let module_group = ModuleGroup::new(module_group_id.clone());
         module_group_graph.add_module_group(module_group);
+        affected_module_groups.insert(module_group_id.clone());
 
         let module_group_ids = {
           let module = module_graph
@@ -276,10 +277,26 @@ pub fn patch_module_group_graph(
     }
   }
 
-  affected_module_groups
+  let affected_module_groups = affected_module_groups
     .into_iter()
     .filter(|g_id| module_group_graph.has(g_id))
-    .collect()
+    .collect::<HashSet<_>>();
+
+  let mut final_affected_module_groups = HashSet::new();
+
+  for module_group_id in affected_module_groups {
+    let module_group = module_group_graph.module_group(&module_group_id).unwrap();
+
+    for module_id in module_group.modules() {
+      let module = module_graph.module(module_id).unwrap();
+
+      for module_group_id in &module.module_groups {
+        final_affected_module_groups.insert(module_group_id.clone());
+      }
+    }
+  }
+
+  final_affected_module_groups
 }
 
 #[cfg(test)]
@@ -420,7 +437,7 @@ mod tests {
     );
     assert_eq!(
       affected_groups,
-      HashSet::from(["A".into(), "B".into(), "F".into()])
+      HashSet::from(["A".into(), "B".into(), "F".into(), "D".into()])
     );
 
     let update_module_group_graph = module_group_graph_from_entries(