feat: remove_es_module

core/modules/map.rs

@@ -248,6 +248,19 @@ impl ModuleMap {
     self.data.borrow().get_id(name, requested_module_type)
   }
 
+  // Removes a module or its alias from the module map.
+  pub(crate) fn remove_id(
+    &self,
+    name: &str,
+    requested_module_type: impl AsRef<RequestedModuleType>,
+    main: bool,
+  ) -> Option<ModuleId> {
+    self
+      .data
+      .borrow_mut()
+      .remove_id(name, requested_module_type, main)
+  }
+
   pub(crate) fn is_main_module(&self, global: &v8::Global<v8::Module>) -> bool {
     self.data.borrow().is_main_module(global)
   }

core/modules/module_map_data.rs

@@ -64,6 +64,16 @@ impl<T> ModuleNameTypeMap<T> {
     map.get(name)
   }
 
+  pub fn remove<Q>(&mut self, ty: &RequestedModuleType, name: &Q) -> Option<T>
+  where
+    ModuleName: std::borrow::Borrow<Q>,
+    Q: std::cmp::Eq + std::hash::Hash + std::fmt::Debug + ?Sized,
+  {
+    let index = self.map_index(ty)?;
+    let map = self.submaps.get_mut(index)?;
+    map.remove(name)
+  }
+
   pub fn insert(
     &mut self,
     module_type: &RequestedModuleType,
@@ -217,6 +227,39 @@ impl ModuleMapData {
     }
   }
 
