Skip to main content

mars_agents/resolve/
context.rs

1use std::collections::HashMap;
2
3use indexmap::IndexMap;
4use semver::Version;
5
6use super::filter::push_filter_constraint;
7use super::{
8    PackageResolutionState, PackageVersions, PendingItem, RegisteredPackage, ResolvedGraph,
9    ResolvedNode, RootedSourceRef, VersionConstraint, VisitedSet,
10};
11use crate::config::FilterMode;
12use crate::source::ResolvedRef;
13use crate::types::{SourceId, SourceName};
14
15/// Mutable resolver state threaded through bottom-up resolution and DFS traversal.
16pub struct ResolverContext {
17    registry: IndexMap<SourceName, RegisteredPackage>,
18    package_states: HashMap<SourceName, PackageResolutionState>,
19    id_index: HashMap<SourceId, SourceName>,
20    version_constraints: HashMap<SourceName, Vec<(String, VersionConstraint)>>,
21    materialization_filters: HashMap<SourceName, Vec<FilterMode>>,
22    stack: Vec<PendingItem>,
23    visited: VisitedSet,
24    package_versions: PackageVersions,
25    /// Version overrides carried from a prior restart pass.
26    ///
27    /// When a restart is triggered because package X would resolve to a different
28    /// version under the full accumulated constraint set, the driver carries the
29    /// correct (new) ref into the fresh context via this map. The first-resolution
30    /// branch in `resolve_package_bottom_up` checks this map and uses the override
31    /// directly, so the same constraint-accumulation pattern does NOT re-trigger a
32    /// restart on the next pass.
33    version_overrides: HashMap<SourceName, (ResolvedRef, RootedSourceRef, Option<Version>)>,
34    /// Pending restart info set by `resolve_package_bottom_up` just before it returns
35    /// `ResolutionRestartNeeded`. The driver reads this before discarding the context.
36    pending_restart: Option<(SourceName, ResolvedRef, RootedSourceRef, Option<Version>)>,
37}
38
39impl Default for ResolverContext {
40    fn default() -> Self {
41        Self::new()
42    }
43}
44
45impl ResolverContext {
46    pub fn new() -> Self {
47        Self {
48            registry: IndexMap::new(),
49            package_states: HashMap::new(),
50            id_index: HashMap::new(),
51            version_constraints: HashMap::new(),
52            materialization_filters: HashMap::new(),
53            stack: Vec::new(),
54            visited: VisitedSet::new(),
55            package_versions: PackageVersions::new(),
56            version_overrides: HashMap::new(),
57            pending_restart: None,
58        }
59    }
60
61    /// Set version overrides from a prior restart pass.
62    /// These are used by `resolve_package_bottom_up` to skip re-resolution for
63    /// packages where the correct version was already computed.
64    pub(super) fn set_version_overrides(
65        &mut self,
66        overrides: HashMap<SourceName, (ResolvedRef, RootedSourceRef, Option<Version>)>,
67    ) {
68        self.version_overrides = overrides;
69    }
70
71    /// Look up an override for the first resolution of `name`.
72    /// Returns the pre-computed (ResolvedRef, RootedSourceRef, latest_version) if present.
73    pub(super) fn version_override(
74        &self,
75        name: &SourceName,
76    ) -> Option<&(ResolvedRef, RootedSourceRef, Option<Version>)> {
77        self.version_overrides.get(name)
78    }
79
80    /// Record the restart info: the package that triggered a restart and the ref
81    /// it should be resolved to on the next pass. Called by `resolve_package_bottom_up`
82    /// just before returning `ResolutionRestartNeeded`.
83    pub(super) fn set_pending_restart(
84        &mut self,
85        package: SourceName,
86        new_ref: ResolvedRef,
87        new_rooted: RootedSourceRef,
88        latest_version: Option<Version>,
89    ) {
90        self.pending_restart = Some((package, new_ref, new_rooted, latest_version));
91    }
92
93    /// Drain the pending restart info. Called by the driver after catching the signal.
94    pub(super) fn take_pending_restart(
95        &mut self,
96    ) -> Option<(SourceName, ResolvedRef, RootedSourceRef, Option<Version>)> {
97        self.pending_restart.take()
98    }
99
100    pub(super) fn registry(&self) -> &IndexMap<SourceName, RegisteredPackage> {
101        &self.registry
102    }
103
104    pub(super) fn registry_mut(&mut self) -> &mut IndexMap<SourceName, RegisteredPackage> {
105        &mut self.registry
106    }
107
108    pub(super) fn package_states(&self) -> &HashMap<SourceName, PackageResolutionState> {
109        &self.package_states
110    }
111
112    pub(super) fn package_states_mut(
113        &mut self,
114    ) -> &mut HashMap<SourceName, PackageResolutionState> {
115        &mut self.package_states
116    }
117
118    pub(super) fn id_index(&self) -> &HashMap<SourceId, SourceName> {
119        &self.id_index
120    }
121
122    pub(super) fn id_index_mut(&mut self) -> &mut HashMap<SourceId, SourceName> {
123        &mut self.id_index
124    }
125
126    pub(super) fn version_constraints(
127        &self,
128    ) -> &HashMap<SourceName, Vec<(String, VersionConstraint)>> {
129        &self.version_constraints
130    }
131
132    pub(super) fn materialization_filters(&self) -> &HashMap<SourceName, Vec<FilterMode>> {
133        &self.materialization_filters
134    }
135
136    pub(super) fn visited(&self) -> &VisitedSet {
137        &self.visited
138    }
139
140    pub(super) fn visited_mut(&mut self) -> &mut VisitedSet {
141        &mut self.visited
142    }
143
144    pub(super) fn package_versions_mut(&mut self) -> &mut PackageVersions {
145        &mut self.package_versions
146    }
147
148    pub fn add_version_constraint(
149        &mut self,
150        package: &SourceName,
151        requester: &str,
152        constraint: VersionConstraint,
153    ) {
154        self.version_constraints
155            .entry(package.clone())
156            .or_default()
157            .push((requester.to_string(), constraint));
158    }
159
160    pub fn add_filter(&mut self, package: &SourceName, filter: FilterMode) {
161        push_filter_constraint(&mut self.materialization_filters, package, &filter);
162    }
163
164    pub fn push_pending(&mut self, item: PendingItem) {
165        self.stack.push(item);
166    }
167
168    pub fn pop_pending(&mut self) -> Option<PendingItem> {
169        self.stack.pop()
170    }
171
172    pub fn into_graph(self) -> ResolvedGraph {
173        let mut nodes: IndexMap<SourceName, ResolvedNode> = IndexMap::new();
174        for (name, package) in self.registry {
175            nodes.insert(name, package.node);
176        }
177
178        let mut order: Vec<SourceName> = nodes.keys().cloned().collect();
179        order.sort();
180
181        ResolvedGraph {
182            nodes,
183            order,
184            filters: self.materialization_filters,
185        }
186    }
187}