Auto merge of #15258 - HKalbasi:mir, r=HKalbasi

Support getrandom syscall

And fix some simd intrinsic bugs and add a memory limit to prevent huge allocations from breaking the main process.

Author: bors

Cargo.lock

@@ -19,6 +19,7 @@ bitflags = "2.1.0"
 smallvec.workspace = true
 ena = "0.14.0"
 either = "1.7.0"
+oorandom = "11.1.3"
 tracing = "0.1.35"
 rustc-hash = "1.1.0"
 scoped-tls = "1.0.0"

crates/hir-ty/Cargo.toml

@@ -2494,6 +2494,28 @@ fn exec_limits() {
     );
 }
 
+#[test]
+fn memory_limit() {
+    check_fail(
+        r#"
+        extern "Rust" {
+            #[rustc_allocator]
+            fn __rust_alloc(size: usize, align: usize) -> *mut u8;
+        }
+
+        const GOAL: u8 = unsafe {
+            __rust_alloc(30_000_000_000, 1); // 30GB
+            2
+        };
+        "#,
+        |e| {
+            e == ConstEvalError::MirEvalError(MirEvalError::Panic(
+                "Memory allocation of 30000000000 bytes failed".to_string(),
+            ))
+        },
+    );
+}
+
 #[test]
 fn type_error() {
     check_fail(

crates/hir-ty/src/consteval/tests.rs

@@ -602,3 +602,41 @@ fn rotate() {
         320192512,
     );
 }
+
+#[test]
+fn simd() {
+    check_number(
+        r#"
+        pub struct i8x16(
+            i8,i8,i8,i8,i8,i8,i8,i8,i8,i8,i8,i8,i8,i8,i8,i8,
+        );
+        extern "platform-intrinsic" {
+            pub fn simd_bitmask<T, U>(x: T) -> U;
+        }
+        const GOAL: u16 = simd_bitmask(i8x16(
+            0, 1, 0, 0, 2, 255, 100, 0, 50, 0, 1, 1, 0, 0, 0, 0
+        ));
+        "#,
+        0b0000110101110010,
+    );
+    check_number(
+        r#"
+        pub struct i8x16(
+            i8,i8,i8,i8,i8,i8,i8,i8,i8,i8,i8,i8,i8,i8,i8,i8,
+        );
+        extern "platform-intrinsic" {
+            pub fn simd_lt<T, U>(x: T, y: T) -> U;
+            pub fn simd_bitmask<T, U>(x: T) -> U;
+        }
+        const GOAL: u16 = simd_bitmask(simd_lt::<i8x16, i8x16>(
+            i8x16(
+                -105, -4, -3, -2, -1, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10
+            ),
+            i8x16(
+                -4, -3, -2, -1, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11
+            ),
+        ));
+        "#,
+        0xFFFF,
+    );
+}

crates/hir-ty/src/consteval/tests/intrinsics.rs

@@ -137,6 +137,7 @@ pub struct Evaluator<'a> {
     /// time of use.
     vtable_map: VTableMap,
     thread_local_storage: TlsData,
+    random_state: oorandom::Rand64,
     stdout: Vec<u8>,
     stderr: Vec<u8>,
     layout_cache: RefCell<FxHashMap<Ty, Arc<Layout>>>,
@@ -147,6 +148,8 @@ pub struct Evaluator<'a> {
     execution_limit: usize,
     /// An additional limit on stack depth, to prevent stack overflow
     stack_depth_limit: usize,
+    /// Maximum count of bytes that heap and stack can grow
+    memory_limit: usize,
 }
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
@@ -520,13 +523,15 @@ impl Evaluator<'_> {
             thread_local_storage: TlsData::default(),
             static_locations: HashMap::default(),
             db,
+            random_state: oorandom::Rand64::new(0),
             trait_env,
             crate_id,
             stdout: vec![],
             stderr: vec![],
             assert_placeholder_ty_is_unused,
             stack_depth_limit: 100,
             execution_limit: 1000_000,
+            memory_limit: 1000_000_000, // 2GB, 1GB for stack and 1GB for heap
             layout_cache: RefCell::new(HashMap::default()),
         }
     }
