use std::collections::BTreeMap;
use brink_format::DefinitionId;
use brink_ir::hir::visit::{self, HirVisitor};
use brink_ir::{
Diagnostic, DiagnosticCode, Expr, FileId, HirFile, Knot, Path as HirPath, ResolutionMap,
Stitch, SymbolIndex, SymbolKind,
};
use rowan::TextRange;
use crate::annotations;
use crate::infer::{InferenceResult, InferredSig, Ty, assignable, ref_assignable};
use crate::resolve::ImportScope;
use crate::structs::{ShapeTable, declared_shapes};
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub struct NodeKey {
pub file: FileId,
pub range: (u32, u32),
}
impl NodeKey {
#[must_use]
pub fn new(file: FileId, range: TextRange) -> Self {
Self {
file,
range: (range.start().into(), range.end().into()),
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SideTable<V> {
entries: BTreeMap<NodeKey, V>,
}
impl<V> Default for SideTable<V> {
fn default() -> Self {
Self {
entries: BTreeMap::new(),
}
}
}
impl<V> SideTable<V> {
#[must_use]
pub fn new() -> Self {
Self::default()
}
pub fn insert(&mut self, key: NodeKey, value: V) -> Option<V> {
self.entries.insert(key, value)
}
#[must_use]
pub fn get(&self, key: NodeKey) -> Option<&V> {
self.entries.get(&key)
}
#[must_use]
pub fn at(&self, file: FileId, range: TextRange) -> Option<&V> {
self.get(NodeKey::new(file, range))
}
pub fn iter(&self) -> impl Iterator<Item = (NodeKey, &V)> {
self.entries.iter().map(|(k, v)| (*k, v))
}
#[must_use]
pub fn len(&self) -> usize {
self.entries.len()
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.entries.is_empty()
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum UfcsVerdict {
FieldCall {
receiver: Ty,
field: String,
field_ty: Ty,
arity_mismatch: Option<UfcsArityMismatch>,
arg_mismatches: Vec<UfcsArgMismatch>,
},
FreeFnAutoRef {
receiver: Ty,
name: String,
target: DefinitionId,
arg_mismatches: Vec<UfcsArgMismatch>,
},
FreeFnDesugar {
receiver: Ty,
name: String,
target: DefinitionId,
arg_mismatches: Vec<UfcsArgMismatch>,
},
PreludeDesugar {
receiver: Ty,
name: String,
arg_mismatches: Vec<UfcsArgMismatch>,
},
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct UfcsArgMismatch {
pub index: usize,
pub expected: Ty,
pub found: Ty,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct UfcsArityMismatch {
pub expected: usize,
pub got: usize,
}
pub type UfcsTable = SideTable<UfcsVerdict>;
#[must_use]
pub fn to_lir_lookup(table: &UfcsTable) -> brink_ir::lir::UfcsLookup {
let entries = table
.iter()
.map(|(key, verdict)| {
let range = TextRange::new(key.range.0.into(), key.range.1.into());
let mirrored = match verdict {
UfcsVerdict::FieldCall { .. } => brink_ir::lir::UfcsVerdict::FieldCall,
UfcsVerdict::FreeFnAutoRef { target, .. } => {
brink_ir::lir::UfcsVerdict::FreeFnAutoRef { target: *target }
}
UfcsVerdict::FreeFnDesugar { target, .. } => {
brink_ir::lir::UfcsVerdict::FreeFnDesugar { target: *target }
}
UfcsVerdict::PreludeDesugar { name, .. } => {
brink_ir::lir::UfcsVerdict::PreludeDesugar { name: name.clone() }
}
};
(key.file, range, mirrored)
})
.collect();
brink_ir::lir::UfcsLookup::from_entries(entries)
}
#[must_use]
pub fn resolve(
files: &[(FileId, &HirFile)],
index: &SymbolIndex,
resolutions: &ResolutionMap,
inference: &InferenceResult,
) -> (UfcsTable, Vec<Diagnostic>) {
let shapes = declared_shapes(files, index);
let globals = crate::infer::collect_globals(files, index, None);
let mut table = UfcsTable::new();
let mut diagnostics = Vec::new();
for &(file, hir) in files {
let resolution_by_range = resolution_index(resolutions, file);
let scope = ImportScope::new(hir.module.as_ref().map(|m| m.name.clone()), &hir.imports);
let mut v = UfcsVisitor {
file,
index,
scope: &scope,
shapes: &shapes,
globals: &globals,
bodies: &inference.bodies,
signatures: &inference.signatures,
resolution_by_range: &resolution_by_range,
current_knot_name: None,
knot_body: None,
stitch_body: None,
lambda_locals: Vec::new(),
table: &mut table,
diagnostics: &mut diagnostics,
};
visit::visit_with_decl_initializers(hir, &mut v);
}
(table, diagnostics)
}
#[must_use]
pub fn check_strict(
files: &[(FileId, &HirFile)],
index: &SymbolIndex,
resolutions: &ResolutionMap,
inference: &InferenceResult,
) -> Vec<Diagnostic> {
if !files.iter().any(|&(_, hir)| project_has_ufcs_call(hir)) {
return Vec::new();
}
let (table, _unconditional) = resolve(files, index, resolutions, inference);
table
.iter()
.flat_map(|(key, verdict)| strict_verdict_diagnostics(key, verdict))
.collect()
}
fn strict_verdict_diagnostics(key: NodeKey, verdict: &UfcsVerdict) -> Vec<Diagnostic> {
match verdict {
UfcsVerdict::PreludeDesugar {
receiver,
name,
arg_mismatches,
} => {
let mut out: Vec<Diagnostic> = arg_mismatches
.iter()
.map(|mismatch| ufcs_arg_mismatch_diagnostic(key, name, mismatch))
.collect();
if let ("remove", Ty::Array(_)) = (name.as_str(), receiver) {
out.push(Diagnostic {
file: key.file,
range: TextRange::new(key.range.0.into(), key.range.1.into()),
message: DiagnosticCode::E149.title().to_owned(),
code: DiagnosticCode::E149,
});
}
out
}
UfcsVerdict::FreeFnDesugar {
name,
arg_mismatches,
..
}
| UfcsVerdict::FreeFnAutoRef {
name,
arg_mismatches,
..
} => arg_mismatches
.iter()
.map(|mismatch| ufcs_arg_mismatch_diagnostic(key, name, mismatch))
.collect(),
UfcsVerdict::FieldCall {
field,
arity_mismatch,
arg_mismatches,
..
} => {
let mut out: Vec<Diagnostic> = Vec::new();
if let Some(arity) = arity_mismatch {
out.push(field_call_arity_diagnostic(key, field, *arity));
}
out.extend(
arg_mismatches
.iter()
.map(|mismatch| field_call_arg_mismatch_diagnostic(key, field, mismatch)),
);
out
}
}
}
fn field_call_arg_mismatch_diagnostic(
key: NodeKey,
field: &str,
mismatch: &UfcsArgMismatch,
) -> Diagnostic {
Diagnostic {
file: key.file,
range: TextRange::new(key.range.0.into(), key.range.1.into()),
message: format!(
"argument {} of call through `{field}` has type `{}` but its known type expects `{}`",
mismatch.index + 1,
mismatch.found.display(),
mismatch.expected.display(),
),
code: DiagnosticCode::E063,
}
}
fn field_call_arity_diagnostic(
key: NodeKey,
field: &str,
mismatch: UfcsArityMismatch,
) -> Diagnostic {
Diagnostic {
file: key.file,
range: TextRange::new(key.range.0.into(), key.range.1.into()),
message: format!(
"call through `{field}` supplies {got} argument(s) but its known type expects \
{expected}",
got = mismatch.got,
expected = mismatch.expected,
),
code: DiagnosticCode::E063,
}
}
fn ufcs_arg_mismatch_diagnostic(
key: NodeKey,
name: &str,
mismatch: &UfcsArgMismatch,
) -> Diagnostic {
Diagnostic {
file: key.file,
range: TextRange::new(key.range.0.into(), key.range.1.into()),
message: format!(
"argument {} of call to `{name}` has type `{}` but its known type expects `{}`",
mismatch.index + 1,
mismatch.found.display(),
mismatch.expected.display(),
),
code: DiagnosticCode::E063,
}
}
#[must_use]
pub fn project_has_ufcs_call(hir: &HirFile) -> bool {
struct Scan {
found: bool,
}
impl HirVisitor for Scan {
fn visit_exprs(&self) -> bool {
true
}
fn enter_expr(&mut self, expr: &Expr) {
if let Expr::Call(path, _) = expr
&& path.segments.len() > 1
{
self.found = true;
}
}
}
let mut scan = Scan { found: false };
visit::visit_with_decl_initializers(hir, &mut scan);
scan.found
}
fn resolution_index(
resolutions: &ResolutionMap,
file: FileId,
) -> BTreeMap<(u32, u32), DefinitionId> {
resolutions
.iter()
.filter(|r| r.file == file)
.map(|r| ((r.range.start().into(), r.range.end().into()), r.target))
.collect()
}
struct UfcsVisitor<'a> {
file: FileId,
index: &'a SymbolIndex,
scope: &'a ImportScope,
shapes: &'a ShapeTable,
globals: &'a BTreeMap<DefinitionId, Ty>,
bodies: &'a BTreeMap<DefinitionId, crate::infer::BodyTypes>,
signatures: &'a BTreeMap<DefinitionId, InferredSig>,
resolution_by_range: &'a BTreeMap<(u32, u32), DefinitionId>,
current_knot_name: Option<String>,
knot_body: Option<&'a crate::infer::BodyTypes>,
stitch_body: Option<&'a crate::infer::BodyTypes>,
lambda_locals: Vec<BTreeMap<String, Ty>>,
table: &'a mut UfcsTable,
diagnostics: &'a mut Vec<Diagnostic>,
}
impl HirVisitor for UfcsVisitor<'_> {
fn visit_exprs(&self) -> bool {
true
}
fn enter_knot(&mut self, knot: &Knot) {
self.current_knot_name = Some(knot.name.text.clone());
self.knot_body =
annotations::def_id_for(self.index, self.file, knot.symbol_kind(), &knot.name.text)
.and_then(|id| self.bodies.get(&id));
}
fn exit_knot(&mut self, _knot: &Knot) {
self.current_knot_name = None;
self.knot_body = None;
}
fn enter_stitch(&mut self, stitch: &Stitch) {
self.stitch_body = self.current_knot_name.as_ref().and_then(|knot_name| {
let qualified = format!("{knot_name}.{}", stitch.name.text);
annotations::def_id_for(self.index, self.file, SymbolKind::Stitch, &qualified)
.and_then(|id| self.bodies.get(&id))
});
}
fn exit_stitch(&mut self, _stitch: &Stitch) {
self.stitch_body = None;
}
fn enter_expr(&mut self, expr: &Expr) {
if let Expr::Call(path, args) = expr {
self.resolve_call(path, args.len());
}
}
fn enter_lambda(&mut self, l: &brink_ir::LambdaExpr) {
let pruned = crate::structs::pruned_locals_for_lambda(l, self.index, self.current_locals());
self.lambda_locals.push(pruned);
}
fn exit_lambda(&mut self, _l: &brink_ir::LambdaExpr) {
self.lambda_locals.pop();
}
}
struct Receiver<'a> {
def: DefinitionId,
segments: &'a [brink_ir::Name],
text: String,
ty: Ty,
}
impl UfcsVisitor<'_> {
fn current_body(&self) -> Option<&crate::infer::BodyTypes> {
self.stitch_body.or(self.knot_body)
}
fn current_locals(&self) -> Option<&BTreeMap<String, Ty>> {
self.lambda_locals
.last()
.or_else(|| self.current_body().map(|b| &b.locals))
}
fn resolve_call(&mut self, path: &HirPath, arg_count: usize) {
let Some((method, receiver_segs)) = path.segments.split_last() else {
return;
};
if receiver_segs.is_empty() {
return;
}
let Some(head_def) = self.value_receiver_def(path) else {
return;
};
let receiver_text = receiver_segs
.iter()
.map(|s| s.text.as_str())
.collect::<Vec<_>>()
.join(".");
let Some(receiver_ty) = self.receiver_ty(head_def, receiver_segs) else {
self.push(
path.range,
DiagnosticCode::E142,
&format!(
"cannot resolve `{receiver_text}.{method}(…)`: the type of `{receiver_text}` \
is not known here, so it is undecidable whether `{method}` is one of its \
fields — annotate the receiver",
method = method.text,
),
);
return;
};
let receiver = Receiver {
def: head_def,
segments: receiver_segs,
text: receiver_text,
ty: receiver_ty,
};
if self.try_field_call(path, method, &receiver, arg_count) {
return;
}
if self.try_free_fn_desugar(path, method, &receiver, arg_count) {
return;
}
self.push(
path.range,
DiagnosticCode::E141,
&format!(
"cannot resolve `{receiver_text}.{method}(…)`: `{recv_ty}` declares no field \
`{method}`, and no function `{method}` is in scope here",
method = method.text,
receiver_text = receiver.text,
recv_ty = receiver.ty.display(),
),
);
}
fn try_field_call(
&mut self,
path: &HirPath,
method: &brink_ir::Name,
receiver: &Receiver<'_>,
arg_count: usize,
) -> bool {
let receiver_ty = &receiver.ty;
let receiver_text = &receiver.text;
let Ty::Struct(shape_name) = receiver_ty else {
return false;
};
let Some(field_ty) = self
.shapes
.resolve(shape_name, self.scope, self.index)
.and_then(|shape| shape.field_ty(&method.text))
else {
return false;
};
if matches!(field_ty, Ty::Fn(..)) {
let (arity_mismatch, arg_mismatches) =
self.check_field_call_args(path.range, field_ty, arg_count);
let verdict = UfcsVerdict::FieldCall {
receiver: receiver_ty.clone(),
field: method.text.clone(),
field_ty: field_ty.clone(),
arity_mismatch,
arg_mismatches,
};
self.table
.insert(NodeKey::new(self.file, path.range), verdict);
} else {
let message = format!(
"field `{field}` on `{shape_name}` is not callable (its type is `{found}`) — \
field access wins over a free function of the same name, so this is never \
re-read as `{field}({receiver_text}, …)`",
field = method.text,
found = field_ty.display(),
);
self.push(path.range, DiagnosticCode::E140, &message);
}
true
}
fn try_free_fn_desugar(
&mut self,
path: &HirPath,
method: &brink_ir::Name,
receiver: &Receiver<'_>,
arg_count: usize,
) -> bool {
let Some(target) = crate::resolve::lookup_by_name(
self.index,
self.scope,
&method.text,
&[SymbolKind::Knot, SymbolKind::External],
) else {
if crate::resolve::is_t1b_stdlib_name(&method.text)
|| crate::resolve::is_builtin_function(&method.text)
{
let arg_mismatches =
self.check_ufcs_prelude_arg_types(path.range, &receiver.ty, &method.text);
let verdict = UfcsVerdict::PreludeDesugar {
receiver: receiver.ty.clone(),
name: method.text.clone(),
arg_mismatches,
};
self.table
.insert(NodeKey::new(self.file, path.range), verdict);
return true;
}
return false;
};
let first_param_is_ref = self
.index
.symbols
.get(&target)
.and_then(|info| info.params.first())
.is_some_and(|p| p.is_ref);
if first_param_is_ref && let Some(cause) = self.auto_ref_fault(receiver) {
let message = format!(
"cannot mutate `{receiver_text}` through `{name}`: `{name}`'s first parameter is \
`ref`, so `{receiver_text}.{name}(…)` auto-refs its receiver (D5) — but {cause}. \
Bind the receiver to a durable cell, or call a by-value function on it",
name = method.text,
receiver_text = receiver.text,
);
self.push(path.range, DiagnosticCode::E143, &message);
return true;
}
let expected = self
.index
.symbols
.get(&target)
.map(|info| info.params.len());
let actual = arg_count + 1;
if let Some(expected) = expected
&& expected != actual
{
let message = format!(
"`{name}` expects {expected} argument(s), got {actual} \
(`{receiver_text}.{name}(…)` desugars to `{name}({receiver_text}, …)`, counting \
the receiver as the first argument)",
name = method.text,
receiver_text = receiver.text,
);
self.push(path.range, DiagnosticCode::E031, &message);
}
let arg_mismatches = self.check_ufcs_arg_types(path.range, target, receiver);
let verdict = if first_param_is_ref {
UfcsVerdict::FreeFnAutoRef {
receiver: receiver.ty.clone(),
name: method.text.clone(),
target,
arg_mismatches,
}
} else {
UfcsVerdict::FreeFnDesugar {
receiver: receiver.ty.clone(),
name: method.text.clone(),
target,
arg_mismatches,
}
};
self.table
.insert(NodeKey::new(self.file, path.range), verdict);
true
}
fn check_ufcs_arg_types(
&self,
range: TextRange,
target: DefinitionId,
receiver: &Receiver<'_>,
) -> Vec<UfcsArgMismatch> {
let Some(sig) = self.signatures.get(&target) else {
return Vec::new();
};
let empty: Vec<Ty> = Vec::new();
let written: &[Ty] = self
.current_body()
.and_then(|b| b.ufcs_call_args.iter().find(|f| f.range == range))
.map_or(empty.as_slice(), |f| f.args.as_slice());
let ref_positions = self.index.symbols.get(&target);
let is_ref_param = |i: usize| {
ref_positions
.and_then(|info| info.params.get(i))
.is_some_and(|p| p.is_ref)
};
let mut mismatches = Vec::new();
if let Some(param_ty) = sig.params.first()
&& !param_ty.is_unresolved()
&& if is_ref_param(0) {
!ref_assignable(param_ty, &receiver.ty)
} else {
!assignable(param_ty, &receiver.ty)
}
{
mismatches.push(UfcsArgMismatch {
index: 0,
expected: param_ty.clone(),
found: receiver.ty.clone(),
});
}
for (i, arg_ty) in written.iter().enumerate() {
if arg_ty.is_unresolved() {
continue;
}
let Some(param_ty) = sig.params.get(i + 1) else {
continue;
};
let ty_disagrees = if is_ref_param(i + 1) {
!ref_assignable(param_ty, arg_ty)
} else {
!assignable(param_ty, arg_ty)
};
if !param_ty.is_unresolved() && ty_disagrees {
mismatches.push(UfcsArgMismatch {
index: i + 1,
expected: param_ty.clone(),
found: arg_ty.clone(),
});
}
}
mismatches
}
fn check_ufcs_prelude_arg_types(
&self,
range: TextRange,
receiver: &Ty,
name: &str,
) -> Vec<UfcsArgMismatch> {
let expected: Vec<Ty> = match (name, receiver) {
("push" | "heap_push" | "index_of" | "contains", Ty::Array(elem)) => {
vec![(**elem).clone()]
}
("contains" | "get" | "remove", Ty::Map(k, _)) => vec![(**k).clone()],
("contains_value", Ty::Map(_, v)) => vec![(**v).clone()],
("insert", Ty::Map(k, v)) => vec![(**k).clone(), (**v).clone()],
_ => return Vec::new(),
};
let empty: Vec<Ty> = Vec::new();
let written: &[Ty] = self
.current_body()
.and_then(|b| b.ufcs_call_args.iter().find(|f| f.range == range))
.map_or(empty.as_slice(), |f| f.args.as_slice());
let mut mismatches = Vec::new();
for (i, expected_ty) in expected.iter().enumerate() {
let Some(arg_ty) = written.get(i) else {
continue;
};
if arg_ty.is_unresolved() {
continue;
}
if !assignable(expected_ty, arg_ty) {
mismatches.push(UfcsArgMismatch {
index: i + 1,
expected: expected_ty.clone(),
found: arg_ty.clone(),
});
}
}
mismatches
}
fn check_field_call_args(
&self,
range: TextRange,
field_ty: &Ty,
arg_count: usize,
) -> (Option<UfcsArityMismatch>, Vec<UfcsArgMismatch>) {
let Ty::Fn(params, _ret, _) = field_ty else {
return (None, Vec::new());
};
let arity_mismatch = (arg_count != params.len()).then_some(UfcsArityMismatch {
expected: params.len(),
got: arg_count,
});
let empty: Vec<Ty> = Vec::new();
let written: &[Ty] = self
.current_body()
.and_then(|b| b.ufcs_call_args.iter().find(|f| f.range == range))
.map_or(empty.as_slice(), |f| f.args.as_slice());
let mut mismatches = Vec::new();
for (i, param_ty) in params.iter().enumerate() {
let Some(arg_ty) = written.get(i) else {
continue;
};
if arg_ty.is_unresolved() || param_ty.is_unresolved() {
continue;
}
if !assignable(param_ty, arg_ty) {
mismatches.push(UfcsArgMismatch {
index: i,
expected: param_ty.clone(),
found: arg_ty.clone(),
});
}
}
(arity_mismatch, mismatches)
}
fn auto_ref_fault(&self, receiver: &Receiver<'_>) -> Option<String> {
let head = receiver.segments.first().map_or("", |s| s.text.as_str());
let frame_local = || {
(receiver.segments.len() > 2).then(|| {
format!(
"`{head}` is a temp/param — a frame-local projection can only reach one \
field level (`{head}.field`); this receiver goes deeper than that"
)
})
};
match self.index.symbols.get(&receiver.def) {
Some(info) => match info.kind {
SymbolKind::Variable => None,
SymbolKind::Constant => Some(format!("`{head}` is a CONST, not a mutable cell")),
SymbolKind::Param | SymbolKind::Temp => frame_local(),
_ => Some(format!(
"`{head}` is not a value that can be written through"
)),
},
None => frame_local(),
}
}
fn value_receiver_def(&self, path: &HirPath) -> Option<DefinitionId> {
let key = (path.range.start().into(), path.range.end().into());
let &target = self.resolution_by_range.get(&key)?;
match self.index.symbols.get(&target) {
Some(info)
if matches!(
info.kind,
SymbolKind::Param
| SymbolKind::Temp
| SymbolKind::Variable
| SymbolKind::Constant
) =>
{
Some(target)
}
None if target.tag() == brink_format::DefinitionTag::LocalVar => Some(target),
Some(_) | None => None,
}
}
fn receiver_ty(&self, head_def: DefinitionId, segments: &[brink_ir::Name]) -> Option<Ty> {
let (head, rest) = segments.split_first()?;
let mut ty = self.head_ty(head_def, head)?;
for seg in rest {
let Ty::Struct(shape_name) = &ty else {
return None;
};
let field = self
.shapes
.resolve(shape_name, self.scope, self.index)?
.field_ty(&seg.text)?
.clone();
ty = field;
}
(!ty.is_unknown() && ty != Ty::Conflicted).then_some(ty)
}
fn head_ty(&self, def: DefinitionId, head: &brink_ir::Name) -> Option<Ty> {
match self.index.symbols.get(&def) {
Some(info) => match info.kind {
SymbolKind::Param | SymbolKind::Temp => {
self.current_locals()?.get(&info.name).cloned()
}
SymbolKind::Variable | SymbolKind::Constant => self.globals.get(&def).cloned(),
_ => None,
},
None => self.current_locals()?.get(&head.text).cloned(),
}
}
fn push(&mut self, range: TextRange, code: DiagnosticCode, detail: &str) {
self.diagnostics.push(Diagnostic {
file: self.file,
range,
message: format!("{}: {detail}", code.title()),
code,
});
}
}