1use crate::ast::*;
2use crate::formatter::{FmtCtx, Formatter, IntoFormatter};
3use crate::ids::Vector;
4use crate::reorder_decls::DeclarationsGroups;
5use derive_generic_visitor::{ControlFlow, Drive, DriveMut};
6use index_vec::Idx;
7use indexmap::IndexSet;
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 #[drive(skip)]
114 pub all_ids: IndexSet<AnyTransId>,
115 #[serde(with = "HashMapToArray::<AnyTransId, Name>")]
120 pub item_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, trans_id: impl Into<AnyTransId>) -> Option<&Name> {
142 self.item_names.get(&trans_id.into())
143 }
144
145 pub fn get_item(&self, trans_id: impl Into<AnyTransId>) -> Option<AnyTransItem<'_>> {
146 match trans_id.into() {
147 AnyTransId::Type(id) => self.type_decls.get(id).map(AnyTransItem::Type),
148 AnyTransId::Fun(id) => self.fun_decls.get(id).map(AnyTransItem::Fun),
149 AnyTransId::Global(id) => self.global_decls.get(id).map(AnyTransItem::Global),
150 AnyTransId::TraitDecl(id) => self.trait_decls.get(id).map(AnyTransItem::TraitDecl),
151 AnyTransId::TraitImpl(id) => self.trait_impls.get(id).map(AnyTransItem::TraitImpl),
152 }
153 }
154
155 pub fn get_item_mut(&mut self, trans_id: AnyTransId) -> Option<AnyTransItemMut<'_>> {
156 match trans_id {
157 AnyTransId::Type(id) => self.type_decls.get_mut(id).map(AnyTransItemMut::Type),
158 AnyTransId::Fun(id) => self.fun_decls.get_mut(id).map(AnyTransItemMut::Fun),
159 AnyTransId::Global(id) => self.global_decls.get_mut(id).map(AnyTransItemMut::Global),
160 AnyTransId::TraitDecl(id) => {
161 self.trait_decls.get_mut(id).map(AnyTransItemMut::TraitDecl)
162 }
163 AnyTransId::TraitImpl(id) => {
164 self.trait_impls.get_mut(id).map(AnyTransItemMut::TraitImpl)
165 }
166 }
167 }
168
169 pub fn all_items(&self) -> impl Iterator<Item = AnyTransItem<'_>> {
170 self.all_items_with_ids().map(|(_, item)| item)
171 }
172 pub fn all_items_with_ids(&self) -> impl Iterator<Item = (AnyTransId, AnyTransItem<'_>)> {
173 self.all_ids
174 .iter()
175 .flat_map(|id| Some((*id, self.get_item(*id)?)))
176 }
177 pub fn all_items_mut(&mut self) -> impl Iterator<Item = AnyTransItemMut<'_>> {
178 self.type_decls
179 .iter_mut()
180 .map(AnyTransItemMut::Type)
181 .chain(self.fun_decls.iter_mut().map(AnyTransItemMut::Fun))
182 .chain(self.global_decls.iter_mut().map(AnyTransItemMut::Global))
183 .chain(self.trait_decls.iter_mut().map(AnyTransItemMut::TraitDecl))
184 .chain(self.trait_impls.iter_mut().map(AnyTransItemMut::TraitImpl))
185 }
186}
187
188impl<'ctx> AnyTransItem<'ctx> {
189 pub fn id(&self) -> AnyTransId {
190 match self {
191 AnyTransItem::Type(d) => d.def_id.into(),
192 AnyTransItem::Fun(d) => d.def_id.into(),
193 AnyTransItem::Global(d) => d.def_id.into(),
194 AnyTransItem::TraitDecl(d) => d.def_id.into(),
195 AnyTransItem::TraitImpl(d) => d.def_id.into(),
196 }
197 }
198
199 pub fn item_meta(&self) -> &'ctx ItemMeta {
200 match self {
201 AnyTransItem::Type(d) => &d.item_meta,
202 AnyTransItem::Fun(d) => &d.item_meta,
203 AnyTransItem::Global(d) => &d.item_meta,
204 AnyTransItem::TraitDecl(d) => &d.item_meta,
205 AnyTransItem::TraitImpl(d) => &d.item_meta,
206 }
207 }
208
209 pub fn generic_params(&self) -> &'ctx GenericParams {
211 match self {
212 AnyTransItem::Type(d) => &d.generics,
213 AnyTransItem::Fun(d) => &d.signature.generics,
214 AnyTransItem::Global(d) => &d.generics,
215 AnyTransItem::TraitDecl(d) => &d.generics,
216 AnyTransItem::TraitImpl(d) => &d.generics,
217 }
218 }
219
220 pub fn identity_args(&self) -> GenericArgs {
222 self.generic_params()
223 .identity_args(GenericsSource::Item(self.id()))
224 }
225
226 pub fn drive<V: VisitAst>(&self, visitor: &mut V) -> ControlFlow<V::Break> {
229 match *self {
230 AnyTransItem::Type(d) => visitor.visit(d),
231 AnyTransItem::Fun(d) => visitor.visit(d),
232 AnyTransItem::Global(d) => visitor.visit(d),
233 AnyTransItem::TraitDecl(d) => visitor.visit(d),
234 AnyTransItem::TraitImpl(d) => visitor.visit(d),
235 }
236 }
237}
238
239impl<'ctx> AnyTransItemMut<'ctx> {
240 pub fn as_ref(&self) -> AnyTransItem<'_> {
241 match self {
242 AnyTransItemMut::Type(d) => AnyTransItem::Type(d),
243 AnyTransItemMut::Fun(d) => AnyTransItem::Fun(d),
244 AnyTransItemMut::Global(d) => AnyTransItem::Global(d),
245 AnyTransItemMut::TraitDecl(d) => AnyTransItem::TraitDecl(d),
246 AnyTransItemMut::TraitImpl(d) => AnyTransItem::TraitImpl(d),
247 }
248 }
249
250 pub fn generic_params(&mut self) -> &mut GenericParams {
252 match self {
253 AnyTransItemMut::Type(d) => &mut d.generics,
254 AnyTransItemMut::Fun(d) => &mut d.signature.generics,
255 AnyTransItemMut::Global(d) => &mut d.generics,
256 AnyTransItemMut::TraitDecl(d) => &mut d.generics,
257 AnyTransItemMut::TraitImpl(d) => &mut d.generics,
258 }
259 }
260
261 pub fn drive_mut<V: VisitAstMut>(&mut self, visitor: &mut V) -> ControlFlow<V::Break> {
264 match self {
265 AnyTransItemMut::Type(d) => visitor.visit(*d),
266 AnyTransItemMut::Fun(d) => visitor.visit(*d),
267 AnyTransItemMut::Global(d) => visitor.visit(*d),
268 AnyTransItemMut::TraitDecl(d) => visitor.visit(*d),
269 AnyTransItemMut::TraitImpl(d) => visitor.visit(*d),
270 }
271 }
272}
273
274impl fmt::Display for TranslatedCrate {
275 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
276 let fmt: FmtCtx = self.into_fmt();
277 match &self.ordered_decls {
278 None => {
279 for d in &self.type_decls {
281 writeln!(f, "{}\n", fmt.format_object(d))?
282 }
283 for d in &self.global_decls {
284 writeln!(f, "{}\n", fmt.format_object(d))?
285 }
286 for d in &self.trait_decls {
287 writeln!(f, "{}\n", fmt.format_object(d))?
288 }
289 for d in &self.trait_impls {
290 writeln!(f, "{}\n", fmt.format_object(d))?
291 }
292 for d in &self.fun_decls {
293 writeln!(f, "{}\n", fmt.format_object(d))?
294 }
295 }
296 Some(ordered_decls) => {
297 for gr in ordered_decls {
298 for id in gr.get_ids() {
299 writeln!(f, "{}\n", fmt.format_decl_id(id))?
300 }
301 }
302 }
303 }
304 fmt::Result::Ok(())
305 }
306}
307
308impl<'tcx, 'ctx, 'a> IntoFormatter for &'a TranslatedCrate {
309 type C = FmtCtx<'a>;
310
311 fn into_fmt(self) -> Self::C {
312 FmtCtx {
313 translated: Some(self),
314 ..Default::default()
315 }
316 }
317}
318
319pub trait HasVectorOf<Id: Idx>: std::ops::Index<Id, Output: Sized> {
320 fn get_vector(&self) -> &Vector<Id, Self::Output>;
321 fn get_vector_mut(&mut self) -> &mut Vector<Id, Self::Output>;
322}
323
324macro_rules! mk_index_impls {
326 ($ty:ident.$field:ident[$idx:ty]: $output:ty) => {
327 impl std::ops::Index<$idx> for $ty {
328 type Output = $output;
329 fn index(&self, index: $idx) -> &Self::Output {
330 &self.$field[index]
331 }
332 }
333 impl std::ops::IndexMut<$idx> for $ty {
334 fn index_mut(&mut self, index: $idx) -> &mut Self::Output {
335 &mut self.$field[index]
336 }
337 }
338 impl HasVectorOf<$idx> for $ty {
339 fn get_vector(&self) -> &Vector<$idx, Self::Output> {
340 &self.$field
341 }
342 fn get_vector_mut(&mut self) -> &mut Vector<$idx, Self::Output> {
343 &mut self.$field
344 }
345 }
346 };
347}
348pub(crate) use mk_index_impls;
349
350mk_index_impls!(TranslatedCrate.type_decls[TypeDeclId]: TypeDecl);
351mk_index_impls!(TranslatedCrate.fun_decls[FunDeclId]: FunDecl);
352mk_index_impls!(TranslatedCrate.global_decls[GlobalDeclId]: GlobalDecl);
353mk_index_impls!(TranslatedCrate.trait_decls[TraitDeclId]: TraitDecl);
354mk_index_impls!(TranslatedCrate.trait_impls[TraitImplId]: TraitImpl);