Skip to content

Commit

Permalink
Rollup merge of rust-lang#123106 - maurer:cfi-closures, r=compiler-er…
Browse files Browse the repository at this point in the history
…rors

CFI: Abstract Closures and Coroutines

This will abstract coroutines in a moment, it's just abstracting closures for now to show `@rcvalle`

This uses the same principal as the methods on traits - figure out the `dyn` type representing the fn trait, instantiate it, and attach that alias set. We're essentially just computing how we would be called in a dynamic context, and attaching that.
  • Loading branch information
matthiaskrgr authored Mar 29, 2024
2 parents 58dcd1f + d243f5f commit 448d4c1
Show file tree
Hide file tree
Showing 5 changed files with 187 additions and 25 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
use rustc_data_structures::base_n;
use rustc_data_structures::fx::FxHashMap;
use rustc_hir as hir;
use rustc_hir::lang_items::LangItem;
use rustc_middle::ty::layout::IntegerExt;
use rustc_middle::ty::TypeVisitableExt;
use rustc_middle::ty::{
Expand Down Expand Up @@ -641,9 +642,7 @@ fn encode_ty<'tcx>(
}

// Function types
ty::FnDef(def_id, args)
| ty::Closure(def_id, args)
| ty::CoroutineClosure(def_id, args) => {
ty::FnDef(def_id, args) | ty::Closure(def_id, args) => {
// u<length><name>[I<element-type1..element-typeN>E], where <element-type> is <subst>,
// as vendor extended type.
let mut s = String::new();
Expand All @@ -654,6 +653,18 @@ fn encode_ty<'tcx>(
typeid.push_str(&s);
}

ty::CoroutineClosure(def_id, args) => {
// u<length><name>[I<element-type1..element-typeN>E], where <element-type> is <subst>,
// as vendor extended type.
let mut s = String::new();
let name = encode_ty_name(tcx, *def_id);
let _ = write!(s, "u{}{}", name.len(), &name);
let parent_args = tcx.mk_args(args.as_coroutine_closure().parent_args());
s.push_str(&encode_args(tcx, parent_args, dict, options));
compress(dict, DictKey::Ty(ty, TyQ::None), &mut s);
typeid.push_str(&s);
}

ty::Coroutine(def_id, args, ..) => {
// u<length><name>[I<element-type1..element-typeN>E], where <element-type> is <subst>,
// as vendor extended type.
Expand Down Expand Up @@ -1142,6 +1153,14 @@ pub fn typeid_for_instance<'tcx>(
instance.args = tcx.mk_args_trait(self_ty, List::empty());
} else if matches!(instance.def, ty::InstanceDef::Virtual(..)) {
instance.args = strip_receiver_auto(tcx, instance.args);
} else if let ty::InstanceDef::VTableShim(def_id) = instance.def
&& let Some(trait_id) = tcx.trait_of_item(def_id)
{
// VTableShims may have a trait method, but a concrete Self. This is not suitable for a vtable,
// as the caller will not know the concrete Self.
let trait_ref = ty::TraitRef::new(tcx, trait_id, instance.args);
let invoke_ty = trait_object_ty(tcx, ty::Binder::dummy(trait_ref));
instance.args = tcx.mk_args_trait(invoke_ty, trait_ref.args.into_iter().skip(1));
}

if !options.contains(EncodeTyOptions::NO_SELF_TYPE_ERASURE)
Expand Down Expand Up @@ -1180,6 +1199,45 @@ pub fn typeid_for_instance<'tcx>(
tcx.mk_args_trait(invoke_ty, trait_ref.args.into_iter().skip(1));
instance.args = instance.args.rebase_onto(tcx, impl_id, abstract_trait_args);
}
} else if tcx.is_closure_like(instance.def_id()) {
// We're either a closure or a coroutine. Our goal is to find the trait we're defined on,
// instantiate it, and take the type of its only method as our own.
let closure_ty = instance.ty(tcx, ty::ParamEnv::reveal_all());
let (trait_id, inputs) = match closure_ty.kind() {
ty::Closure(..) => {
let closure_args = instance.args.as_closure();
let trait_id = tcx.fn_trait_kind_to_def_id(closure_args.kind()).unwrap();
let tuple_args =
tcx.instantiate_bound_regions_with_erased(closure_args.sig()).inputs()[0];
(trait_id, tuple_args)
}
ty::Coroutine(..) => (
tcx.require_lang_item(LangItem::Coroutine, None),
instance.args.as_coroutine().resume_ty(),
),
ty::CoroutineClosure(..) => (
tcx.require_lang_item(LangItem::FnOnce, None),
tcx.instantiate_bound_regions_with_erased(
instance.args.as_coroutine_closure().coroutine_closure_sig(),
)
.tupled_inputs_ty,
),
x => bug!("Unexpected type kind for closure-like: {x:?}"),
};
let trait_ref = ty::TraitRef::new(tcx, trait_id, [closure_ty, inputs]);
let invoke_ty = trait_object_ty(tcx, ty::Binder::dummy(trait_ref));
let abstract_args = tcx.mk_args_trait(invoke_ty, trait_ref.args.into_iter().skip(1));
// There should be exactly one method on this trait, and it should be the one we're
// defining.
let call = tcx
.associated_items(trait_id)
.in_definition_order()
.find(|it| it.kind == ty::AssocKind::Fn)
.expect("No call-family function on closure-like Fn trait?")
.def_id;

instance.def = ty::InstanceDef::Virtual(call, 0);
instance.args = abstract_args;
}