+  // Removes a module or its alias from the module map.
+  pub fn remove_id(
+    &mut self,
+    name: &str,
+    requested_module_type: impl AsRef<RequestedModuleType>,
+    main: bool,
+  ) -> Option<ModuleId> {
+    let map = &mut self.by_name;
+    let first_symbolic_module =
+      map.remove(requested_module_type.as_ref(), name)?;
+    let mod_id = match first_symbolic_module {
+      SymbolicModule::Mod(mod_id) => mod_id,
+      SymbolicModule::Alias(mut mod_name) => loop {
+        let symbolic_module =
+          map.remove(requested_module_type.as_ref(), &mod_name)?;
+        match symbolic_module {
+          SymbolicModule::Alias(target) => {
+            debug_assert_ne!(mod_name, target);
+            mod_name = target;
+          }
+          SymbolicModule::Mod(mod_id) => break mod_id,
+        }
+      },
+    };
+
+    if main {
+      if let Some(main_id) = self.main_module_id.take() {
+        debug_assert_eq!(main_id, mod_id);
+      }
+    }
+    Some(mod_id)
+  }
+
   pub fn is_registered(
     &self,
     specifier: &str,

core/modules/tests.rs

@@ -1476,6 +1476,168 @@ fn main_and_side_module() {
     .unwrap();
 }
 
+#[test]
+#[should_panic]
+fn load_main_module_twice() {
+  let main_specifier = resolve_url("file:///main_module.js").unwrap();
+
+  let loader = StaticModuleLoader::with(
+    main_specifier.clone(),
+    ascii_str!("if (!import.meta.main) throw Error();"),
+  );
+
+  let mut runtime = JsRuntime::new(RuntimeOptions {
+    module_loader: Some(Rc::new(loader)),
+    ..Default::default()
+  });
+
+  let main_id_fut = runtime.load_main_es_module(&main_specifier).boxed_local();
+  let main_id = futures::executor::block_on(main_id_fut).unwrap();
+
+  #[allow(clippy::let_underscore_future)]
+  let _ = runtime.mod_evaluate(main_id);
+  futures::executor::block_on(runtime.run_event_loop(Default::default()))
+    .unwrap();
+
+  let main_id_fut = runtime.load_main_es_module(&main_specifier).boxed_local();
+  futures::executor::block_on(main_id_fut).unwrap_err();
+
+  #[allow(clippy::let_underscore_future)]
+  // Try loading the same module - should panic.
+  let _ = runtime.mod_evaluate(main_id);
+  futures::executor::block_on(runtime.run_event_loop(Default::default()))
+    .unwrap();
+}
+
+#[test]
+fn load_main_module_twice_with_remove() {
+  let main_specifier = resolve_url("file:///main_module.js").unwrap();
+
+  let loader = StaticModuleLoader::with(
+    main_specifier.clone(),
+    ascii_str!("if (!import.meta.main) throw Error();"),
+  );
+
+  let mut runtime = JsRuntime::new(RuntimeOptions {
+    module_loader: Some(Rc::new(loader)),
+    ..Default::default()
+  });
+  let module_map_rc = runtime.module_map();
+
+  let main_id_fut = runtime.load_main_es_module(&main_specifier).boxed_local();
+  let main_id = futures::executor::block_on(main_id_fut).unwrap();
+  // Ensure that the loaded module is marked as the main module.
+  // Main module is `main_id`.
+  assert!(module_map_rc.is_main_module_id(main_id));
+
+  #[allow(clippy::let_underscore_future)]
+  let _ = runtime.mod_evaluate(main_id);
+  futures::executor::block_on(runtime.run_event_loop(Default::default()))
+    .unwrap();
+
+  // And now remove id from module_map and try to load module again.
+  let main_id_fut =
+    runtime.remove_main_es_module(&main_specifier).boxed_local();
+
+  let old_main_id = futures::executor::block_on(main_id_fut).unwrap();
+  // Ensure that the module is removed by checking that it
+  // no longer has a main module id.
+  assert!(!module_map_rc.is_main_module_id(main_id));
+
+  let main_id_fut = runtime.load_main_es_module(&main_specifier).boxed_local();
+  let updated_main_id = futures::executor::block_on(main_id_fut).unwrap();
+  // Since the old main_id was removed, a new module is loaded
+  // with an incremented module id.
+  assert_eq!(updated_main_id, main_id + 1);
+  assert_eq!(old_main_id, main_id);
+  assert!(module_map_rc.is_main_module_id(updated_main_id));
+
+  // Evaluate the newly loaded module.
+  #[allow(clippy::let_underscore_future)]
+  let _ = runtime.mod_evaluate(updated_main_id);
+  futures::executor::block_on(runtime.run_event_loop(Default::default()))
+    .unwrap();
+}
+
+#[test]
+#[should_panic]
+fn remove_main_module_before_load() {
+  let side_specifier = resolve_url("file:///side_module.js").unwrap();
+
+  let loader = StaticModuleLoader::with(
+    side_specifier.clone(),
+    ascii_str!("if (!import.meta.main) throw Error();"),
+  );
+
+  let mut runtime = JsRuntime::new(RuntimeOptions {
+    module_loader: Some(Rc::new(loader)),
+    ..Default::default()
+  });
+
+  // Try to remove module id before loading module - should panic.
+  let side_id_fut =
+    runtime.remove_main_es_module(&side_specifier).boxed_local();
+  futures::executor::block_on(side_id_fut).unwrap();
+}
+
+#[test]
+fn load_side_module_twice_with_remove() {
+  let main_specifier = resolve_url("file:///main_module.js").unwrap();
+  let side_specifier = resolve_url("file:///side_module.js").unwrap();
+
+  let loader = StaticModuleLoader::new([
+    (
+      main_specifier.clone(),
+      ascii_str!("if (!import.meta.main) throw Error();"),
+    ),
+    (
+      side_specifier.clone(),
+      ascii_str!("if (import.meta.main) throw Error();"),
+    ),
+  ]);
+
+  let mut runtime = JsRuntime::new(RuntimeOptions {
+    module_loader: Some(Rc::new(loader)),
+    ..Default::default()
+  });
+  let module_map_rc = runtime.module_map();
+
+  let main_id_fut = runtime.load_main_es_module(&main_specifier).boxed_local();
+  let main_id = futures::executor::block_on(main_id_fut).unwrap();
+
+  let side_id_fut = runtime.load_side_es_module(&side_specifier).boxed_local();
+  let side_id = futures::executor::block_on(side_id_fut).unwrap();
+
+  // Main module is main_id
+  assert!(module_map_rc.is_main_module_id(main_id));
+
+  #[allow(clippy::let_underscore_future)]
+  let _ = runtime.mod_evaluate(main_id);
+  futures::executor::block_on(runtime.run_event_loop(Default::default()))
+    .unwrap();
+
+  let side_id_fut =
+    runtime.remove_side_es_module(&side_specifier).boxed_local();
+  let old_side_id = futures::executor::block_on(side_id_fut).unwrap();
+
+  // Main module has not changed
+  assert!(module_map_rc.is_main_module_id(main_id));
+
+  // Ensure that the module is removed by checking that it
+  // no longer has a main module id.
+  let side_id_fut = runtime.load_side_es_module(&side_specifier).boxed_local();
+  let updated_side_id = futures::executor::block_on(side_id_fut).unwrap();
+
+  assert_eq!(old_side_id, side_id);
+  assert_eq!(updated_side_id, side_id + 1);
+
+  // Evaluate the newly loaded module.
+  #[allow(clippy::let_underscore_future)]
+  let _ = runtime.mod_evaluate(updated_side_id);
+  futures::executor::block_on(runtime.run_event_loop(Default::default()))
+    .unwrap();
+}
+
 #[test]
 fn dynamic_imports_snapshot() {
   //TODO: Once the issue with the ModuleNamespaceEntryGetter is fixed, we can maintain a reference to the module

core/runtime/jsrealm.rs

@@ -16,7 +16,7 @@ use crate::ops::OpCtx;
 use crate::stats::RuntimeActivityTraces;
 use crate::tasks::V8TaskSpawnerFactory;
 use crate::web_timeout::WebTimers;
-use crate::GetErrorClassFn;
+use crate::{GetErrorClassFn, RequestedModuleType};
 use anyhow::Error;
 use futures::stream::StreamExt;
 use std::cell::Cell;
@@ -413,6 +413,26 @@ impl JsRealm {
     Ok(root_id)
   }
 
+  /// Removes the specified main module from the module map.
+  ///
+  /// This method will panic if the module is not found.
+  #[allow(clippy::unused_async)]
+  pub(crate) async fn remove_main_es_module(
+    &self,
+    specifier: &ModuleSpecifier,
+  ) -> Result<ModuleId, Error> {
+    let module_map_rc = self.0.module_map();
+    let removed_id = module_map_rc.remove_id(
+      specifier.as_str(),
+      RequestedModuleType::None,
+      true,
+    );
+
+    let module_id =
+      removed_id.expect("Module id was not found in the module map");
+    Ok(module_id)
+  }
+
   /// Asynchronously load specified ES module and all of its dependencies.
   ///
   /// This method is meant to be used when loading some utility code that
@@ -459,6 +479,26 @@ impl JsRealm {
     Ok(root_id)
   }
 
+  /// Removes the specified side module from the module map.
+  ///
+  /// This method will panic if the module is not found.
+  #[allow(clippy::unused_async)]
+  pub(crate) async fn remove_side_es_module(
+    &self,
+    specifier: &ModuleSpecifier,
+  ) -> Result<ModuleId, Error> {
+    let module_map_rc = self.0.module_map();
+    let removed_id = module_map_rc.remove_id(
+      specifier.as_str(),
+      RequestedModuleType::None,
+      false,
+    );
+
+    let module_id =
+      removed_id.expect("Module id was not found in the module map");
+    Ok(module_id)
+  }
+
   /// Load and evaluate an ES module provided the specifier and source code.
   ///
   /// The module should not have Top-Level Await (that is, it should be

core/runtime/jsruntime.rs

@@ -2364,6 +2364,18 @@ impl JsRuntime {
       .await
   }
 
+  /// Remove the specified main module from the module map.
+  ///
+  /// This method is useful when you need to update a previously loaded main
+  /// module and potentially replace it with a different version.
+  /// It returns the `ModuleId` of the old module that was removed from the module map.
+  pub async fn remove_main_es_module(
+    &mut self,
+    specifier: &ModuleSpecifier,
+  ) -> Result<ModuleId, Error> {
+    self.inner.main_realm.remove_main_es_module(specifier).await
+  }
+
   /// Asynchronously load specified ES module and all of its dependencies from the
   /// provided source.
   ///
@@ -2414,6 +2426,18 @@ impl JsRuntime {
       .await
   }
 
+  /// Remove the specified side module from the module map.
+  ///
+  /// This method is useful when you need to update a previously loaded side
+  /// module and potentially replace it with a different version.
+  /// It returns the `ModuleId` of the old module that was removed from the module map.
+  pub async fn remove_side_es_module(
+    &mut self,
+    specifier: &ModuleSpecifier,
+  ) -> Result<ModuleId, Error> {
+    self.inner.main_realm.remove_side_es_module(specifier).await
+  }
+
   /// Load and evaluate an ES module provided the specifier and source code.
   ///
   /// The module should not have Top-Level Await (that is, it should be