ryo_analysis/query/

var_id.rs

//! VarId - Lightweight variable identifier for DataFlowGraph.
//!
//! VarId is a local identifier for variables within DataFlowGraph,
//! providing O(1) access while maintaining a mapping to SymbolId
//! for integration with the broader symbol system.
//!
//! Uses SlotMap for stable identifiers that support deletion.

use crate::symbol::SymbolId;
use slotmap::{new_key_type, SecondaryMap, SlotMap};

new_key_type! {
    /// Lightweight variable identifier for DataFlowGraph internal use.
    ///
    /// VarId is a SlotMap key optimized for:
    /// - O(1) comparison and hashing
    /// - Stable identity across insertions/deletions
    /// - Safe deletion without index invalidation
    ///
    /// For full symbol information, convert to SymbolId via VarSymbolMapping.
    pub struct VarId;
}

/// Bidirectional mapping between VarId and SymbolId.
///
/// Uses SlotMap for stable VarId generation that supports deletion.
///
/// This provides:
/// - VarId → SymbolId: O(1) via SlotMap
/// - SymbolId → VarId: O(1) via SecondaryMap
///
/// # Usage
///
/// ```ignore
/// let mut mapping = VarSymbolMapping::new();
///
/// // Register a variable
/// let var_id = mapping.register(symbol_id);
///
/// // Lookup
/// let sym = mapping.to_symbol(var_id);
/// let var = mapping.to_var(symbol_id);
///
/// // Remove (for incremental updates)
/// mapping.remove(var_id);
/// ```
#[derive(Debug, Clone, Default)]
pub struct VarSymbolMapping {
    /// VarId → `Option<SymbolId>` (forward mapping via SlotMap)
    /// None for local variables without a SymbolId in the registry.
    var_to_symbol: SlotMap<VarId, Option<SymbolId>>,

    /// SymbolId → VarId (reverse mapping, only for registered variables)
    symbol_to_var: SecondaryMap<SymbolId, VarId>,
}

impl VarSymbolMapping {
    /// Create a new empty mapping.
    pub fn new() -> Self {
        Self::default()
    }

    /// Create with pre-allocated capacity.
    pub fn with_capacity(capacity: usize) -> Self {
        Self {
            var_to_symbol: SlotMap::with_capacity_and_key(capacity),
            symbol_to_var: SecondaryMap::with_capacity(capacity),
        }
    }

    /// Register a named variable (with SymbolId) and return its VarId.
    ///
    /// If the SymbolId is already registered, returns the existing VarId.
    pub fn register(&mut self, symbol_id: SymbolId) -> VarId {
        // Check if already registered
        if let Some(&var_id) = self.symbol_to_var.get(symbol_id) {
            // Verify forward mapping still exists (wasn't removed)
            if self.var_to_symbol.contains_key(var_id) {
                return var_id;
            }
            // Stale reverse mapping, will be overwritten below
        }

        // Create new VarId
        let var_id = self.var_to_symbol.insert(Some(symbol_id));
        self.symbol_to_var.insert(symbol_id, var_id);
        var_id
    }

    /// Allocate a fresh VarId for a local variable without a SymbolId.
    ///
    /// Always creates a new unique VarId. No reverse lookup is registered.
    pub fn allocate(&mut self) -> VarId {
        self.var_to_symbol.insert(None)
    }

    /// Remove a variable by VarId.
    ///
    /// Returns the associated SymbolId if it existed.
    pub fn remove(&mut self, var_id: VarId) -> Option<SymbolId> {
        if let Some(Some(symbol_id)) = self.var_to_symbol.remove(var_id) {
            self.symbol_to_var.remove(symbol_id);
            return Some(symbol_id);
        }
        None
    }

    /// Remove a variable by SymbolId.
    ///
    /// Returns the VarId if it existed.
    pub fn remove_by_symbol(&mut self, symbol_id: SymbolId) -> Option<VarId> {
        if let Some(var_id) = self.symbol_to_var.remove(symbol_id) {
            self.var_to_symbol.remove(var_id);
            Some(var_id)
        } else {
            None
        }
    }

