use alloc::{collections::BTreeMap, vec::Vec};
use hir_lang::{DefId, ItemKind, Name, Ns, Res, Shape, Symbol, Vis};
use crate::policy::Namespace;
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
#[non_exhaustive]
pub enum DefKind {
Module,
Fn,
Const,
Global,
Record {
unit: bool,
},
Sum,
Class {
mixin: bool,
},
Interface,
Alias,
AssocType,
Variant {
unit: bool,
},
Extern,
Prim,
Local,
Err,
}
impl DefKind {
#[must_use]
pub const fn of_item(kind: &ItemKind) -> Option<Self> {
Some(match kind {
ItemKind::Fn(_) => Self::Fn,
ItemKind::Record(r) => Self::Record {
unit: matches!(r.shape, Shape::Unit),
},
ItemKind::Sum(_) => Self::Sum,
ItemKind::Class(c) => Self::Class { mixin: c.mixin },
ItemKind::Interface(_) => Self::Interface,
ItemKind::Alias { .. } => Self::Alias,
ItemKind::AssocType { .. } => Self::AssocType,
ItemKind::Const { .. } => Self::Const,
ItemKind::Global { .. } => Self::Global,
ItemKind::Module { .. } => Self::Module,
ItemKind::Err => Self::Err,
ItemKind::Impl(_) | ItemKind::Import { .. } | ItemKind::MixinUse(_) => return None,
})
}
#[must_use]
pub const fn fits(self, ns: Ns) -> bool {
match ns {
Ns::Value => matches!(
self,
Self::Fn
| Self::Const
| Self::Global
| Self::Record { .. }
| Self::Class { .. }
| Self::Err
| Self::Variant { .. }
| Self::Extern
),
Ns::Type => matches!(
self,
Self::Record { .. }
| Self::Sum
| Self::Class { .. }
| Self::Interface
| Self::Alias
| Self::AssocType
| Self::Err
| Self::Variant { .. }
| Self::Prim
| Self::Extern
),
Ns::Pattern => matches!(
self,
Self::Const
| Self::Record { .. }
| Self::Class { .. }
| Self::Err
| Self::Variant { .. }
| Self::Extern
),
Ns::Region => false,
Ns::Import => !matches!(self, Self::Prim),
}
}
#[must_use]
pub const fn is_prefix(self) -> bool {
matches!(
self,
Self::Module
| Self::Record { .. }
| Self::Sum
| Self::Class { .. }
| Self::Interface
| Self::Alias
| Self::AssocType
| Self::Err
| Self::Prim
| Self::Extern
)
}
#[must_use]
pub const fn is_pattern_constant(self) -> bool {
matches!(
self,
Self::Const | Self::Variant { unit: true } | Self::Record { unit: true }
)
}
#[must_use]
pub const fn name(self) -> &'static str {
match self {
Self::Module => "module",
Self::Fn => "function",
Self::Const => "constant",
Self::Global => "global",
Self::Record { .. } => "record",
Self::Sum => "sum type",
Self::Class { mixin: false } => "class",
Self::Class { mixin: true } => "mixin",
Self::Interface => "interface",
Self::Alias => "type alias",
Self::AssocType => "associated type",
Self::Variant { .. } => "variant",
Self::Extern => "extern",
Self::Prim => "primitive type",
Self::Local => "local",
Self::Err => "error item",
}
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Export {
pub res: Res,
pub kind: DefKind,
pub vis: Vis,
}
impl Export {
#[must_use]
pub const fn new(res: Res, kind: DefKind) -> Self {
Self {
res,
kind,
vis: Vis::Public,
}
}
#[must_use]
pub const fn with_vis(mut self, vis: Vis) -> Self {
self.vis = vis;
self
}
}
pub trait Env {
fn root(&self, name: Name) -> Option<Export> {
let _ = name;
None
}
fn prelude(&self, name: Name, ns: Namespace) -> Option<Export> {
let _ = (name, ns);
None
}
fn member(&self, container: Res, name: Name, ns: Namespace) -> Option<Export> {
let _ = (container, name, ns);
None
}
fn for_each_member(&self, container: Res, f: &mut dyn FnMut(Name, Namespace, Export)) {
let _ = (container, f);
}
fn for_each_root(&self, f: &mut dyn FnMut(Name)) {
let _ = f;
}
fn for_each_prelude(&self, f: &mut dyn FnMut(Name, Namespace)) {
let _ = f;
}
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct NoEnv;
impl Env for NoEnv {}
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
enum ContainerKey {
Def(DefId),
Extern(Symbol),
}
impl ContainerKey {
fn of(res: Res) -> Option<Self> {
match res {
Res::Def(d) => Some(Self::Def(d)),
Res::Extern(s) => Some(Self::Extern(s)),
_ => None,
}
}
}
#[derive(Clone, Debug, Default)]
pub struct MapEnv {
roots: BTreeMap<Name, Export>,
prelude: BTreeMap<(Name, Namespace), Export>,
members: BTreeMap<ContainerKey, BTreeMap<(Name, Namespace), Export>>,
}
impl MapEnv {
#[must_use]
pub fn new() -> Self {
Self::default()
}
#[must_use]
pub fn with_root(mut self, name: Name, export: Export) -> Self {
let _replaced = self.roots.insert(name, export);
self
}
#[must_use]
pub fn with_prelude(mut self, name: Name, ns: Namespace, export: Export) -> Self {
let _replaced = self.prelude.insert((name, ns), export);
self
}
#[must_use]
pub fn with_member(
mut self,
container: Res,
name: Name,
ns: Namespace,
export: Export,
) -> Self {
if let Some(key) = ContainerKey::of(container) {
let _replaced = self
.members
.entry(key)
.or_default()
.insert((name, ns), export);
}
self
}
}
impl Env for MapEnv {
fn root(&self, name: Name) -> Option<Export> {
self.roots.get(&name).copied()
}
fn prelude(&self, name: Name, ns: Namespace) -> Option<Export> {
self.prelude.get(&(name, ns)).copied()
}
fn member(&self, container: Res, name: Name, ns: Namespace) -> Option<Export> {
let key = ContainerKey::of(container)?;
self.members.get(&key)?.get(&(name, ns)).copied()
}
fn for_each_member(&self, container: Res, f: &mut dyn FnMut(Name, Namespace, Export)) {
let Some(members) = ContainerKey::of(container).and_then(|k| self.members.get(&k)) else {
return;
};
for (&(name, ns), export) in members {
f(name, ns, *export);
}
}
fn for_each_root(&self, f: &mut dyn FnMut(Name)) {
for name in self.roots.keys() {
f(*name);
}
}
fn for_each_prelude(&self, f: &mut dyn FnMut(Name, Namespace)) {
for (name, ns) in self.prelude.keys() {
f(*name, *ns);
}
}
}
pub(crate) fn env_members(env: &dyn Env, container: Res) -> Vec<(Name, Namespace, Export)> {
let mut out = Vec::new();
env.for_each_member(container, &mut |n, ns, e| out.push((n, ns, e)));
out
}
#[cfg(test)]
#[allow(clippy::unwrap_used)]
mod tests {
use super::*;
use hir_lang::{Def, ItemId, Prim, UnitId};
fn sym(i: u32) -> Symbol {
Symbol::from_u32(i).unwrap()
}
#[test]
fn test_map_env_enumerates_members_of_one_container_only() {
let a = Res::Extern(sym(1));
let b = Res::Extern(sym(2));
let e = Export::new(Res::Extern(sym(9)), DefKind::Extern);
let env = MapEnv::new()
.with_member(a, Name::new(sym(5)), Namespace::Value, e)
.with_member(a, Name::new(sym(3)), Namespace::Type, e)
.with_member(b, Name::new(sym(4)), Namespace::Value, e);
let names: Vec<u32> = env_members(&env, a)
.iter()
.map(|(n, _, _)| n.sym.as_u32())
.collect();
assert_eq!(names, [3, 5]);
assert!(env_members(&env, Res::Prim(Prim::I32)).is_empty());
}
#[test]
fn test_def_kind_matches_hir_namespace_table() {
assert!(DefKind::Class { mixin: false }.fits(Ns::Pattern));
assert!(!DefKind::Sum.fits(Ns::Value));
assert!(!DefKind::Fn.fits(Ns::Region));
assert!(DefKind::Extern.fits(Ns::Import));
assert!(!DefKind::Prim.fits(Ns::Import));
}
#[test]
fn test_member_on_non_container_is_ignored() {
let d = DefId::foreign(UnitId::new(1), Def::Item(ItemId::from_index(0).unwrap()));
let n = Name::new(sym(1));
let env = MapEnv::new().with_member(
Res::Err,
n,
Namespace::Value,
Export::new(Res::Def(d), DefKind::Fn),
);
assert!(env.member(Res::Err, n, Namespace::Value).is_none());
let env = env.with_member(
Res::Def(d),
n,
Namespace::Value,
Export::new(Res::Def(d), DefKind::Fn),
);
assert!(env.member(Res::Def(d), n, Namespace::Value).is_some());
}
}