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