rapx/analysis/api_dependency/
mono.rs

1#![allow(warnings, unused)]
2
3use super::graph::TyWrapper;
4use super::utils::{self, fn_sig_with_generic_args};
5use crate::helpers::def_path::path_str_def_id;
6use crate::{rap_debug, rap_trace};
7use rand::Rng;
8use rand::seq::SliceRandom;
9use rustc_hir::LangItem;
10use rustc_hir::def_id::DefId;
11use rustc_infer::infer::DefineOpaqueTypes;
12use rustc_infer::infer::{InferCtxt, TyCtxtInferExt};
13use rustc_infer::traits::{ImplSource, Obligation, ObligationCause};
14use rustc_middle::ty::{self, GenericArgsRef, Ty, TyCtxt, TypeVisitableExt, TypingEnv};
15use rustc_span::DUMMY_SP;
16use rustc_trait_selection::traits::query::evaluate_obligation::InferCtxtExt as _;
17use std::collections::HashSet;
18
19static MAX_STEP_SET_SIZE: usize = 1000;
20
21#[derive(Clone, Debug, Hash, PartialEq, Eq)]
22pub struct Mono<'tcx> {
23    pub value: Vec<ty::GenericArg<'tcx>>,
24}
25
26impl<'tcx> FromIterator<ty::GenericArg<'tcx>> for Mono<'tcx> {
27    fn from_iter<T>(iter: T) -> Self
28    where
29        T: IntoIterator<Item = ty::GenericArg<'tcx>>,
30    {
31        Mono {
32            value: iter.into_iter().collect(),
33        }
34    }
35}
36
37impl<'tcx> Mono<'tcx> {
38    pub fn new(identity: &[ty::GenericArg<'tcx>]) -> Self {
39        Mono {
40            value: Vec::from(identity),
41        }
42    }
43
44    fn has_infer_types(&self) -> bool {
45        self.value.iter().any(|arg| match arg.kind() {
46            ty::GenericArgKind::Type(ty) => ty.has_infer_types(),
47            _ => false,
48        })
49    }
50
51    fn mut_arg_at(&mut self, idx: usize) -> &mut ty::GenericArg<'tcx> {
52        &mut self.value[idx]
53    }
54
55    fn merge(&self, other: &Mono<'tcx>, tcx: TyCtxt<'tcx>) -> Option<Mono<'tcx>> {
56        assert!(self.value.len() == other.value.len());
57        let mut res = Vec::new();
58        for i in 0..self.value.len() {
59            let arg = self.value[i];
60            let other_arg = other.value[i];
61            let new_arg = if let Some(ty) = arg.as_type() {
62                let other_ty = other_arg.expect_ty();
63                if ty.is_ty_var() && other_ty.is_ty_var() {
64                    arg
65                } else if ty.is_ty_var() {
66                    other_arg
67                } else if other_ty.is_ty_var() {
68                    arg
69                } else if utils::is_ty_eq(ty, other_ty, tcx) {
70                    arg
71                } else {
72                    return None;
73                }
74            } else {
75                arg
76            };
77            res.push(new_arg);
78        }
79        Some(Mono { value: res })
80    }
81
82    fn fill_unbound_var(&self, tcx: TyCtxt<'tcx>) -> Vec<Mono<'tcx>> {
83        let candidates = get_unbound_generic_candidates(tcx);
84        let mut res = vec![self.clone()];
85        rap_trace!("fill unbound: {:?}", self);
86
87        for (i, arg) in self.value.iter().enumerate() {
88            if let Some(ty) = arg.as_type() {
89                if ty.is_ty_var() {
90                    let mut last = Vec::new();
91                    std::mem::swap(&mut res, &mut last);
92                    last.into_iter().for_each(|mono| {
93                        for candidate in &candidates {
94                            let mut new_mono = mono.clone();
95                            *new_mono.mut_arg_at(i) = (*candidate).into();
96                            res.push(new_mono);
97                        }
98                    });
99                }
100            }
101        }
102        res
103    }
104}
105
106#[derive(Clone, Debug, Default)]
107pub struct MonoSet<'tcx> {
108    pub monos: Vec<Mono<'tcx>>,
109}
110
111impl<'tcx> MonoSet<'tcx> {
112    pub fn all(identity: &[ty::GenericArg<'tcx>]) -> MonoSet<'tcx> {
113        MonoSet {
114            monos: vec![Mono::new(identity)],
115        }
116    }
117
118    pub fn empty() -> MonoSet<'tcx> {
119        MonoSet { monos: Vec::new() }
120    }
121
122    pub fn count(&self) -> usize {
123        self.monos.len()
124    }
125
126    pub fn at(&self, no: usize) -> &Mono<'tcx> {
127        &self.monos[no]
128    }
129
130    pub fn is_empty(&self) -> bool {
131        self.monos.is_empty()
132    }
133
134    pub fn new() -> MonoSet<'tcx> {
135        MonoSet { monos: Vec::new() }
136    }
137
138    pub fn insert(&mut self, mono: Mono<'tcx>) {
139        self.monos.push(mono);
140    }
141
142    pub fn merge(&mut self, other: &MonoSet<'tcx>, tcx: TyCtxt<'tcx>) -> MonoSet<'tcx> {
143        let mut res = MonoSet::new();
144
145        for args in self.monos.iter() {
146            for other_args in other.monos.iter() {
147                let merged = args.merge(other_args, tcx);
148                if let Some(mono) = merged {
149                    res.insert(mono);
150                }
151            }
152        }
153        res
154    }
155
156    fn filter_unbound_solution(mut self) -> Self {
157        self.monos.retain(|mono| mono.has_infer_types());
158        self
159    }
160
161    // if the unbound generic type is still exist (this could happen
162    // if `T` has no trait bounds at all)
163    // we substitute the unbound generic type with predefined type candidates
164    fn instantiate_unbound(&self, tcx: TyCtxt<'tcx>) -> Self {
165        let mut res = MonoSet::new();
166        for mono in &self.monos {
167            let filled = mono.fill_unbound_var(tcx);
168            res.monos.extend(filled);
169        }
170        res
171    }
172
173    fn erase_region_var(&mut self, tcx: TyCtxt<'tcx>) {
174        for mono in &mut self.monos {
175            mono.value
176                .iter_mut()
177                .for_each(|arg| *arg = tcx.erase_and_anonymize_regions(*arg))
178        }
179    }
180
181    pub fn filter(mut self, f: impl Fn(&Mono<'tcx>) -> bool) -> Self {
182        self.monos.retain(|args| f(args));
183        self
184    }
185
186    pub fn random_sample<R: Rng>(&mut self, rng: &mut R) {
187        if self.monos.len() <= MAX_STEP_SET_SIZE {
188            return;
189        }
190        self.monos.shuffle(rng);
191        self.monos.truncate(MAX_STEP_SET_SIZE);
192    }
193}
194
195/// try to unfiy lhs = rhs,
196/// e.g.,
197/// try_unify(Vec<T>, Vec<i32>, ...) = Some(i32)
198/// try_unify(Vec<T>, i32, ...) = None
199fn unify_ty<'tcx>(
200    lhs: Ty<'tcx>,
201    rhs: Ty<'tcx>,
202    identity: &[ty::GenericArg<'tcx>],
203    infcx: &InferCtxt<'tcx>,
204    cause: &ObligationCause<'tcx>,
205    param_env: ty::ParamEnv<'tcx>,
206) -> Option<Mono<'tcx>> {
207    // rap_info!("check {} = {}", lhs, rhs);
208    infcx.probe(|_| {
209        match infcx
210            .at(cause, param_env)
211            .eq(DefineOpaqueTypes::Yes, lhs, rhs)
212        {
213            Ok(_infer_ok) => {
214                // rap_trace!("[infer_ok] {} = {} : {:?}", lhs, rhs, infer_ok);
215                let mono = identity
216                    .iter()
217                    .map(|arg| match arg.kind() {
218                        ty::GenericArgKind::Lifetime(region) => {
219                            infcx.resolve_vars_if_possible(region).into()
220                        }
221                        ty::GenericArgKind::Type(ty) => infcx.resolve_vars_if_possible(ty).into(),
222                        ty::GenericArgKind::Const(ct) => infcx.resolve_vars_if_possible(ct).into(),
223                    })
224                    .collect();
225                Some(mono)
226            }
227            Err(_e) => {
228                // rap_trace!("[infer_err] {} = {} : {:?}", lhs, rhs, e);
229                None
230            }
231        }
232    })
233}
234
235fn is_args_fit_trait_bound<'tcx>(
236    fn_did: DefId,
237    args: &[ty::GenericArg<'tcx>],
238    tcx: TyCtxt<'tcx>,
239) -> bool {
240    let args = tcx.mk_args(args);
241    rap_trace!(
242        "fn: {:?} args: {:?} identity: {:?}",
243        fn_did,
244        args,
245        ty::GenericArgs::identity_for_item(tcx, fn_did)
246    );
247    let infcx = tcx.infer_ctxt().build(ty::TypingMode::PostAnalysis);
248    let pred = tcx.predicates_of(fn_did);
249    let inst_pred = pred.instantiate(tcx, args);
250    let param_env = tcx.param_env(fn_did);
251    rap_trace!(
252        "[trait bound] check {}",
253        tcx.def_path_str_with_args(fn_did, args)
254    );
255
256    for pred in inst_pred.predicates.iter() {
257        #[cfg(rapx_rustc_ge_198)]
258        let pred = pred.skip_norm_wip();
259        let obligation = Obligation::new(
260            tcx,
261            ObligationCause::dummy(),
262            param_env,
263            pred.as_predicate(),
264        );
265        rap_trace!("[trait bound] check pred: {:?}", pred);
266
267        let res = infcx.evaluate_obligation(&obligation);
268        match res {
269            Ok(eva) => {
270                if !eva.may_apply() {
271                    rap_trace!("[trait bound] check fail for {pred:?}");
272                    return false;
273                }
274            }
275            Err(_) => {
276                rap_trace!("[trait bound] check fail for {pred:?}");
277                return false;
278            }
279        }
280    }
281    rap_trace!("[trait bound] check succ");
282    true
283}
284
285fn is_fn_solvable<'tcx>(fn_did: DefId, tcx: TyCtxt<'tcx>) -> bool {
286    for pred in tcx
287        .predicates_of(fn_did)
288        .instantiate_identity(tcx)
289        .predicates
290    {
291        #[cfg(rapx_rustc_ge_198)]
292        let pred = pred.skip_norm_wip();
293        if let Some(pred) = pred.as_trait_clause() {
294            let trait_did = pred.skip_binder().trait_ref.def_id;
295            if tcx.is_lang_item(trait_did, LangItem::Fn)
296                || tcx.is_lang_item(trait_did, LangItem::FnMut)
297                || tcx.is_lang_item(trait_did, LangItem::FnOnce)
298            {
299                return false;
300            }
301        }
302    }
303    true
304}
305
306fn get_mono_set<'tcx>(
307    fn_did: DefId,
308    available_ty: &HashSet<TyWrapper<'tcx>>,
309    tcx: TyCtxt<'tcx>,
310) -> MonoSet<'tcx> {
311    let mut rng = rand::rng();
312
313    // sample from reachable types
314    rap_debug!("[get_mono_set] solve {}", tcx.def_path_str(fn_did));
315    let identity = ty::GenericArgs::identity_for_item(tcx, fn_did);
316    let infcx = tcx
317        .infer_ctxt()
318        .ignoring_regions()
319        .build(ty::TypingMode::PostAnalysis);
320    let param_env = tcx.param_env(fn_did);
321    let dummy_cause = ObligationCause::dummy();
322    let fresh_args = infcx.fresh_args_for_item(DUMMY_SP, fn_did);
323    // this replace generic types in fn_sig to infer var, e.g. fn(Vec<T>, i32) => fn(Vec<?0>, i32)
324    let fn_sig = fn_sig_with_generic_args(fn_did, fresh_args, tcx);
325    let identity_fnsig = fn_sig_with_generic_args(fn_did, identity, tcx);
326    let generics = tcx.generics_of(fn_did);
327
328    // print fresh_args for debugging
329    for i in 0..fresh_args.len() {
330        rap_trace!(
331            "[get_mono_set] arg#{}: {:?} -> {:?}",
332            i,
333            generics.param_at(i, tcx).name,
334            fresh_args[i]
335        );
336    }
337
338    let mut s = MonoSet::all(&fresh_args);
339
340    rap_trace!("[get_mono_set] initialize s: {:?}", s);
341
342    for (no, input_ty) in fn_sig.inputs().iter().enumerate() {
343        if !input_ty.has_infer_types() {
344            continue;
345        }
346        rap_debug!(
347            "[get_mono_set] input_ty#{}: {}",
348            no,
349            identity_fnsig.inputs()[no]
350        );
351
352        let reachable_set = available_ty
353            .iter()
354            .fold(MonoSet::new(), |mut reachable_set, ty| {
355                if let Some(mono) = unify_ty(
356                    *input_ty,
357                    (*ty).into(),
358                    &fresh_args,
359                    &infcx,
360                    &dummy_cause,
361                    param_env,
362                ) {
363                    reachable_set.insert(mono);
364                }
365                reachable_set
366            });
367        // reachable_set.random_sample(&mut rng);
368        rap_debug!(
369            "[get_mono_set] size of s: {}, size of input: {}",
370            s.count(),
371            reachable_set.count()
372        );
373        rap_trace!("[get_mono_set] input = {:?}", reachable_set);
374        s = s.merge(&reachable_set, tcx);
375        s.random_sample(&mut rng);
376        rap_trace!("[get_mono_set] after merge s = {:?}", reachable_set);
377    }
378
379    rap_debug!(
380        "[get_mono_set] after input filter, size of s: {}",
381        s.count()
382    );
383
384    let mut res = MonoSet::new();
385
386    for mono in s.monos {
387        solve_unbound_type_generics(
388            fn_did,
389            mono,
390            &mut res,
391            // &fresh_args,
392            &infcx,
393            &dummy_cause,
394            param_env,
395            tcx,
396        );
397    }
398
399    // erase infer region var
400    res.erase_region_var(tcx);
401
402    // if there is still unbound generic type, we try to instantiate it with predefined candidates
403    res.instantiate_unbound(tcx)
404}
405
406fn solve_unbound_type_generics<'tcx>(
407    did: DefId,
408    mono: Mono<'tcx>,
409    res: &mut MonoSet<'tcx>,
410    infcx: &InferCtxt<'tcx>,
411    cause: &ObligationCause<'tcx>,
412    param_env: ty::ParamEnv<'tcx>,
413    tcx: TyCtxt<'tcx>,
414) {
415    if !mono.has_infer_types() {
416        res.insert(mono);
417        return;
418    }
419    let args = tcx.mk_args(&mono.value);
420    let preds = tcx.predicates_of(did).instantiate(tcx, args);
421    let mut mset = MonoSet::all(args);
422    rap_debug!("[solve_unbound] did = {did:?}, mset={mset:?}");
423    for pred in preds.predicates.iter() {
424        rap_debug!("[solve_unbound] pred = {:?}", pred);
425        #[cfg(rapx_rustc_ge_198)]
426        let pred = pred.skip_norm_wip();
427        if let Some(trait_pred) = pred.as_trait_clause() {
428            let trait_pred = trait_pred.skip_binder();
429
430            rap_trace!("[solve_unbound] pred: {:?}", trait_pred);
431
432            let trait_def_id = trait_pred.trait_ref.def_id;
433            // ignore Sized trait
434            if tcx.is_lang_item(trait_def_id, LangItem::Sized)
435                || tcx.is_lang_item(trait_def_id, LangItem::Copy)
436            {
437                continue;
438            }
439
440            let mut p = MonoSet::new();
441
442            for impl_did in tcx.all_impls(trait_def_id)
443            // .chain(tcx.inherent_impls(trait_def_id).iter().map(|did| *did))
444            {
445                // format: <arg0 as Trait<arg1, arg2>>
446                #[cfg(rapx_rustc_ge_193)]
447                let impl_trait_ref = tcx.impl_trait_ref(impl_did).skip_binder();
448                #[cfg(not(rapx_rustc_ge_193))]
449                let impl_trait_ref = tcx
450                    .impl_trait_ref(impl_did)
451                    .expect("impl must have trait ref")
452                    .skip_binder();
453
454                // filter irrelevant implementation. We only consider implementation that:
455                // 1. it is local
456                // 2. it is not local, but its' self_ty is a primitive
457                if !impl_did.is_local() && !impl_trait_ref.self_ty().is_primitive() {
458                    continue;
459                }
460
461                if let Some(mono) = unify_trait(
462                    trait_pred.trait_ref,
463                    impl_trait_ref,
464                    args,
465                    &infcx,
466                    &cause,
467                    param_env,
468                    tcx,
469                ) {
470                    p.insert(mono);
471                }
472            }
473            mset = mset.merge(&p, tcx);
474            rap_trace!("[solve_unbound] mset: {:?}", mset);
475        }
476    }
477
478    rap_trace!("[solve_unbound] (final) mset: {:?}", mset);
479    for mono in mset.monos {
480        res.insert(mono);
481    }
482}
483
484/// only handle the case that rhs does not have any infer types
485/// e.g., `<T as Into<U>> == <Foo as Into<Bar>> => Some(T=Foo, U=Bar))`
486fn unify_trait<'tcx>(
487    lhs: ty::TraitRef<'tcx>,
488    rhs: ty::TraitRef<'tcx>,
489    identity: &[ty::GenericArg<'tcx>],
490    infcx: &InferCtxt<'tcx>,
491    cause: &ObligationCause<'tcx>,
492    param_env: ty::ParamEnv<'tcx>,
493    tcx: TyCtxt<'tcx>,
494) -> Option<Mono<'tcx>> {
495    rap_trace!("[unify_trait] lhs: {:?}, rhs: {:?}", lhs, rhs);
496    if lhs.def_id != rhs.def_id {
497        return None;
498    }
499
500    assert!(lhs.args.len() == rhs.args.len());
501    let mut s = Mono::new(identity);
502    for (lhs_arg, rhs_arg) in lhs.args.iter().zip(rhs.args.iter()) {
503        if let (Some(lhs_ty), Some(rhs_ty)) = (lhs_arg.as_type(), rhs_arg.as_type()) {
504            if rhs_ty.has_infer_types() || rhs_ty.has_param() {
505                // if rhs has infer types, we cannot unify it with lhs
506                return None;
507            }
508            let mono = unify_ty(lhs_ty, rhs_ty, identity, infcx, cause, param_env)?;
509            rap_trace!("[unify_trait] unified mono: {:?}", mono);
510            s = s.merge(&mono, tcx)?;
511        }
512    }
513    Some(s)
514}
515
516pub fn resolve_mono_apis<'tcx>(
517    fn_did: DefId,
518    available_ty: &HashSet<TyWrapper<'tcx>>,
519    tcx: TyCtxt<'tcx>,
520) -> MonoSet<'tcx> {
521    // 1. check solvable condition
522    if !is_fn_solvable(fn_did, tcx) {
523        return MonoSet::empty();
524    }
525
526    // 2. get mono set from available types
527    let ret = get_mono_set(fn_did, &available_ty, tcx);
528
529    // 3. check trait bound & ty is stable
530    let ret = ret.filter(|mono| {
531        is_args_fit_trait_bound(fn_did, &mono.value, tcx)
532            && mono.value.iter().all(|arg| {
533                if let Some(ty) = arg.as_type() {
534                    !utils::is_ty_unstable(ty, tcx)
535                } else {
536                    true
537                }
538            })
539    });
540
541    rap_debug!(
542        "[resolve_mono_apis] fn_did: {:?}, size of mono: {:?}",
543        fn_did,
544        ret.count()
545    );
546
547    ret
548}
549
550pub fn add_transform_tys<'tcx>(available_ty: &mut HashSet<TyWrapper<'tcx>>, tcx: TyCtxt<'tcx>) {
551    let mut new_tys = Vec::new();
552    available_ty.iter().for_each(|ty| {
553        new_tys.push(
554            Ty::new_ref(
555                tcx,
556                tcx.lifetimes.re_erased,
557                (*ty).into(),
558                ty::Mutability::Not,
559            )
560            .into(),
561        );
562        new_tys.push(Ty::new_ref(
563            tcx,
564            tcx.lifetimes.re_erased,
565            (*ty).into(),
566            ty::Mutability::Mut,
567        ));
568        new_tys.push(Ty::new_ref(
569            tcx,
570            tcx.lifetimes.re_erased,
571            Ty::new_slice(tcx, (*ty).into()),
572            ty::Mutability::Not,
573        ));
574        new_tys.push(Ty::new_ref(
575            tcx,
576            tcx.lifetimes.re_erased,
577            Ty::new_slice(tcx, (*ty).into()),
578            ty::Mutability::Mut,
579        ));
580    });
581
582    new_tys.into_iter().for_each(|ty| {
583        available_ty.insert(ty.into());
584    });
585}
586
587pub fn eliminate_infer_var<'tcx>(
588    fn_did: DefId,
589    args: &[ty::GenericArg<'tcx>],
590    tcx: TyCtxt<'tcx>,
591) -> Vec<ty::GenericArg<'tcx>> {
592    let mut res = Vec::new();
593    let identity = ty::GenericArgs::identity_for_item(tcx, fn_did);
594    for (i, arg) in args.iter().enumerate() {
595        if let Some(ty) = arg.as_type() {
596            if ty.is_ty_var() {
597                res.push(identity[i]);
598            } else {
599                res.push(*arg);
600            }
601        } else {
602            res.push(*arg);
603        }
604    }
605    res
606}
607
608/// if type parameter is unbound, e.g., `T` in `fn foo<T>()`,
609/// we use some predefined types to substitute it
610pub fn get_unbound_generic_candidates<'tcx>(tcx: TyCtxt<'tcx>) -> Vec<ty::Ty<'tcx>> {
611    vec![
612        tcx.types.bool,
613        tcx.types.char,
614        tcx.types.u8,
615        tcx.types.i8,
616        tcx.types.i32,
617        tcx.types.u32,
618        // tcx.types.i64,
619        // tcx.types.u64,
620        tcx.types.f32,
621        // tcx.types.f64,
622        Ty::new_imm_ref(
623            tcx,
624            tcx.lifetimes.re_erased,
625            Ty::new_slice(tcx, tcx.types.u8),
626        ),
627        Ty::new_mut_ref(
628            tcx,
629            tcx.lifetimes.re_erased,
630            Ty::new_slice(tcx, tcx.types.u8),
631        ),
632    ]
633}
634
635// calculate the complexity of monomorphic solution,
636// complexity = sum of complexity of each type argument
637pub fn get_mono_complexity<'tcx>(args: &GenericArgsRef<'tcx>) -> usize {
638    args.iter().fold(0, |acc, arg| {
639        if let Some(ty) = arg.as_type() {
640            acc + utils::ty_complexity(ty)
641        } else {
642            acc
643        }
644    })
645}
646
647pub fn get_impls<'tcx>(
648    tcx: TyCtxt<'tcx>,
649    fn_did: DefId,
650    args: GenericArgsRef<'tcx>,
651) -> HashSet<DefId> {
652    rap_debug!(
653        "get impls for fn: {:?} args: {:?}",
654        tcx.def_path_str_with_args(fn_did, args),
655        args
656    );
657    let mut impls = HashSet::new();
658    let preds = tcx.predicates_of(fn_did).instantiate(tcx, args);
659    for (pred, _) in preds {
660        #[cfg(rapx_rustc_ge_198)]
661        let pred = pred.skip_norm_wip();
662        if let Some(trait_pred) = pred.as_trait_clause() {
663            let trait_ref: rustc_type_ir::TraitRef<TyCtxt<'tcx>> = tcx
664                .liberate_late_bound_regions(fn_did, trait_pred)
665                .trait_ref;
666
667            let res = tcx.codegen_select_candidate(
668                TypingEnv::fully_monomorphized().as_query_input(trait_ref),
669            );
670            if let Ok(source) = res {
671                match source {
672                    ImplSource::UserDefined(data) => {
673                        if data.impl_def_id.is_local() {
674                            impls.insert(data.impl_def_id);
675                        }
676                    }
677                    _ => {}
678                }
679            }
680            // rap_debug!("{:?} => {:?}", trait_ref, res);
681        }
682    }
683    rap_trace!("fn: {:?} args: {:?} impls: {:?}", fn_did, args, impls);
684    impls
685}
rapx/analysis/api_dependency/mono.rs

rapx/analysis/api_dependency/
mono.rs
