Implement memory initialization state copy functionality (#3350)

This PR adds support of copying memory initialization state without
checks in-between. Every time a copy is performed, the tracked byte is
non-deterministically switched.

Resolves #3347

By submitting this pull request, I confirm that my contribution is made
under the terms of the Apache 2.0 and MIT licenses.

kani-compiler/src/kani_middle/transform/check_uninit/mod.rs

@@ -178,6 +178,9 @@ impl<'a> UninitInstrumenter<'a> {
             | MemoryInitOp::SetRef { .. } => {
                 self.build_set(tcx, body, source, operation, pointee_ty_info, skip_first)
             }
+            MemoryInitOp::Copy { .. } => {
+                self.build_copy(tcx, body, source, operation, pointee_ty_info, skip_first)
+            }
             MemoryInitOp::Unsupported { .. } | MemoryInitOp::TriviallyUnsafe { .. } => {
                 unreachable!()
             }
@@ -397,6 +400,38 @@ impl<'a> UninitInstrumenter<'a> {
         };
     }
 
+    /// Copy memory initialization state from one pointer to the other.
+    fn build_copy(
+        &mut self,
+        tcx: TyCtxt,
+        body: &mut MutableBody,
+        source: &mut SourceInstruction,
+        operation: MemoryInitOp,
+        pointee_info: PointeeInfo,
+        skip_first: &mut HashSet<usize>,
+    ) {
+        let ret_place = Place {
+            local: body.new_local(Ty::new_tuple(&[]), source.span(body.blocks()), Mutability::Not),
+            projection: vec![],
+        };
+        let PointeeLayout::Sized { layout } = pointee_info.layout() else { unreachable!() };
+        let copy_init_state_instance = resolve_mem_init_fn(
+            get_mem_init_fn_def(tcx, "KaniCopyInitState", &mut self.mem_init_fn_cache),
+            layout.len(),
+            *pointee_info.ty(),
+        );
+        collect_skipped(&operation, body, skip_first);
+        let position = operation.position();
+        let MemoryInitOp::Copy { from, to, count } = operation else { unreachable!() };
+        body.insert_call(
+            &copy_init_state_instance,
+            source,
+            position,
+            vec![from, to, count],
+            ret_place.clone(),
+        );
+    }
+
     fn inject_assert_false(
         &self,
         tcx: TyCtxt,

kani-compiler/src/kani_middle/transform/check_uninit/ptr_uninit/uninit_visitor.rs

@@ -88,18 +88,13 @@ impl MirVisitor for CheckUninitVisitor {
             match &stmt.kind {
                 StatementKind::Intrinsic(NonDivergingIntrinsic::CopyNonOverlapping(copy)) => {
                     self.super_statement(stmt, location);
-                    // Source is a *const T and it must be initialized.
-                    self.push_target(MemoryInitOp::CheckSliceChunk {
-                        operand: copy.src.clone(),
+                    // The copy is untyped, so we should copy memory initialization state from `src`
+                    // to `dst`.
+                    self.push_target(MemoryInitOp::Copy {
+                        from: copy.src.clone(),
+                        to: copy.dst.clone(),
                         count: copy.count.clone(),
                     });
-                    // Destination is a *mut T so it gets initialized.
-                    self.push_target(MemoryInitOp::SetSliceChunk {
-                        operand: copy.dst.clone(),
-                        count: copy.count.clone(),
-                        value: true,
-                        position: InsertPosition::After,
-                    });
                 }
                 StatementKind::Assign(place, rvalue) => {
                     // First check rvalue.
@@ -219,29 +214,24 @@ impl MirVisitor for CheckUninitVisitor {
                                     });
                                 }
                                 Intrinsic::Copy => {
-                                    self.push_target(MemoryInitOp::CheckSliceChunk {
-                                        operand: args[0].clone(),
+                                    // The copy is untyped, so we should copy memory
+                                    // initialization state from `src` to `dst`.
+                                    self.push_target(MemoryInitOp::Copy {
+                                        from: args[0].clone(),
+                                        to: args[1].clone(),
                                         count: args[2].clone(),
                                     });
-                                    self.push_target(MemoryInitOp::SetSliceChunk {
-                                        operand: args[1].clone(),
-                                        count: args[2].clone(),
-                                        value: true,
-                                        position: InsertPosition::After,
-                                    });
                                 }
                                 Intrinsic::VolatileCopyMemory
                                 | Intrinsic::VolatileCopyNonOverlappingMemory => {
-                                    self.push_target(MemoryInitOp::CheckSliceChunk {
-                                        operand: args[1].clone(),
+                                    // The copy is untyped, so we should copy initialization state
+                                    // from `src` to `dst`. Note that the `dst` comes before `src`
+                                    // in this case.
+                                    self.push_target(MemoryInitOp::Copy {
+                                        from: args[1].clone(),
+                                        to: args[0].clone(),
                                         count: args[2].clone(),
                                     });
-                                    self.push_target(MemoryInitOp::SetSliceChunk {
-                                        operand: args[0].clone(),
-                                        count: args[2].clone(),
-                                        value: true,
-                                        position: InsertPosition::After,
-                                    });
                                 }
                                 Intrinsic::TypedSwap => {
                                     self.push_target(MemoryInitOp::Check {

kani-compiler/src/kani_middle/transform/check_uninit/relevant_instruction.rs

@@ -5,7 +5,10 @@
 //! character of instrumentation needed.
 
 use crate::kani_middle::transform::body::{InsertPosition, MutableBody, SourceInstruction};
-use stable_mir::mir::{Mutability, Operand, Place, Rvalue};
+use stable_mir::{
+    mir::{Mutability, Operand, Place, Rvalue},
+    ty::RigidTy,
+};
 use strum_macros::AsRefStr;
 
 /// Memory initialization operations: set or get memory initialization state for a given pointer.
@@ -33,6 +36,8 @@ pub enum MemoryInitOp {
     Unsupported { reason: String },
     /// Operation that trivially accesses uninitialized memory, results in injecting `assert!(false)`.
     TriviallyUnsafe { reason: String },
+    /// Operation that copies memory initialization state over to another operand.
+    Copy { from: Operand, to: Operand, count: Operand },
 }
 
 impl MemoryInitOp {
@@ -62,15 +67,31 @@ impl MemoryInitOp {
                 })
             }
             MemoryInitOp::Unsupported { .. } | MemoryInitOp::TriviallyUnsafe { .. } => {
-                unreachable!()
+                unreachable!("operands do not exist for this operation")
+            }
+            MemoryInitOp::Copy { from, to, .. } => {
+                // It does not matter which operand to return for layout generation, since both of
+                // them have the same pointee type, so we assert that.
+                let from_kind = from.ty(body.locals()).unwrap().kind();
+                let to_kind = to.ty(body.locals()).unwrap().kind();
+
+                let RigidTy::RawPtr(from_pointee_ty, _) = from_kind.rigid().unwrap().clone() else {
+                    unreachable!()
+                };
+                let RigidTy::RawPtr(to_pointee_ty, _) = to_kind.rigid().unwrap().clone() else {
+                    unreachable!()
+                };
+                assert!(from_pointee_ty == to_pointee_ty);
+                from.clone()
             }
         }
     }
 
     pub fn expect_count(&self) -> Operand {
         match self {
             MemoryInitOp::CheckSliceChunk { count, .. }
-            | MemoryInitOp::SetSliceChunk { count, .. } => count.clone(),
+            | MemoryInitOp::SetSliceChunk { count, .. }
+            | MemoryInitOp::Copy { count, .. } => count.clone(),
             MemoryInitOp::Check { .. }
             | MemoryInitOp::Set { .. }
             | MemoryInitOp::CheckRef { .. }
@@ -89,7 +110,8 @@ impl MemoryInitOp {
             | MemoryInitOp::CheckSliceChunk { .. }
             | MemoryInitOp::CheckRef { .. }
             | MemoryInitOp::Unsupported { .. }
-            | MemoryInitOp::TriviallyUnsafe { .. } => unreachable!(),
+            | MemoryInitOp::TriviallyUnsafe { .. }
+            | MemoryInitOp::Copy { .. } => unreachable!(),
         }
     }
 
@@ -103,6 +125,7 @@ impl MemoryInitOp {
             | MemoryInitOp::CheckRef { .. }
             | MemoryInitOp::Unsupported { .. }
             | MemoryInitOp::TriviallyUnsafe { .. } => InsertPosition::Before,
+            MemoryInitOp::Copy { .. } => InsertPosition::After,
         }
     }
 }

library/kani/src/mem_init.rs

@@ -91,6 +91,33 @@ impl MemoryInitializationState {
         }
     }
 
+    /// Copy memory initialization state by non-deterministically switching the tracked object and
+    /// adjusting the tracked offset.
+    pub fn copy<const LAYOUT_SIZE: usize>(
+        &mut self,
+        from_ptr: *const u8,
+        to_ptr: *const u8,
+        num_elts: usize,
+    ) {
+        let from_obj = crate::mem::pointer_object(from_ptr);
+        let from_offset = crate::mem::pointer_offset(from_ptr);
+
+        let to_obj = crate::mem::pointer_object(to_ptr);
+        let to_offset = crate::mem::pointer_offset(to_ptr);
+
+        if self.tracked_object_id == from_obj
+            && self.tracked_offset >= from_offset
+            && self.tracked_offset < from_offset + num_elts * LAYOUT_SIZE
+        {
+            let should_reset: bool = crate::any();
+            if should_reset {
+                self.tracked_object_id = to_obj;
+                self.tracked_offset += to_offset - from_offset;
+                // Note that this preserves the value.
+            }
+        }
+    }
+
     /// Return currently tracked memory initialization state if `ptr` points to the currently
     /// tracked object and the tracked offset lies within `LAYOUT_SIZE * num_elts` bytes of `ptr`.
     /// Return `true` otherwise.
@@ -281,3 +308,16 @@ fn set_str_ptr_initialized<const LAYOUT_SIZE: usize>(
         MEM_INIT_STATE.set_slice(ptr as *const u8, layout, num_elts, value);
     }
 }
+
+/// Copy initialization state of `size_of::<T> * num_elts` bytes from one pointer to the other.
+#[rustc_diagnostic_item = "KaniCopyInitState"]
+fn copy_init_state<const LAYOUT_SIZE: usize, T>(from: *const T, to: *const T, num_elts: usize) {
+    if LAYOUT_SIZE == 0 {
+        return;
+    }
+    let (from_ptr, _) = from.to_raw_parts();
+    let (to_ptr, _) = to.to_raw_parts();
+    unsafe {
+        MEM_INIT_STATE.copy::<LAYOUT_SIZE>(from_ptr as *const u8, to_ptr as *const u8, num_elts);
+    }
+}

tests/expected/uninit/copy/copy_without_padding.expected

@@ -0,0 +1 @@
+Complete - 1 successfully verified harnesses, 0 failures, 1 total.

tests/expected/uninit/copy/copy_without_padding.rs

@@ -0,0 +1,23 @@
+// Copyright Kani Contributors
+// SPDX-License-Identifier: Apache-2.0 OR MIT
+// kani-flags: -Z uninit-checks
+
+#[repr(C)]
+#[derive(kani::Arbitrary)]
+struct S(u32, u8); // 5 bytes of data + 3 bytes of padding.
+
+#[kani::proof]
+/// This checks that reading copied initialized bytes verifies correctly.
+unsafe fn copy_without_padding() {
+    let from: S = kani::any();
+    let mut to: u64 = kani::any();
+
+    let from_ptr = &from as *const S;
+    let to_ptr = &mut to as *mut u64;
+
+    // This should not cause UB since `copy` is untyped.
+    std::ptr::copy(from_ptr as *const u8, to_ptr as *mut u8, std::mem::size_of::<u32>());
+
+    // Since the previous copy only copied 4 bytes, no padding was copied, so no padding is read.
+    let data: u64 = std::ptr::read(to_ptr);
+}

tests/expected/uninit/copy/expose_padding_via_copy.expected

@@ -0,0 +1,11 @@
+std::ptr::read::<u64>.assertion.1\
+	 - Status: FAILURE\
+	 - Description: "Undefined Behavior: Reading from an uninitialized pointer of type `*const u64`"\
+
+std::ptr::read::<u64>.assertion.2\
+	 - Status: FAILURE\
+	 - Description: "Undefined Behavior: Reading from an uninitialized pointer of type `u64`"\
+
+Summary:
+Verification failed for - expose_padding_via_copy
+Complete - 0 successfully verified harnesses, 1 failures, 1 total.

tests/expected/uninit/copy/expose_padding_via_copy.rs

@@ -0,0 +1,23 @@
+// Copyright Kani Contributors
+// SPDX-License-Identifier: Apache-2.0 OR MIT
+// kani-flags: -Z uninit-checks
+
+#[repr(C)]
+#[derive(kani::Arbitrary)]
+struct S(u32, u8); // 5 bytes of data + 3 bytes of padding.
+
+/// This checks that reading copied uninitialized bytes fails an assertion.
+#[kani::proof]
+unsafe fn expose_padding_via_copy() {
+    let from: S = kani::any();
+    let mut to: u64 = kani::any();
+
+    let from_ptr = &from as *const S;
+    let to_ptr = &mut to as *mut u64;
+
+    // This should not cause UB since `copy` is untyped.
+    std::ptr::copy(from_ptr as *const u8, to_ptr as *mut u8, std::mem::size_of::<S>());
+
+    // This reads uninitialized bytes, which is UB.
+    let padding: u64 = std::ptr::read(to_ptr);
+}

tests/expected/uninit/copy/expose_padding_via_copy_convoluted.expected

@@ -0,0 +1,11 @@
+std::ptr::read::<u64>.assertion.1\
+	 - Status: FAILURE\
+	 - Description: "Undefined Behavior: Reading from an uninitialized pointer of type `*const u64`"\
+
+std::ptr::read::<u64>.assertion.2\
+	 - Status: FAILURE\
+	 - Description: "Undefined Behavior: Reading from an uninitialized pointer of type `u64`"\
+
+Summary:
+Verification failed for - expose_padding_via_copy_convoluted
+Complete - 0 successfully verified harnesses, 1 failures, 1 total.

tests/expected/uninit/copy/expose_padding_via_copy_convoluted.rs

@@ -0,0 +1,42 @@
+// Copyright Kani Contributors
+// SPDX-License-Identifier: Apache-2.0 OR MIT
+// kani-flags: -Z uninit-checks
+
+#[repr(C)]
+#[derive(kani::Arbitrary)]
+struct S(u32, u8); // 5 bytes of data + 3 bytes of padding.
+
+/// This checks that reading copied uninitialized bytes fails an assertion even if pointer are
+/// passed around different functions.
+#[kani::proof]
+unsafe fn expose_padding_via_copy_convoluted() {
+    unsafe fn copy_and_read_helper(from_ptr: *const S, to_ptr: *mut u64) -> u64 {
+        // This should not cause UB since `copy` is untyped.
+        std::ptr::copy(from_ptr as *const u8, to_ptr as *mut u8, std::mem::size_of::<S>());
+        // This reads uninitialized bytes, which is UB.
+        let padding: u64 = std::ptr::read(to_ptr);
+        padding
+    }
+
+    unsafe fn partial_copy_and_read_helper(from_ptr: *const S, to_ptr: *mut u64) -> u32 {
+        // This should not cause UB since `copy` is untyped.
+        std::ptr::copy(from_ptr as *const u8, to_ptr as *mut u8, std::mem::size_of::<u32>());
+        // This does not read uninitialized bytes.
+        let not_padding: u32 = std::ptr::read(to_ptr as *mut u32);
+        not_padding
+    }
+
+    let flag: bool = kani::any();
+
+    let from: S = kani::any();
+    let mut to: u64 = kani::any();
+
+    let from_ptr = &from as *const S;
+    let to_ptr = &mut to as *mut u64;
+
+    if flag {
+        copy_and_read_helper(from_ptr, to_ptr);
+    } else {
+        partial_copy_and_read_helper(from_ptr, to_ptr);
+    }
+}

tests/expected/uninit/copy/expose_padding_via_non_byte_copy.expected

@@ -0,0 +1,11 @@
+std::ptr::read::<u64>.assertion.1\
+	 - Status: FAILURE\
+	 - Description: "Undefined Behavior: Reading from an uninitialized pointer of type `*const u64`"\
+
+std::ptr::read::<u64>.assertion.2\
+	 - Status: FAILURE\
+	 - Description: "Undefined Behavior: Reading from an uninitialized pointer of type `u64`"\
+
+Summary:
+Verification failed for - expose_padding_via_non_byte_copy
+Complete - 0 successfully verified harnesses, 1 failures, 1 total.

tests/expected/uninit/copy/expose_padding_via_non_byte_copy.rs

@@ -0,0 +1,23 @@
+// Copyright Kani Contributors
+// SPDX-License-Identifier: Apache-2.0 OR MIT
+// kani-flags: -Z uninit-checks
+
+#[repr(C)]
+#[derive(kani::Arbitrary)]
+struct S(u32, u8); // 5 bytes of data + 3 bytes of padding.
+
+/// This checks that reading copied uninitialized bytes after a multi-byte copy fails an assertion.
+#[kani::proof]
+unsafe fn expose_padding_via_non_byte_copy() {
+    let from: S = kani::any();
+    let mut to: u64 = kani::any();
+
+    let from_ptr = &from as *const S;
+    let to_ptr = &mut to as *mut u64;
+
+    // This should not cause UB since `copy` is untyped.
+    std::ptr::copy(from_ptr as *const u64, to_ptr as *mut u64, 1);
+
+    // This reads uninitialized bytes, which is UB.
+    let padding: u64 = std::ptr::read(to_ptr);
+}

tests/expected/uninit/copy/non_byte_copy_without_padding.expected

@@ -0,0 +1 @@
+Complete - 1 successfully verified harnesses, 0 failures, 1 total.