1use flux_common::{bug, cache::QueryCache, dbg, iter::IterExt, result::ResultExt};
2use flux_config as config;
3use flux_errors::FluxSession;
4use flux_infer::fixpoint_encoding::FixQueryCache;
5use flux_metadata::CStore;
6use flux_middle::{
7 Specs, fhir,
8 global_env::GlobalEnv,
9 queries::{Providers, QueryResult},
10 timings,
11};
12use flux_refineck as refineck;
13use itertools::Itertools;
14use rustc_borrowck::consumers::ConsumerOptions;
15use rustc_driver::{Callbacks, Compilation};
16use rustc_errors::ErrorGuaranteed;
17use rustc_hir::{
18 def::DefKind,
19 def_id::{DefId, LocalDefId},
20};
21use rustc_interface::interface::Compiler;
22use rustc_middle::{query, ty::TyCtxt};
23use rustc_session::config::OutputType;
24use rustc_span::FileName;
25
26use crate::{DEFAULT_LOCALE_RESOURCES, collector::SpecCollector};
27
28#[derive(Default)]
29pub struct FluxCallbacks;
30
31impl Callbacks for FluxCallbacks {
32 fn config(&mut self, config: &mut rustc_interface::interface::Config) {
33 assert!(config.override_queries.is_none());
34
35 config.override_queries = Some(|_, local| {
36 local.mir_borrowck = mir_borrowck;
37 });
38 }
39
40 fn after_analysis(&mut self, compiler: &Compiler, tcx: TyCtxt<'_>) -> Compilation {
41 timings::enter(tcx, || self.verify(compiler, tcx));
42 if config::full_compilation() { Compilation::Continue } else { Compilation::Stop }
43 }
44}
45
46impl FluxCallbacks {
47 fn verify(&self, compiler: &Compiler, tcx: TyCtxt<'_>) {
48 if compiler.sess.dcx().has_errors().is_some() {
49 return;
50 }
51
52 let sess = FluxSession::new(
53 &tcx.sess.opts,
54 tcx.sess.psess.clone_source_map(),
55 rustc_errors::fallback_fluent_bundle(DEFAULT_LOCALE_RESOURCES.to_vec(), false),
56 );
57
58 let mut providers = Providers::default();
59 flux_desugar::provide(&mut providers);
60 flux_fhir_analysis::provide(&mut providers);
61 providers.collect_specs = collect_specs;
62
63 let cstore = CStore::load(tcx, &sess);
64 let arena = fhir::Arena::new();
65 GlobalEnv::enter(tcx, &sess, Box::new(cstore), &arena, providers, |genv| {
66 if check_crate(genv).is_ok() {
67 encode_and_save_metadata(genv);
68 }
69 });
70 sess.finish_diagnostics();
71 }
72}
73
74fn check_crate(genv: GlobalEnv) -> Result<(), ErrorGuaranteed> {
75 tracing::info_span!("check_crate").in_scope(move || {
76 tracing::info!("Callbacks::check_wf");
77
78 flux_fhir_analysis::check_crate_wf(genv)?;
79
80 let mut ck = CrateChecker::new(genv);
81
82 let crate_items = genv.tcx().hir_crate_items(());
83
84 let dups = crate_items.definitions().duplicates().collect_vec();
85 if !dups.is_empty() {
86 bug!("TODO: {dups:#?}");
87 }
88 let result = crate_items
89 .definitions()
90 .try_for_each_exhaust(|def_id| ck.check_def_catching_bugs(def_id));
91
92 ck.cache.save().unwrap_or(());
93
94 tracing::info!("Callbacks::check_crate");
95
96 result
97 })
98}
99
100fn collect_specs(genv: GlobalEnv) -> Specs {
101 match SpecCollector::collect(genv.tcx(), genv.sess()) {
102 Ok(specs) => specs,
103 Err(err) => {
104 genv.sess().abort(err);
105 }
106 }
107}
108
109fn encode_and_save_metadata(genv: GlobalEnv) {
110 let tcx = genv.tcx();
114 if tcx
115 .output_filenames(())
116 .outputs
117 .contains_key(&OutputType::Metadata)
118 {
119 let path = flux_metadata::filename_for_metadata(tcx);
120 flux_metadata::encode_metadata(genv, path.as_path());
121 }
122}
123
124struct CrateChecker<'genv, 'tcx> {
125 genv: GlobalEnv<'genv, 'tcx>,
126 cache: FixQueryCache,
127}
128
129impl<'genv, 'tcx> CrateChecker<'genv, 'tcx> {
130 fn new(genv: GlobalEnv<'genv, 'tcx>) -> Self {
131 CrateChecker { genv, cache: QueryCache::load() }
132 }
133
134 fn matches_check_def(&self, def_id: DefId) -> bool {
135 let def_path = self.genv.tcx().def_path_str(def_id);
136 def_path.contains(config::check_def())
137 }
138
139 fn matches_check_file(&self, def_id: LocalDefId) -> bool {
140 let tcx = self.genv.tcx();
141 let span = tcx.def_span(def_id);
142 let sm = tcx.sess.source_map();
143 let current_dir = tcx.sess.opts.working_dir.clone();
144 if let FileName::Real(file_name) = sm.span_to_filename(span) {
145 if let Some(file_path) = file_name.local_path()
146 && let Some(current_dir_path) = current_dir.local_path()
147 {
148 let file = current_dir_path
149 .join(file_path)
150 .to_string_lossy()
151 .to_string();
152 return config::is_checked_file(&file);
153 }
154 }
155 true
156 }
157
158 fn check_def_catching_bugs(&mut self, def_id: LocalDefId) -> Result<(), ErrorGuaranteed> {
159 let mut this = std::panic::AssertUnwindSafe(self);
160 let msg = format!("def_id: {:?}, span: {:?}", def_id, this.genv.tcx().def_span(def_id));
161 flux_common::bug::catch_bugs(&msg, move || this.check_def(def_id))?
162 }
163
164 fn check_def(&mut self, def_id: LocalDefId) -> Result<(), ErrorGuaranteed> {
165 let def_id = self.genv.maybe_extern_id(def_id);
166
167 if !self.matches_check_def(def_id.resolved_id()) {
168 return Ok(());
169 }
170 if self.genv.ignored(def_id.local_id()) || self.genv.is_dummy(def_id.local_id()) {
171 return Ok(());
172 }
173 if !self.matches_check_file(def_id.local_id()) {
174 return Ok(());
175 }
176
177 match self.genv.def_kind(def_id) {
178 DefKind::Fn | DefKind::AssocFn => {
179 force_conv(self.genv, def_id.resolved_id()).emit(&self.genv)?;
181 let Some(local_id) = def_id.as_local() else { return Ok(()) };
182 refineck::check_fn(self.genv, &mut self.cache, local_id)
183 }
184 DefKind::Enum => {
185 let adt_def = self.genv.adt_def(def_id).emit(&self.genv)?;
186 let _ = self.genv.variants_of(def_id).emit(&self.genv)?;
187 let enum_def = self
188 .genv
189 .map()
190 .expect_item(def_id.local_id())
191 .emit(&self.genv)?
192 .expect_enum();
193 refineck::invariants::check_invariants(
194 self.genv,
195 &mut self.cache,
196 def_id,
197 enum_def.invariants,
198 &adt_def,
199 )
200 }
201 DefKind::Struct => {
202 let _adt_def = self.genv.adt_def(def_id).emit(&self.genv)?;
206 let _ = self.genv.variants_of(def_id).emit(&self.genv)?;
207 let _struct_def = self
208 .genv
209 .map()
210 .expect_item(def_id.local_id())
211 .emit(&self.genv)?
212 .expect_struct();
213 Ok(())
214 }
215 DefKind::Impl { of_trait } => {
216 if of_trait {
217 refineck::compare_impl_item::check_impl_against_trait(self.genv, def_id)
218 .emit(&self.genv)?;
219 }
220 Ok(())
221 }
222 DefKind::TyAlias => {
223 self.genv.type_of(def_id).emit(&self.genv)?;
224 Ok(())
225 }
226 _ => Ok(()),
227 }
228 }
229}
230
231fn force_conv(genv: GlobalEnv, def_id: DefId) -> QueryResult {
232 genv.generics_of(def_id)?;
233 genv.refinement_generics_of(def_id)?;
234 genv.predicates_of(def_id)?;
235 genv.fn_sig(def_id)?;
236 Ok(())
237}
238
239#[expect(clippy::needless_lifetimes, reason = "we want to be explicit about lifetimes here")]
240fn mir_borrowck<'tcx>(
241 tcx: TyCtxt<'tcx>,
242 def_id: LocalDefId,
243) -> query::queries::mir_borrowck::ProvidedValue<'tcx> {
244 let body_with_facts = rustc_borrowck::consumers::get_body_with_borrowck_facts(
245 tcx,
246 def_id,
247 ConsumerOptions::RegionInferenceContext,
248 );
249
250 if config::dump_mir() {
251 rustc_middle::mir::pretty::write_mir_fn(
252 tcx,
253 &body_with_facts.body,
254 &mut |_, _| Ok(()),
255 &mut dbg::writer_for_item(tcx, def_id.to_def_id(), "mir").unwrap(),
256 rustc_middle::mir::pretty::PrettyPrintMirOptions::from_cli(tcx),
257 )
258 .unwrap();
259 }
260
261 unsafe {
264 flux_common::mir_storage::store_mir_body(tcx, def_id, body_with_facts);
265 }
266 let mut providers = query::Providers::default();
267 rustc_borrowck::provide(&mut providers);
268 let original_mir_borrowck = providers.mir_borrowck;
269 original_mir_borrowck(tcx, def_id)
270}