rspack_core 0.101.5

use std::{
  fmt::Debug,
  ops::{BitOr, BitOrAssign},
  sync::Arc,
};

use async_trait::async_trait;
use rspack_cacheable::cacheable;
use rspack_collections::Identifier;
use rspack_error::Result;
use rspack_hash::{RspackHash, RspackHashDigest, RspackHasher};
use rspack_sources::{BoxSource, OriginalSource, RawStringSource, Source, SourceExt};
use rspack_util::source_map::SourceMapKind;
use tokio::sync::OnceCell;

use crate::{
  ChunkUkey, CodeGenerationResult, Compilation, Module, ModuleCodeGenerationContext,
  RuntimeCodeTemplate, RuntimeGlobals, RuntimeSpec, RuntimeTemplate, SourceType,
  runtime_mode::RuntimeMode,
};

pub struct RuntimeModuleGenerateContext<'a> {
  pub compilation: &'a Compilation,
  pub runtime_template: &'a RuntimeCodeTemplate,
}

pub fn runtime_module_owned_define_fields(
  compilation: &Compilation,
  chunk_ukey: &ChunkUkey,
) -> RuntimeGlobals {
  compilation
    .build_chunk_graph_artifact
    .chunk_graph
    .get_chunk_runtime_modules_iterable(chunk_ukey)
    .fold(RuntimeGlobals::default(), |fields, runtime_module_id| {
      let runtime_module = compilation
        .runtime_modules
        .get(runtime_module_id)
        .expect("should have runtime module");
      if runtime_module.get_custom_source().is_some()
        || runtime_module.get_constructor_name() == "RuntimeModuleFromJs"
      {
        return fields;
      }
      let runtime_requirements = runtime_module.runtime_requirements(compilation);
      let define_fields = runtime_requirements
        .define
        .difference(RuntimeGlobals::STARTUP | runtime_requirements.force_context);
      fields | define_fields
    })
}

#[derive(Debug, Default, Clone, Copy)]
pub struct RuntimeModuleRuntimeRequirements {
  pub dependencies: RuntimeGlobals,
  pub weak: RuntimeGlobals,
  pub define: RuntimeGlobals,
  pub force_context: RuntimeGlobals,
}

impl RuntimeModuleRuntimeRequirements {
  pub fn lexical_requirements(&self) -> RuntimeGlobals {
    self.dependencies | self.weak | self.define | self.force_context
  }
}

impl BitOr for RuntimeModuleRuntimeRequirements {
  type Output = Self;

  fn bitor(self, rhs: Self) -> Self::Output {
    Self {
      dependencies: self.dependencies | rhs.dependencies,
      weak: self.weak | rhs.weak,
      define: self.define | rhs.define,
      force_context: self.force_context | rhs.force_context,
    }
  }
}

impl BitOrAssign for RuntimeModuleRuntimeRequirements {
  fn bitor_assign(&mut self, rhs: Self) {
    self.dependencies.insert(rhs.dependencies);
    self.weak.insert(rhs.weak);
    self.define.insert(rhs.define);
    self.force_context.insert(rhs.force_context);
  }
}

#[cacheable]
#[derive(Debug, Default, Clone)]
pub struct RuntimeModuleCommon {
  pub id: Identifier,
  pub chunk: Option<ChunkUkey>,
  pub source_map_kind: SourceMapKind,
  pub custom_source: Option<String>,
  #[cacheable(with=rspack_cacheable::with::Skip)]
  pub cached_generated_code: Arc<OnceCell<BoxSource>>,
}

impl RuntimeModuleCommon {
  pub fn with_default(runtime_template: &RuntimeTemplate, type_name: &str) -> Self {
    Self {
      source_map_kind: SourceMapKind::empty(),
      custom_source: None,
      cached_generated_code: Default::default(),
      chunk: None,
      id: runtime_template.create_runtime_module_identifier(type_name),
    }
  }

  pub fn with_name(runtime_template: &RuntimeTemplate, name: &str) -> Self {
    Self {
      source_map_kind: SourceMapKind::empty(),
      custom_source: None,
      cached_generated_code: Default::default(),
      chunk: None,
      id: runtime_template.create_custom_runtime_module_identifier(name),
    }
  }

  pub fn attach(&mut self, chunk: ChunkUkey) {
    self.chunk = Some(chunk);
  }

  pub fn name(&self) -> Identifier {
    self.id
  }

  pub fn id(&self) -> &Identifier {
    &self.id
  }

  pub fn chunk(&self) -> Option<ChunkUkey> {
    self.chunk
  }

  pub fn set_custom_source(&mut self, source: String) {
    self.custom_source = Some(source);
  }

  pub fn get_custom_source(&self) -> Option<String> {
    self.custom_source.clone()
  }

  pub fn get_source_map_kind(&self) -> &SourceMapKind {
    &self.source_map_kind
  }

  pub fn set_source_map_kind(&mut self, source_map_kind: SourceMapKind) {
    self.source_map_kind = source_map_kind;
  }

  pub fn size(&self) -> f64 {
    self
      .cached_generated_code
      .get()
      .map_or(0f64, |cached_generated_code| {
        cached_generated_code.size() as f64
      })
  }
}

