use code_moniker_core::core::kinds::REF_REEXPORTS;
use code_moniker_core::core::moniker::query::bare_callable_name;
use code_moniker_core::core::moniker::{Moniker, MonikerBuilder};
use code_moniker_core::lang::kinds;
use rustc_hash::{FxHashMap, FxHashSet};
use crate::linkage::binding::{
ExternalOrigin, ReferenceLinkageDecision, ResolutionDecision, ResolutionScope,
};
use crate::linkage::catalog::{CandidateCatalog, SymbolSet};
use crate::snapshot::{
CandidateReason, DynamicReason, RecordTable, ReferenceRecord, ResolutionEvidence,
};
use crate::source::CodeIndexMaterial;
pub(in crate::linkage) struct PythonBindings {
aliases: FxHashMap<Moniker, FxHashMap<Vec<u8>, BindingTarget>>,
export_policies: FxHashMap<Moniker, ExportPolicy>,
wildcard_imports: Vec<WildcardImport>,
dynamic_wildcard_owners: FxHashSet<Moniker>,
}
impl PythonBindings {
pub(in crate::linkage) fn build(
material: &CodeIndexMaterial,
candidates: &CandidateCatalog,
decisions: &[ReferenceLinkageDecision],
references: &RecordTable<ReferenceRecord>,
decision_indices: &[usize],
) -> Self {
let mut bindings = Self {
aliases: FxHashMap::default(),
export_policies: FxHashMap::default(),
wildcard_imports: Vec::new(),
dynamic_wildcard_owners: FxHashSet::default(),
};
let mut pending_bindings = Vec::new();
bindings.seed_module_definitions(material, candidates);
for &decision_idx in decision_indices {
let decision = &decisions[decision_idx];
bindings.record_decision(
material,
decision,
&references[decision.reference_idx()],
&mut pending_bindings,
);
}
bindings.propagate_bindings(&pending_bindings);
bindings
}
fn record_decision(
&mut self,
material: &CodeIndexMaterial,
decision: &ReferenceLinkageDecision,
reference: &ReferenceRecord,
pending_bindings: &mut Vec<PendingBinding>,
) {
let Some(raw_owner) = material.symbol_moniker(&reference.source_symbol) else {
return;
};
let is_reexport = reference.kind.as_bytes() == REF_REEXPORTS;
let receiver = reference.receiver.as_deref();
let python_module = is_python_module(raw_owner);
if !is_reexport && !python_module {
return;
}
let owners = if python_module {
binding_owners(raw_owner)
} else {
vec![raw_owner.clone()]
};
if is_reexport && reference.alias.as_deref().is_none_or(str::is_empty) {
match receiver {
Some("python_all_replace") => {
for owner in owners {
self.export_policies
.insert(owner, ExportPolicy::Static(FxHashSet::default()));
}
return;
}
Some("python_all_extend") => {
for owner in owners {
self.export_policies
.entry(owner)
.or_insert(ExportPolicy::Dynamic);
}
return;
}
Some("python_all_dynamic") => {
for owner in owners {
self.export_policies.insert(owner, ExportPolicy::Dynamic);
}
return;
}
_ => {}
}
}
if python_module
&& reference.kind.as_bytes() == kinds::IMPORTS_MODULE
&& reference.alias.as_deref() == Some("*")
{
if receiver == Some("python_conditional_import") {
self.dynamic_wildcard_owners.extend(owners);
return;
}
let Some(target) = material.reference_target(&reference.id).cloned() else {
return;
};
let external = match decision {
ReferenceLinkageDecision::External { origin, .. } => Some(*origin),
_ => None,
};
for owner in owners {
self.wildcard_imports.push(WildcardImport {
owner,
target: target.clone(),
external,
});
}
return;
}
if !is_reexport
&& !matches!(
reference.kind.as_bytes(),
kinds::IMPORTS_SYMBOL | kinds::IMPORTS_MODULE
) {
return;
}
let Some(name) = binding_name(material, reference) else {
return;
};
if is_reexport {
for owner in &owners {
let policy = self
.export_policies
.entry(owner.clone())
.or_insert_with(|| ExportPolicy::Static(FxHashSet::default()));
if let ExportPolicy::Static(names) = policy {
names.insert(name.clone());
}
}
}
let fallback = material.reference_target(&reference.id).cloned();
let Some(mut target) = BindingTarget::from_decision(decision, fallback) else {
if let Some((target_owner, target_name)) = Self::target_key(material, reference) {
pending_bindings.push(PendingBinding {
owners,
name,
target_owner,
target_name,
conditional: receiver == Some("python_conditional_import"),
});
}
return;
};
if receiver == Some("python_conditional_import") {
target = target.into_dynamic();
}
for owner in owners {
self.merge_alias((owner, name.clone()), target.clone());
}
}
fn seed_module_definitions(
&mut self,
material: &CodeIndexMaterial,
candidates: &CandidateCatalog,
) {
for (file_idx, file) in material.files.iter().enumerate() {
if file.lang != code_moniker_core::lang::Lang::Python {
continue;
}
for (def_idx, def) in file.graph.defs().enumerate() {
let Some(owner) = def.moniker.parent().filter(is_python_module) else {
continue;
};
let Some(segment) = def.moniker.as_view().segments().last() else {
continue;
};
let Some(symbol) = candidates.symbol_at(file_idx, def_idx) else {
continue;
};
let target = BindingTarget::Workspace {
scope: ResolutionScope::Global,
evidence: ResolutionEvidence::ExactBinding,
targets: SymbolSet::from_symbol(symbol),
candidate_reason: None,
};
for owner in binding_owners(&owner) {
self.merge_alias(
(owner, bare_callable_name(segment.name).to_vec()),
target.clone(),
);
}
}
}
}
pub(in crate::linkage) fn target_key(
material: &CodeIndexMaterial,
reference: &ReferenceRecord,
) -> Option<(Moniker, Vec<u8>)> {
let target = material.reference_target(&reference.id)?;
let name = target
.as_view()
.segments()
.last()
.map(|segment| bare_callable_name(segment.name).to_vec())
.or_else(|| {
reference
.call_name
.as_deref()
.map(|name| name.as_bytes().to_vec())
})?;
Some((target.parent()?, name))
}
pub(in crate::linkage) fn alias(&self, owner: &Moniker, name: &[u8]) -> Option<&BindingTarget> {
self.aliases.get(owner)?.get(name)
}
pub(in crate::linkage) fn has_dynamic_wildcards(&self, owner: &Moniker) -> bool {
self.dynamic_wildcard_owners.contains(owner)
}
pub(in crate::linkage) fn external_wildcards(
&self,
owner: &Moniker,
) -> impl Iterator<Item = (ExternalOrigin, &Moniker)> {
self.wildcard_imports
.iter()
.filter(move |wildcard| &wildcard.owner == owner)
.filter_map(|wildcard| wildcard.external.map(|origin| (origin, &wildcard.target)))
}
#[cfg(test)]
pub(in crate::linkage) fn with_external_wildcard(
owner: Moniker,
target: Moniker,
origin: ExternalOrigin,
) -> Self {
Self {
aliases: FxHashMap::default(),
export_policies: FxHashMap::default(),
wildcard_imports: vec![WildcardImport {
owner,
target,
external: Some(origin),
}],
dynamic_wildcard_owners: FxHashSet::default(),
}
}
fn propagate_wildcards(&mut self) {
loop {
let mut changed = false;
for wildcard in self.wildcard_imports.clone() {
if wildcard.external.is_some() {
continue;
}
if matches!(
self.export_policies.get(&wildcard.target),
Some(ExportPolicy::Dynamic)
) || self.dynamic_wildcard_owners.contains(&wildcard.target)
{
changed |= self.dynamic_wildcard_owners.insert(wildcard.owner.clone());
continue;
}
let exports = self.exported_aliases(&wildcard.target);
for (name, target) in exports {
changed |= self.merge_alias((wildcard.owner.clone(), name), target);
}
}
if !changed {
break;
}
}
}
fn propagate_bindings(&mut self, pending_bindings: &[PendingBinding]) {
loop {
self.propagate_wildcards();
let mut changed = false;
for pending in pending_bindings {
let Some(mut target) = self
.alias(&pending.target_owner, &pending.target_name)
.cloned()
else {
continue;
};
if pending.conditional {
target = target.into_dynamic();
}
for owner in &pending.owners {
changed |=
self.merge_alias((owner.clone(), pending.name.clone()), target.clone());
}
}
if !changed {
break;
}
}
}
fn exported_aliases(&self, owner: &Moniker) -> Vec<(Vec<u8>, BindingTarget)> {
let policy = self.export_policies.get(owner);
self.aliases
.get(owner)
.into_iter()
.flat_map(|aliases| aliases.iter())
.filter(|(name, _)| match policy {
None => !name.starts_with(b"_"),
Some(ExportPolicy::Static(names)) => names.contains(*name),
Some(ExportPolicy::Dynamic) => false,
})
.map(|(name, target)| (name.clone(), target.clone()))
.collect()
}
fn merge_alias(&mut self, key: (Moniker, Vec<u8>), target: BindingTarget) -> bool {
let (owner, name) = key;
let aliases = self.aliases.entry(owner).or_default();
let Some(existing) = aliases.remove(&name) else {
aliases.insert(name, target);
return true;
};
let merged = existing.clone().merge(target);
let changed = merged != existing;
aliases.insert(name, merged);
changed
}
}
enum ExportPolicy {
Static(FxHashSet<Vec<u8>>),
Dynamic,
}
#[derive(Clone)]
struct WildcardImport {
owner: Moniker,
target: Moniker,
external: Option<ExternalOrigin>,
}
#[derive(Clone)]
struct PendingBinding {
owners: Vec<Moniker>,
name: Vec<u8>,
target_owner: Moniker,
target_name: Vec<u8>,
conditional: bool,
}
#[derive(Clone, Eq, PartialEq)]
pub(in crate::linkage) enum BindingTarget {
Workspace {
scope: ResolutionScope,
evidence: ResolutionEvidence,
targets: SymbolSet,
candidate_reason: Option<CandidateReason>,
},
External {
origin: ExternalOrigin,
target: Moniker,
},
Dynamic {
candidates: SymbolSet,
},
}
impl BindingTarget {
fn into_dynamic(self) -> Self {
match self {
Self::Workspace { targets, .. }
| Self::Dynamic {
candidates: targets,
} => Self::Dynamic {
candidates: targets,
},
Self::External { .. } => Self::Dynamic {
candidates: SymbolSet::new(),
},
}
}
pub(in crate::linkage) fn workspace_candidates(&self) -> SymbolSet {
match self {
Self::Workspace { targets, .. } => targets.clone(),
Self::Dynamic { candidates } => candidates.clone(),
Self::External { .. } => SymbolSet::new(),
}
}
fn from_decision(
decision: &ReferenceLinkageDecision,
fallback_external_target: Option<Moniker>,
) -> Option<Self> {
match decision {
ReferenceLinkageDecision::Unique { resolution } => {
Some(Self::from_workspace_resolution(resolution, None))
}
ReferenceLinkageDecision::Candidate { reason, resolution } => {
Some(Self::from_workspace_resolution(resolution, Some(*reason)))
}
ReferenceLinkageDecision::External { origin, target, .. } => Some(Self::External {
origin: *origin,
target: target.clone().or(fallback_external_target)?,
}),
ReferenceLinkageDecision::Dynamic { candidates, .. } => Some(Self::Dynamic {
candidates: candidates.clone(),
}),
_ => None,
}
}
fn from_workspace_resolution(
resolution: &ResolutionDecision,
candidate_reason: Option<CandidateReason>,
) -> Self {
Self::Workspace {
scope: resolution.scope,
evidence: resolution.evidence,
targets: resolution.targets.clone(),
candidate_reason,
}
}
fn merge(self, other: Self) -> Self {
match (self, other) {
(
Self::Workspace {
scope,
evidence,
targets,
candidate_reason,
},
Self::Workspace {
scope: other_scope,
evidence: other_evidence,
targets: other_targets,
candidate_reason: other_reason,
},
) => {
let merged = targets.union(&other_targets);
let candidate_reason = candidate_reason
.or(other_reason)
.or_else(|| (merged.len() > 1).then_some(CandidateReason::MultipleTargets));
Self::Workspace {
scope: if scope == other_scope {
scope
} else {
ResolutionScope::Global
},
evidence: if evidence == other_evidence {
evidence
} else {
ResolutionEvidence::GlobalBinding
},
targets: merged,
candidate_reason,
}
}
(Self::External { origin, target }, Self::External { target: other, .. })
if target == other =>
{
Self::External { origin, target }
}
(Self::Dynamic { candidates }, Self::Workspace { targets, .. }) => Self::Dynamic {
candidates: candidates.union(&targets),
},
(Self::Workspace { targets, .. }, Self::Dynamic { candidates }) => Self::Dynamic {
candidates: targets.union(&candidates),
},
(Self::Dynamic { candidates }, Self::Dynamic { candidates: other }) => Self::Dynamic {
candidates: candidates.union(&other),
},
(Self::Workspace { targets, .. }, Self::External { .. }) => Self::Dynamic {
candidates: targets,
},
(Self::External { .. }, Self::Workspace { targets, .. }) => Self::Dynamic {
candidates: targets,
},
(Self::Dynamic { candidates }, Self::External { .. }) => Self::Dynamic { candidates },
(Self::External { .. }, Self::Dynamic { candidates }) => Self::Dynamic { candidates },
(Self::External { .. }, Self::External { .. }) => Self::Dynamic {
candidates: SymbolSet::new(),
},
}
}
pub(in crate::linkage) fn to_decision(
&self,
reference_idx: usize,
reference: &ReferenceRecord,
requested_target: Option<&Moniker>,
) -> ReferenceLinkageDecision {
match self {
Self::Workspace {
scope,
evidence,
targets,
candidate_reason,
} => {
let resolution = ResolutionDecision::new(
*scope,
*evidence,
reference.id,
reference_idx,
targets.clone(),
);
if targets.len() == 1 && candidate_reason.is_none() {
ReferenceLinkageDecision::resolved(resolution)
} else {
ReferenceLinkageDecision::candidate(
candidate_reason.unwrap_or(CandidateReason::MultipleTargets),
resolution,
)
}
}
Self::External { origin, target } => ReferenceLinkageDecision::external_target(
*origin,
reference_idx,
reference.id,
reexport_external_target(target, requested_target),
),
Self::Dynamic { candidates } => ReferenceLinkageDecision::dynamic(
DynamicReason::RuntimeImport,
reference_idx,
reference.id,
candidates.clone(),
),
}
}
}
fn is_python_module(owner: &Moniker) -> bool {
let segments = owner.as_view().segments().collect::<Vec<_>>();
segments
.first()
.is_some_and(|segment| segment.kind == kinds::LANG && segment.name == b"python")
&& segments
.last()
.is_some_and(|segment| segment.kind == kinds::MODULE)
}
fn binding_owners(owner: &Moniker) -> Vec<Moniker> {
let mut owners = vec![owner.clone()];
for alternate in [collapsed_init_owner(owner), folded_init_import_owner(owner)]
.into_iter()
.flatten()
{
if !owners.contains(&alternate) {
owners.push(alternate);
}
}
owners
}
fn collapsed_init_owner(owner: &Moniker) -> Option<Moniker> {
let segments = owner.as_view().segments().collect::<Vec<_>>();
let [first, .., package, module] = segments.as_slice() else {
return None;
};
if first.kind != kinds::LANG
|| first.name != b"python"
|| module.kind != kinds::MODULE
|| module.name != b"__init__"
|| package.kind != kinds::PACKAGE
{
return None;
}
let package_name = package.name.to_vec();
let prefix = owner.parent()?.parent()?;
Some(
MonikerBuilder::from_view(prefix.as_view())
.segment(kinds::MODULE, &package_name)
.build(),
)
}
fn folded_init_import_owner(owner: &Moniker) -> Option<Moniker> {
let view = owner.as_view();
let segments = view.segments().collect::<Vec<_>>();
let [lang, packages @ .., module] = segments.as_slice() else {
return None;
};
if lang.kind != kinds::LANG
|| lang.name != b"python"
|| packages.is_empty()
|| packages
.iter()
.any(|segment| segment.kind != kinds::PACKAGE)
|| module.kind != kinds::MODULE
|| module.name != b"__init__"
{
return None;
}
let mut builder = MonikerBuilder::new();
builder
.project(view.project())
.segment(lang.kind, lang.name);
if packages.len() == 1 {
builder.segment(kinds::MODULE, packages[0].name);
} else {
builder
.segment(kinds::PACKAGE, packages[0].name)
.segment(kinds::MODULE, packages[1].name);
for package in &packages[2..] {
builder.segment(kinds::PATH, package.name);
}
}
Some(builder.build())
}
fn binding_name(material: &CodeIndexMaterial, reference: &ReferenceRecord) -> Option<Vec<u8>> {
if let Some(alias) = reference.alias.as_deref().filter(|alias| !alias.is_empty()) {
return Some(alias.as_bytes().to_vec());
}
let target = material.reference_target(&reference.id)?;
if reference.kind.as_bytes() == kinds::IMPORTS_MODULE {
let mut segments = target.as_view().segments();
let first = segments.next()?;
let binding = if matches!(first.kind, kinds::LANG | kinds::SDK) {
segments.next()?
} else {
first
};
return Some(bare_callable_name(binding.name).to_vec());
}
let last = target.as_view().segments().last()?;
(last.kind == kinds::PATH).then(|| bare_callable_name(last.name).to_vec())
}
fn reexport_external_target(alias_target: &Moniker, requested_target: Option<&Moniker>) -> Moniker {
let Some(requested_target) = requested_target else {
return alias_target.clone();
};
let Some(alias_last) = alias_target.as_view().segments().last() else {
return alias_target.clone();
};
let Some(requested_last) = requested_target.as_view().segments().last() else {
return alias_target.clone();
};
if bare_callable_name(alias_last.name) != bare_callable_name(requested_last.name) {
return alias_target.clone();
}
let Some(owner) = alias_target.parent() else {
return alias_target.clone();
};
MonikerBuilder::from_view(owner.as_view())
.segment(requested_last.kind, requested_last.name)
.build()
}