rolldown_plugin_vite_resolve 0.1.1

use std::{
  ffi::OsString,
  fs,
  path::{self, Path, PathBuf},
  sync::Arc,
};

use oxc_resolver::{PackageJson, ResolveOptions, TsconfigDiscovery};
use rolldown_common::side_effects::HookSideEffects;
use rolldown_plugin::{HookResolveIdOutput, HookResolveIdReturn};
use rolldown_utils::{dashmap::FxDashSet, url::clean_url};
use rustc_hash::FxHashSet;
use sugar_path::SugarPath;

use crate::{
  builtin::BuiltinChecker,
  package_json_cache::{PackageJsonCache, PackageJsonWithOptionalPeerDependencies},
  utils::{
    BROWSER_EXTERNAL_ID, OPTIONAL_PEER_DEP_ID, can_externalize_file, get_extension,
    get_npm_package_name, is_bare_import, is_deep_import, normalize_path,
  },
};

#[derive(Debug, Clone)]
pub struct BaseOptions<'a> {
  pub main_fields: &'a Vec<String>,
  pub conditions: &'a Vec<String>,
  pub extensions: &'a Vec<String>,
  pub is_production: bool,
  pub try_index: bool,
  pub try_prefix: &'a Option<String>,
  pub as_src: bool,
  pub root: PathBuf,
  pub preserve_symlinks: bool,
  pub tsconfig_paths: bool,
  pub yarn_pnp: bool,
}

const ADDITIONAL_OPTIONS_FIELD_COUNT: u8 = 2;
const RESOLVER_COUNT: u8 = 2_u8.pow(ADDITIONAL_OPTIONS_FIELD_COUNT as u32);

const DEV_PROD_CONDITION: &str = "development|production";

#[derive(Debug, Clone)]
pub struct AdditionalOptions {
  is_require: bool,
  prefer_relative: bool,
}

impl AdditionalOptions {
  pub fn new(is_require: bool, prefer_relative: bool) -> Self {
    Self { is_require, prefer_relative }
  }

  fn as_bools(&self) -> [bool; ADDITIONAL_OPTIONS_FIELD_COUNT as usize] {
    [self.is_require, self.prefer_relative]
  }

  fn as_u8(&self) -> u8 {
    bools_to_u8(self.as_bools())
  }
}

impl From<[bool; RESOLVER_COUNT as usize]> for AdditionalOptions {
  fn from(value: [bool; RESOLVER_COUNT as usize]) -> Self {
    Self { is_require: value[0], prefer_relative: value[1] }
  }
}

impl From<u8> for AdditionalOptions {
  fn from(value: u8) -> Self {
    u8_to_bools(value).into()
  }
}

#[derive(Debug, Default)]
struct ResolverCaches {
  package_json: PackageJsonCache,
  importer_exists: FxDashSet<String>,
}

#[derive(Debug)]
pub struct Resolvers {
  resolvers: [Resolver; RESOLVER_COUNT as usize],
  external_resolver: Arc<Resolver>,
  resolver_caches: Arc<ResolverCaches>,
  lock: Arc<ResolverLock>,
}

impl Resolvers {
  pub fn new(
    base_options: &BaseOptions,
    external_conditions: &Vec<String>,
    builtin_checker: Arc<BuiltinChecker>,
  ) -> Self {
    let resolver_caches = Arc::new(ResolverCaches::default());

    let resolver_lock = Arc::new(ResolverLock::new());

    let base_resolver = oxc_resolver::Resolver::new(oxc_resolver::ResolveOptions {
      // NOTE: yarn_pnp option affects the underlying fs cache, so it should be consistent for all resolvers
      yarn_pnp: base_options.yarn_pnp,
      ..oxc_resolver::ResolveOptions::default()
    });

    let resolvers = (0..RESOLVER_COUNT)
      .map(|v| {
        Resolver::new(
          &base_resolver,
          base_options,
          v.into(),
          external_conditions,
          Arc::clone(&resolver_lock),
          Arc::clone(&builtin_checker),
          Arc::clone(&resolver_caches),
        )
      })
      .collect::<Vec<_>>()
      .try_into()
      .unwrap();

    let external_resolver = Resolver::new(
      &base_resolver,
      &BaseOptions {
        is_production: false,
        conditions: external_conditions,
        ..(*base_options).clone()
      },
      AdditionalOptions { is_require: false, prefer_relative: false },
      external_conditions,
      Arc::clone(&resolver_lock),
      Arc::clone(&builtin_checker),
      Arc::clone(&resolver_caches),
    );

    Self {
      resolvers,
      external_resolver: Arc::new(external_resolver),
      resolver_caches,
      lock: resolver_lock,
    }
  }

  pub fn get(&self, additional_options: AdditionalOptions) -> &Resolver {
    &self.resolvers[additional_options.as_u8() as usize]
  }