@@ -938,6 +943,11 @@ impl Evaluator<'_> {
         };
         locals.ptr = locals_ptr;
         let prev_stack_pointer = self.stack.len();
+        if stack_size > self.memory_limit {
+            return Err(MirEvalError::Panic(format!(
+                "Stack overflow. Tried to grow stack to {stack_size} bytes"
+            )));
+        }
         self.stack.extend(iter::repeat(0).take(stack_size));
         Ok((locals, prev_stack_pointer))
     }
@@ -1180,7 +1190,7 @@ impl Evaluator<'_> {
                 let Some((size, align)) = self.size_align_of(ty, locals)? else {
                     not_supported!("unsized box initialization");
                 };
-                let addr = self.heap_allocate(size, align);
+                let addr = self.heap_allocate(size, align)?;
                 Owned(addr.to_bytes())
             }
             Rvalue::CopyForDeref(_) => not_supported!("copy for deref"),
@@ -1565,7 +1575,7 @@ impl Evaluator<'_> {
             ConstScalar::Bytes(v, memory_map) => {
                 let mut v: Cow<'_, [u8]> = Cow::Borrowed(v);
                 let patch_map = memory_map.transform_addresses(|b, align| {
-                    let addr = self.heap_allocate(b.len(), align);
+                    let addr = self.heap_allocate(b.len(), align)?;
                     self.write_memory(addr, b)?;
                     Ok(addr.to_usize())
                 })?;
@@ -1580,7 +1590,7 @@ impl Evaluator<'_> {
                         return Err(MirEvalError::InvalidConst(konst.clone()));
                     }
                 }
-                let addr = self.heap_allocate(size, align);
+                let addr = self.heap_allocate(size, align)?;
                 self.write_memory(addr, &v)?;
                 self.patch_addresses(&patch_map, &memory_map.vtable, addr, ty, locals)?;
                 Interval::new(addr, size)
@@ -1683,13 +1693,19 @@ impl Evaluator<'_> {
         }
     }
 
-    fn heap_allocate(&mut self, size: usize, align: usize) -> Address {
+    fn heap_allocate(&mut self, size: usize, align: usize) -> Result<Address> {
+        if !align.is_power_of_two() || align > 10000 {
+            return Err(MirEvalError::UndefinedBehavior(format!("Alignment {align} is invalid")));
+        }
         while self.heap.len() % align != 0 {
             self.heap.push(0);
         }
+        if size.checked_add(self.heap.len()).map_or(true, |x| x > self.memory_limit) {
+            return Err(MirEvalError::Panic(format!("Memory allocation of {size} bytes failed")));
+        }
         let pos = self.heap.len();
         self.heap.extend(iter::repeat(0).take(size));
-        Address::Heap(pos)
+        Ok(Address::Heap(pos))
     }
 
     fn detect_fn_trait(&self, def: FunctionId) -> Option<FnTrait> {
@@ -2200,7 +2216,7 @@ impl Evaluator<'_> {
                     )?;
                     // FIXME: there is some leak here
                     let size = layout.size.bytes_usize();
-                    let addr = self.heap_allocate(size, layout.align.abi.bytes() as usize);
+                    let addr = self.heap_allocate(size, layout.align.abi.bytes() as usize)?;
                     self.write_memory(addr, &result)?;
                     IntervalAndTy { interval: Interval { addr, size }, ty }
                 };
@@ -2235,10 +2251,10 @@ impl Evaluator<'_> {
             let Some((size, align)) = self.size_align_of(&ty, locals)? else {
                 not_supported!("unsized extern static");
             };
-            let addr = self.heap_allocate(size, align);
+            let addr = self.heap_allocate(size, align)?;
             Interval::new(addr, size)
         };
-        let addr = self.heap_allocate(self.ptr_size(), self.ptr_size());
+        let addr = self.heap_allocate(self.ptr_size(), self.ptr_size())?;
         self.write_memory(addr, &result.addr.to_bytes())?;
         self.static_locations.insert(st, addr);
         Ok(addr)
@@ -2398,11 +2414,11 @@ pub fn render_const_using_debug_impl(
         not_supported!("core::fmt::Debug::fmt not found");
     };
     // a1 = &[""]
