use std::collections::HashMap;
use crate::codegen::ModuleInfo;
use crate::ir::SymbolTable;
use crate::ir::hir::{ResolveCtx, ResolvedFnDef, ResolvedTopLevel, resolve_fn_def_external};
#[derive(Debug, Clone, Default)]
pub struct ResolvedModuleFns {
pub prefix: String,
pub fn_defs: Vec<ResolvedFnDef>,
}
#[derive(Debug, Clone, Default)]
pub struct ResolvedProgramView {
pub entry_items: Vec<ResolvedTopLevel>,
pub modules: Vec<ResolvedModuleFns>,
fn_index: HashMap<crate::ir::FnId, FnIndexEntry>,
}
#[derive(Debug, Clone, Copy)]
struct FnIndexEntry {
module: Option<usize>,
pos: usize,
}
impl ResolvedProgramView {
pub fn build(
entry_items: Vec<ResolvedTopLevel>,
modules: &[ModuleInfo],
symbol_table: &SymbolTable,
) -> Self {
let module_views: Vec<ResolvedModuleFns> = modules
.iter()
.map(|module| {
let mut rctx = ResolveCtx::new(symbol_table);
rctx.current_module = Some(module.prefix.clone());
let fn_defs = module
.fn_defs
.iter()
.filter_map(|fd| resolve_fn_def_external(&rctx, fd))
.collect();
ResolvedModuleFns {
prefix: module.prefix.clone(),
fn_defs,
}
})
.collect();
let mut fn_index: HashMap<crate::ir::FnId, FnIndexEntry> = HashMap::new();
for (pos, item) in entry_items.iter().enumerate() {
if let ResolvedTopLevel::FnDef(rfd) = item {
fn_index.insert(rfd.fn_id, FnIndexEntry { module: None, pos });
}
}
for (i, m) in module_views.iter().enumerate() {
for (pos, rfd) in m.fn_defs.iter().enumerate() {
fn_index.insert(
rfd.fn_id,
FnIndexEntry {
module: Some(i),
pos,
},
);
}
}
Self {
entry_items,
modules: module_views,
fn_index,
}
}
pub fn entry_fns(&self) -> impl Iterator<Item = &ResolvedFnDef> + '_ {
self.entry_items.iter().filter_map(|item| match item {
ResolvedTopLevel::FnDef(rfd) => Some(rfd),
_ => None,
})
}
pub fn module_fns(&self, prefix: &str) -> impl Iterator<Item = &ResolvedFnDef> + '_ {
self.modules
.iter()
.find(|m| m.prefix == prefix)
.into_iter()
.flat_map(|m| m.fn_defs.iter())
}
pub fn module(&self, prefix: &str) -> Option<&ResolvedModuleFns> {
self.modules.iter().find(|m| m.prefix == prefix)
}
pub fn fn_by_id(&self, id: crate::ir::FnId) -> Option<&ResolvedFnDef> {
let entry = *self.fn_index.get(&id)?;
match entry.module {
None => {
let item = self.entry_items.get(entry.pos)?;
match item {
ResolvedTopLevel::FnDef(rfd) => Some(rfd),
_ => None,
}
}
Some(i) => self.modules.get(i)?.fn_defs.get(entry.pos),
}
}
pub fn fn_by_name(&self, name: &str) -> Option<&ResolvedFnDef> {
for fd in self.entry_fns() {
if fd.name == name {
return Some(fd);
}
}
for m in &self.modules {
for fd in &m.fn_defs {
if fd.name == name {
return Some(fd);
}
}
}
None
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::ast::{FnBody, FnDef, Module, Spanned, TopLevel};
use crate::codegen::ModuleInfo;
use crate::ir::SymbolTable;
fn mk_fn(name: &str) -> FnDef {
FnDef {
name: name.to_string(),
line: 1,
params: vec![],
return_type: "Int".to_string(),
effects: vec![],
desc: None,
body: std::sync::Arc::new(FnBody::Block(vec![crate::ast::Stmt::Expr(Spanned::bare(
crate::ast::Expr::Literal(crate::ast::Literal::Int(0)),
))])),
resolution: None,
}
}
fn mk_module(prefix: &str, fn_names: &[&str]) -> ModuleInfo {
ModuleInfo {
prefix: prefix.to_string(),
depends: vec![],
type_defs: vec![],
fn_defs: fn_names.iter().map(|n| mk_fn(n)).collect(),
analysis: None,
}
}
#[test]
fn view_indexes_entry_fns_by_fn_id() {
let entry_items = vec![TopLevel::FnDef(mk_fn("foo")), TopLevel::FnDef(mk_fn("bar"))];
let symbol_table = SymbolTable::build(&entry_items, &[]);
let resolved = crate::ir::hir::resolve_program(&symbol_table, &entry_items);
let view = ResolvedProgramView::build(resolved, &[], &symbol_table);
assert_eq!(view.entry_fns().count(), 2);
let foo_id = symbol_table
.fn_id_of(&crate::ir::FnKey::entry("foo"))
.expect("foo entry FnId");
let bar_id = symbol_table
.fn_id_of(&crate::ir::FnKey::entry("bar"))
.expect("bar entry FnId");
assert_eq!(view.fn_by_id(foo_id).map(|f| f.name.as_str()), Some("foo"));
assert_eq!(view.fn_by_id(bar_id).map(|f| f.name.as_str()), Some("bar"));
}
#[test]
fn cross_module_same_bare_name_disambiguates_by_fn_id() {
let entry_items: Vec<TopLevel> = vec![TopLevel::Module(Module {
name: "Entry".to_string(),
line: 1,
depends: vec!["A".to_string(), "B".to_string()],
exposes: vec![],
exposes_opaque: vec![],
exposes_line: None,
intent: "Test".to_string(),
effects: None,
effects_line: None,
})];
let modules = vec![mk_module("A", &["walker"]), mk_module("B", &["walker"])];
let symbol_table = SymbolTable::build(&entry_items, &modules);
let resolved = crate::ir::hir::resolve_program(&symbol_table, &entry_items);
let view = ResolvedProgramView::build(resolved, &modules, &symbol_table);
let a_id = symbol_table
.fn_id_of(&crate::ir::FnKey::in_module("A".to_string(), "walker"))
.expect("A.walker FnId");
let b_id = symbol_table
.fn_id_of(&crate::ir::FnKey::in_module("B".to_string(), "walker"))
.expect("B.walker FnId");
assert_ne!(a_id, b_id, "FnIds must be distinct across modules");
let a_walker = view.fn_by_id(a_id).expect("A.walker present");
let b_walker = view.fn_by_id(b_id).expect("B.walker present");
assert_eq!(a_walker.name, "walker");
assert_eq!(b_walker.name, "walker");
assert_eq!(a_walker.fn_id, a_id);
assert_eq!(b_walker.fn_id, b_id);
assert_eq!(view.module_fns("A").count(), 1);
assert_eq!(view.module_fns("B").count(), 1);
}
}