pub async fn runtime_module_get_generated_code(
  module: &dyn RuntimeModule,
  common: &RuntimeModuleCommon,
  compilation: &Compilation,
) -> Result<Arc<dyn Source>> {
  let result: Result<&BoxSource> = common
    .cached_generated_code
    .get_or_try_init(|| async {
      let runtime_template = compilation
        .runtime_template
        .create_runtime_module_code_template();
      let context = RuntimeModuleGenerateContext {
        compilation,
        runtime_template: &runtime_template,
      };
      let source_str = module.generate_with_custom(&context).await?;
      let source_map_kind = module.get_source_map_kind();
      Ok(if source_map_kind.enabled() {
        OriginalSource::new(source_str, module.identifier().as_str()).boxed()
      } else {
        RawStringSource::from(source_str).boxed()
      })
    })
    .await;
  let source = result?.clone();
  Ok(source)
}

pub async fn runtime_module_code_generation(
  module: &dyn RuntimeModule,
  common: &RuntimeModuleCommon,
  ctx: &mut ModuleCodeGenerationContext<'_>,
) -> Result<CodeGenerationResult> {
  let mut result = CodeGenerationResult::default();
  let source = runtime_module_get_generated_code(module, common, ctx.compilation).await?;
  result.add(SourceType::Runtime, source);
  Ok(result)
}

pub async fn runtime_module_get_runtime_hash(
  module: &dyn RuntimeModule,
  common: &RuntimeModuleCommon,
  compilation: &Compilation,
  _runtime: Option<&RuntimeSpec>,
) -> Result<RspackHashDigest> {
  let mut hasher = RspackHasher::from(&compilation.options.output);
  module.name().hash(&mut hasher);
  module.stage().hash(&mut hasher);
  if module.full_hash() || module.dependent_hash() {
    let runtime_template = compilation
      .runtime_template
      .create_runtime_module_code_template();
    let context = RuntimeModuleGenerateContext {
      compilation,
      runtime_template: &runtime_template,
    };
    RspackHash::hash(&module.generate_with_custom(&context).await?, &mut hasher);
  } else {
    use std::hash::Hash;

    runtime_module_get_generated_code(module, common, compilation)
      .await?
      .hash(&mut hasher);
  }
  Ok(hasher.digest(&compilation.options.output.hash_digest))
}

#[async_trait]
pub trait RuntimeModule:
  Module + CustomSourceRuntimeModule + AttachableRuntimeModule + NamedRuntimeModule
{
  fn stage(&self) -> RuntimeModuleStage {
    RuntimeModuleStage::Normal
  }
  fn full_hash(&self) -> bool {
    false
  }
  fn dependent_hash(&self) -> bool {
    false
  }
  // if wrap iife
  fn should_isolate(&self, runtime_mode: RuntimeMode) -> bool {
    matches!(runtime_mode, RuntimeMode::Webpack)
  }
  fn template(&self) -> Vec<(String, String)> {
    vec![]
  }
  async fn generate(
    &self,
    context: &RuntimeModuleGenerateContext<'_>,
  ) -> rspack_error::Result<String>;
  async fn generate_with_custom(
    &self,
    context: &RuntimeModuleGenerateContext<'_>,
  ) -> rspack_error::Result<String> {
    if let Some(custom_source) = self.get_custom_source() {
      Ok(custom_source)
    } else {
      self.generate(context).await
    }
  }
  fn runtime_requirements(&self, _compilation: &Compilation) -> RuntimeModuleRuntimeRequirements {
    Default::default()
  }
}

pub trait AttachableRuntimeModule {
  fn attach(&mut self, chunk: ChunkUkey);
}

pub trait NamedRuntimeModule {
  fn name(&self) -> Identifier;
}

pub trait CustomSourceRuntimeModule {
  fn set_custom_source(&mut self, source: String);
  fn get_custom_source(&self) -> Option<String>;
  fn get_constructor_name(&self) -> String;
}

pub type BoxRuntimeModule = Box<dyn RuntimeModule>;

#[cacheable]
#[derive(Debug, Default, Clone, PartialEq, Eq, PartialOrd, Ord)]
pub enum RuntimeModuleStage {
  #[default]
  Normal, // Runtime modules without any dependencies to other runtime modules
  Basic,   // Runtime modules with simple dependencies on other runtime modules
  Attach,  // Runtime modules which attach to handlers of other runtime modules
  Trigger, // Runtime modules which trigger actions on bootstrap
}

impl From<u32> for RuntimeModuleStage {
  fn from(stage: u32) -> Self {
    match stage {
      0 => RuntimeModuleStage::Normal,
      5 => RuntimeModuleStage::Basic,
      10 => RuntimeModuleStage::Attach,
      20 => RuntimeModuleStage::Trigger,
      _ => RuntimeModuleStage::Normal,
    }
  }
}

impl From<RuntimeModuleStage> for u32 {
  fn from(value: RuntimeModuleStage) -> Self {
    match value {
      RuntimeModuleStage::Normal => 0,
      RuntimeModuleStage::Basic => 5,
      RuntimeModuleStage::Attach => 10,
      RuntimeModuleStage::Trigger => 20,
    }
  }
}

impl RspackHash for RuntimeModuleStage {
  fn hash(&self, state: &mut RspackHasher) {
    let stage: u32 = self.clone().into();
    stage.hash(state);
  }
}

pub trait RuntimeModuleExt {
  fn boxed(self) -> Box<dyn RuntimeModule>;
}

impl<T: RuntimeModule + 'static> RuntimeModuleExt for T {
  fn boxed(self) -> Box<dyn RuntimeModule> {
    Box::new(self)
  }
}