macroforge_ts_syn 0.1.79

//! # Import Registry
//!
//! Unified import registry built during IR lowering, used throughout macro expansion.
//! All insertions are idempotent — duplicates are impossible by design.
//!
//! This module lives in `macroforge_ts_syn` so that [`TsStream`](crate::TsStream) can
//! register imports directly instead of pushing raw patches.

use std::cell::RefCell;
use std::collections::HashMap;

use indexmap::IndexMap;

#[cfg(feature = "swc")]
use swc_core::ecma::ast::{
    ImportDecl, ImportSpecifier, Module, ModuleDecl, ModuleExportName, ModuleItem,
};

// ============================================================================
// Core types
// ============================================================================

/// A single import from the user's source file.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct SourceImport {
    /// The module this identifier is imported from (e.g., `"effect"`, `"./types"`)
    pub source_module: String,
    /// The original exported name, if the import uses an alias.
    /// For `import { Option as EffectOption }`, this is `Some("Option")`.
    pub original_name: Option<String>,
    /// Whether this import is type-only (`import type { ... }` or `import { type ... }`).
    pub is_type_only: bool,
}

/// Serializable version of [`SourceImport`] for cross-process transfer.
/// Used in [`MacroContextIR`] to pass source imports to external macros.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct SourceImportEntry {
    pub local_name: String,
    pub source_module: String,
    pub original_name: Option<String>,
    pub is_type_only: bool,
}

/// An import generated by macro expansion (to be emitted at the top of the file).
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct GeneratedImport {
    /// The local name used in the generated code.
    pub local_name: String,
    /// The module to import from.
    pub source_module: String,
    /// The original name if different from local (for `import { X as Y }` patterns).
    pub original_name: Option<String>,
    /// Whether this should be emitted as `import type`.
    pub is_type_only: bool,
}

/// Unified import registry. Built during IR lowering, used throughout expansion.
///
/// All insertions are idempotent — duplicates are impossible by design.
/// The registry has three sections:
/// 1. **Source imports** — pre-populated from the user's AST during lowering. Read-only after creation.
/// 2. **Config imports** — from `macroforge.config.ts` import statements. Name → module.
/// 3. **Generated imports** — accumulated during macro expansion via `request_*` methods.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct ImportRegistry {
    /// Source file imports. Key: local_name → entry.
    /// Pre-populated from AST during lowering. Never modified after initial construction.
    source_imports: HashMap<String, SourceImport>,

    /// Config file imports (macroforge.config.ts `import` statements). Name → module.
    pub config_imports: HashMap<String, String>,

    /// Generated imports accumulated during macro expansion.
    /// Key: local_name. Insertion order preserved for deterministic output.
    generated: IndexMap<String, GeneratedImport>,
}

impl Default for ImportRegistry {
    fn default() -> Self {
        Self::new()
    }
}

impl ImportRegistry {
    /// Create an empty registry.
    pub fn new() -> Self {
        Self {
            source_imports: HashMap::new(),
            config_imports: HashMap::new(),
            generated: IndexMap::new(),
        }
    }

    /// Build from AST — extracts all import declarations from the module.
    /// Called during IR lowering in `prepare_expansion_context`.
    #[cfg(feature = "swc")]
    pub fn from_module(module: &Module, source: &str) -> Self {
        let mut source_imports = HashMap::new();

        // Also extract imports from JSDoc `@import macro` comments
        source_imports.extend(collect_macro_import_comments(source).into_iter().map(
            |(name, module_src)| {
                (
                    name,
                    SourceImport {
                        source_module: module_src,
                        original_name: None,
                        is_type_only: false,
                    },
                )
            },
        ));

        for item in &module.body {
            if let ModuleItem::ModuleDecl(ModuleDecl::Import(ImportDecl {
                specifiers,
                src,
                type_only: import_type_only,
                ..
            })) = item
            {
                let module_source = src.value.to_string_lossy().to_string();

                for specifier in specifiers {
                    match specifier {
                        ImportSpecifier::Named(named) => {
                            let local_name = named.local.sym.to_string();
                            let is_type_only = *import_type_only || named.is_type_only;

                            let original_name = named.imported.as_ref().and_then(|imported| {
                                let orig = match imported {
                                    ModuleExportName::Ident(ident) => ident.sym.to_string(),
                                    ModuleExportName::Str(s) => {
                                        String::from_utf8_lossy(s.value.as_bytes()).to_string()
                                    }
                                };
                                if orig != local_name { Some(orig) } else { None }
                            });

                            source_imports.insert(
                                local_name,
                                SourceImport {
                                    source_module: module_source.clone(),
                                    original_name,
                                    is_type_only,
                                },
                            );
                        }
                        ImportSpecifier::Default(default) => {
                            let local_name = default.local.sym.to_string();
                            source_imports.insert(
                                local_name,
                                SourceImport {
                                    source_module: module_source.clone(),
                                    original_name: None,
                                    is_type_only: *import_type_only,
                                },
                            );
                        }
                        ImportSpecifier::Namespace(ns) => {
                            let local_name = ns.local.sym.to_string();
                            source_imports.insert(
                                local_name,
                                SourceImport {
                                    source_module: module_source.clone(),
                                    original_name: None,
                                    is_type_only: *import_type_only,
                                },
                            );
                        }
                    }
                }
            }
        }

        Self {
            source_imports,
            config_imports: HashMap::new(),
            generated: IndexMap::new(),
        }
    }

