erg/crates/erg_compiler/module/index.rs

use std::collections::hash_map::{Iter, Keys, Values};
use std::fmt;

use erg_common::dict::Dict;
use erg_common::pathutil::NormalizedPathBuf;
use erg_common::set;
use erg_common::set::Set;
use erg_common::shared::{MappedRwLockReadGuard, RwLockReadGuard, Shared};
use erg_common::Str;

use crate::varinfo::{AbsLocation, VarInfo};

pub struct Members<'a>(MappedRwLockReadGuard<'a, Dict<AbsLocation, ModuleIndexValue>>);

impl Members<'_> {
    pub fn iter(&self) -> Iter<AbsLocation, ModuleIndexValue> {
        self.0.iter()
    }

    pub fn keys(&self) -> Keys<AbsLocation, ModuleIndexValue> {
        self.0.keys()
    }

    pub fn values(&self) -> Values<AbsLocation, ModuleIndexValue> {
        self.0.values()
    }
}

#[derive(Debug, Clone)]
pub struct ModuleIndexValue {
    pub name: Str,
    pub vi: VarInfo,
    pub referrers: Set<AbsLocation>,
}

impl fmt::Display for ModuleIndexValue {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "{{ name: {}, vi: {}, referrers: {} }}",
            self.name, self.vi, self.referrers
        )
    }
}

impl ModuleIndexValue {
    pub const fn new(name: Str, vi: VarInfo, referrers: Set<AbsLocation>) -> Self {
        Self {
            name,
            vi,
            referrers,
        }
    }

    pub fn push_ref(&mut self, referrer: AbsLocation) {
        self.referrers.insert(referrer);
    }
}

#[derive(Debug, Clone, Default)]
pub struct ModuleIndex {
    members: Dict<AbsLocation, ModuleIndexValue>,
}

impl fmt::Display for ModuleIndex {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        self.members.fmt(f)
    }
}

impl ModuleIndex {
    pub fn new() -> Self {
        Self {
            members: Dict::new(),
        }
    }

    pub fn inc_ref(&mut self, name: &Str, vi: &VarInfo, referrer: AbsLocation) {
        let referee = vi.def_loc.clone();
        if let Some(referrers) = self.members.get_mut(&referee) {
            referrers.push_ref(referrer);
        } else {
            let value = ModuleIndexValue::new(name.clone(), vi.clone(), set! {referrer});
            self.members.insert(referee, value);
        }
    }

    pub fn register(&mut self, name: Str, vi: &VarInfo) {
        if self.members.contains_key(&vi.def_loc) {
            return;
        }
        let referee = vi.def_loc.clone();
        let value = ModuleIndexValue::new(name, vi.clone(), set! {});
        self.members.insert(referee, value);
    }

    pub fn get_refs(&self, referee: &AbsLocation) -> Option<&ModuleIndexValue> {
        self.members.get(referee)
    }

    pub fn initialize(&mut self) {
        self.members.clear();
    }

    pub fn remove_path(&mut self, path: &NormalizedPathBuf) {
        self.members.retain(|loc, value| {
            value
                .referrers
                .retain(|ref_loc| ref_loc.module.as_deref() != Some(path));
            loc.module.as_deref() != Some(path)
        });
    }

    pub fn rename_path(&mut self, old: &NormalizedPathBuf, new: NormalizedPathBuf) {
        let mut new_members = Dict::new();
        for (loc, mut value) in std::mem::take(&mut self.members) {
            if value.vi.def_loc.module.as_deref() == Some(old) {
                value.vi.def_loc.module = Some(new.clone());
            }
            let mut new_referrers = set! {};
            for referee in value.referrers.into_iter() {
                if referee.module.as_deref() == Some(old) {
                    new_referrers.insert(AbsLocation {
                        module: Some(new.clone()),
                        ..referee
                    });
                } else {
                    new_referrers.insert(referee);
                }
            }
            value.referrers = new_referrers;
            if loc.module.as_deref() != Some(old) {
                new_members.insert(loc.clone(), value.clone());
            } else {
                new_members.insert(
                    AbsLocation {
                        module: Some(new.clone()),
                        ..loc
                    },
                    value,
                );
            }
        }
        self.members = new_members;
    }
}

#[derive(Debug, Clone, Default)]
pub struct SharedModuleIndex(Shared<ModuleIndex>);

impl fmt::Display for SharedModuleIndex {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        self.0.borrow().fmt(f)
    }
}

impl SharedModuleIndex {
    pub fn new() -> Self {
        Self(Shared::new(ModuleIndex::new()))
    }

    pub fn inc_ref(&self, name: &Str, vi: &VarInfo, referrer: AbsLocation) {
        self.0.borrow_mut().inc_ref(name, vi, referrer);
    }

    pub fn register(&self, name: Str, vi: &VarInfo) {
        self.0.borrow_mut().register(name, vi);
    }

    pub fn get_refs(
        &self,
        referee: &AbsLocation,
    ) -> Option<MappedRwLockReadGuard<ModuleIndexValue>> {
        RwLockReadGuard::try_map(self.0.borrow(), |index| index.get_refs(referee)).ok()
    }

    pub fn members(&self) -> Members {
        Members(RwLockReadGuard::map(self.0.borrow(), |mi| &mi.members))
    }

    pub fn initialize(&self) {
        self.0.borrow_mut().initialize();
    }

    pub fn remove_path(&self, path: &NormalizedPathBuf) {
        self.0.borrow_mut().remove_path(path);
    }

    pub fn rename_path(&self, old: &NormalizedPathBuf, new: NormalizedPathBuf) {
        self.0.borrow_mut().rename_path(old, new);
    }
}