Skip to main content

rapx/analysis/heap_ownership/
default.rs

1use rustc_abi::VariantIdx;
2use rustc_middle::{
3    mir::{
4        BasicBlock, BasicBlockData, Body, Local, LocalDecl, Operand, TerminatorKind,
5        visit::{TyContext, Visitor},
6    },
7    ty::{
8        self, EarlyBinder, GenericArgKind, InstanceKind::Item, Ty, TyCtxt, TyKind,
9        TypeSuperVisitable, TypeVisitable, TypeVisitor,
10    },
11};
12use rustc_span::def_id::DefId;
13use std::{collections::HashMap, ops::ControlFlow};
14
15use super::*;
16
17pub struct HeapOwnershipAnalyzer<'tcx> {
18    tcx: TyCtxt<'tcx>,
19    adt_heap: HeapOwnershipResultMap,
20    fn_set: HashSet<DefId>,
21    ty_map: HashMap<Ty<'tcx>, String>,
22    adt_recorder: HashSet<DefId>,
23}
24
25impl<'tcx> Analysis for HeapOwnershipAnalyzer<'tcx> {
26    fn run(&mut self) {
27        self.start();
28    }
29}
30
31impl<'tcx> HeapOwnershipAnalysis for HeapOwnershipAnalyzer<'tcx> {
32    fn get_all_items(&self) -> HeapOwnershipResultMap {
33        self.adt_heap.clone()
34    }
35}
36
37// This function is aiming at resolving problems due to 'TyContext' not implementing 'Clone' trait,
38// thus we call function 'copy_ty_context' to simulate 'self.clone()'.
39#[inline(always)]
40pub(crate) fn copy_ty_context(tc: &TyContext) -> TyContext {
41    match tc {
42        TyContext::LocalDecl { local, source_info } => TyContext::LocalDecl {
43            local: local.clone(),
44            source_info: source_info.clone(),
45        },
46        _ => unreachable!(),
47    }
48}
49
50impl<'tcx> HeapOwnershipAnalyzer<'tcx> {
51    pub fn new(tcx: TyCtxt<'tcx>) -> Self {
52        Self {
53            tcx,
54            adt_heap: HashMap::default(),
55            fn_set: HashSet::new(),
56            ty_map: HashMap::new(),
57            adt_recorder: HashSet::new(),
58        }
59    }
60
61    pub fn ty_map(&self) -> &HashMap<Ty<'tcx>, String> {
62        &self.ty_map
63    }
64
65    pub fn ty_map_mut(&mut self) -> &mut HashMap<Ty<'tcx>, String> {
66        &mut self.ty_map
67    }
68
69    pub fn fn_set(&self) -> &HashSet<DefId> {
70        &self.fn_set
71    }
72
73    pub fn fn_set_mut(&mut self) -> &mut HashSet<DefId> {
74        &mut self.fn_set
75    }
76
77    pub fn adt_recorder(&self) -> &HashSet<DefId> {
78        &self.adt_recorder
79    }
80
81    pub fn adt_recorder_mut(&mut self) -> &mut HashSet<DefId> {
82        &mut self.adt_recorder
83    }
84
85    pub fn adt_heap(&self) -> &HeapOwnershipResultMap {
86        &self.adt_heap
87    }
88
89    pub fn adt_heap_mut(&mut self) -> &mut HeapOwnershipResultMap {
90        &mut self.adt_heap
91    }
92
93    pub fn format_heap_unit(unit: &(HeapOwnership, Vec<bool>)) -> String {
94        let (heap, flags) = unit;
95        let vec_str = flags
96            .iter()
97            .map(|&b| if b { "1" } else { "0" })
98            .collect::<Vec<_>>()
99            .join(",");
100        format!("({}, [{}])", heap, vec_str)
101    }
102
103    pub fn output(&mut self) {
104        for elem in self.adt_heap() {
105            let name = format!("{:?}", EarlyBinder::skip_binder(self.tcx.type_of(*elem.0)));
106            let owning = elem
107                .1
108                .iter()
109                .map(Self::format_heap_unit)
110                .collect::<Vec<_>>()
111                .join(", ");
112            rap_info!("{} {}", name, owning);
113        }
114    }
115
116    // From the top-down method of our approach, this 'visitor' is the set of several sub-phases
117    // which means it contains multiple sub-visitors to make whole method 'self.visitor()' work.
118    //
119    // For example, given an adtef (like Vec<T>), the result of 'visitor' contains two parts:
120    //
121    //     pt1 Enum:  {True / UnTrue} indicates whether it will directly have a heap data
122    //     pt2 Array: [bool;N] indicates whether each generic parameter will have a raw param
123    //
124    // Those 2 parts can accelerate heap-heap inference in the data-flow analysis.
125    pub fn start(&mut self) {
126        #[inline(always)]
127        fn start_channel<M>(mut method: M, v_did: &Vec<DefId>)
128        where
129            M: FnMut(DefId) -> (),
130        {
131            for did in v_did {
132                method(*did);
133            }
134        }
135
136        #[inline(always)]
137        fn show_heap(ref_type_analysis: &mut HeapOwnershipAnalyzer) {
138            for elem in ref_type_analysis.adt_heap() {
139                let name = format!(
140                    "{:?}",
141                    EarlyBinder::skip_binder(ref_type_analysis.tcx.type_of(*elem.0))
142                );
143                let owning = format!("{:?}", elem.1);
144                rap_debug!("ADT analysis: {} {}", name, owning);
145            }
146        }
147
148        // Get the Global TyCtxt from rustc
149        // Grasp all mir Keys defined in current crate
150        let tcx = self.tcx;
151        let mir_keys = tcx.mir_keys(());
152
153        for each_mir in mir_keys {
154            // Get the defid of current crate and get mir Body through this id
155            let def_id = each_mir.to_def_id();
156            let body = tcx.instance_mir(Item(def_id));
157
158            // Insert the defid to hashset if is not existed and visit the body
159            if self.fn_set_mut().insert(def_id) {
160                self.visit_body(body);
161            } else {
162                continue;
163            }
164        }
165
166        let dids: Vec<DefId> = self.adt_recorder.iter().map(|did| *did).collect();
167
168        start_channel(|did| self.extract_raw_generic(did), &dids);
169        start_channel(|did| self.extract_raw_generic_prop(did), &dids);
170        start_channel(|did| self.extract_phantom_unit(did), &dids);
171        start_channel(|did| self.extract_heap_prop(did), &dids);
172
173        show_heap(self);
174    }
175
176    // Extract params in adt types, the 'param' means one generic parameter acting like 'T', 'A', etc...
177    // In the sub-visitor RawGeneric, it will visit the given type recursively, and extract all params.
178    //
179    // Note that RAPx is only interested in 'raw' params ('T' not like '*mut T').
180    // It lies in 'one-entire field' | recursive in tuple | recursive in array | mixed before
181    //
182    // Given a struct Example<A, B, T, S>:
183    //
184    // struct Example<A, B, T, S> {
185    //     a: A,
186    //     b: (i32, (f64, B)),
187    //     c: [[(S) ; 1] ; 2],
188    //     d: Vec<T>,
189    // }
190    //
191    // the final result for <A, B, T, S> is <true, true, false, true>.
192    #[inline(always)]
193    fn extract_raw_generic(&mut self, did: DefId) {
194        // Get the definition and subset reference from adt did
195        let ty = EarlyBinder::skip_binder(self.tcx.type_of(did));
196        let (adt_def, substs) = match ty.kind() {
197            TyKind::Adt(adt_def, substs) => (adt_def, substs),
198            _ => unreachable!(),
199        };
200
201        let mut v_res = Vec::new();
202
203        for variant in adt_def.variants().iter() {
204            let mut raw_generic = IsolatedParam::new(substs.len());
205
206            for field in &variant.fields {
207                let field_ty = crate::helpers::mir_utils::field_ty(self.tcx, field, substs);
208                let _ = field_ty.visit_with(&mut raw_generic);
209            }
210            v_res.push((HeapOwnership::False, raw_generic.record_mut().clone()));
211        }
212
213        self.adt_heap_mut().insert(did, v_res);
214    }
215
216    // Extract all params in the adt types like param 'T' and then propagate from the bottom to top.
217    // This procedural is the successor of `extract_raw_generic`, and the main idea of RawGenericPropagation
218    // is to propagate params from bottom adt to the top as well as updating Analysis Context.
219    //
220    // Note that it will thorough consider mono-morphization existed in adt-def.
221    // That means the type 'Vec<T>', 'Vec<Vec<T>>' and 'Vec<i32>' are totally different!!!!
222    //
223    // Given a struct Example<A, B, T, S>:
224    //
225    // struct X<A> {
226    //     a: A,
227    // }
228    // the final result for <A> is <true>.
229    //
230    // struct Y1<B> {
231    //     a: (i32, (f64, B)),
232    //     b: X<i32>,
233    // }
234    // the final result for <B> is <true>.
235    //
236    // struct Example<A, B, T, S> {
237    //     a: X<A>,
238    //     b: (i32, (f64, B)),
239    //     c: [[(S) ; 1] ; 2],
240    //     d: Vec<T>,
241    // }
242    //
243    // the final result for <A, B, T, S> is <true, true, false, true>.
244    #[inline(always)]
245    fn extract_raw_generic_prop(&mut self, did: DefId) {
246        // Get the definition and subset reference from adt did
247        let ty = EarlyBinder::skip_binder(self.tcx.type_of(did));
248        let (adt_def, substs) = match ty.kind() {
249            TyKind::Adt(adt_def, substs) => (adt_def, substs),
250            _ => unreachable!(),
251        };
252
253        let source_enum = adt_def.is_enum();
254
255        let mut v_res = self.adt_heap_mut().get_mut(&did).unwrap().clone();
256
257        for (variant_index, variant) in adt_def.variants().iter().enumerate() {
258            let res = v_res[variant_index ].clone();
259
260            let mut raw_generic_prop = IsolatedParamPropagation::new(
261                self.tcx,
262                res.1.clone(),
263                source_enum,
264                self.adt_heap(),
265            );
266
267            for field in &variant.fields {
268                let field_ty = crate::helpers::mir_utils::field_ty(self.tcx, field, substs);
269                let _ = field_ty.visit_with(&mut raw_generic_prop);
270            }
271            v_res[variant_index] =
272                (HeapOwnership::False, raw_generic_prop.record_mut().clone());
273        }
274
275        self.adt_heap_mut().insert(did, v_res);
276    }
277
278    // Extract all types that include PhantomData<T> which T must be a raw Param
279    // Consider these types as a unit to guide the traversal over adt types
280    #[inline(always)]
281    fn extract_phantom_unit(&mut self, did: DefId) {
282        // Get ty from defid and the ty is made up with generic type
283        let ty = EarlyBinder::skip_binder(self.tcx.type_of(did));
284        let (adt_def, substs) = match ty.kind() {
285            TyKind::Adt(adt_def, substs) => (adt_def, substs),
286            _ => unreachable!(),
287        };
288
289        // As for one heap-allocation unit, only struct will contains the information that we want
290        // Example:
291        // struct Foo<T> {
292        //     NonNull<T>,      // this indicates a pointer
293        //     PhantomData<T>,  // this indicates a heap
294        // }
295        if adt_def.is_struct() {
296            let mut res = self.adt_heap_mut().get_mut(&did).unwrap()[0].clone();
297            // Extract all fields in one given struct
298            for field in adt_def.all_fields() {
299                let field_ty = crate::helpers::mir_utils::field_ty(self.tcx, field, substs);
300                match field_ty.kind() {
301                    // Filter the field which is also a struct due to PhantomData<T> is struct
302                    TyKind::Adt(field_adt_def, field_substs) => {
303                        if field_adt_def.is_phantom_data() {
304                            // Extract all generic args in the type
305                            for generic_arg in *field_substs {
306                                match generic_arg.kind() {
307                                    GenericArgKind::Type(g_ty) => {
308                                        let mut raw_generic_field_subst =
309                                            IsolatedParamFieldSubst::new();
310                                        let _ = g_ty.visit_with(&mut raw_generic_field_subst);
311                                        if raw_generic_field_subst.contains_param() {
312                                            {
313                                                // To enhance the soundness of phantom unit, the struct should have a
314                                                // pointer to store T
315                                                let mut has_ptr = false;
316                                                for field in adt_def.all_fields() {
317                                                    let field_ty =
318                                                        crate::helpers::mir_utils::field_ty(
319                                                            self.tcx, field, substs,
320                                                        );
321                                                    let mut find_ptr = FindPtr::new(self.tcx);
322                                                    let _ = field_ty.visit_with(&mut find_ptr);
323                                                    if find_ptr.has_ptr() {
324                                                        has_ptr = true;
325                                                        break;
326                                                    }
327                                                }
328                                                if !has_ptr {
329                                                    return;
330                                                }
331                                            }
332
333                                            res.0 = HeapOwnership::True;
334                                            self.adt_heap_mut().insert(did, vec![res.clone()]);
335                                            return;
336                                        }
337                                    }
338                                    GenericArgKind::Lifetime(..) => {
339                                        return;
340                                    }
341                                    GenericArgKind::Const(..) => {
342                                        return;
343                                    }
344                                }
345                            }
346                        }
347                    }
348                    _ => continue,
349                }
350            }
351        }
352    }
353
354    #[inline(always)]
355    fn extract_heap_prop(&mut self, did: DefId) {
356        // Get the definition and subset reference from adt did
357        let ty = EarlyBinder::skip_binder(self.tcx.type_of(did));
358        let (adt_def, substs) = match ty.kind() {
359            TyKind::Adt(adt_def, substs) => (adt_def, substs),
360            _ => unreachable!(),
361        };
362
363        let mut v_res = self.adt_heap_mut().get_mut(&did).unwrap().clone();
364
365        for (variant_index, variant) in adt_def.variants().iter().enumerate() {
366            let res = v_res[variant_index ].clone();
367
368            let mut heap_prop = HeapPropagation::new(self.tcx, res.0, self.adt_heap());
369
370            for field in &variant.fields {
371                let field_ty = crate::helpers::mir_utils::field_ty(self.tcx, field, substs);
372                let _ = field_ty.visit_with(&mut heap_prop);
373            }
374            v_res[variant_index].0 = heap_prop.heap();
375        }
376
377        self.adt_heap_mut().insert(did, v_res);
378    }
379}
380
381impl<'tcx> Visitor<'tcx> for HeapOwnershipAnalyzer<'tcx> {
382    fn visit_body(&mut self, body: &Body<'tcx>) {
383        for (local, local_decl) in body.local_decls.iter().enumerate() {
384            self.visit_local_decl(Local::from(local), local_decl);
385        }
386
387        for (block, data) in body.basic_blocks.iter().enumerate() {
388            self.visit_basic_block_data(BasicBlock::from(block), data);
389        }
390    }
391
392    fn visit_basic_block_data(&mut self, _block: BasicBlock, data: &BasicBlockData<'tcx>) {
393        let term = data.terminator();
394        if let TerminatorKind::Call { func, .. } = &term.kind { if let Operand::Constant(constant) = func { if let ty::FnDef(def_id, ..) = constant.ty().kind() {
395            if self.tcx.is_mir_available(*def_id) && self.fn_set_mut().insert(*def_id) {
396                let body = self.tcx.instance_mir(Item(*def_id));
397                self.visit_body(body);
398            }
399        } } }
400    }
401
402    fn visit_ty(&mut self, ty: Ty<'tcx>, ty_context: TyContext) {
403        match ty.kind() {
404            TyKind::Adt(adtdef, substs) => {
405                if self.ty_map().get(&ty).is_some() {
406                    return;
407                }
408                self.ty_map_mut().insert(ty, format!("{:?}", ty));
409                self.adt_recorder_mut().insert(adtdef.did());
410
411                for field in adtdef.all_fields() {
412                    let fty = crate::helpers::mir_utils::field_ty(self.tcx, field, substs);
413                    self.visit_ty(fty, copy_ty_context(&ty_context))
414                }
415
416                for ty in substs.types() {
417                    self.visit_ty(ty, copy_ty_context(&ty_context));
418                }
419            }
420            TyKind::Array(ty, ..) => {
421                self.visit_ty(*ty, ty_context);
422            }
423            TyKind::Slice(ty) => {
424                self.visit_ty(*ty, ty_context);
425            }
426            TyKind::RawPtr(ty, _) => {
427                self.visit_ty(*ty, ty_context);
428            }
429            TyKind::Ref(_, ty, ..) => {
430                self.visit_ty(*ty, ty_context);
431            }
432            TyKind::Tuple(tuple_fields) => {
433                for field in tuple_fields.iter() {
434                    self.visit_ty(field, copy_ty_context(&ty_context));
435                }
436            }
437            _ => (),
438        }
439    }
440
441    fn visit_local_decl(&mut self, local: Local, local_decl: &LocalDecl<'tcx>) {
442        let ty_context = TyContext::LocalDecl {
443            local,
444            source_info: local_decl.source_info,
445        };
446        self.visit_ty(local_decl.ty, ty_context);
447    }
448}
449
450impl<'tcx> TypeVisitor<TyCtxt<'tcx>> for IsolatedParam {
451    type Result = ControlFlow<()>;
452    fn visit_ty(&mut self, ty: Ty<'tcx>) -> Self::Result {
453        match ty.kind() {
454            TyKind::Array(..) => ty.super_visit_with(self),
455            TyKind::Tuple(..) => ty.super_visit_with(self),
456            TyKind::Param(param_ty) => {
457                self.record_mut()[param_ty.index as usize] = true;
458                ControlFlow::Continue(())
459            }
460            _ => ControlFlow::Continue(()),
461        }
462    }
463}
464
465impl<'tcx> TypeVisitor<TyCtxt<'tcx>> for IsolatedParamFieldSubst {
466    type Result = ControlFlow<()>;
467    #[inline(always)]
468    fn visit_ty(&mut self, ty: Ty<'tcx>) -> Self::Result {
469        match ty.kind() {
470            TyKind::Array(..) => ty.super_visit_with(self),
471            TyKind::Tuple(..) => ty.super_visit_with(self),
472            TyKind::Adt(..) => ty.super_visit_with(self),
473            TyKind::Param(param_ty) => {
474                self.parameters_mut().insert(param_ty.index as usize);
475                ControlFlow::Continue(())
476            }
477            _ => ControlFlow::Continue(()),
478        }
479    }
480}
481
482impl<'tcx, 'a> TypeVisitor<TyCtxt<'tcx>> for IsolatedParamPropagation<'tcx, 'a> {
483    // #[inline(always)]
484    // fn tcx_for_anon_const_substs(&self) -> Option<TyCtxt<'tcx>> {
485    //     Some(self.tcx)
486    // }
487    type Result = ControlFlow<()>;
488
489    #[inline(always)]
490    fn visit_ty(&mut self, ty: Ty<'tcx>) -> Self::Result {
491        match ty.kind() {
492            TyKind::Adt(adtdef, substs) => {
493                if substs.len() == 0 {
494                    return ControlFlow::Break(());
495                }
496
497                if !self.source_enum() && adtdef.is_enum() {
498                    return ControlFlow::Break(());
499                }
500
501                if !self.unique_mut().insert(adtdef.did()) {
502                    return ControlFlow::Continue(());
503                }
504
505                let mut map_raw_generic_field_subst = HashMap::new();
506                for (index, subst) in substs.iter().enumerate() {
507                    match subst.kind() {
508                        GenericArgKind::Lifetime(..) => continue,
509                        GenericArgKind::Const(..) => continue,
510                        GenericArgKind::Type(g_ty) => {
511                            let mut raw_generic_field_subst = IsolatedParamFieldSubst::new();
512                            let _ = g_ty.visit_with(&mut raw_generic_field_subst);
513                            if !raw_generic_field_subst.contains_param() {
514                                continue;
515                            }
516                            map_raw_generic_field_subst
517                                .insert(index, raw_generic_field_subst);
518                        }
519                    }
520                }
521                if map_raw_generic_field_subst.is_empty() {
522                    return ControlFlow::Break(());
523                }
524
525                let get_ans = self.heap().get(&adtdef.did()).unwrap();
526                if get_ans.len() == 0 {
527                    return ControlFlow::Break(());
528                }
529                let get_ans = get_ans[0].clone();
530
531                for (index, flag) in get_ans.1.iter().enumerate() {
532                    if *flag && map_raw_generic_field_subst.contains_key(&index) {
533                        for elem in map_raw_generic_field_subst
534                            .get(&index)
535                            .unwrap()
536                            .parameters()
537                        {
538                            self.record[*elem] = true;
539                        }
540                    }
541                }
542
543                for field in adtdef.all_fields() {
544                    let field_ty = crate::helpers::mir_utils::field_ty(self.tcx, field, substs);
545                    let _ = field_ty.visit_with(self);
546                }
547
548                self.unique_mut().remove(&adtdef.did());
549
550                ty.super_visit_with(self)
551            }
552            TyKind::Array(..) => ty.super_visit_with(self),
553            TyKind::Tuple(..) => ty.super_visit_with(self),
554            _ => ControlFlow::Continue(()),
555        }
556    }
557}
558
559impl<'tcx, 'a> TypeVisitor<TyCtxt<'tcx>> for HeapPropagation<'tcx, 'a> {
560    // #[inline(always)]
561    // fn tcx_for_anon_const_substs(&self) -> Option<TyCtxt<'tcx>> {
562    //     Some(self.tcx)
563    // }
564    type Result = ControlFlow<()>;
565    #[inline(always)]
566    fn visit_ty(&mut self, ty: Ty<'tcx>) -> Self::Result {
567        match ty.kind() {
568            TyKind::Adt(adtdef, substs) => {
569                if !self.unique_mut().insert(adtdef.did()) {
570                    return ControlFlow::Continue(());
571                }
572
573                if adtdef.is_enum() {
574                    return ControlFlow::Break(());
575                }
576
577                let get_ans = self.heap_res().get(&adtdef.did()).unwrap();
578                if get_ans.len() == 0 {
579                    return ControlFlow::Break(());
580                }
581                let get_ans = get_ans[0].clone();
582
583                if get_ans.0 == HeapOwnership::True {
584                    self.heap = HeapOwnership::True;
585                    return ControlFlow::Break(());
586                };
587
588                for field in adtdef.all_fields() {
589                    let field_ty = crate::helpers::mir_utils::field_ty(self.tcx, field, substs);
590                    let _ = field_ty.visit_with(self);
591                }
592
593                self.unique_mut().remove(&adtdef.did());
594
595                ty.super_visit_with(self)
596            }
597            TyKind::Array(..) => ty.super_visit_with(self),
598            TyKind::Tuple(..) => ty.super_visit_with(self),
599            _ => ControlFlow::Continue(()),
600        }
601    }
602}
603
604impl<'tcx> TypeVisitor<TyCtxt<'tcx>> for FindPtr<'tcx> {
605    type Result = ControlFlow<()>;
606    #[inline(always)]
607    fn visit_ty(&mut self, ty: Ty<'tcx>) -> Self::Result {
608        match ty.kind() {
609            TyKind::Adt(adtdef, substs) => {
610                if adtdef.is_struct() {
611                    if !self.unique_mut().insert(adtdef.did()) {
612                        return ControlFlow::Continue(());
613                    }
614
615                    for field in adtdef.all_fields() {
616                        let field_ty = crate::helpers::mir_utils::field_ty(self.tcx, field, substs);
617                        let _ = field_ty.visit_with(self);
618                    }
619                    self.unique_mut().remove(&adtdef.did());
620                }
621                ControlFlow::Continue(())
622            }
623            TyKind::Tuple(..) => ty.super_visit_with(self),
624            TyKind::RawPtr(..) => {
625                self.set_ptr(true);
626                ControlFlow::Break(())
627            }
628            TyKind::Ref(..) => {
629                self.set_ptr(true);
630                ControlFlow::Break(())
631            }
632            #[cfg(rapx_ge_99)]
633            TyKind::Pat(..) => {
634                self.set_ptr(true);
635                ControlFlow::Break(())
636            }
637            _ => ControlFlow::Continue(()),
638        }
639    }
640}
641
642impl<'tcx, 'a> TypeVisitor<TyCtxt<'tcx>> for DefaultOwnership<'tcx, 'a> {
643    // #[inline(always)]
644    // fn tcx_for_anon_const_substs(&self) -> Option<TyCtxt<'tcx>> {
645    //     Some(self.tcx)
646    // }
647    type Result = ControlFlow<()>;
648    #[inline(always)]
649    fn visit_ty(&mut self, ty: Ty<'tcx>) -> Self::Result {
650        match ty.kind() {
651            TyKind::Adt(adtdef, substs) => {
652                if adtdef.is_enum() {
653                    return ControlFlow::Break(());
654                }
655
656                if !self.unique_mut().insert(adtdef.did()) {
657                    return ControlFlow::Continue(());
658                }
659
660                let get_ans = self.heap().get(&adtdef.did()).unwrap();
661
662                // handle the secene of Zero Sized Types
663                if get_ans.len() == 0 {
664                    return ControlFlow::Break(());
665                }
666                let (unit_res, generic_list) = get_ans[0].clone();
667
668                match unit_res {
669                    HeapOwnership::True => {
670                        self.set_res(HeapOwnership::True);
671                        return ControlFlow::Break(());
672                    }
673                    HeapOwnership::False => {
674                        for (index, each_generic) in generic_list.iter().enumerate() {
675                            if !*each_generic {
676                                continue;
677                            } else {
678                                let subset_ty = substs[index].expect_ty();
679                                self.unique_mut().remove(&adtdef.did());
680                                let _ = subset_ty.visit_with(self);
681                            }
682                        }
683                    }
684                    _ => {
685                        unreachable!();
686                    }
687                }
688                ControlFlow::Continue(())
689            }
690            TyKind::Array(..) => ty.super_visit_with(self),
691            TyKind::Tuple(..) => ty.super_visit_with(self),
692            TyKind::Param(..) => {
693                self.set_param(true);
694                self.set_res(HeapOwnership::True);
695                ControlFlow::Break(())
696            }
697            TyKind::RawPtr(..) => {
698                self.set_ptr(true);
699                ControlFlow::Continue(())
700            }
701            TyKind::Ref(..) => {
702                self.set_ptr(true);
703                ControlFlow::Continue(())
704            }
705            _ => ControlFlow::Continue(()),
706        }
707    }
708}
709
710#[derive(Debug, Clone, Hash, Eq, PartialEq, Default)]
711pub struct TyWithIndex<'tcx>(pub Option<(usize, &'tcx TyKind<'tcx>, Option<usize>, bool)>);
712
713impl<'tcx> TyWithIndex<'tcx> {
714    pub fn new(ty: Ty<'tcx>, vidx: Option<VariantIdx>) -> Self {
715        match &ty.kind() {
716            TyKind::Tuple(list) => TyWithIndex(Some((list.len(), ty.kind(), None, true))),
717            TyKind::Adt(adtdef, ..) => {
718                if adtdef.is_enum() {
719                    if vidx.is_none() {
720                        return TyWithIndex(None);
721                    }
722                    let idx = vidx.unwrap();
723                    let len = adtdef.variants()[idx].fields.len();
724                    TyWithIndex(Some((len, ty.kind(), Some(idx.index()), true)))
725                } else {
726                    let len = adtdef.variants()[VariantIdx::from_usize(0)].fields.len();
727                    TyWithIndex(Some((len, ty.kind(), None, true)))
728                }
729            }
730            TyKind::Array(..) | TyKind::Param(..) | TyKind::RawPtr(..) | TyKind::Ref(..) => {
731                TyWithIndex(Some((1, ty.kind(), None, true)))
732            }
733            TyKind::Bool
734            | TyKind::Char
735            | TyKind::Int(..)
736            | TyKind::Uint(..)
737            | TyKind::Float(..)
738            | TyKind::Str
739            | TyKind::Slice(..) => TyWithIndex(Some((1, ty.kind(), None, false))),
740            _ => TyWithIndex(None),
741        }
742    }
743
744    // 0->unsupported, 1->trivial, 2-> needed
745    pub fn get_priority(&self) -> usize {
746        if self.0.is_none() {
747            return 0;
748        }
749        match self.0.unwrap().0 {
750            0 => 1,
751            _ => match self.0.unwrap().3 {
752                true => 2,
753                false => 1,
754            },
755        }
756    }
757}
758
759#[derive(Copy, Clone, Debug)]
760pub struct Encoder;
761
762impl<'tcx> Encoder {
763    pub fn encode(
764        tcx: TyCtxt<'tcx>,
765        ty: Ty<'tcx>,
766        adt_heap: HeapOwnershipResultMap,
767        variant: Option<VariantIdx>,
768    ) -> OwnershipLayoutResult {
769        match ty.kind() {
770            TyKind::Array(..) => {
771                let mut res = OwnershipLayoutResult::new();
772                let mut default_heap = DefaultOwnership::new(tcx, &adt_heap);
773
774                let _ = ty.visit_with(&mut default_heap);
775                res.update_from_default_heap_visitor(&mut default_heap);
776
777                res
778            }
779            TyKind::Tuple(tuple_ty_list) => {
780                let mut res = OwnershipLayoutResult::new();
781
782                for tuple_ty in tuple_ty_list.iter() {
783                    let mut default_heap = DefaultOwnership::new(tcx, &adt_heap);
784
785                    let _ = tuple_ty.visit_with(&mut default_heap);
786                    res.update_from_default_heap_visitor(&mut default_heap);
787                }
788
789                res
790            }
791            TyKind::Adt(adtdef, substs) => {
792                // check the ty is or is not an enum and the variant of this enum is or is not given
793                if adtdef.is_enum() && variant.is_none() {
794                    return OwnershipLayoutResult::new();
795                }
796
797                let mut res = OwnershipLayoutResult::new();
798
799                // check the ty if it is a struct or union
800                if adtdef.is_struct() || adtdef.is_union() {
801                    for field in adtdef.all_fields() {
802                        let field_ty = crate::helpers::mir_utils::field_ty(tcx, field, substs);
803
804                        let mut default_heap = DefaultOwnership::new(tcx, &adt_heap);
805
806                        let _ = field_ty.visit_with(&mut default_heap);
807                        res.update_from_default_heap_visitor(&mut default_heap);
808                    }
809                }
810                // check the ty which is an enum with a exact variant idx
811                else if adtdef.is_enum() {
812                    let vidx = variant.unwrap();
813
814                    for field in &adtdef.variants()[vidx].fields {
815                        let field_ty = crate::helpers::mir_utils::field_ty(tcx, field, substs);
816
817                        let mut default_heap = DefaultOwnership::new(tcx, &adt_heap);
818
819                        let _ = field_ty.visit_with(&mut default_heap);
820                        res.update_from_default_heap_visitor(&mut default_heap);
821                    }
822                }
823                res
824            }
825            TyKind::Param(..) => {
826                let mut res = OwnershipLayoutResult::new();
827                res.set_requirement(true);
828                res.set_param(true);
829                res.set_owned(true);
830                res.layout_mut().push(HeapOwnership::True);
831                res
832            }
833            TyKind::RawPtr(..) => {
834                let mut res = OwnershipLayoutResult::new();
835                res.set_requirement(true);
836                res.layout_mut().push(HeapOwnership::False);
837                res
838            }
839            TyKind::Ref(..) => {
840                let mut res = OwnershipLayoutResult::new();
841                res.set_requirement(true);
842                res.layout_mut().push(HeapOwnership::False);
843                res
844            }
845            _ => OwnershipLayoutResult::new(),
846        }
847    }
848}
849
850#[derive(Clone)]
851struct IsolatedParamFieldSubst {
852    parameters: HashSet<usize>,
853}
854
855impl<'tcx> IsolatedParamFieldSubst {
856    pub fn new() -> Self {
857        Self {
858            parameters: HashSet::new(),
859        }
860    }
861
862    pub fn parameters(&self) -> &HashSet<usize> {
863        &self.parameters
864    }
865
866    pub fn parameters_mut(&mut self) -> &mut HashSet<usize> {
867        &mut self.parameters
868    }
869
870    pub fn contains_param(&self) -> bool {
871        !self.parameters.is_empty()
872    }
873}
874
875#[derive(Clone)]
876struct IsolatedParamPropagation<'tcx, 'a> {
877    tcx: TyCtxt<'tcx>,
878    record: Vec<bool>,
879    unique: HashSet<DefId>,
880    source_enum: bool,
881    ref_adt_heap: &'a HeapOwnershipResultMap,
882}
883
884impl<'tcx, 'a> IsolatedParamPropagation<'tcx, 'a> {
885    pub fn new(
886        tcx: TyCtxt<'tcx>,
887        record: Vec<bool>,
888        source_enum: bool,
889        ref_adt_heap: &'a HeapOwnershipResultMap,
890    ) -> Self {
891        Self {
892            tcx,
893            record,
894            unique: HashSet::new(),
895            source_enum,
896            ref_adt_heap,
897        }
898    }
899
900    pub fn record_mut(&mut self) -> &mut Vec<bool> {
901        &mut self.record
902    }
903
904    pub fn unique_mut(&mut self) -> &mut HashSet<DefId> {
905        &mut self.unique
906    }
907
908    pub fn source_enum(&mut self) -> bool {
909        self.source_enum
910    }
911
912    pub fn heap(&self) -> &'a HeapOwnershipResultMap {
913        self.ref_adt_heap
914    }
915}
916
917#[derive(Clone)]
918struct HeapPropagation<'tcx, 'a> {
919    tcx: TyCtxt<'tcx>,
920    heap: HeapOwnership,
921    unique: HashSet<DefId>,
922    heap_res: &'a HeapOwnershipResultMap,
923}
924
925impl<'tcx, 'a> HeapPropagation<'tcx, 'a> {
926    pub fn new(
927        tcx: TyCtxt<'tcx>,
928        heap: HeapOwnership,
929        heap_res: &'a HeapOwnershipResultMap,
930    ) -> Self {
931        Self {
932            tcx,
933            heap,
934            unique: HashSet::new(),
935            heap_res,
936        }
937    }
938
939    pub fn heap(&self) -> HeapOwnership {
940        self.heap
941    }
942
943    pub fn unique_mut(&mut self) -> &mut HashSet<DefId> {
944        &mut self.unique
945    }
946
947    pub fn heap_res(&self) -> &'a HeapOwnershipResultMap {
948        self.heap_res
949    }
950}
951
952#[derive(Clone)]
953struct IsolatedParam {
954    record: Vec<bool>,
955}
956
957impl IsolatedParam {
958    pub fn new(len: usize) -> Self {
959        Self {
960            record: vec![false; len],
961        }
962    }
963
964    pub fn record_mut(&mut self) -> &mut Vec<bool> {
965        &mut self.record
966    }
967}
968
969#[derive(Clone)]
970pub struct DefaultOwnership<'tcx, 'a> {
971    tcx: TyCtxt<'tcx>,
972    unique: HashSet<DefId>,
973    ref_adt_heap: &'a HeapOwnershipResultMap,
974    res: HeapOwnership,
975    param: bool,
976    ptr: bool,
977}
978
979impl<'tcx, 'a> DefaultOwnership<'tcx, 'a> {
980    pub fn new(tcx: TyCtxt<'tcx>, ref_adt_heap: &'a HeapOwnershipResultMap) -> Self {
981        Self {
982            tcx,
983            unique: HashSet::new(),
984            ref_adt_heap,
985            res: HeapOwnership::False,
986            param: false,
987            ptr: false,
988        }
989    }
990
991    pub fn tcx(&self) -> TyCtxt<'tcx> {
992        self.tcx
993    }
994
995    pub fn unique(&self) -> &HashSet<DefId> {
996        &self.unique
997    }
998
999    pub fn unique_mut(&mut self) -> &mut HashSet<DefId> {
1000        &mut self.unique
1001    }
1002
1003    pub fn get_res(&self) -> HeapOwnership {
1004        self.res
1005    }
1006
1007    pub fn set_res(&mut self, res: HeapOwnership) {
1008        self.res = res;
1009    }
1010
1011    pub fn is_owning_true(&self) -> bool {
1012        self.res == HeapOwnership::True
1013    }
1014
1015    pub fn get_param(&self) -> bool {
1016        self.param
1017    }
1018
1019    pub fn set_param(&mut self, p: bool) {
1020        self.param = p;
1021    }
1022
1023    pub fn is_param_true(&self) -> bool {
1024        self.param
1025    }
1026
1027    pub fn get_ptr(&self) -> bool {
1028        self.ptr
1029    }
1030
1031    pub fn set_ptr(&mut self, p: bool) {
1032        self.ptr = p;
1033    }
1034
1035    pub fn is_ptr_true(&self) -> bool {
1036        self.ptr
1037    }
1038
1039    pub fn heap(&self) -> &'a HeapOwnershipResultMap {
1040        self.ref_adt_heap
1041    }
1042}
1043
1044#[derive(Clone)]
1045pub struct FindPtr<'tcx> {
1046    tcx: TyCtxt<'tcx>,
1047    unique: HashSet<DefId>,
1048    ptr: bool,
1049}
1050
1051impl<'tcx> FindPtr<'tcx> {
1052    pub fn new(tcx: TyCtxt<'tcx>) -> Self {
1053        Self {
1054            tcx,
1055            unique: HashSet::<DefId>::default(),
1056            ptr: false,
1057        }
1058    }
1059
1060    pub fn tcx(&self) -> TyCtxt<'tcx> {
1061        self.tcx
1062    }
1063
1064    pub fn unique(&self) -> &HashSet<DefId> {
1065        &self.unique
1066    }
1067
1068    pub fn unique_mut(&mut self) -> &mut HashSet<DefId> {
1069        &mut self.unique
1070    }
1071
1072    pub fn has_ptr(&self) -> bool {
1073        self.ptr
1074    }
1075
1076    pub fn set_ptr(&mut self, ptr: bool) {
1077        self.ptr = ptr;
1078    }
1079}
1080
1081pub fn is_display_verbose() -> bool {
1082    env::var_os("ADT_DISPLAY").is_some()
1083}
1084
1085#[derive(Debug, Clone, Hash, Eq, PartialEq, Default)]
1086pub struct IndexedTy<'tcx>(pub Option<(usize, &'tcx TyKind<'tcx>, Option<usize>, bool)>);
1087
1088impl<'tcx> IndexedTy<'tcx> {
1089    pub fn new(ty: Ty<'tcx>, vidx: Option<VariantIdx>) -> Self {
1090        match &ty.kind() {
1091            TyKind::Tuple(list) => IndexedTy(Some((list.len(), ty.kind(), None, true))),
1092            TyKind::Adt(adtdef, ..) => {
1093                if adtdef.is_enum() {
1094                    if vidx.is_none() {
1095                        return IndexedTy(None);
1096                    }
1097                    let idx = vidx.unwrap();
1098                    let len = adtdef.variants()[idx].fields.len();
1099                    IndexedTy(Some((len, ty.kind(), Some(idx.index()), true)))
1100                } else {
1101                    let len = adtdef.variants()[VariantIdx::from_usize(0)].fields.len();
1102                    IndexedTy(Some((len, ty.kind(), None, true)))
1103                }
1104            }
1105            TyKind::Array(..) | TyKind::Param(..) | TyKind::RawPtr(..) | TyKind::Ref(..) => {
1106                IndexedTy(Some((1, ty.kind(), None, true)))
1107            }
1108            TyKind::Bool
1109            | TyKind::Char
1110            | TyKind::Int(..)
1111            | TyKind::Uint(..)
1112            | TyKind::Float(..)
1113            | TyKind::Str
1114            | TyKind::Slice(..) => IndexedTy(Some((1, ty.kind(), None, false))),
1115            _ => IndexedTy(None),
1116        }
1117    }
1118
1119    // 0->unsupported, 1->trivial, 2-> needed
1120    pub fn get_priority(&self) -> usize {
1121        if self.0.is_none() {
1122            return 0;
1123        }
1124        match self.0.unwrap().0 {
1125            0 => 1,
1126            _ => match self.0.unwrap().3 {
1127                true => 2,
1128                false => 1,
1129            },
1130        }
1131    }
1132}
1133
1134#[derive(Clone, Debug)]
1135pub struct OwnershipLayoutResult {
1136    layout: Vec<HeapOwnership>,
1137    param: bool,
1138    requirement: bool,
1139    owned: bool,
1140}
1141
1142impl OwnershipLayoutResult {
1143    pub fn new() -> Self {
1144        Self {
1145            layout: Vec::new(),
1146            param: false,
1147            requirement: false,
1148            owned: false,
1149        }
1150    }
1151
1152    pub fn layout(&self) -> &Vec<HeapOwnership> {
1153        &self.layout
1154    }
1155
1156    pub fn layout_mut(&mut self) -> &mut Vec<HeapOwnership> {
1157        &mut self.layout
1158    }
1159
1160    pub fn get_param(&self) -> bool {
1161        self.param
1162    }
1163
1164    pub fn set_param(&mut self, p: bool) {
1165        self.param = p;
1166    }
1167
1168    pub fn is_param_true(&self) -> bool {
1169        self.param
1170    }
1171
1172    pub fn get_requirement(&self) -> bool {
1173        self.requirement
1174    }
1175
1176    pub fn set_requirement(&mut self, r: bool) {
1177        self.requirement = r;
1178    }
1179
1180    pub fn is_requirement_true(&self) -> bool {
1181        self.requirement
1182    }
1183
1184    pub fn is_empty(&self) -> bool {
1185        self.layout.is_empty()
1186    }
1187
1188    pub fn is_owned(&self) -> bool {
1189        self.owned
1190    }
1191
1192    pub fn set_owned(&mut self, o: bool) {
1193        self.owned = o;
1194    }
1195
1196    pub fn update_from_default_heap_visitor<'tcx, 'a>(
1197        &mut self,
1198        default_heap: &mut DefaultOwnership<'tcx, 'a>,
1199    ) {
1200        if default_heap.is_owning_true() || default_heap.is_ptr_true() {
1201            self.set_requirement(true);
1202        }
1203
1204        if default_heap.is_owning_true() {
1205            self.set_owned(true);
1206        }
1207
1208        self.layout_mut().push(default_heap.get_res());
1209
1210        self.set_param(default_heap.get_param());
1211    }
1212}