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