rolldown_plugin 0.1.1

use std::sync::Arc;

use crate::{
  HookBuildEndArgs, HookLoadArgs, HookLoadReturn, HookNoopReturn, HookResolveIdArgs,
  HookResolveIdReturn, HookTransformArgs, LoadPluginContext, PluginContext, PluginDriver,
  TransformPluginContext,
  pluginable::HookTransformAstReturn,
  types::{
    hook_resolve_id_skipped::HookResolveIdSkipped, hook_transform_ast_args::HookTransformAstArgs,
  },
};
use anyhow::{Context, Result};
use arcstr::ArcStr;
use rolldown_common::{
  ModuleInfo, ModuleType, NormalModule, PluginIdx, SharedNormalizedBundlerOptions,
  SourcemapChainElement, side_effects::HookSideEffects,
};
use rolldown_devtools::{action, trace_action};
use rolldown_error::CausedPlugin;
use rolldown_sourcemap::SourceMap;
use rolldown_utils::{IndexBitSet, unique_arc::UniqueArc};
use string_wizard::{MagicString, SourceMapOptions};
use tracing::{Instrument, debug_span};

impl PluginDriver {
  #[tracing::instrument(
    level = "trace",
    target = "rolldown_plugin::plugin_driver::build_hooks::total::build_start",
    skip_all
  )]
  pub async fn build_start(&self, opts: &SharedNormalizedBundlerOptions) -> HookNoopReturn {
    for (plugin_idx, plugin, ctx) in
      self.iter_plugin_with_context_by_order(&self.order_by_build_start_meta)
    {
      let start = self.start_timing();
      let result = plugin.call_build_start(ctx, &crate::HookBuildStartArgs { options: opts }).await;
      self.record_timing(plugin_idx, start);
      result.with_context(|| CausedPlugin::new(plugin.call_name()))?;
    }

    Ok(())
  }

  #[inline]
  fn get_resolve_call_skipped_plugins(
    specifier: &str,
    importer: Option<&str>,
    skipped_resolve_calls: Option<&Vec<Arc<HookResolveIdSkipped>>>,
    plugin_count: usize,
  ) -> IndexBitSet<PluginIdx> {
    let mut skipped_plugins = IndexBitSet::new(plugin_count);
    if let Some(skipped_resolve_calls) = skipped_resolve_calls {
      for skip_resolve_call in skipped_resolve_calls {
        if skip_resolve_call.specifier == specifier
          && skip_resolve_call.importer.as_deref() == importer
        {
          skipped_plugins.set_bit(skip_resolve_call.plugin_idx);
        }
      }
    }
    skipped_plugins
  }

  #[tracing::instrument(
    level = "trace",
    target = "rolldown_plugin::plugin_driver::build_hooks::total::resolve_id",
    skip_all
  )]
  pub async fn resolve_id(
    &self,
    args: &HookResolveIdArgs<'_>,
    skipped_resolve_calls: Option<&Vec<Arc<HookResolveIdSkipped>>>,
  ) -> HookResolveIdReturn {
    let skipped_plugins = Self::get_resolve_call_skipped_plugins(
      args.specifier,
      args.importer,
      skipped_resolve_calls,
      self.plugins.len(),
    );
    for (plugin_idx, plugin, ctx) in
      self.iter_plugin_with_context_by_order(&self.order_by_resolve_id_meta)
    {
      if skipped_plugins.has_bit(plugin_idx) {
        continue;
      }
      let ret = async {
        trace_action!(action::HookResolveIdCallStart {
          action: "HookResolveIdCallStart",
          importer: args.importer.map(ToString::to_string),
          module_request: args.specifier.to_string(),
          import_kind: args.kind.to_string(),
          plugin_name: plugin.call_name().to_string(),
          plugin_id: plugin_idx.raw(),
          trigger: "${hook_resolve_id_trigger}",
          call_id: "${call_id}",
        });
        let start = self.start_timing();
        let result = plugin
          .call_resolve_id(
            &skipped_resolve_calls.map_or_else(
              || ctx.clone(),
              |skipped_resolve_calls| {
                PluginContext::fork_with_skipped_resolve_calls(ctx, skipped_resolve_calls.clone())
              },
            ),
            args,
          )
          .await;
        self.record_timing(plugin_idx, start);
        if let Some(r) = result? {
          trace_action!(action::HookResolveIdCallEnd {
            action: "HookResolveIdCallEnd",
            resolved_id: Some(r.id.to_string()),
            is_external: r.external.map(|v| v.is_external()),
            plugin_name: plugin.call_name().to_string(),
            plugin_id: plugin_idx.raw(),
            trigger: "${hook_resolve_id_trigger}",
            call_id: "${call_id}",
          });
          anyhow::Ok(Some(r))
        } else {
          trace_action!(action::HookResolveIdCallEnd {
            action: "HookResolveIdCallEnd",
            resolved_id: None,
            is_external: None,
            plugin_name: plugin.call_name().to_string(),
            plugin_id: plugin_idx.raw(),
            trigger: "${hook_resolve_id_trigger}",
            call_id: "${call_id}",
          });
          Ok(None)
        }
      }
      .instrument(tracing::trace_span!(
        "HookResolveIdCall",
        CONTEXT_call_id = rolldown_utils::uuid::uuid_v4()
      ))
      .await
      .with_context(|| CausedPlugin::new(plugin.call_name()))?;
      if ret.is_some() {
        return Ok(ret);
      }
    }
    Ok(None)
  }

  // Only for rollup compatibility
  #[expect(deprecated)]
  #[tracing::instrument(
    level = "trace",
    target = "rolldown_plugin::plugin_driver::build_hooks::total::resolve_dynamic_import",
    skip_all
  )]
  pub async fn resolve_dynamic_import(
    &self,
    args: &HookResolveIdArgs<'_>,
    skipped_resolve_calls: Option<&Vec<Arc<HookResolveIdSkipped>>>,
  ) -> HookResolveIdReturn {
    let skipped_plugins = Self::get_resolve_call_skipped_plugins(
      args.specifier,
      args.importer,
      skipped_resolve_calls,
      self.plugins.len(),
    );
    for (plugin_idx, plugin, ctx) in
      self.iter_plugin_with_context_by_order(&self.order_by_resolve_dynamic_import_meta)
    {
      if skipped_plugins.has_bit(plugin_idx) {
        continue;
      }
      let start = self.start_timing();
      let result = plugin
        .call_resolve_dynamic_import(
          &skipped_resolve_calls.map_or_else(
            || ctx.clone(),
            |skipped_resolve_calls| {
              PluginContext::fork_with_skipped_resolve_calls(ctx, skipped_resolve_calls.clone())
            },
          ),
          args,
        )
        .await;
      self.record_timing(plugin_idx, start);
      if let Some(r) = result.with_context(|| CausedPlugin::new(plugin.call_name()))? {
        return Ok(Some(r));
      }
    }
    Ok(None)
  }

  #[tracing::instrument(
    level = "trace",
    target = "rolldown_plugin::plugin_driver::build_hooks::total::load",
    skip_all
  )]
  pub async fn load(&self, args: &HookLoadArgs<'_>) -> HookLoadReturn {
    for (plugin_idx, plugin, ctx) in
      self.iter_plugin_with_context_by_order(&self.order_by_load_meta)
    {
      let ret = async {
        trace_action!(action::HookLoadCallStart {
          action: "HookLoadCallStart",
          module_id: args.id.to_string(),
          plugin_name: plugin.call_name().to_string(),
          plugin_id: plugin_idx.raw(),
          call_id: "${call_id}",
        });
        let load_ctx = Arc::new(LoadPluginContext::new(ctx.clone(), args.module_idx));
        let start = self.start_timing();
        let result = plugin.call_load(load_ctx, args).await;
        self.record_timing(plugin_idx, start);
        if let Some(r) = result? {
          trace_action!(action::HookLoadCallEnd {
            action: "HookLoadCallEnd",
            module_id: args.id.to_string(),
            content: Some(r.code.to_string()),
            plugin_name: plugin.call_name().to_string(),
            plugin_id: plugin_idx.raw(),
            call_id: "${call_id}",
          });
          anyhow::Ok(Some(r))
        } else {
          trace_action!(action::HookLoadCallEnd {
            action: "HookLoadCallEnd",
            module_id: args.id.to_string(),
            content: None,
            plugin_name: plugin.call_name().to_string(),
            plugin_id: plugin_idx.raw(),
            call_id: "${call_id}",
          });
          Ok(None)
        }
      }
      .instrument(tracing::trace_span!(
        "HookLoadCall",
        CONTEXT_call_id = rolldown_utils::uuid::uuid_v4()
      ))
      .await
      .with_context(|| CausedPlugin::new(plugin.call_name()))?;
      if ret.is_some() {
        return Ok(ret);
      }
    }
    Ok(None)
  }

  #[tracing::instrument(
    level = "trace",
    target = "rolldown_plugin::plugin_driver::build_hooks::total::transform",
    skip_all
  )]
  #[expect(clippy::too_many_arguments)]
  pub async fn transform(
    &self,
    id: &str,
    module_idx: rolldown_common::ModuleIdx,
    original_code: String,
    sourcemap_chain: &mut Vec<SourcemapChainElement>,
    side_effects: &mut Option<HookSideEffects>,
    module_type: &mut ModuleType,
    magic_string_tx: Option<Arc<std::sync::mpsc::Sender<rolldown_common::SourceMapGenMsg>>>,
    code_changed_by_plugins: &mut Option<Vec<String>>,
  ) -> Result<String> {
    let mut code = original_code;
    let mut code_arc: Option<ArcStr> = None;
    let mut original_sourcemap_chain = std::mem::take(sourcemap_chain);
    let mut plugin_sourcemap_chain = UniqueArc::new(original_sourcemap_chain);
    for (plugin_idx, plugin, ctx) in
      self.iter_plugin_with_context_by_order(&self.order_by_transform_meta)
    {
      let call_id = tracing::enabled!(tracing::Level::TRACE).then(rolldown_utils::uuid::uuid_v4);
      let code_arc_ref = code_arc.get_or_insert_with(|| ArcStr::from(code.as_str()));

      trace_action!(action::HookTransformCallStart {
        action: "HookTransformCallStart",
        module_id: id.to_string(),
        content: code.clone(),
        plugin_name: plugin.call_name().to_string(),
        plugin_id: plugin_idx.raw(),
        call_id: call_id.clone().unwrap_or_default(),
      });
      let start = self.start_timing();
      let result = plugin
        .call_transform(
          Arc::new(TransformPluginContext::new(
            ctx.clone(),
            plugin_sourcemap_chain.weak_ref(),
            code_arc_ref.clone(),
            id.into(),
            module_idx,
            plugin_idx,
            magic_string_tx.clone(),
          )),
          &HookTransformArgs { id, code: code_arc_ref, module_type: &*module_type },
        )
        .await;
      self.record_timing(plugin_idx, start);
      if let Some(r) = result.with_context(|| CausedPlugin::new(plugin.call_name()))? {
        original_sourcemap_chain = plugin_sourcemap_chain.into_inner();
        if let Some(map) =
          self.normalize_transform_sourcemap(r.map.into_sourcemap(), id, &code, r.code.as_ref())
        {
          original_sourcemap_chain.push(SourcemapChainElement::Transform((plugin_idx, map)));
        }
        plugin_sourcemap_chain = UniqueArc::new(original_sourcemap_chain);
        if let Some(v) = r.side_effects {
          *side_effects = Some(v);
        }
        if let Some(v) = r.code {
          if let Some(changed_by) = code_changed_by_plugins {
            if v != code {
              changed_by.push(plugin.call_name().to_string());
            }
          }
          code = v;
          code_arc = None;
          trace_action!(action::HookTransformCallEnd {
            action: "HookTransformCallEnd",
            module_id: id.to_string(),
            content: Some(code.clone()),
            plugin_name: plugin.call_name().to_string(),
            plugin_id: plugin_idx.raw(),
            call_id: call_id.unwrap_or_default()
          });
        }
        if let Some(ty) = r.module_type {
          *module_type = ty;
        }
      } else {
        trace_action!(action::HookTransformCallEnd {
          action: "HookTransformCallEnd",
          module_id: id.to_string(),
          content: Some(code.clone()),
          plugin_name: plugin.call_name().to_string(),
          plugin_id: plugin_idx.raw(),
          call_id: call_id.unwrap_or_default()
        });
      }
    }
    *sourcemap_chain = plugin_sourcemap_chain.into_inner();
    Ok(code)
  }

  #[inline]
  fn normalize_transform_sourcemap(
    &self,
    map: Option<SourceMap>,
    id: &str,
    original_code: &str,
    code: Option<&String>,
  ) -> Option<SourceMap> {
    if let Some(mut map) = map {
      // If sourcemap  hasn't `sources`, using original id to fill it.
      let source = map.get_source(0);
      if source.is_none_or(|s| s.is_empty())
        || (map.get_sources().count() == 1 && (source.map(AsRef::as_ref) != Some(id)))
      {
        map.set_sources(vec![id]);
      }
      // If sourcemap hasn't `sourcesContent`, using original code to fill it.
      if map.get_source_content(0).is_none_or(|s| s.is_empty()) {
        map.set_source_contents(vec![Some(original_code)]);
      }
      Some(map)
    } else if let Some(code) = code {
      if original_code == code {
        None
      } else {
        // If sourcemap is empty and code has changed, need to create one remapping original code.
        let magic_string = MagicString::new(original_code);
        Some(magic_string.source_map(SourceMapOptions {
          hires: string_wizard::Hires::Boundary,
          include_content: true,
          source: id.into(),
        }))
      }
    } else {
      None
    }
  }

  #[tracing::instrument(
    level = "trace",
    target = "rolldown_plugin::plugin_driver::build_hooks::total::transform_ast",
    skip_all
  )]
  pub async fn transform_ast(&self, mut args: HookTransformAstArgs<'_>) -> HookTransformAstReturn {
    for (_, plugin, ctx) in
      self.iter_plugin_with_context_by_order(&self.order_by_transform_ast_meta)
    {
      // Reconstructing the struct is necessary because `args.ast` is moved and reassigned each iteration
      #[expect(clippy::unnecessary_struct_initialization)]
      let transform_args = HookTransformAstArgs {
        cwd: args.cwd,
        ast: args.ast,
        id: args.id,
        stable_id: args.stable_id,
        is_user_defined_entry: args.is_user_defined_entry,
        module_type: args.module_type,
      };
      args.ast = plugin
        .call_transform_ast(ctx, transform_args)
        .instrument(debug_span!("transform_ast_hook", plugin_name = plugin.call_name().as_ref()))
        .await
        .with_context(|| CausedPlugin::new(plugin.call_name()))?;
    }
    Ok(args.ast)
  }

  #[tracing::instrument(
    level = "trace",
    target = "rolldown_plugin::plugin_driver::build_hooks::total::module_parsed",
    skip_all
  )]
  pub async fn module_parsed(
    &self,
    module_info: Arc<ModuleInfo>,
    normal_module: &NormalModule,
  ) -> HookNoopReturn {
    for (plugin_idx, plugin, ctx) in
      self.iter_plugin_with_context_by_order(&self.order_by_module_parsed_meta)
    {
      let start = self.start_timing();
      let result = plugin.call_module_parsed(ctx, Arc::clone(&module_info), normal_module).await;
      self.record_timing(plugin_idx, start);
      result.with_context(|| CausedPlugin::new(plugin.call_name()))?;
    }
    Ok(())
  }

  #[tracing::instrument(
    level = "trace",
    target = "rolldown_plugin::plugin_driver::build_hooks::total::build_end",
    skip_all
  )]
  pub async fn build_end(&self, args: Option<&HookBuildEndArgs<'_>>) -> HookNoopReturn {
    for (plugin_idx, plugin, ctx) in
      self.iter_plugin_with_context_by_order(&self.order_by_build_end_meta)
    {
      let start = self.start_timing();
      let result = plugin.call_build_end(ctx, args).await;
      self.record_timing(plugin_idx, start);
      result.with_context(|| CausedPlugin::new(plugin.call_name()))?;
    }
    Ok(())
  }
}