    /// Check if an identifier is available (in source OR already generated).
    pub fn is_available(&self, name: &str) -> bool {
        self.source_imports.contains_key(name) || self.generated.contains_key(name)
    }

    /// Check if a source import is type-only.
    /// Returns false if the import is not found.
    pub fn is_type_only(&self, name: &str) -> bool {
        self.source_imports
            .get(name)
            .map(|si| si.is_type_only)
            .unwrap_or(false)
    }

    /// Get import module for a source import.
    pub fn get_source(&self, name: &str) -> Option<&str> {
        self.source_imports
            .get(name)
            .map(|si| si.source_module.as_str())
    }

    /// Resolve alias to original name.
    /// For `import { Option as EffectOption }`, `resolve_alias("EffectOption")` → `Some("Option")`.
    pub fn resolve_alias(&self, name: &str) -> Option<&str> {
        self.source_imports
            .get(name)
            .and_then(|si| si.original_name.as_deref())
    }

    /// Get all source imports as a reference to the internal HashMap.
    pub fn source_map(&self) -> &HashMap<String, SourceImport> {
        &self.source_imports
    }

    /// Get source imports as a flat name→module HashMap (for backward-compat callers).
    pub fn source_modules(&self) -> HashMap<String, String> {
        self.source_imports
            .iter()
            .map(|(name, si)| (name.clone(), si.source_module.clone()))
            .collect()
    }

    /// Get import aliases as a flat local→original HashMap (for backward-compat callers).
    pub fn aliases(&self) -> HashMap<String, String> {
        self.source_imports
            .iter()
            .filter_map(|(name, si)| {
                si.original_name
                    .as_ref()
                    .map(|orig| (name.clone(), orig.clone()))
            })
            .collect()
    }

    /// Export source imports as serializable entries for cross-process transfer.
    pub fn source_import_entries(&self) -> Vec<SourceImportEntry> {
        self.source_imports
            .iter()
            .map(|(name, si)| SourceImportEntry {
                local_name: name.clone(),
                source_module: si.source_module.clone(),
                original_name: si.original_name.clone(),
                is_type_only: si.is_type_only,
            })
            .collect()
    }

    /// Install source imports from serializable entries (for external macro processes).
    pub fn install_source_imports(&mut self, entries: Vec<SourceImportEntry>) {
        for entry in entries {
            self.source_imports.insert(
                entry.local_name,
                SourceImport {
                    source_module: entry.source_module,
                    original_name: entry.original_name,
                    is_type_only: entry.is_type_only,
                },
            );
        }
    }

    /// Get type-only tracking as a flat name→bool HashMap (for backward-compat callers).
    pub fn type_only_map(&self) -> HashMap<String, bool> {
        self.source_imports
            .iter()
            .map(|(name, si)| (name.clone(), si.is_type_only))
            .collect()
    }

    /// General-purpose import request. Idempotent — skips if `local_name` already
    /// exists in source or generated imports.
    ///
    /// This is the method TsStream import helpers delegate to.
    pub fn request_import(
        &mut self,
        local_name: &str,
        original_name: Option<&str>,
        module: &str,
        is_type_only: bool,
    ) {
        // Reject dotted names
        if local_name.contains('.') {
            return;
        }

        // Already in source — skip
        if self.source_imports.contains_key(local_name) {
            return;
        }

        // Already generated — skip
        if self.generated.contains_key(local_name) {
            return;
        }

        self.generated.insert(
            local_name.to_string(),
            GeneratedImport {
                local_name: local_name.to_string(),
                source_module: module.to_string(),
                original_name: original_name.map(|s| s.to_string()),
                is_type_only,
            },
        );
    }

    /// Request a namespace import (e.g., `import { DateTime as __mf_DateTime } from "effect"`).
    ///
    /// No-op if `namespace` is already a value import in source.
    /// No-op if `alias` is already in generated.
    pub fn request_namespace_import(&mut self, namespace: &str, module: &str, alias: &str) {
        // If namespace is already available as a non-type-only import, skip
        if let Some(si) = self.source_imports.get(namespace)
            && !si.is_type_only
        {
            return;
        }

        // If alias is already generated, skip
        if self.generated.contains_key(alias) {
            return;
        }

        self.generated.insert(
            alias.to_string(),
            GeneratedImport {
                local_name: alias.to_string(),
                source_module: module.to_string(),
                original_name: Some(namespace.to_string()),
                is_type_only: false,
            },
        );
    }

    /// Request a type-only import (e.g., `import type { Utc } from "effect/DateTime"`).
    ///
    /// No-op if `name` is already in source or generated.
    /// Rejects dotted names (logs warning, returns early).
    pub fn request_type_import(&mut self, name: &str, module: &str) {
        self.request_import(name, None, module, true);
    }

