Skip to content

Commit b11e564

Browse files
committed
Check for occupied niches
1 parent 94885bc commit b11e564

File tree

23 files changed

+479
-15
lines changed

23 files changed

+479
-15
lines changed

compiler/rustc_codegen_ssa/src/mir/block.rs

+12-1
Original file line numberDiff line numberDiff line change
@@ -1297,6 +1297,17 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
12971297
) -> MergingSucc {
12981298
debug!("codegen_terminator: {:?}", terminator);
12991299

1300+
if bx.tcx().may_insert_niche_checks() {
1301+
if let mir::TerminatorKind::Return = terminator.kind {
1302+
let op = mir::Operand::Copy(mir::Place::return_place());
1303+
let ty = op.ty(self.mir, bx.tcx());
1304+
let ty = self.monomorphize(ty);
1305+
if let Some(niche) = bx.layout_of(ty).largest_niche {
1306+
self.codegen_niche_check(bx, op, niche, terminator.source_info);
1307+
}
1308+
}
1309+
}
1310+
13001311
let helper = TerminatorCodegenHelper { bb, terminator };
13011312

13021313
let mergeable_succ = || {
@@ -1598,7 +1609,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
15981609
tuple.layout.fields.count()
15991610
}
16001611

1601-
fn get_caller_location(
1612+
pub fn get_caller_location(
16021613
&mut self,
16031614
bx: &mut Bx,
16041615
source_info: mir::SourceInfo,

compiler/rustc_codegen_ssa/src/mir/mod.rs

+1
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ pub mod coverageinfo;
2020
pub mod debuginfo;
2121
mod intrinsic;
2222
mod locals;
23+
mod niche_check;
2324
pub mod operand;
2425
pub mod place;
2526
mod rvalue;
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,283 @@
1+
use rustc_hir::LangItem;
2+
use rustc_middle::mir;
3+
use rustc_middle::mir::visit::{NonMutatingUseContext, PlaceContext, Visitor};
4+
use rustc_middle::ty::{Mutability, Ty, TyCtxt};
5+
use rustc_span::def_id::LOCAL_CRATE;
6+
use rustc_span::Span;
7+
use rustc_target::abi::{Float, Integer, Niche, Primitive, Size};
8+
9+
use super::FunctionCx;
10+
use crate::mir::place::PlaceValue;
11+
use crate::mir::OperandValue;
12+
use crate::traits::*;
13+
use crate::{base, common};
14+
15+
pub struct NicheFinder<'s, 'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> {
16+
pub fx: &'s mut FunctionCx<'a, 'tcx, Bx>,
17+
pub bx: &'s mut Bx,
18+
pub places: Vec<(mir::Operand<'tcx>, Niche)>,
19+
}
20+
21+
impl<'s, 'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> Visitor<'tcx> for NicheFinder<'s, 'a, 'tcx, Bx> {
22+
fn visit_rvalue(&mut self, rvalue: &mir::Rvalue<'tcx>, location: mir::Location) {
23+
match rvalue {
24+
mir::Rvalue::Cast(mir::CastKind::Transmute, op, ty) => {
25+
let ty = self.fx.monomorphize(*ty);
26+
if let Some(niche) = self.bx.layout_of(ty).largest_niche {
27+
self.places.push((op.clone(), niche));
28+
}
29+
}
30+
_ => self.super_rvalue(rvalue, location),
31+
}
32+
}
33+
34+
fn visit_terminator(&mut self, terminator: &mir::Terminator<'tcx>, _location: mir::Location) {
35+
if let mir::TerminatorKind::Return = terminator.kind {
36+
let op = mir::Operand::Copy(mir::Place::return_place());
37+
let ty = op.ty(self.fx.mir, self.bx.tcx());
38+
let ty = self.fx.monomorphize(ty);
39+
if let Some(niche) = self.bx.layout_of(ty).largest_niche {
40+
self.places.push((op, niche));
41+
}
42+
}
43+
}
44+
45+
fn visit_place(
46+
&mut self,
47+
place: &mir::Place<'tcx>,
48+
context: PlaceContext,
49+
_location: mir::Location,
50+
) {
51+
match context {
52+
PlaceContext::NonMutatingUse(
53+
NonMutatingUseContext::Copy | NonMutatingUseContext::Move,
54+
) => {}
55+
_ => {
56+
return;
57+
}
58+
}
59+
60+
let ty = place.ty(self.fx.mir, self.bx.tcx()).ty;
61+
let ty = self.fx.monomorphize(ty);
62+
if let Some(niche) = self.bx.layout_of(ty).largest_niche {
63+
self.places.push((mir::Operand::Copy(*place), niche));
64+
};
65+
}
66+
}
67+
68+
use rustc_target::abi::{Abi, Scalar, WrappingRange};
69+
impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
70+
fn value_in_niche(
71+
&mut self,
72+
bx: &mut Bx,
73+
op: crate::mir::OperandRef<'tcx, Bx::Value>,
74+
niche: Niche,
75+
) -> Option<Bx::Value> {
76+
let niche_ty = niche.ty(bx.tcx());
77+
let niche_layout = bx.layout_of(niche_ty);
78+
79+
let (imm, from_scalar, from_backend_ty) = match op.val {
80+
OperandValue::Immediate(imm) => {
81+
let Abi::Scalar(from_scalar) = op.layout.abi else { unreachable!() };
82+
let from_backend_ty = bx.backend_type(op.layout);
83+
(imm, from_scalar, from_backend_ty)
84+
}
85+
OperandValue::Pair(first, second) => {
86+
let Abi::ScalarPair(first_scalar, second_scalar) = op.layout.abi else {
87+
unreachable!()
88+
};
89+
if niche.offset == Size::ZERO {
90+
(first, first_scalar, bx.scalar_pair_element_backend_type(op.layout, 0, true))
91+
} else {
92+
// yolo
93+
(second, second_scalar, bx.scalar_pair_element_backend_type(op.layout, 1, true))
94+
}
95+
}
96+
OperandValue::ZeroSized => unreachable!(),
97+
OperandValue::Ref(PlaceValue { llval: ptr, .. }) => {
98+
// General case: Load the niche primitive via pointer arithmetic.
99+
let niche_ptr_ty = Ty::new_ptr(bx.tcx(), niche_ty, Mutability::Not);
100+
let ptr = bx.pointercast(ptr, bx.backend_type(bx.layout_of(niche_ptr_ty)));
101+
102+
let offset = niche.offset.bytes() / niche_layout.size.bytes();
103+
let niche_backend_ty = bx.backend_type(bx.layout_of(niche_ty));
104+
let ptr = bx.inbounds_gep(niche_backend_ty, ptr, &[bx.const_usize(offset)]);
105+
let value = bx.load(niche_backend_ty, ptr, rustc_target::abi::Align::ONE);
106+
return Some(value);
107+
}
108+
};
109+
110+
// Any type whose ABI is a Scalar bool is turned into an i1, so it cannot contain a value
111+
// outside of its niche.
112+
if from_scalar.is_bool() {
113+
return None;
114+
}
115+
116+
let to_scalar = Scalar::Initialized {
117+
value: niche.value,
118+
valid_range: WrappingRange::full(niche.size(bx.tcx())),
119+
};
120+
let to_backend_ty = bx.backend_type(niche_layout);
121+
if from_backend_ty == to_backend_ty {
122+
return Some(imm);
123+
}
124+
let value = self.transmute_immediate(
125+
bx,
126+
imm,
127+
from_scalar,
128+
from_backend_ty,
129+
to_scalar,
130+
to_backend_ty,
131+
);
132+
Some(value)
133+
}
134+
135+
#[instrument(level = "debug", skip(self, bx))]
136+
pub fn codegen_niche_check(
137+
&mut self,
138+
bx: &mut Bx,
139+
mir_op: mir::Operand<'tcx>,
140+
niche: Niche,
141+
source_info: mir::SourceInfo,
142+
) {
143+
let tcx = bx.tcx();
144+
let op_ty = self.monomorphize(mir_op.ty(self.mir, tcx));
145+
if op_ty == tcx.types.bool {
146+
return;
147+
}
148+
149+
let op = self.codegen_operand(bx, &mir_op);
150+
151+
let Some(value_in_niche) = self.value_in_niche(bx, op, niche) else {
152+
return;
153+
};
154+
let size = niche.size(tcx);
155+
156+
let start = niche.scalar(niche.valid_range.start, bx);
157+
let end = niche.scalar(niche.valid_range.end, bx);
158+
159+
let binop_le = base::bin_op_to_icmp_predicate(mir::BinOp::Le.to_hir_binop(), false);
160+
let binop_ge = base::bin_op_to_icmp_predicate(mir::BinOp::Ge.to_hir_binop(), false);
161+
let is_valid = if niche.valid_range.start == 0 {
162+
bx.icmp(binop_le, value_in_niche, end)
163+
} else if niche.valid_range.end == (u128::MAX >> 128 - size.bits()) {
164+
bx.icmp(binop_ge, value_in_niche, start)
165+
} else {
166+
// We need to check if the value is within a *wrapping* range. We could do this:
167+
// (niche >= start) && (niche <= end)
168+
// But what we're going to actually do is this:
169+
// max = end - start
170+
// (niche - start) <= max
171+
// The latter is much more complicated conceptually, but is actually less operations
172+
// because we can compute max in codegen.
173+
let mut max = niche.valid_range.end.wrapping_sub(niche.valid_range.start);
174+
let size = niche.size(tcx);
175+
if size.bits() < 128 {
176+
let mask = (1 << size.bits()) - 1;
177+
max &= mask;
178+
}
179+
let max_adjusted_allowed_value = niche.scalar(max, bx);
180+
181+
let biased = bx.sub(value_in_niche, start);
182+
bx.icmp(binop_le, biased, max_adjusted_allowed_value)
183+
};
184+
185+
// Create destination blocks, branching on is_valid
186+
let panic = bx.append_sibling_block("panic");
187+
let success = bx.append_sibling_block("success");
188+
bx.cond_br(is_valid, success, panic);
189+
190+
// Switch to the failure block and codegen a call to the panic intrinsic
191+
bx.switch_to_block(panic);
192+
self.set_debug_loc(bx, source_info);
193+
let location = self.get_caller_location(bx, source_info).immediate();
194+
self.codegen_panic(
195+
bx,
196+
niche.lang_item(),
197+
&[value_in_niche, start, end, location],
198+
source_info.span,
199+
);
200+
201+
// Continue codegen in the success block.
202+
bx.switch_to_block(success);
203+
self.set_debug_loc(bx, source_info);
204+
}
205+
206+
#[instrument(level = "debug", skip(self, bx))]
207+
fn codegen_panic(&mut self, bx: &mut Bx, lang_item: LangItem, args: &[Bx::Value], span: Span) {
208+
if bx.tcx().is_compiler_builtins(LOCAL_CRATE) {
209+
bx.abort()
210+
} else {
211+
let (fn_abi, fn_ptr, instance) = common::build_langcall(bx, Some(span), lang_item);
212+
let fn_ty = bx.fn_decl_backend_type(&fn_abi);
213+
let fn_attrs = if bx.tcx().def_kind(self.instance.def_id()).has_codegen_attrs() {
214+
Some(bx.tcx().codegen_fn_attrs(self.instance.def_id()))
215+
} else {
216+
None
217+
};
218+
bx.call(fn_ty, fn_attrs, Some(&fn_abi), fn_ptr, args, None, Some(instance));
219+
}
220+
bx.unreachable();
221+
}
222+
}
223+
224+
pub trait NicheExt {
225+
fn ty<'tcx>(&self, tcx: TyCtxt<'tcx>) -> Ty<'tcx>;
226+
fn size<'tcx>(&self, tcx: TyCtxt<'tcx>) -> Size;
227+
fn scalar<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(&self, val: u128, bx: &mut Bx) -> Bx::Value;
228+
fn lang_item(&self) -> LangItem;
229+
}
230+
231+
impl NicheExt for Niche {
232+
fn lang_item(&self) -> LangItem {
233+
match self.value {
234+
Primitive::Int(Integer::I8, _) => LangItem::PanicOccupiedNicheU8,
235+
Primitive::Int(Integer::I16, _) => LangItem::PanicOccupiedNicheU16,
236+
Primitive::Int(Integer::I32, _) => LangItem::PanicOccupiedNicheU32,
237+
Primitive::Int(Integer::I64, _) => LangItem::PanicOccupiedNicheU64,
238+
Primitive::Int(Integer::I128, _) => LangItem::PanicOccupiedNicheU128,
239+
Primitive::Pointer(_) => LangItem::PanicOccupiedNichePtr,
240+
Primitive::Float(Float::F16) => LangItem::PanicOccupiedNicheU16,
241+
Primitive::Float(Float::F32) => LangItem::PanicOccupiedNicheU32,
242+
Primitive::Float(Float::F64) => LangItem::PanicOccupiedNicheU64,
243+
Primitive::Float(Float::F128) => LangItem::PanicOccupiedNicheU128,
244+
}
245+
}
246+
247+
fn ty<'tcx>(&self, tcx: TyCtxt<'tcx>) -> Ty<'tcx> {
248+
let types = &tcx.types;
249+
match self.value {
250+
Primitive::Int(Integer::I8, _) => types.u8,
251+
Primitive::Int(Integer::I16, _) => types.u16,
252+
Primitive::Int(Integer::I32, _) => types.u32,
253+
Primitive::Int(Integer::I64, _) => types.u64,
254+
Primitive::Int(Integer::I128, _) => types.u128,
255+
Primitive::Pointer(_) => Ty::new_ptr(tcx, types.unit, Mutability::Not),
256+
Primitive::Float(Float::F16) => types.u16,
257+
Primitive::Float(Float::F32) => types.u32,
258+
Primitive::Float(Float::F64) => types.u64,
259+
Primitive::Float(Float::F128) => types.u128,
260+
}
261+
}
262+
263+
fn size<'tcx>(&self, tcx: TyCtxt<'tcx>) -> Size {
264+
self.value.size(&tcx)
265+
}
266+
267+
fn scalar<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(&self, val: u128, bx: &mut Bx) -> Bx::Value {
268+
use rustc_middle::mir::interpret::{Pointer, Scalar};
269+
270+
let tcx = bx.tcx();
271+
let niche_ty = self.ty(tcx);
272+
let value = if niche_ty.is_any_ptr() {
273+
Scalar::from_maybe_pointer(Pointer::from_addr_invalid(val as u64), &tcx)
274+
} else {
275+
Scalar::from_uint(val, self.size(tcx))
276+
};
277+
let layout = rustc_target::abi::Scalar::Initialized {
278+
value: self.value,
279+
valid_range: WrappingRange::full(self.size(tcx)),
280+
};
281+
bx.scalar_to_backend(value, layout, bx.backend_type(bx.layout_of(self.ty(tcx))))
282+
}
283+
}