-    let a1 = evaluator.heap_allocate(evaluator.ptr_size() * 2, evaluator.ptr_size());
+    let a1 = evaluator.heap_allocate(evaluator.ptr_size() * 2, evaluator.ptr_size())?;
     // a2 = &[::core::fmt::ArgumentV1::new(&(THE_CONST), ::core::fmt::Debug::fmt)]
     // FIXME: we should call the said function, but since its name is going to break in the next rustc version
     // and its ABI doesn't break yet, we put it in memory manually.
-    let a2 = evaluator.heap_allocate(evaluator.ptr_size() * 2, evaluator.ptr_size());
+    let a2 = evaluator.heap_allocate(evaluator.ptr_size() * 2, evaluator.ptr_size())?;
     evaluator.write_memory(a2, &data.addr.to_bytes())?;
     let debug_fmt_fn_ptr = evaluator.vtable_map.id(TyKind::FnDef(
         db.intern_callable_def(debug_fmt_fn.into()).into(),
@@ -2412,7 +2428,7 @@ pub fn render_const_using_debug_impl(
     evaluator.write_memory(a2.offset(evaluator.ptr_size()), &debug_fmt_fn_ptr.to_le_bytes())?;
     // a3 = ::core::fmt::Arguments::new_v1(a1, a2)
     // FIXME: similarly, we should call function here, not directly working with memory.
-    let a3 = evaluator.heap_allocate(evaluator.ptr_size() * 6, evaluator.ptr_size());
+    let a3 = evaluator.heap_allocate(evaluator.ptr_size() * 6, evaluator.ptr_size())?;
     evaluator.write_memory(a3.offset(2 * evaluator.ptr_size()), &a1.to_bytes())?;
     evaluator.write_memory(a3.offset(3 * evaluator.ptr_size()), &[1])?;
     evaluator.write_memory(a3.offset(4 * evaluator.ptr_size()), &a2.to_bytes())?;

crates/hir-ty/src/mir/eval.rs

@@ -140,7 +140,7 @@ impl Evaluator<'_> {
                 };
                 let size = from_bytes!(usize, size.get(self)?);
                 let align = from_bytes!(usize, align.get(self)?);
-                let result = self.heap_allocate(size, align);
+                let result = self.heap_allocate(size, align)?;
                 destination.write_from_bytes(self, &result.to_bytes())?;
             }
             "rustc_deallocator" => { /* no-op for now */ }
@@ -155,7 +155,7 @@ impl Evaluator<'_> {
                 } else {
                     let ptr = Address::from_bytes(ptr.get(self)?)?;
                     let align = from_bytes!(usize, align.get(self)?);
-                    let result = self.heap_allocate(new_size, align);
+                    let result = self.heap_allocate(new_size, align)?;
                     Interval { addr: result, size: old_size }
                         .write_from_interval(self, Interval { addr: ptr, size: old_size })?;
                     destination.write_from_bytes(self, &result.to_bytes())?;
@@ -239,14 +239,42 @@ impl Evaluator<'_> {
         }
     }
 
+    fn exec_syscall(
+        &mut self,
+        id: i64,
+        args: &[IntervalAndTy],
+        destination: Interval,
+        _locals: &Locals,
+        _span: MirSpan,
+    ) -> Result<()> {
+        match id {
+            318 => {
+                // SYS_getrandom
+                let [buf, len, _flags] = args else {
+                    return Err(MirEvalError::TypeError("SYS_getrandom args are not provided"));
+                };
+                let addr = Address::from_bytes(buf.get(self)?)?;
+                let size = from_bytes!(usize, len.get(self)?);
+                for i in 0..size {
+                    let rand_byte = self.random_state.rand_u64() as u8;
+                    self.write_memory(addr.offset(i), &[rand_byte])?;
+                }
+                destination.write_from_interval(self, len.interval)
+            }
+            _ => {
+                not_supported!("Unknown syscall id {id:?}")
+            }
+        }
+    }
+
     fn exec_extern_c(
         &mut self,
         as_str: &str,
         args: &[IntervalAndTy],
         _generic_args: &Substitution,
         destination: Interval,
         locals: &Locals,
-        _span: MirSpan,
+        span: MirSpan,
     ) -> Result<()> {
         match as_str {
             "memcmp" => {
@@ -343,6 +371,15 @@ impl Evaluator<'_> {
                 destination.write_from_bytes(self, &0u64.to_le_bytes()[0..destination.size])?;
                 Ok(())
             }
+            "syscall" => {
+                let Some((id, rest)) = args.split_first() else {
+                    return Err(MirEvalError::TypeError(
+                        "syscall arg1 is not provided",
+                    ));
+                };
+                let id = from_bytes!(i64, id.get(self)?);
+                self.exec_syscall(id, rest, destination, locals, span)
+            }
             _ => not_supported!("unknown external function {as_str}"),
         }
     }

