Add tests to allocation methods and fix is_null()

3 files changed, 114 insertions, 7 deletions

diff --git a/compiler/stable_mir/src/mir/body.rs b/compiler/stable_mir/src/mir/body.rs
index d03ded20e9b..663275d9a0f 100644
--- a/compiler/stable_mir/src/mir/body.rs
+++ b/compiler/stable_mir/src/mir/body.rs
@@ -21,7 +21,7 @@ pub struct Body {
     pub(super) arg_count: usize,
 
     /// Debug information pertaining to user variables, including captures.
-    pub(super) var_debug_info: Vec<VarDebugInfo>,
+    pub var_debug_info: Vec<VarDebugInfo>,
 }
 
 pub type BasicBlockIdx = usize;
@@ -616,6 +616,24 @@ pub struct VarDebugInfo {
     pub argument_index: Option<u16>,
 }
 
+impl VarDebugInfo {
+    /// Return a local variable if this info is related to one.
+    pub fn local(&self) -> Option<Local> {
+        match &self.value {
+            VarDebugInfoContents::Place(place) if place.projection.is_empty() => Some(place.local),
+            VarDebugInfoContents::Place(_) | VarDebugInfoContents::Const(_) => None,
+        }
+    }
+
+    /// Return a constant if this info is related to one.
+    pub fn constant(&self) -> Option<&ConstOperand> {
+        match &self.value {
+            VarDebugInfoContents::Place(_) => None,
+            VarDebugInfoContents::Const(const_op) => Some(const_op),
+        }
+    }
+}
+
 pub type SourceScope = u32;
 
 #[derive(Clone, Debug, Eq, PartialEq)]
diff --git a/compiler/stable_mir/src/ty.rs b/compiler/stable_mir/src/ty.rs
index 4caac575c1e..bea7702bd34 100644
--- a/compiler/stable_mir/src/ty.rs
+++ b/compiler/stable_mir/src/ty.rs
@@ -901,6 +901,7 @@ impl Allocation {
         read_target_uint(&raw)
     }
 