compiler/rustc_codegen_ssa/src/mir/rvalue.rs

+3-3
Original file line numberDiff line numberDiff line change
@@ -159,7 +159,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
159159
}
160160
}
161161

162-
fn codegen_transmute(
162+
pub fn codegen_transmute(
163163
&mut self,
164164
bx: &mut Bx,
165165
src: OperandRef<'tcx, Bx::Value>,
@@ -194,7 +194,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
194194
///
195195
/// Returns `None` for cases that can't work in that framework, such as for
196196
/// `Immediate`->`Ref` that needs an `alloc` to get the location.
197-
fn codegen_transmute_operand(
197+
pub fn codegen_transmute_operand(
198198
&mut self,
199199
bx: &mut Bx,
200200
operand: OperandRef<'tcx, Bx::Value>,
@@ -336,7 +336,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
336336
///
337337
/// `to_backend_ty` must be the *non*-immediate backend type (so it will be
338338
/// `i8`, not `i1`, for `bool`-like types.)
339-
fn transmute_immediate(
339+
pub fn transmute_immediate(
340340
&self,
341341
bx: &mut Bx,
342342
mut imm: Bx::Value,

compiler/rustc_codegen_ssa/src/mir/statement.rs

+21-2
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,34 @@
1-
use rustc_middle::mir::{self, NonDivergingIntrinsic};
2-
use rustc_middle::span_bug;
1+
use rustc_middle::mir::visit::Visitor;
2+
use rustc_middle::mir::NonDivergingIntrinsic;
3+
use rustc_middle::{mir, span_bug};
34
use rustc_session::config::OptLevel;
45
use tracing::instrument;
56

67
use super::{FunctionCx, LocalRef};
8+
use crate::mir::niche_check::NicheFinder;
79
use crate::traits::*;
810

911
impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
12+
fn niches_to_check(
13+
&mut self,
14+
bx: &mut Bx,
15+
statement: &mir::Statement<'tcx>,
16+
) -> Vec<(mir::Operand<'tcx>, rustc_target::abi::Niche)> {
17+
let mut finder = NicheFinder { fx: self, bx, places: Vec::new() };
18+
finder.visit_statement(statement, rustc_middle::mir::Location::START);
19+
finder.places
20+
}
21+
1022
#[instrument(level = "debug", skip(self, bx))]
1123
pub fn codegen_statement(&mut self, bx: &mut Bx, statement: &mir::Statement<'tcx>) {
1224
self.set_debug_loc(bx, statement.source_info);
25+
26+
if bx.tcx().may_insert_niche_checks() {
27+
for (op, niche) in self.niches_to_check(bx, statement) {
28+
self.codegen_niche_check(bx, op, niche, statement.source_info);
29+
}
30+
}
31+
1332
match statement.kind {
1433
mir::StatementKind::Assign(box (ref place, ref rvalue)) => {
1534
if let Some(index) = place.as_local() {

compiler/rustc_hir/src/lang_items.rs

+6
Original file line numberDiff line numberDiff line change
@@ -282,6 +282,12 @@ language_item_table! {
282282
ConstPanicFmt, sym::const_panic_fmt, const_panic_fmt, Target::Fn, GenericRequirement::None;
283283
PanicBoundsCheck, sym::panic_bounds_check, panic_bounds_check_fn, Target::Fn, GenericRequirement::Exact(0);
284284
PanicMisalignedPointerDereference, sym::panic_misaligned_pointer_dereference, panic_misaligned_pointer_dereference_fn, Target::Fn, GenericRequirement::Exact(0);
285+
PanicOccupiedNicheU8, sym::panic_occupied_niche_u8, panic_occupied_niche_u8, Target::Fn, GenericRequirement::None;
286+
PanicOccupiedNicheU16, sym::panic_occupied_niche_u16, panic_occupied_niche_u16, Target::Fn, GenericRequirement::None;
287+
PanicOccupiedNicheU32, sym::panic_occupied_niche_u32, panic_occupied_niche_u32, Target::Fn, GenericRequirement::None;
288+
PanicOccupiedNicheU64, sym::panic_occupied_niche_u64, panic_occupied_niche_u64, Target::Fn, GenericRequirement::None;
289+
PanicOccupiedNicheU128, sym::panic_occupied_niche_u128, panic_occupied_niche_u128, Target::Fn, GenericRequirement::None;
290+
PanicOccupiedNichePtr, sym::panic_occupied_niche_ptr, panic_occupied_niche_ptr, Target::Fn, GenericRequirement::None;
285291
PanicInfo, sym::panic_info, panic_info, Target::Struct, GenericRequirement::None;
286292
PanicLocation, sym::panic_location, panic_location, Target::Struct, GenericRequirement::None;
287293
PanicImpl, sym::panic_impl, panic_impl, Target::Fn, GenericRequirement::None;

0 commit comments

Comments
 (0)