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

rapx/analysis/api_dependency/
mono.rs
