Implement memory initialization state copy functionality

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

@@ -183,6 +183,9 @@ impl UninitPass {
             | 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!()
             }
@@ -394,6 +397,38 @@ impl UninitPass {
         };
     }
 
+    /// 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.add_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/uninit_visitor.rs

@@ -39,6 +39,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 {
@@ -66,15 +68,31 @@ impl MemoryInitOp {
                 projection: vec![],
             }),
             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::SetRef { .. }
@@ -91,7 +109,8 @@ impl MemoryInitOp {
             MemoryInitOp::Check { .. }
             | MemoryInitOp::CheckSliceChunk { .. }
             | MemoryInitOp::Unsupported { .. }
-            | MemoryInitOp::TriviallyUnsafe { .. } => unreachable!(),
+            | MemoryInitOp::TriviallyUnsafe { .. }
+            | MemoryInitOp::Copy { .. } => unreachable!(),
         }
     }
 
@@ -104,6 +123,7 @@ impl MemoryInitOp {
             | MemoryInitOp::CheckSliceChunk { .. }
             | MemoryInitOp::Unsupported { .. }
             | MemoryInitOp::TriviallyUnsafe { .. } => InsertPosition::Before,
+            MemoryInitOp::Copy { .. } => InsertPosition::After,
         }
     }
 }
@@ -181,18 +201,12 @@ impl<'a> MirVisitor for CheckUninitVisitor<'a> {
             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(),
+                    // 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(),
                     });
-                    // Destimation 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.
@@ -327,9 +341,7 @@ impl<'a> MirVisitor for CheckUninitVisitor<'a> {
                                         count: args[2].clone(),
                                     });
                                 }
-                                "copy"
-                                | "volatile_copy_memory"
-                                | "volatile_copy_nonoverlapping_memory" => {
+                                "copy" => {
                                     assert_eq!(
                                         args.len(),
                                         3,
@@ -343,15 +355,33 @@ impl<'a> MirVisitor for CheckUninitVisitor<'a> {
                                         args[1].ty(self.locals).unwrap().kind(),
                                         TyKind::RigidTy(RigidTy::RawPtr(_, Mutability::Mut))
                                     ));
-                                    self.push_target(MemoryInitOp::CheckSliceChunk {
-                                        operand: args[0].clone(),
+                                    // 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(),
+                                }
+                                // Here, `dst` is arg[0] and `src` is arg[1], so need to swap the order.
+                                "volatile_copy_memory" | "volatile_copy_nonoverlapping_memory" => {
+                                    assert_eq!(
+                                        args.len(),
+                                        3,
+                                        "Unexpected number of arguments for `volatile_copy`"
+                                    );
+                                    assert!(matches!(
+                                        args[0].ty(self.locals).unwrap().kind(),
+                                        TyKind::RigidTy(RigidTy::RawPtr(_, Mutability::Mut))
+                                    ));
+                                    assert!(matches!(
+                                        args[1].ty(self.locals).unwrap().kind(),
+                                        TyKind::RigidTy(RigidTy::RawPtr(_, Mutability::Not))
+                                    ));
+                                    // Copy memory initialization state from `src` to `dst`.
+                                    self.push_target(MemoryInitOp::Copy {
+                                        from: args[1].clone(),
+                                        to: args[0].clone(),
                                         count: args[2].clone(),
-                                        value: true,
-                                        position: InsertPosition::After,
                                     });
                                 }
                                 "typed_swap" => {

library/kani/src/mem_init.rs

@@ -89,6 +89,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.
@@ -268,3 +295,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.rs

@@ -0,0 +1,87 @@
+// 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 uninitialized bytes fails an assertion.
+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);
+}
+
+#[kani::proof]
+/// This checks that reading copied uninitialized bytes after a multi-byte copy fails an assertion.
+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);
+}
+
+#[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);
+}
+
+#[kani::proof]
+/// This checks that reading copied initialized bytes after a multi-byte copy verifies correctly.
+unsafe fn non_byte_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 u32, to_ptr as *mut u32, 1);
+
+    // 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);
+}
+
+#[kani::proof]
+/// This checks that reading uninitialized bytes fails an assertion even after copy.
+unsafe fn read_after_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::<u64>());
+
+    // Reading padding from the previous place should be UB even after copy.
+    let data: u64 = std::ptr::read(from_ptr as *const u64);
+}

tests/expected/uninit/copy/expected

@@ -0,0 +1,31 @@
+Checking harness read_after_copy...
+
+Failed Checks: Undefined Behavior: Reading from an uninitialized pointer of type `*const u64`
+
+VERIFICATION:- FAILED
+
+Checking harness non_byte_copy_without_padding...
+
+VERIFICATION:- SUCCESSFUL
+
+Checking harness copy_without_padding...
+
+VERIFICATION:- SUCCESSFUL
+
+Checking harness expose_padding_via_non_byte_copy...
+
+Failed Checks: Undefined Behavior: Reading from an uninitialized pointer of type `*const u64`
+
+VERIFICATION:- FAILED
+
+Checking harness expose_padding_via_copy...
+
+Failed Checks: Undefined Behavior: Reading from an uninitialized pointer of type `*const u64`
+
+VERIFICATION:- FAILED
+
+Summary:
+Verification failed for - read_after_copy
+Verification failed for - expose_padding_via_non_byte_copy
+Verification failed for - expose_padding_via_copy
+Complete - 2 successfully verified harnesses, 3 failures, 5 total.

tests/expected/uninit/intrinsics/expected

@@ -46,6 +46,10 @@ VERIFICATION:- SUCCESSFUL
 
 Checking harness check_copy...
 
+VERIFICATION:- SUCCESSFUL
+
+Checking harness check_copy_read...
+
 Failed Checks: Undefined Behavior: Reading from an uninitialized pointer of type `*const u8`
 
 VERIFICATION:- FAILED
@@ -56,6 +60,10 @@ VERIFICATION:- SUCCESSFUL
 
 Checking harness check_copy_nonoverlapping...
 
+VERIFICATION:- SUCCESSFUL
+
+Checking harness check_copy_nonoverlapping_read...
+
 Failed Checks: Undefined Behavior: Reading from an uninitialized pointer of type `*const u8`
 
 VERIFICATION:- FAILED
@@ -65,6 +73,6 @@ Verification failed for - check_typed_swap_safe
 Verification failed for - check_typed_swap
 Verification failed for - check_volatile_load
 Verification failed for - check_compare_bytes
-Verification failed for - check_copy
-Verification failed for - check_copy_nonoverlapping
-Complete - 5 successfully verified harnesses, 6 failures, 11 total.
+Verification failed for - check_copy_read
+Verification failed for - check_copy_nonoverlapping_read
+Complete - 7 successfully verified harnesses, 6 failures, 13 total.