use syntax::ast::{Attribute, Expression, Span, VariantFields};
use syntax::program::{Definition, DefinitionBody, Visibility};
use syntax::types::{CompoundKind, Symbol, Type};
use super::TaskState;
use crate::store::Store;
impl TaskState<'_> {
pub(super) fn register_iterables(&mut self, store: &mut Store, items: &[Expression]) {
let module_id = self.cursor.module_id.clone();
let is_d_lis = self.is_d_lis(store);
let mut candidates = Vec::new();
for item in items {
collect_iterable_candidates(item, is_d_lis, &mut candidates);
}
for candidate in candidates {
self.process_iterable_candidate(store, &module_id, candidate);
}
}
pub(super) fn register_module_iterables(&mut self, store: &mut Store, module_id: &str) {
let candidates = {
let module = store.get_module(module_id).expect("module must exist");
let mut candidates = Vec::new();
for file in module.files.values() {
let is_d_lis = file.is_d_lis();
for item in &file.items {
collect_iterable_candidates(item, is_d_lis, &mut candidates);
}
}
candidates
};
for candidate in candidates {
self.process_iterable_candidate(store, module_id, candidate);
}
}
fn process_iterable_candidate(
&mut self,
store: &mut Store,
module_id: &str,
candidate: IterableCandidate,
) {
let IterableCandidate {
attribute_span,
kind,
} = candidate;
let EnumCandidate {
name,
name_span,
is_generic,
is_d_lis,
payload_variant_span,
} = match kind {
CandidateKind::NotAnEnum => {
self.sink.push(diagnostics::attribute::iterable_not_an_enum(
&attribute_span,
));
return;
}
CandidateKind::Enum(enum_candidate) => enum_candidate,
};
if is_d_lis {
self.sink
.push(diagnostics::attribute::iterable_in_typedef(&attribute_span));
return;
}
if is_generic {
self.sink
.push(diagnostics::attribute::iterable_generic_enum(
&attribute_span,
));
return;
}
if let Some(variant_span) = payload_variant_span {
self.sink
.push(diagnostics::attribute::iterable_non_unit_variant(
&attribute_span,
&variant_span,
));
return;
}
let qualified = Symbol::from_parts(module_id, &name);
let variants_key = qualified.with_segment("variants");
let existing_span = store
.get_definition(variants_key.as_str())
.and_then(|definition| definition.name_span);
let (has_instance_variants, visibility) = match store.get_definition(qualified.as_str()) {
Some(definition) => (
matches!(&definition.body, DefinitionBody::Enum { methods, .. } if methods.contains_key("variants")),
definition.visibility().clone(),
),
None => (false, Visibility::Private),
};
if existing_span.is_some() || has_instance_variants {
self.sink
.push(diagnostics::attribute::iterable_variants_conflict(
&attribute_span,
existing_span.as_ref(),
));
return;
}
let Some(enum_ty) = store.get_type(qualified.as_str()).cloned() else {
return;
};
let slice_ty = Type::Compound {
kind: CompoundKind::Slice,
args: vec![enum_ty],
};
let fn_ty = Type::function(vec![], vec![], Default::default(), Box::new(slice_ty));
let module = store.get_module_mut(module_id).expect("module must exist");
module.definitions.insert(
variants_key,
Definition {
visibility,
ty: fn_ty,
name: None,
name_span: Some(name_span),
doc: None,
body: DefinitionBody::Value {
allowed_lints: vec![],
go_hints: vec![],
go_name: None,
},
},
);
}
}
struct IterableCandidate {
attribute_span: Span,
kind: CandidateKind,
}
enum CandidateKind {
Enum(EnumCandidate),
NotAnEnum,
}
struct EnumCandidate {
name: String,
name_span: Span,
is_generic: bool,
is_d_lis: bool,
payload_variant_span: Option<Span>,
}
fn iterable_attribute_span(attributes: &[Attribute]) -> Option<Span> {
attributes
.iter()
.find(|a| a.name == "iterable")
.map(|a| a.span)
}
fn misplaced_candidate(span: Span) -> IterableCandidate {
IterableCandidate {
attribute_span: span,
kind: CandidateKind::NotAnEnum,
}
}
fn collect_method_attributes(methods: &[Expression], out: &mut Vec<IterableCandidate>) {
for method in methods {
if let Expression::Function { attributes, .. } = method {
out.extend(iterable_attribute_span(attributes).map(misplaced_candidate));
}
}
}
fn collect_iterable_candidates(
item: &Expression,
is_d_lis: bool,
out: &mut Vec<IterableCandidate>,
) {
match item {
Expression::Enum {
attributes,
name,
name_span,
generics,
variants,
..
} => {
if let Some(attribute_span) = iterable_attribute_span(attributes) {
let payload_variant_span = variants
.iter()
.find(|v| !matches!(v.fields, VariantFields::Unit))
.map(|v| v.name_span);
out.push(IterableCandidate {
attribute_span,
kind: CandidateKind::Enum(EnumCandidate {
name: name.to_string(),
name_span: *name_span,
is_generic: !generics.is_empty(),
is_d_lis,
payload_variant_span,
}),
});
}
}
Expression::Struct {
attributes, fields, ..
} => {
out.extend(iterable_attribute_span(attributes).map(misplaced_candidate));
for field in fields {
out.extend(iterable_attribute_span(&field.attributes).map(misplaced_candidate));
}
}
Expression::Function { attributes, .. } => {
out.extend(iterable_attribute_span(attributes).map(misplaced_candidate));
}
Expression::ImplBlock { methods, .. } => collect_method_attributes(methods, out),
Expression::Interface {
method_signatures, ..
} => collect_method_attributes(method_signatures, out),
_ => {}
}
}