1use crate::ast::*;
2use crate::formatter::{FmtCtx, IntoFormatter};
3use crate::ids::Vector;
4use crate::pretty::FmtWithCtx;
5use crate::reorder_decls::DeclarationsGroups;
6use derive_generic_visitor::{ControlFlow, Drive, DriveMut};
7use index_vec::Idx;
8use indexmap::IndexSet;
9use macros::{EnumAsGetters, EnumIsA, VariantIndexArity, VariantName};
10use serde::{Deserialize, Serialize};
11use serde_map_to_array::HashMapToArray;
12use std::cmp::{Ord, PartialOrd};
13use std::collections::HashMap;
14use std::fmt;
15
16generate_index_type!(FunDeclId, "Fun");
17generate_index_type!(TypeDeclId, "Adt");
18generate_index_type!(GlobalDeclId, "Global");
19generate_index_type!(TraitDeclId, "TraitDecl");
20generate_index_type!(TraitImplId, "TraitImpl");
21
22#[derive(
24 Copy,
25 Clone,
26 Debug,
27 PartialOrd,
28 Ord,
29 PartialEq,
30 Eq,
31 Hash,
32 EnumIsA,
33 EnumAsGetters,
34 VariantName,
35 VariantIndexArity,
36 Serialize,
37 Deserialize,
38 Drive,
39 DriveMut,
40)]
41#[charon::rename("AnyDeclId")]
42#[charon::variants_prefix("Id")]
43pub enum AnyTransId {
44 Type(TypeDeclId),
45 Fun(FunDeclId),
46 Global(GlobalDeclId),
47 TraitDecl(TraitDeclId),
48 TraitImpl(TraitImplId),
49}
50
51macro_rules! wrap_unwrap_enum {
53 ($enum:ident::$variant:ident($variant_ty:ident)) => {
54 impl TryFrom<$enum> for $variant_ty {
55 type Error = ();
56 fn try_from(x: $enum) -> Result<Self, Self::Error> {
57 match x {
58 $enum::$variant(x) => Ok(x),
59 _ => Err(()),
60 }
61 }
62 }
63
64 impl From<$variant_ty> for $enum {
65 fn from(x: $variant_ty) -> Self {
66 $enum::$variant(x)
67 }
68 }
69 };
70}
71
72wrap_unwrap_enum!(AnyTransId::Fun(FunDeclId));
73wrap_unwrap_enum!(AnyTransId::Global(GlobalDeclId));
74wrap_unwrap_enum!(AnyTransId::Type(TypeDeclId));
75wrap_unwrap_enum!(AnyTransId::TraitDecl(TraitDeclId));
76wrap_unwrap_enum!(AnyTransId::TraitImpl(TraitImplId));
77
78#[derive(
80 Debug, Clone, Copy, EnumIsA, EnumAsGetters, VariantName, VariantIndexArity, Drive, DriveMut,
81)]
82pub enum AnyTransItem<'ctx> {
83 Type(&'ctx TypeDecl),
84 Fun(&'ctx FunDecl),
85 Global(&'ctx GlobalDecl),
86 TraitDecl(&'ctx TraitDecl),
87 TraitImpl(&'ctx TraitImpl),
88}
89
90#[derive(Debug, EnumIsA, EnumAsGetters, VariantName, VariantIndexArity, Drive, DriveMut)]
92pub enum AnyTransItemMut<'ctx> {
93 Type(&'ctx mut TypeDecl),
94 Fun(&'ctx mut FunDecl),
95 Global(&'ctx mut GlobalDecl),
96 TraitDecl(&'ctx mut TraitDecl),
97 TraitImpl(&'ctx mut TraitImpl),
98}
99
100#[derive(Default, Clone, Drive, DriveMut, Serialize, Deserialize)]
102pub struct TranslatedCrate {
103 #[drive(skip)]
105 pub crate_name: String,
106
107 #[drive(skip)]
111 pub options: crate::options::CliOpts,
112
113 #[drive(skip)]
115 pub all_ids: IndexSet<AnyTransId>,
116 #[serde(with = "HashMapToArray::<AnyTransId, Name>")]
121 pub item_names: HashMap<AnyTransId, Name>,
122 #[serde(with = "HashMapToArray::<AnyTransId, Name>")]
124 pub short_names: HashMap<AnyTransId, Name>,
125
126 #[drive(skip)]
128 pub files: Vector<FileId, File>,
129 pub type_decls: Vector<TypeDeclId, TypeDecl>,
131 pub fun_decls: Vector<FunDeclId, FunDecl>,
133 pub global_decls: Vector<GlobalDeclId, GlobalDecl>,
135 pub trait_decls: Vector<TraitDeclId, TraitDecl>,
137 pub trait_impls: Vector<TraitImplId, TraitImpl>,
139 #[drive(skip)]
141 pub ordered_decls: Option<DeclarationsGroups>,
142}
143
144impl TranslatedCrate {
145 pub fn item_name(&self, id: impl Into<AnyTransId>) -> Option<&Name> {
146 self.item_names.get(&id.into())
147 }
148
149 pub fn item_short_name(&self, id: impl Into<AnyTransId>) -> Option<&Name> {
150 let id = id.into();
151 self.short_names.get(&id).or_else(|| self.item_name(id))
152 }
153
154 pub fn get_item(&self, trans_id: impl Into<AnyTransId>) -> Option<AnyTransItem<'_>> {
155 match trans_id.into() {
156 AnyTransId::Type(id) => self.type_decls.get(id).map(AnyTransItem::Type),
157 AnyTransId::Fun(id) => self.fun_decls.get(id).map(AnyTransItem::Fun),
158 AnyTransId::Global(id) => self.global_decls.get(id).map(AnyTransItem::Global),
159 AnyTransId::TraitDecl(id) => self.trait_decls.get(id).map(AnyTransItem::TraitDecl),
160 AnyTransId::TraitImpl(id) => self.trait_impls.get(id).map(AnyTransItem::TraitImpl),
161 }
162 }
163
164 pub fn get_item_mut(&mut self, trans_id: AnyTransId) -> Option<AnyTransItemMut<'_>> {
165 match trans_id {
166 AnyTransId::Type(id) => self.type_decls.get_mut(id).map(AnyTransItemMut::Type),
167 AnyTransId::Fun(id) => self.fun_decls.get_mut(id).map(AnyTransItemMut::Fun),
168 AnyTransId::Global(id) => self.global_decls.get_mut(id).map(AnyTransItemMut::Global),
169 AnyTransId::TraitDecl(id) => {
170 self.trait_decls.get_mut(id).map(AnyTransItemMut::TraitDecl)
171 }
172 AnyTransId::TraitImpl(id) => {
173 self.trait_impls.get_mut(id).map(AnyTransItemMut::TraitImpl)
174 }
175 }
176 }
177
178 pub fn all_items(&self) -> impl Iterator<Item = AnyTransItem<'_>> {
179 self.all_items_with_ids().map(|(_, item)| item)
180 }
181 pub fn all_items_with_ids(&self) -> impl Iterator<Item = (AnyTransId, AnyTransItem<'_>)> {
182 self.all_ids
183 .iter()
184 .flat_map(|id| Some((*id, self.get_item(*id)?)))
185 }
186 pub fn all_items_mut(&mut self) -> impl Iterator<Item = AnyTransItemMut<'_>> {
187 self.type_decls
188 .iter_mut()
189 .map(AnyTransItemMut::Type)
190 .chain(self.fun_decls.iter_mut().map(AnyTransItemMut::Fun))
191 .chain(self.global_decls.iter_mut().map(AnyTransItemMut::Global))
192 .chain(self.trait_decls.iter_mut().map(AnyTransItemMut::TraitDecl))
193 .chain(self.trait_impls.iter_mut().map(AnyTransItemMut::TraitImpl))
194 }
195}
196
197impl<'ctx> AnyTransItem<'ctx> {
198 pub fn id(&self) -> AnyTransId {
199 match self {
200 AnyTransItem::Type(d) => d.def_id.into(),
201 AnyTransItem::Fun(d) => d.def_id.into(),
202 AnyTransItem::Global(d) => d.def_id.into(),
203 AnyTransItem::TraitDecl(d) => d.def_id.into(),
204 AnyTransItem::TraitImpl(d) => d.def_id.into(),
205 }
206 }
207
208 pub fn item_meta(&self) -> &'ctx ItemMeta {
209 match self {
210 AnyTransItem::Type(d) => &d.item_meta,
211 AnyTransItem::Fun(d) => &d.item_meta,
212 AnyTransItem::Global(d) => &d.item_meta,
213 AnyTransItem::TraitDecl(d) => &d.item_meta,
214 AnyTransItem::TraitImpl(d) => &d.item_meta,
215 }
216 }
217
218 pub fn generic_params(&self) -> &'ctx GenericParams {
220 match self {
221 AnyTransItem::Type(d) => &d.generics,
222 AnyTransItem::Fun(d) => &d.signature.generics,
223 AnyTransItem::Global(d) => &d.generics,
224 AnyTransItem::TraitDecl(d) => &d.generics,
225 AnyTransItem::TraitImpl(d) => &d.generics,
226 }
227 }
228
229 pub fn parent_info(&self) -> &'ctx ItemKind {
231 match self {
232 AnyTransItem::Fun(d) => &d.kind,
233 AnyTransItem::Global(d) => &d.kind,
234 AnyTransItem::Type(_) | AnyTransItem::TraitDecl(_) | AnyTransItem::TraitImpl(_) => {
235 &ItemKind::TopLevel
236 }
237 }
238 }
239
240 pub fn identity_args(&self) -> GenericArgs {
242 self.generic_params()
243 .identity_args(GenericsSource::Item(self.id()))
244 }
245
246 pub fn drive<V: VisitAst>(&self, visitor: &mut V) -> ControlFlow<V::Break> {
249 match *self {
250 AnyTransItem::Type(d) => visitor.visit(d),
251 AnyTransItem::Fun(d) => visitor.visit(d),
252 AnyTransItem::Global(d) => visitor.visit(d),
253 AnyTransItem::TraitDecl(d) => visitor.visit(d),
254 AnyTransItem::TraitImpl(d) => visitor.visit(d),
255 }
256 }
257
258 pub fn dyn_visit<T: AstVisitable>(&self, f: impl FnMut(&T)) {
260 match *self {
261 AnyTransItem::Type(d) => d.dyn_visit(f),
262 AnyTransItem::Fun(d) => d.dyn_visit(f),
263 AnyTransItem::Global(d) => d.dyn_visit(f),
264 AnyTransItem::TraitDecl(d) => d.dyn_visit(f),
265 AnyTransItem::TraitImpl(d) => d.dyn_visit(f),
266 }
267 }
268}
269
270impl<'ctx> AnyTransItemMut<'ctx> {
271 pub fn as_ref(&self) -> AnyTransItem<'_> {
272 match self {
273 AnyTransItemMut::Type(d) => AnyTransItem::Type(d),
274 AnyTransItemMut::Fun(d) => AnyTransItem::Fun(d),
275 AnyTransItemMut::Global(d) => AnyTransItem::Global(d),
276 AnyTransItemMut::TraitDecl(d) => AnyTransItem::TraitDecl(d),
277 AnyTransItemMut::TraitImpl(d) => AnyTransItem::TraitImpl(d),
278 }
279 }
280
281 pub fn generic_params(&mut self) -> &mut GenericParams {
283 match self {
284 AnyTransItemMut::Type(d) => &mut d.generics,
285 AnyTransItemMut::Fun(d) => &mut d.signature.generics,
286 AnyTransItemMut::Global(d) => &mut d.generics,
287 AnyTransItemMut::TraitDecl(d) => &mut d.generics,
288 AnyTransItemMut::TraitImpl(d) => &mut d.generics,
289 }
290 }
291
292 pub fn drive_mut<V: VisitAstMut>(&mut self, visitor: &mut V) -> ControlFlow<V::Break> {
295 match self {
296 AnyTransItemMut::Type(d) => visitor.visit(*d),
297 AnyTransItemMut::Fun(d) => visitor.visit(*d),
298 AnyTransItemMut::Global(d) => visitor.visit(*d),
299 AnyTransItemMut::TraitDecl(d) => visitor.visit(*d),
300 AnyTransItemMut::TraitImpl(d) => visitor.visit(*d),
301 }
302 }
303
304 pub fn dyn_visit_mut<T: AstVisitable>(&mut self, f: impl FnMut(&mut T)) {
306 match self {
307 AnyTransItemMut::Type(d) => d.dyn_visit_mut(f),
308 AnyTransItemMut::Fun(d) => d.dyn_visit_mut(f),
309 AnyTransItemMut::Global(d) => d.dyn_visit_mut(f),
310 AnyTransItemMut::TraitDecl(d) => d.dyn_visit_mut(f),
311 AnyTransItemMut::TraitImpl(d) => d.dyn_visit_mut(f),
312 }
313 }
314}
315
316impl fmt::Display for TranslatedCrate {
317 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
318 let fmt: &FmtCtx = &self.into_fmt();
319 match &self.ordered_decls {
320 None => {
321 for d in &self.type_decls {
323 writeln!(f, "{}\n", d.with_ctx(fmt))?
324 }
325 for d in &self.global_decls {
326 writeln!(f, "{}\n", d.with_ctx(fmt))?
327 }
328 for d in &self.trait_decls {
329 writeln!(f, "{}\n", d.with_ctx(fmt))?
330 }
331 for d in &self.trait_impls {
332 writeln!(f, "{}\n", d.with_ctx(fmt))?
333 }
334 for d in &self.fun_decls {
335 writeln!(f, "{}\n", d.with_ctx(fmt))?
336 }
337 }
338 Some(ordered_decls) => {
339 for gr in ordered_decls {
340 for id in gr.get_ids() {
341 writeln!(f, "{}\n", fmt.format_decl_id(id))?
342 }
343 }
344 }
345 }
346 fmt::Result::Ok(())
347 }
348}
349
350impl<'tcx, 'ctx, 'a> IntoFormatter for &'a TranslatedCrate {
351 type C = FmtCtx<'a>;
352
353 fn into_fmt(self) -> Self::C {
354 FmtCtx {
355 translated: Some(self),
356 ..Default::default()
357 }
358 }
359}
360
361pub trait HasVectorOf<Id: Idx>: std::ops::Index<Id, Output: Sized> {
362 fn get_vector(&self) -> &Vector<Id, Self::Output>;
363 fn get_vector_mut(&mut self) -> &mut Vector<Id, Self::Output>;
364}
365
366macro_rules! mk_index_impls {
368 ($ty:ident.$field:ident[$idx:ty]: $output:ty) => {
369 impl std::ops::Index<$idx> for $ty {
370 type Output = $output;
371 fn index(&self, index: $idx) -> &Self::Output {
372 &self.$field[index]
373 }
374 }
375 impl std::ops::IndexMut<$idx> for $ty {
376 fn index_mut(&mut self, index: $idx) -> &mut Self::Output {
377 &mut self.$field[index]
378 }
379 }
380 impl HasVectorOf<$idx> for $ty {
381 fn get_vector(&self) -> &Vector<$idx, Self::Output> {
382 &self.$field
383 }
384 fn get_vector_mut(&mut self) -> &mut Vector<$idx, Self::Output> {
385 &mut self.$field
386 }
387 }
388 };
389}
390pub(crate) use mk_index_impls;
391
392mk_index_impls!(TranslatedCrate.type_decls[TypeDeclId]: TypeDecl);
393mk_index_impls!(TranslatedCrate.fun_decls[FunDeclId]: FunDecl);
394mk_index_impls!(TranslatedCrate.global_decls[GlobalDeclId]: GlobalDecl);
395mk_index_impls!(TranslatedCrate.trait_decls[TraitDeclId]: TraitDecl);
396mk_index_impls!(TranslatedCrate.trait_impls[TraitImplId]: TraitImpl);