+    /// Read this allocation and try to convert it to an unassigned integer.
     pub fn read_uint(&self) -> Result<u128, Error> {
         if self.bytes.len() > 16 {
             return Err(error!("Allocation is bigger than largest integer"));
@@ -909,6 +910,7 @@ impl Allocation {
         read_target_uint(&raw)
     }
 
+    /// Read this allocation and try to convert it to a signed integer.
     pub fn read_int(&self) -> Result<i128, Error> {
         if self.bytes.len() > 16 {
             return Err(error!("Allocation is bigger than largest integer"));
@@ -917,6 +919,7 @@ impl Allocation {
         read_target_int(&raw)
     }
 
+    /// Read this allocation and try to convert it to a boolean.
     pub fn read_bool(&self) -> Result<bool, Error> {
         match self.read_int()? {
             0 => Ok(false),
@@ -925,13 +928,14 @@ impl Allocation {
         }
     }
 
+    /// Read this allocation as a pointer and return whether it represents a `null` pointer.
     pub fn is_null(&self) -> Result<bool, Error> {
         let len = self.bytes.len();
         let ptr_len = MachineInfo::target_pointer_width().bytes();
         if len != ptr_len {
             return Err(error!("Expected width of pointer (`{ptr_len}`), but found: `{len}`"));
         }
-        Ok(self.read_uint()? == 0)
+        Ok(self.read_uint()? == 0 && self.provenance.ptrs.is_empty())
     }
 }
 
diff --git a/tests/ui-fulldeps/stable-mir/check_allocation.rs b/tests/ui-fulldeps/stable-mir/check_allocation.rs
index 5051d7b7d82..d2bc58be4c1 100644
--- a/tests/ui-fulldeps/stable-mir/check_allocation.rs
+++ b/tests/ui-fulldeps/stable-mir/check_allocation.rs
@@ -23,12 +23,16 @@ extern crate stable_mir;
 
 use rustc_middle::ty::TyCtxt;
 use rustc_smir::rustc_internal;
-use stable_mir::{CrateItem, CrateItems, ItemKind};
 use stable_mir::crate_def::CrateDef;
 use stable_mir::mir::alloc::GlobalAlloc;
-use stable_mir::mir::mono::StaticDef;
+use stable_mir::mir::mono::{Instance, StaticDef};
+use stable_mir::mir::Body;
+use stable_mir::ty::{Allocation, ConstantKind};
+use stable_mir::{CrateItem, CrateItems, ItemKind};
 use std::ascii::Char;
 use std::assert_matches::assert_matches;
+use std::cmp::{max, min};
+use std::collections::HashMap;
 use std::io::Write;
 use std::ops::ControlFlow;
 
@@ -41,6 +45,7 @@ fn test_stable_mir(_tcx: TyCtxt<'_>) -> ControlFlow<()> {
     check_foo(*get_item(&items, (ItemKind::Static, "FOO")).unwrap());
     check_bar(*get_item(&items, (ItemKind::Static, "BAR")).unwrap());
     check_len(*get_item(&items, (ItemKind::Static, "LEN")).unwrap());
+    check_other_consts(*get_item(&items, (ItemKind::Fn, "other_consts")).unwrap());
     ControlFlow::Continue(())
 }
 
@@ -80,6 +85,73 @@ fn check_bar(item: CrateItem) {
     assert_eq!(std::str::from_utf8(&allocation.raw_bytes().unwrap()), Ok("Bar"));
 }
 
+/// Check the allocation data for constants used in `other_consts` function.
+fn check_other_consts(item: CrateItem) {
+    // Instance body will force constant evaluation.
+    let body = Instance::try_from(item).unwrap().body().unwrap();
+    let assigns = collect_consts(&body);
+    assert_eq!(assigns.len(), 9);
+    for (name, alloc) in assigns {
+        match name.as_str() {
+            "_max_u128" => {
+                assert_eq!(alloc.read_uint(), Ok(u128::MAX), "Failed parsing allocation: {alloc:?}")
+            }
+            "_min_i128" => {
+                assert_eq!(alloc.read_int(), Ok(i128::MIN), "Failed parsing allocation: {alloc:?}")
+            }
+            "_max_i8" => {
+                assert_eq!(
+                    alloc.read_int().unwrap() as i8,
+                    i8::MAX,
+                    "Failed parsing allocation: {alloc:?}"
+                )
+            }
+            "_char" => {
+                assert_eq!(
+                    char::from_u32(alloc.read_uint().unwrap() as u32),
+                    Some('x'),
+                    "Failed parsing allocation: {alloc:?}"
+                )
+            }
+            "_false" => {
+                assert_eq!(alloc.read_bool(), Ok(false), "Failed parsing allocation: {alloc:?}")
+            }
+            "_true" => {
+                assert_eq!(alloc.read_bool(), Ok(true), "Failed parsing allocation: {alloc:?}")
+            }
+            "_ptr" => {
+                assert_eq!(alloc.is_null(), Ok(false), "Failed parsing allocation: {alloc:?}")
+            }
+            "_null_ptr" => {
+                assert_eq!(alloc.is_null(), Ok(true), "Failed parsing allocation: {alloc:?}")
+            }
+            "_tuple" => {
+                // The order of fields is not guaranteed.
+                let first = alloc.read_partial_uint(0..4).unwrap();
+                let second = alloc.read_partial_uint(4..8).unwrap();
+                assert_eq!(max(first, second) as u32, u32::MAX);
+                assert_eq!(min(first, second), 10);
+            }
+            _ => {
+                unreachable!("{name} -- {alloc:?}")
+            }
+        }
+    }
+}
+
+/// Collects all the constant assignments.
+pub fn collect_consts(body: &Body) -> HashMap<String, &Allocation> {
+    body.var_debug_info
+        .iter()
+        .filter_map(|info| {
+            info.constant().map(|const_op| {
+                let ConstantKind::Allocated(alloc) = const_op.const_.kind() else { unreachable!() };
+                (info.name.clone(), alloc)
+            })
+        })
+        .collect::<HashMap<_, _>>()
+}
+
 /// Check the allocation data for `LEN`.
 ///
 /// ```no_run
@@ -97,9 +169,7 @@ fn get_item<'a>(
     items: &'a CrateItems,
     item: (ItemKind, &str),
 ) -> Option<&'a stable_mir::CrateItem> {
-    items.iter().find(|crate_item| {
-        (item.0 == crate_item.kind()) && crate_item.name() == item.1
-    })
+    items.iter().find(|crate_item| (item.0 == crate_item.kind()) && crate_item.name() == item.1)
 }
 
 /// This test will generate and analyze a dummy crate using the stable mir.
@@ -126,10 +196,25 @@ fn generate_input(path: &str) -> std::io::Result<()> {
     static LEN: usize = 2;
     static FOO: [&str; 2] = ["hi", "there"];
     static BAR: &str = "Bar";
+    const NULL: *const u8 = std::ptr::null();
+    const TUPLE: (u32, u32) = (10, u32::MAX);
+
+    fn other_consts() {{
+        let _max_u128 = u128::MAX;
+        let _min_i128 = i128::MIN;
+        let _max_i8 = i8::MAX;
+        let _char = 'x';
+        let _false = false;
+        let _true = true;
+        let _ptr = &BAR;
+        let _null_ptr: *const u8 = NULL;
+        let _tuple = TUPLE;
+    }}
 
     pub fn main() {{
         println!("{{FOO:?}}! {{BAR}}");
         assert_eq!(FOO.len(), LEN);
+        other_consts();
     }}"#
     )?;
     Ok(())