let fn_abi = tcx
Expand Down
32 changes: 32 additions & 0 deletions tests/ui/sanitizer/cfi-async-closures.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,32 @@
// Check various forms of dynamic closure calls

//@ edition: 2021
//@ revisions: cfi kcfi
// FIXME(#122848) Remove only-linux once OSX CFI binaries work
//@ only-linux
//@ [cfi] needs-sanitizer-cfi
//@ [kcfi] needs-sanitizer-kcfi
//@ compile-flags: -C target-feature=-crt-static
//@ [cfi] compile-flags: -C codegen-units=1 -C lto -C prefer-dynamic=off -C opt-level=0
//@ [cfi] compile-flags: -Z sanitizer=cfi
//@ [kcfi] compile-flags: -Z sanitizer=kcfi
//@ run-pass

#![feature(async_closure)]
#![feature(async_fn_traits)]

use std::ops::AsyncFn;

#[inline(never)]
fn identity<T>(x: T) -> T { x }

// We can't actually create a `dyn AsyncFn()`, because it's not object-safe, but we should check
// that we don't bug out when we encounter one.

fn main() {
let f = identity(async || ());
let _ = f.async_call(());
let _ = f();
let g: Box<dyn FnOnce() -> _> = Box::new(f) as _;
let _ = g();
}
22 changes: 0 additions & 22 deletions tests/ui/sanitizer/cfi-closure-fn-ptr-cast.rs

This file was deleted.

65 changes: 65 additions & 0 deletions tests/ui/sanitizer/cfi-closures.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,65 @@
// Check various forms of dynamic closure calls

//@ revisions: cfi kcfi
// FIXME(#122848) Remove only-linux once OSX CFI binaries work
//@ only-linux
//@ [cfi] needs-sanitizer-cfi
//@ [kcfi] needs-sanitizer-kcfi
//@ compile-flags: -C target-feature=-crt-static
//@ [cfi] compile-flags: -C codegen-units=1 -C lto -C prefer-dynamic=off -C opt-level=0
//@ [cfi] compile-flags: -Z sanitizer=cfi
//@ [kcfi] compile-flags: -Z sanitizer=kcfi
//@ compile-flags: --test
//@ run-pass

#![feature(fn_traits)]

fn foo<'a, T>() -> Box<dyn Fn(&'a T) -> &'a T> {
Box::new(|x| x)
}

#[test]
fn dyn_fn_with_params() {
let x = 3;
let f = foo();
f(&x);
// FIXME remove once drops are working.
std::mem::forget(f);
}

#[test]
fn call_fn_trait() {
let f: &(dyn Fn()) = &(|| {}) as _;
f.call(());
}

#[test]
fn fn_ptr_cast() {
let f: &fn() = &((|| ()) as _);
f();
}

fn use_fnmut<F: FnMut()>(mut f: F) {
f()
}

#[test]
fn fn_to_fnmut() {
let f: &(dyn Fn()) = &(|| {}) as _;
use_fnmut(f);
}

fn hrtb_helper(f: &dyn for<'a> Fn(&'a usize)) {
f(&10)
}

#[test]
fn hrtb_fn() {
hrtb_helper((&|x: &usize| println!("{}", *x)) as _)
}

#[test]
fn fnonce() {
let f: Box<dyn FnOnce()> = Box::new(|| {}) as _;
f();
}
29 changes: 29 additions & 0 deletions tests/ui/sanitizer/cfi-coroutine.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,29 @@
// Verifies that we can call dynamic coroutines

//@ revisions: cfi kcfi
// FIXME(#122848) Remove only-linux once OSX CFI binaries work
//@ only-linux
//@ [cfi] needs-sanitizer-cfi
//@ [kcfi] needs-sanitizer-kcfi
//@ compile-flags: -C target-feature=-crt-static
//@ [cfi] compile-flags: -C codegen-units=1 -C lto -C prefer-dynamic=off -C opt-level=0
//@ [cfi] compile-flags: -Z sanitizer=cfi
//@ [kcfi] compile-flags: -Z sanitizer=kcfi
//@ compile-flags: --test
//@ run-pass

#![feature(coroutines)]
#![feature(coroutine_trait)]

use std::ops::{Coroutine, CoroutineState};
use std::pin::{pin, Pin};

fn main() {
let mut coro = |x: i32| {
yield x;
"done"
};
let mut abstract_coro: Pin<&mut dyn Coroutine<i32,Yield=i32,Return=&'static str>> = pin!(coro);
assert_eq!(abstract_coro.as_mut().resume(2), CoroutineState::Yielded(2));
assert_eq!(abstract_coro.as_mut().resume(0), CoroutineState::Complete("done"));
}

0 comments on commit 448d4c1

Please sign in to comment.