Skip to main content

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