  pub fn get_for_external(&self) -> Arc<Resolver> {
    Arc::clone(&self.external_resolver)
  }

  pub fn get_nearest_package_json(&self, p: &str) -> Option<Arc<PackageJson>> {
    self.resolvers[0].get_nearest_package_json(p)
  }

  pub fn clear_cache(&self) {
    let _guard = self.lock.lock_for_clear();
    self.resolvers.iter().for_each(|v| v.clear_cache());
    self.external_resolver.clear_cache();
    self.resolver_caches.importer_exists.clear();
  }
}

fn get_resolve_options(
  base_options: &BaseOptions,
  additional_options: AdditionalOptions,
) -> oxc_resolver::ResolveOptions {
  let mut extensions = base_options.extensions.clone();
  extensions.push(".node".to_string());

  let main_fields = base_options.main_fields.clone();

  oxc_resolver::ResolveOptions {
    tsconfig: base_options.tsconfig_paths.then_some(TsconfigDiscovery::Auto),
    alias_fields: if base_options.main_fields.iter().any(|field| field == "browser") {
      vec![vec!["browser".to_string()]]
    } else {
      vec![]
    },
    condition_names: get_conditions(
      base_options.conditions,
      base_options.is_production,
      &additional_options,
    ),
    extensions,
    // The first alias is the original extension to make sure that `foo.js` will be resolved to `foo.js` if `foo.js` exists.
    extension_alias: vec![
      (".js".to_string(), vec![".js".to_string(), ".ts".to_string(), ".tsx".to_string()]),
      (".jsx".to_string(), vec![".jsx".to_string(), ".ts".to_string(), ".tsx".to_string()]),
      (".mjs".to_string(), vec![".mjs".to_string(), ".mts".to_string()]),
      (".cjs".to_string(), vec![".cjs".to_string(), ".cts".to_string()]),
    ],
    main_fields,
    main_files: if !base_options.try_index {
      vec![]
    } else if let Some(try_prefix) = &base_options.try_prefix {
      vec!["index".to_string(), format!("{try_prefix}index")]
    } else {
      vec!["index".to_string()]
    },
    node_path: false,
    prefer_relative: additional_options.prefer_relative,
    roots: if base_options.as_src { vec![base_options.root.clone()] } else { vec![] },
    symlinks: !base_options.preserve_symlinks,
    yarn_pnp: base_options.yarn_pnp,
    // This is not part of the spec, but required to align with rollup based vite.
    allow_package_exports_in_directory_resolve: true,
    ..Default::default()
  }
}

fn get_conditions(
  condition_names: &[String],
  is_production: bool,
  additional_options: &AdditionalOptions,
) -> Vec<String> {
  let mut conditions = condition_names
    .iter()
    .map(|c| {
      if c == DEV_PROD_CONDITION {
        if is_production { "production" } else { "development" }
      } else {
        c
      }
    })
    .map(|c| c.to_string())
    .collect::<Vec<_>>();

  if additional_options.is_require {
    conditions.push("require".to_string());
  } else {
    conditions.push("import".to_string());
  }

  conditions
}

fn bools_to_u8<const N: usize>(bools: [bool; N]) -> u8 {
  bools.iter().enumerate().map(|(i, v)| if *v { 1 << i } else { 0 }).sum()
}

fn u8_to_bools<const N: usize>(n: u8) -> [bool; N] {
  let mut ret = [false; N];
  ret.iter_mut().enumerate().for_each(|(i, v)| *v = n & (1 << i) != 0);
  ret
}

#[derive(Debug)]
pub struct Resolver {
  inner: oxc_resolver::Resolver,
  inner_for_external: oxc_resolver::Resolver,
  lock: Arc<ResolverLock>,
  built_in_checker: Arc<BuiltinChecker>,
  resolver_caches: Arc<ResolverCaches>,
  root: PathBuf,
  try_prefix: Option<String>,
}

impl Resolver {
  fn new(
    base_resolver: &oxc_resolver::Resolver,
    base_options: &BaseOptions,
    additional_options: AdditionalOptions,
    external_conditions: &[String],
    resolver_lock: Arc<ResolverLock>,
    built_in_checker: Arc<BuiltinChecker>,
    resolver_caches: Arc<ResolverCaches>,
  ) -> Self {
    let external_condition_names =
      get_conditions(external_conditions, base_options.is_production, &additional_options);
    let inner =
      base_resolver.clone_with_options(get_resolve_options(base_options, additional_options));
    let inner_for_external = inner.clone_with_options(ResolveOptions {
      condition_names: external_condition_names,
      ..inner.options().clone()
    });
    Self {
      inner,
      inner_for_external,
      lock: resolver_lock,
      built_in_checker,
      resolver_caches,
      root: base_options.root.clone(),
      try_prefix: base_options.try_prefix.clone(),
    }
  }