crates/hir-ty/src/mir/eval/shim.rs

@@ -22,21 +22,29 @@ macro_rules! not_supported {
 }
 
 impl Evaluator<'_> {
-    fn detect_simd_ty(&self, ty: &Ty) -> Result<usize> {
+    fn detect_simd_ty(&self, ty: &Ty) -> Result<(usize, Ty)> {
         match ty.kind(Interner) {
             TyKind::Adt(id, subst) => {
                 let len = match subst.as_slice(Interner).get(1).and_then(|it| it.constant(Interner))
                 {
                     Some(len) => len,
                     _ => {
                         if let AdtId::StructId(id) = id.0 {
-                            return Ok(self.db.struct_data(id).variant_data.fields().len());
+                            let struct_data = self.db.struct_data(id);
+                            let fields = struct_data.variant_data.fields();
+                            let Some((first_field, _)) = fields.iter().next() else {
+                                not_supported!("simd type with no field");
+                            };
+                            let field_ty = self.db.field_types(id.into())[first_field]
+                                .clone()
+                                .substitute(Interner, subst);
+                            return Ok((fields.len(), field_ty));
                         }
                         return Err(MirEvalError::TypeError("simd type with no len param"));
                     }
                 };
                 match try_const_usize(self.db, len) {
-                    Some(it) => Ok(it as usize),
+                    Some(_) => not_supported!("array like simd type"),
                     None => Err(MirEvalError::TypeError("simd type with unevaluatable len param")),
                 }
             }
@@ -75,7 +83,8 @@ impl Evaluator<'_> {
                 let [left, right] = args else {
                     return Err(MirEvalError::TypeError("simd args are not provided"));
                 };
-                let len = self.detect_simd_ty(&left.ty)?;
+                let (len, ty) = self.detect_simd_ty(&left.ty)?;
+                let is_signed = matches!(ty.as_builtin(), Some(BuiltinType::Int(_)));
                 let size = left.interval.size / len;
                 let dest_size = destination.size / len;
                 let mut destination_bytes = vec![];
@@ -89,6 +98,13 @@ impl Evaluator<'_> {
                             break;
                         }
                     }
+                    if is_signed {
+                        if let Some((&l, &r)) = l.iter().zip(r).rev().next() {
+                            if l != r {
+                                result = (l as i8).cmp(&(r as i8));
+                            }
+                        }
+                    }
                     let result = match result {
                         Ordering::Less => ["lt", "le", "ne"].contains(&name),
                         Ordering::Equal => ["ge", "le", "eq"].contains(&name),
@@ -102,9 +118,9 @@ impl Evaluator<'_> {
             }
             "bitmask" => {
                 let [op] = args else {
-                    return Err(MirEvalError::TypeError("simd_shuffle args are not provided"));
+                    return Err(MirEvalError::TypeError("simd_bitmask args are not provided"));
                 };
-                let op_len = self.detect_simd_ty(&op.ty)?;
+                let (op_len, _) = self.detect_simd_ty(&op.ty)?;
                 let op_count = op.interval.size / op_len;
                 let mut result: u64 = 0;
                 for (i, val) in op.get(self)?.chunks(op_count).enumerate() {
@@ -131,7 +147,7 @@ impl Evaluator<'_> {
                         ))
                     }
                 };
-                let left_len = self.detect_simd_ty(&left.ty)?;
+                let (left_len, _) = self.detect_simd_ty(&left.ty)?;
                 let left_size = left.interval.size / left_len;
                 let vector =
                     left.get(self)?.chunks(left_size).chain(right.get(self)?.chunks(left_size));