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