  pub fn resolve_raw(
    &self,
    specifier: &str,
    importer: Option<&str>,
    external: bool,
  ) -> Result<oxc_resolver::Resolution, oxc_resolver::ResolveError> {
    let _guard = self.lock.lock_for_update();

    let inner_resolver = if external { &self.inner_for_external } else { &self.inner };
    let result = if let Some(importer) = importer {
      // check if `is_absolute` to avoid extra `join` overhead
      if Path::new(importer).is_absolute() {
        inner_resolver.resolve_file(importer, specifier)
      } else {
        inner_resolver.resolve_file(self.root.join(importer), specifier)
      }
    } else {
      inner_resolver.resolve(&self.root, specifier)
    };

    match result {
      Err(
        oxc_resolver::ResolveError::Ignored(_)
        | oxc_resolver::ResolveError::TsconfigNotFound(_)
        | oxc_resolver::ResolveError::TsconfigSelfReference(_)
        | oxc_resolver::ResolveError::TsconfigCircularExtend(_)
        | oxc_resolver::ResolveError::PathNotSupported(_)
        | oxc_resolver::ResolveError::FailedToFindYarnPnpManifest(_),
      ) => return result,
      Ok(_) => return result,
      Err(_) => {}
    }

    // try with prefix if exists
    let Some(path_with_prefix) = self.try_prefix.as_ref().and_then(|try_prefix| {
      let mut path = Path::new(specifier).components();
      let filename = path.next_back()?;
      let path::Component::Normal(filename) = filename else {
        return None;
      };
      let mut filename_with_prefix = OsString::with_capacity(try_prefix.len() + filename.len());
      filename_with_prefix.push(try_prefix);
      filename_with_prefix.push(filename);

      Some(path.as_path().join(filename_with_prefix))
    }) else {
      return result;
    };
    let Some(path_with_prefix) = path_with_prefix.to_str() else {
      return result;
    };

    // this allows resolving `@pkg/pkg/foo.scss` to `@pkg/pkg/_foo.scss`, which is probably not allowed by sass's resolver
    // but that's an edge case so we ignore it here
    if let Some(importer) = importer {
      // check if `is_absolute` to avoid extra `join` overhead
      if Path::new(importer).is_absolute() {
        inner_resolver.resolve_file(importer, path_with_prefix)
      } else {
        inner_resolver.resolve_file(self.root.join(importer), path_with_prefix)
      }
    } else {
      inner_resolver.resolve(&self.root, path_with_prefix)
    }
  }

  pub fn normalize_oxc_resolver_result(
    &self,
    id: &str,
    importer: Option<&str>,
    dedupe: &FxHashSet<String>,
    legacy_inconsistent_cjs_interop: bool,
    result: &Result<oxc_resolver::Resolution, oxc_resolver::ResolveError>,
  ) -> Result<Option<HookResolveIdOutput>, oxc_resolver::ResolveError> {
    match result {
      Ok(result) => {
        let raw_path = result.full_path().to_str().unwrap().to_string();
        let path = normalize_path(&raw_path);

        let side_effects = result
          .package_json()
          .and_then(|pkg_json| {
            // the glob expr is based on parent path of package.json, which is package path
            // so we should use the relative path of the module to package path
            let module_path_relative_to_package =
              raw_path.as_path().relative(pkg_json.realpath.parent()?);
            self
              .resolver_caches
              .package_json
              .cached_package_json_side_effects(pkg_json)
              .check_side_effects_for(module_path_relative_to_package.to_str()?)
          })
          .map(
            |side_effects| {
              if side_effects { HookSideEffects::True } else { HookSideEffects::False }
            },
          );
        Ok(Some(HookResolveIdOutput {
          id: path.into(),
          side_effects,
          package_json_path: (!legacy_inconsistent_cjs_interop)
            .then(|| result.package_json().map(|pj| pj.realpath().to_str().unwrap().to_string()))
            .flatten(),
          ..Default::default()
        }))
      }
      Err(
        oxc_resolver::ResolveError::YarnPnpError(
          pnp::Error::BadSpecifier(_)
          | pnp::Error::MissingPeerDependency(_)
          | pnp::Error::UndeclaredDependency(_)
          | pnp::Error::MissingDependency(_),
        )
        | oxc_resolver::ResolveError::NotFound(_),
      ) => {
        // if import can't be found, check if it's an optional peer dep.
        // if so, we can resolve to a special id that errors only when imported.
        if is_bare_import(id)
          && !self.built_in_checker.is_builtin(id)
          && !id.contains('\0')
          && let Some(pkg_name) = get_npm_package_name(id)
        {
          let resolved_importer_dir = self
            .should_use_importer(id, importer, dedupe)
            .then(|| {
              importer.map(|importer| {
                Path::new(importer).parent().map(|i| i.to_str().unwrap()).unwrap_or(importer)
              })
            })
            .flatten();
          if resolved_importer_dir.is_some_and(|dir| dir != self.root.to_str().unwrap())
            && let Some(package_json) =
              self.get_nearest_package_json_optional_peer_deps(importer.unwrap())
            && package_json.optional_peer_dependencies.contains(pkg_name)
          {
            return Ok(Some(HookResolveIdOutput {
              id: format!("{OPTIONAL_PEER_DEP_ID}:{id}:{}", package_json.name).into(),
              ..Default::default()
            }));
          }
        }
        Ok(None)
      }
      Err(oxc_resolver::ResolveError::Ignored(_)) => Ok(Some(HookResolveIdOutput {
        id: arcstr::literal!(BROWSER_EXTERNAL_ID),
        ..Default::default()
      })),
      Err(err) => Err(err.to_owned()),
    }
  }

