Skip to main content

mars_agents/resolve/
context.rs

1use std::collections::HashMap;
2
3use indexmap::IndexMap;
4
5use super::filter::push_filter_constraint;
6use super::{
7    PackageResolutionState, PackageVersions, PendingItem, RegisteredPackage, ResolvedGraph,
8    ResolvedNode, VersionConstraint, VisitedSet,
9};
10use crate::config::FilterMode;
11use crate::types::{SourceId, SourceName};
12
13/// Mutable resolver state threaded through bottom-up resolution and DFS traversal.
14pub struct ResolverContext {
15    registry: IndexMap<SourceName, RegisteredPackage>,
16    package_states: HashMap<SourceName, PackageResolutionState>,
17    id_index: HashMap<SourceId, SourceName>,
18    version_constraints: HashMap<SourceName, Vec<(String, VersionConstraint)>>,
19    materialization_filters: HashMap<SourceName, Vec<FilterMode>>,
20    stack: Vec<PendingItem>,
21    visited: VisitedSet,
22    package_versions: PackageVersions,
23}
24
25impl Default for ResolverContext {
26    fn default() -> Self {
27        Self::new()
28    }
29}
30
31impl ResolverContext {
32    pub fn new() -> Self {
33        Self {
34            registry: IndexMap::new(),
35            package_states: HashMap::new(),
36            id_index: HashMap::new(),
37            version_constraints: HashMap::new(),
38            materialization_filters: HashMap::new(),
39            stack: Vec::new(),
40            visited: VisitedSet::new(),
41            package_versions: PackageVersions::new(),
42        }
43    }
44
45    pub(super) fn registry(&self) -> &IndexMap<SourceName, RegisteredPackage> {
46        &self.registry
47    }
48
49    pub(super) fn registry_mut(&mut self) -> &mut IndexMap<SourceName, RegisteredPackage> {
50        &mut self.registry
51    }
52
53    pub(super) fn package_states(&self) -> &HashMap<SourceName, PackageResolutionState> {
54        &self.package_states
55    }
56
57    pub(super) fn package_states_mut(
58        &mut self,
59    ) -> &mut HashMap<SourceName, PackageResolutionState> {
60        &mut self.package_states
61    }
62
63    pub(super) fn id_index(&self) -> &HashMap<SourceId, SourceName> {
64        &self.id_index
65    }
66
67    pub(super) fn id_index_mut(&mut self) -> &mut HashMap<SourceId, SourceName> {
68        &mut self.id_index
69    }
70
71    pub(super) fn version_constraints(
72        &self,
73    ) -> &HashMap<SourceName, Vec<(String, VersionConstraint)>> {
74        &self.version_constraints
75    }
76
77    pub(super) fn materialization_filters(&self) -> &HashMap<SourceName, Vec<FilterMode>> {
78        &self.materialization_filters
79    }
80
81    pub(super) fn visited(&self) -> &VisitedSet {
82        &self.visited
83    }
84
85    pub(super) fn visited_mut(&mut self) -> &mut VisitedSet {
86        &mut self.visited
87    }
88
89    pub(super) fn package_versions_mut(&mut self) -> &mut PackageVersions {
90        &mut self.package_versions
91    }
92
93    pub fn add_version_constraint(
94        &mut self,
95        package: &SourceName,
96        requester: &str,
97        constraint: VersionConstraint,
98    ) {
99        self.version_constraints
100            .entry(package.clone())
101            .or_default()
102            .push((requester.to_string(), constraint));
103    }
104
105    pub fn add_filter(&mut self, package: &SourceName, filter: FilterMode) {
106        push_filter_constraint(&mut self.materialization_filters, package, &filter);
107    }
108
109    pub fn push_pending(&mut self, item: PendingItem) {
110        self.stack.push(item);
111    }
112
113    pub fn pop_pending(&mut self) -> Option<PendingItem> {
114        self.stack.pop()
115    }
116
117    pub fn into_graph(self) -> ResolvedGraph {
118        let mut nodes: IndexMap<SourceName, ResolvedNode> = IndexMap::new();
119        for (name, package) in self.registry {
120            nodes.insert(name, package.node);
121        }
122
123        let mut order: Vec<SourceName> = nodes.keys().cloned().collect();
124        order.sort();
125
126        ResolvedGraph {
127            nodes,
128            order,
129            filters: self.materialization_filters,
130        }
131    }
132}