    /// Get an iterator over generated imports.
    pub fn generated_imports(&self) -> impl Iterator<Item = &GeneratedImport> {
        self.generated.values()
    }

    /// Take a snapshot of generated imports and clear them from the registry.
    /// Used by `TsStream::into_result()` to capture imports into `MacroResult`
    /// so they survive serialization across process boundaries.
    pub fn take_generated_imports(&mut self) -> Vec<GeneratedImport> {
        std::mem::take(&mut self.generated).into_values().collect()
    }

    /// Merge a list of generated imports back into the registry.
    /// Used to restore imports from a `MacroResult` (e.g., from an external macro).
    pub fn merge_imports(&mut self, imports: Vec<GeneratedImport>) {
        for import in imports {
            if !self.source_imports.contains_key(&import.local_name)
                && !self.generated.contains_key(&import.local_name)
            {
                self.generated.insert(import.local_name.clone(), import);
            }
        }
    }

    /// Format all generated imports as TypeScript import lines.
    pub fn emit_generated_imports(&self) -> String {
        if self.generated.is_empty() {
            return String::new();
        }

        let mut lines = String::new();

        for (_alias, import) in &self.generated {
            let keyword = if import.is_type_only {
                "import type"
            } else {
                "import"
            };

            let specifier = if let Some(ref original) = import.original_name {
                format!("{} as {}", original, import.local_name)
            } else {
                import.local_name.clone()
            };

            lines.push_str(&format!(
                "{} {{ {} }} from \"{}\";\n",
                keyword, specifier, import.source_module
            ));
        }

        lines
    }
}

// ============================================================================
// Thread-local registry — single source of truth for all import state
// ============================================================================

thread_local! {
    /// Single thread-local holding the unified import registry.
    static IMPORT_REGISTRY: RefCell<ImportRegistry> = RefCell::new(ImportRegistry::new());
}

/// Read from the registry (immutable access).
pub fn with_registry<R>(f: impl FnOnce(&ImportRegistry) -> R) -> R {
    IMPORT_REGISTRY.with(|r| f(&r.borrow()))
}

/// Write to the registry (mutable access).
pub fn with_registry_mut<R>(f: impl FnOnce(&mut ImportRegistry) -> R) -> R {
    IMPORT_REGISTRY.with(|r| f(&mut r.borrow_mut()))
}

/// Install a pre-built registry into the thread-local.
pub fn install_registry(registry: ImportRegistry) {
    IMPORT_REGISTRY.with(|r| {
        *r.borrow_mut() = registry;
    });
}

/// Take the registry out of the thread-local, replacing it with an empty one.
pub fn take_registry() -> ImportRegistry {
    IMPORT_REGISTRY.with(|r| std::mem::take(&mut *r.borrow_mut()))
}

/// Clear the registry after expansion.
pub fn clear_registry() {
    IMPORT_REGISTRY.with(|r| {
        *r.borrow_mut() = ImportRegistry::new();
    });
}

// ============================================================================
// JSDoc @import macro comment parsing
// ============================================================================

/// Extracts import information from JSDoc `@import macro` comments.
fn collect_macro_import_comments(source: &str) -> HashMap<String, String> {
    let mut out = HashMap::new();
    let mut search_start = 0usize;

    while let Some(idx) = source[search_start..].find("/**") {
        let abs_idx = search_start + idx;
        let remaining = &source[abs_idx + 3..];
        let Some(end_rel) = remaining.find("*/") else {
            break;
        };
        let body = &remaining[..end_rel];
        let normalized = normalize_macro_import_body(body);
        let normalized_lower = normalized.to_ascii_lowercase();

        if normalized_lower.contains("import macro")
            && let (Some(open_brace), Some(close_brace)) =
                (normalized.find('{'), normalized.find('}'))
            && close_brace > open_brace
            && let Some(from_idx) = normalized_lower[close_brace..].find("from")
        {
            let names_src = normalized[open_brace + 1..close_brace].trim();
            let from_section = &normalized[close_brace + from_idx + "from".len()..];
            if let Some(module_src) = extract_quoted_string(from_section) {
                for name in names_src.split(',') {
                    let trimmed = name.trim();
                    if !trimmed.is_empty() {
                        out.insert(trimmed.to_string(), module_src.clone());
                    }
                }
            }
        }

        search_start = abs_idx + 3 + end_rel + 2;
    }

    out
}

fn normalize_macro_import_body(body: &str) -> String {
    let mut normalized = String::new();
    for line in body.lines() {
        let mut trimmed = line.trim();
        if let Some(stripped) = trimmed.strip_prefix('*') {
            trimmed = stripped.trim();
        }
        if trimmed.is_empty() {
            continue;
        }
        if !normalized.is_empty() {
            normalized.push(' ');
        }
        normalized.push_str(trimmed);
    }
    normalized
}

fn extract_quoted_string(input: &str) -> Option<String> {
    for (idx, ch) in input.char_indices() {
        if ch == '"' || ch == '\'' {
            let start = idx + 1;
            let rest = &input[start..];
            if let Some(end) = rest.find(ch) {
                return Some(rest[..end].trim().to_string());
            }
            break;
        }
    }
    None
}