  pub fn get_nearest_package_json(&self, p: &str) -> Option<Arc<PackageJson>> {
    let _guard = self.lock.lock_for_update();

    let specifier = Path::new(p).absolutize();
    let Ok(result) = self.inner.resolve(
      /* actually this can be anything, as the specifier is absolute path */ &self.root,
      specifier.to_str().unwrap_or(p),
    ) else {
      // Errors when p is a virtual module
      return None;
    };

    result.package_json().map(Arc::clone)
  }

  fn get_nearest_package_json_optional_peer_deps(
    &self,
    p: &str,
  ) -> Option<Arc<PackageJsonWithOptionalPeerDependencies>> {
    let pj = self.get_nearest_package_json(p)?;
    Some(self.resolver_caches.package_json.cached_package_json_optional_peer_dep(&pj))
  }

  pub fn resolve_bare_import(
    &self,
    specifier: &str,
    importer: Option<&str>,
    external: bool,
    dedupe: &FxHashSet<String>,
    legacy_inconsistent_cjs_interop: bool,
  ) -> HookResolveIdReturn {
    let oxc_resolved_result = self.resolve_raw(
      specifier,
      if self.should_use_importer(specifier, importer, dedupe) { importer } else { None },
      external,
    );
    let resolved = self.normalize_oxc_resolver_result(
      specifier,
      importer,
      dedupe,
      legacy_inconsistent_cjs_interop,
      &oxc_resolved_result,
    )?;
    if let Some(mut resolved) = resolved {
      if !external || !can_externalize_file(&resolved.id) {
        return Ok(Some(resolved));
      }

      let id = specifier;
      let mut resolved_id = id;
      if is_deep_import(id)
        && get_extension(id) != get_extension(&resolved.id)
        && let Some(pkg_json) = oxc_resolved_result.unwrap().package_json()
      {
        let has_exports_field = pkg_json.exports().is_some();
        if !has_exports_field {
          // id date-fns/locale
          // resolve.id ...date-fns/esm/locale/index.js
          if let Some(index) = resolved.id.find(id) {
            resolved_id = &resolved.id[index..];
          }
        }
      }
      resolved.id = resolved_id.into();
      resolved.external = Some(true.into());

      return Ok(Some(resolved));
    }
    Ok(None)
  }

  pub fn clear_cache(&self) {
    self.inner.clear_cache();
  }

  fn should_use_importer(
    &self,
    specifier: &'_ str,
    importer: Option<&'_ str>,
    dedupe: &FxHashSet<String>,
  ) -> bool {
    if should_dedupe(specifier, dedupe) {
      return false;
    }

    if let Some(importer) = importer {
      let imp = Path::new(importer);
      if imp.is_absolute() {
        // css processing appends `*` for importer
        if importer.ends_with('*') {
          return true;
        }

        let importer = clean_url(importer);
        if self.resolver_caches.importer_exists.contains(importer) {
          return true;
        }

        let exists = fs::exists(importer).unwrap_or(false);
        if exists {
          self.resolver_caches.importer_exists.insert(importer.to_string());
        }
        return exists;
      }
    }
    false
  }
}

fn should_dedupe(specifier: &str, dedupe: &FxHashSet<String>) -> bool {
  if dedupe.is_empty() {
    return false;
  }

  let pkg_id = get_npm_package_name(specifier).unwrap_or(clean_url(specifier));
  dedupe.contains(pkg_id)
}

#[derive(Debug)]
pub struct ResolverLock(
  // we should use parking_lot instead of std::sync to avoid write starvation
  parking_lot::RwLock<()>,
);

impl ResolverLock {
  pub fn new() -> Self {
    Self(parking_lot::RwLock::new(()))
  }

  pub fn lock_for_update(&self) -> parking_lot::RwLockReadGuard<'_, ()> {
    self.0.read()
  }

  pub fn lock_for_clear(&self) -> parking_lot::RwLockWriteGuard<'_, ()> {
    self.0.write()
  }
}