    /// Get the SymbolId for a VarId.
    ///
    /// Returns None if the VarId doesn't exist or has no associated SymbolId.
    #[inline]
    pub fn to_symbol(&self, var_id: VarId) -> Option<SymbolId> {
        self.var_to_symbol.get(var_id).copied().flatten()
    }

    /// Get the VarId for a SymbolId.
    ///
    /// Returns None if the SymbolId is not registered.
    #[inline]
    pub fn to_var(&self, symbol_id: SymbolId) -> Option<VarId> {
        self.symbol_to_var.get(symbol_id).copied()
    }

    /// Check if a SymbolId is registered.
    #[inline]
    pub fn contains_symbol(&self, symbol_id: SymbolId) -> bool {
        self.symbol_to_var.contains_key(symbol_id)
    }

    /// Check if a VarId is valid.
    #[inline]
    pub fn contains_var(&self, var_id: VarId) -> bool {
        self.var_to_symbol.contains_key(var_id)
    }

    /// Get the number of registered variables.
    #[inline]
    pub fn len(&self) -> usize {
        self.var_to_symbol.len()
    }

    /// Check if empty.
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.var_to_symbol.is_empty()
    }

    /// Iterate over all `(VarId, Option<SymbolId>)` pairs.
    pub fn iter(&self) -> impl Iterator<Item = (VarId, Option<SymbolId>)> + '_ {
        self.var_to_symbol.iter().map(|(id, &sym)| (id, sym))
    }

    /// Clear all mappings.
    pub fn clear(&mut self) {
        self.var_to_symbol.clear();
        self.symbol_to_var.clear();
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use slotmap::SlotMap;

    #[test]
    fn test_var_symbol_mapping_register() {
        let mut symbols: SlotMap<SymbolId, &str> = SlotMap::with_key();
        let sym1 = symbols.insert("x");
        let sym2 = symbols.insert("y");

        let mut mapping = VarSymbolMapping::new();

        let var1 = mapping.register(sym1);
        let var2 = mapping.register(sym2);

        assert_ne!(var1, var2);
        assert_eq!(mapping.len(), 2);
    }

    #[test]
    fn test_var_symbol_mapping_idempotent() {
        let mut symbols: SlotMap<SymbolId, &str> = SlotMap::with_key();
        let sym = symbols.insert("x");

        let mut mapping = VarSymbolMapping::new();

        let var1 = mapping.register(sym);
        let var2 = mapping.register(sym); // Same symbol again

        assert_eq!(var1, var2);
        assert_eq!(mapping.len(), 1);
    }

    #[test]
    fn test_var_symbol_mapping_bidirectional() {
        let mut symbols: SlotMap<SymbolId, &str> = SlotMap::with_key();
        let sym = symbols.insert("x");

        let mut mapping = VarSymbolMapping::new();
        let var = mapping.register(sym);

        // Forward lookup
        assert_eq!(mapping.to_symbol(var), Some(sym));

        // Reverse lookup
        assert_eq!(mapping.to_var(sym), Some(var));
    }

    #[test]
    fn test_var_symbol_mapping_contains() {
        let mut symbols: SlotMap<SymbolId, &str> = SlotMap::with_key();
        let sym1 = symbols.insert("x");
        let sym2 = symbols.insert("y");

        let mut mapping = VarSymbolMapping::new();
        let var = mapping.register(sym1);

        assert!(mapping.contains_symbol(sym1));
        assert!(!mapping.contains_symbol(sym2));
        assert!(mapping.contains_var(var));
    }

    #[test]
    fn test_var_symbol_mapping_remove() {
        let mut symbols: SlotMap<SymbolId, &str> = SlotMap::with_key();
        let sym = symbols.insert("x");

        let mut mapping = VarSymbolMapping::new();
        let var = mapping.register(sym);

        assert_eq!(mapping.len(), 1);
        assert!(mapping.contains_var(var));

        // Remove by VarId
        let removed = mapping.remove(var);
        assert_eq!(removed, Some(sym));
        assert_eq!(mapping.len(), 0);
        assert!(!mapping.contains_var(var));
        assert!(!mapping.contains_symbol(sym));
    }

    #[test]
    fn test_var_symbol_mapping_remove_by_symbol() {
        let mut symbols: SlotMap<SymbolId, &str> = SlotMap::with_key();
        let sym = symbols.insert("x");

        let mut mapping = VarSymbolMapping::new();
        let var = mapping.register(sym);

        // Remove by SymbolId
        let removed = mapping.remove_by_symbol(sym);
        assert_eq!(removed, Some(var));
        assert_eq!(mapping.len(), 0);
    }

    #[test]
    fn test_var_symbol_mapping_iter() {
        let mut symbols: SlotMap<SymbolId, &str> = SlotMap::with_key();
        let sym1 = symbols.insert("x");
        let sym2 = symbols.insert("y");

        let mut mapping = VarSymbolMapping::new();
        mapping.register(sym1);
        mapping.register(sym2);

        let pairs: Vec<_> = mapping.iter().collect();
        assert_eq!(pairs.len(), 2);
    }

    #[test]
    fn test_var_symbol_mapping_clear() {
        let mut symbols: SlotMap<SymbolId, &str> = SlotMap::with_key();
        let sym = symbols.insert("x");

        let mut mapping = VarSymbolMapping::new();
        mapping.register(sym);
        assert!(!mapping.is_empty());

        mapping.clear();
        assert!(mapping.is_empty());
        assert!(!mapping.contains_symbol(sym));
    }

    #[test]
    fn test_var_symbol_mapping_reregister_after_remove() {
        let mut symbols: SlotMap<SymbolId, &str> = SlotMap::with_key();
        let sym = symbols.insert("x");

        let mut mapping = VarSymbolMapping::new();
        let var1 = mapping.register(sym);
        mapping.remove(var1);

        // Re-register same symbol
        let var2 = mapping.register(sym);
        assert_ne!(var1, var2); // Should get a new VarId
        assert_eq!(mapping.len(), 1);
        assert!(mapping.contains_var(var2));
    }

    #[test]
    fn test_allocate_produces_distinct_ids() {
        let mut mapping = VarSymbolMapping::new();

        let v1 = mapping.allocate();
        let v2 = mapping.allocate();
        let v3 = mapping.allocate();

        assert_ne!(v1, v2);
        assert_ne!(v2, v3);
        assert_ne!(v1, v3);
        assert_eq!(mapping.len(), 3);

        // Allocated vars have no SymbolId
        assert_eq!(mapping.to_symbol(v1), None);
        assert_eq!(mapping.to_symbol(v2), None);
    }

    #[test]
    fn test_allocate_and_register_coexist() {
        let mut symbols: SlotMap<SymbolId, &str> = SlotMap::with_key();
        let sym = symbols.insert("x");

        let mut mapping = VarSymbolMapping::new();
        let anon = mapping.allocate();
        let named = mapping.register(sym);

        assert_ne!(anon, named);
        assert_eq!(mapping.len(), 2);
        assert_eq!(mapping.to_symbol(anon), None);
        assert_eq!(mapping.to_symbol(named), Some(sym));
        assert_eq!(mapping.to_var(sym), Some(named));
    }

    #[test]
    fn test_remove_allocated_var() {
        let mut mapping = VarSymbolMapping::new();
        let v = mapping.allocate();
        assert_eq!(mapping.len(), 1);

        let removed = mapping.remove(v);
        assert_eq!(removed, None); // No SymbolId associated
        assert_eq!(mapping.len(), 0);
        assert!(!mapping.contains_var(v));
    }
}