Basic support for memory initialization checks for unions

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

@@ -164,8 +164,14 @@ impl<'a, 'tcx> UninitInstrumenter<'a, 'tcx> {
             }
             MemoryInitOp::SetSliceChunk { .. }
             | MemoryInitOp::Set { .. }
-            | MemoryInitOp::SetRef { .. } => self.build_set(body, source, operation, pointee_info),
+            | MemoryInitOp::SetRef { .. }
+            | MemoryInitOp::CreateUnion { .. } => {
+                self.build_set(body, source, operation, pointee_info)
+            }
             MemoryInitOp::Copy { .. } => self.build_copy(body, source, operation, pointee_info),
+            MemoryInitOp::AssignUnion { .. } => {
+                self.build_assign_union(body, source, operation, pointee_info)
+            }
             MemoryInitOp::Unsupported { .. } | MemoryInitOp::TriviallyUnsafe { .. } => {
                 unreachable!()
             }
@@ -266,6 +272,14 @@ impl<'a, 'tcx> UninitInstrumenter<'a, 'tcx> {
                 self.inject_assert_false(self.tcx, body, source, operation.position(), reason);
                 return;
             }
+            PointeeLayout::Union { .. } => {
+                // Here we are reading from a pointer to a union.
+                // TODO: we perhaps need to check that the union at least contains an intersection
+                // of all layouts initialized.
+                let reason = "Interaction between raw pointers and unions is not yet supported.";
+                self.inject_assert_false(self.tcx, body, source, operation.position(), reason);
+                return;
+            }
         };
 
         // Construct the basic block and insert it into the body.
@@ -409,6 +423,62 @@ impl<'a, 'tcx> UninitInstrumenter<'a, 'tcx> {
             PointeeLayout::TraitObject => {
                 unreachable!("Cannot change the initialization state of a trait object directly.");
             }
+            PointeeLayout::Union { field_layouts } => {
+                // Writing union data, which could be either creating a union from scratch or
+                // performing some pointer operations with it.
+
+                // TODO: If we don't have a union field, we are either creating a pointer to a union
+                // or assigning to one. In the former case, it is safe to return from this function,
+                // since the union must be already tracked (on creation and update). In the latter
+                // case, we should have been using union assignment instead. Nevertheless, this is
+                // currently mitigated by injecting `assert!(false)`.
+                let union_field = match operation.union_field() {
+                    Some(field) => field,
+                    None => {
+                        let reason =
+                            "Interaction between raw pointers and unions is not yet supported.";
+                        self.inject_assert_false(
+                            self.tcx,
+                            body,
+                            source,
+                            operation.position(),
+                            reason,
+                        );
+                        return;
+                    }
+                };
+                let layout = &field_layouts[union_field];
+                let layout_operand = mk_layout_operand(body, &mut statements, source, layout);
+                let diagnostic = "KaniSetPtrInitialized";
+                let args = vec![
+                    ptr_operand,
+                    layout_operand,
+                    Operand::Constant(ConstOperand {
+                        span: source.span(body.blocks()),
+                        user_ty: None,
+                        const_: MirConst::from_bool(value),
+                    }),
+                ];
+                let set_ptr_initialized_instance = resolve_mem_init_fn(
+                    get_mem_init_fn_def(self.tcx, diagnostic, &mut self.mem_init_fn_cache),
+                    layout.len(),
+                    *pointee_info.ty(),
+                );
+                Terminator {
+                    kind: TerminatorKind::Call {
+                        func: Operand::Copy(Place::from(body.new_local(
+                            set_ptr_initialized_instance.ty(),
+                            source.span(body.blocks()),
+                            Mutability::Not,
+                        ))),
+                        args,
+                        destination: ret_place.clone(),
+                        target: Some(0), // this will be overriden in add_bb
+                        unwind: UnwindAction::Terminate,
+                    },
+                    span: source.span(body.blocks()),
+                }
+            }
         };
         // Construct the basic block and insert it into the body.
         body.insert_bb(BasicBlock { statements, terminator }, source, operation.position());
@@ -426,14 +496,15 @@ impl<'a, 'tcx> UninitInstrumenter<'a, 'tcx> {
             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 layout_size = pointee_info.layout().maybe_size().unwrap();
         let copy_init_state_instance = resolve_mem_init_fn(
             get_mem_init_fn_def(self.tcx, "KaniCopyInitState", &mut self.mem_init_fn_cache),
-            layout.len(),
+            layout_size,
             *pointee_info.ty(),
         );
         let position = operation.position();
-        let MemoryInitOp::Copy { from, to, count } = operation else { unreachable!() };
+        let (from, to) = operation.expect_copy_operands();
+        let count = operation.expect_count();
         body.insert_call(
             &copy_init_state_instance,
             source,
@@ -443,6 +514,45 @@ impl<'a, 'tcx> UninitInstrumenter<'a, 'tcx> {
         );
     }
 
+    /// Copy memory initialization state from one union variable to another.
+    fn build_assign_union(
+        &mut self,
+        body: &mut MutableBody,
+        source: &mut SourceInstruction,
+        operation: MemoryInitOp,
+        pointee_info: PointeeInfo,
+    ) {
+        let ret_place = Place {
+            local: body.new_local(Ty::new_tuple(&[]), source.span(body.blocks()), Mutability::Not),
+            projection: vec![],
+        };
+        let mut statements = vec![];
+        let layout_size = pointee_info.layout().maybe_size().unwrap();
+        let copy_init_state_instance = resolve_mem_init_fn(
+            get_mem_init_fn_def(self.tcx, "KaniCopyInitStateSingle", &mut self.mem_init_fn_cache),
+            layout_size,
+            *pointee_info.ty(),
+        );
+        let (from, to) = operation.expect_assign_union_operands(body, &mut statements, source);
+        let terminator = Terminator {
+            kind: TerminatorKind::Call {
+                func: Operand::Copy(Place::from(body.new_local(
+                    copy_init_state_instance.ty(),
+                    source.span(body.blocks()),
+                    Mutability::Not,
+                ))),
+                args: vec![from, to],
+                destination: ret_place.clone(),
+                target: Some(0), // this will be overriden in add_bb
+                unwind: UnwindAction::Terminate,
+            },
+            span: source.span(body.blocks()),
+        };
+
+        // Construct the basic block and insert it into the body.
+        body.insert_bb(BasicBlock { statements, terminator }, source, operation.position());
+    }
+
     fn inject_assert_false(
         &self,
         tcx: TyCtxt,

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

@@ -19,10 +19,11 @@ use stable_mir::{
         alloc::GlobalAlloc,
         mono::{Instance, InstanceKind},
         visit::{Location, PlaceContext},
-        CastKind, LocalDecl, MirVisitor, NonDivergingIntrinsic, Operand, Place, PointerCoercion,
-        ProjectionElem, Rvalue, Statement, StatementKind, Terminator, TerminatorKind,
+        AggregateKind, CastKind, LocalDecl, MirVisitor, NonDivergingIntrinsic, Operand, Place,
+        PointerCoercion, ProjectionElem, Rvalue, Statement, StatementKind, Terminator,
+        TerminatorKind,
     },
-    ty::{ConstantKind, RigidTy, TyKind},
+    ty::{AdtKind, ConstantKind, RigidTy, TyKind},
 };
 
 pub struct CheckUninitVisitor {
@@ -112,7 +113,7 @@ impl MirVisitor for CheckUninitVisitor {
                     }
                 }
                 // Check whether Rvalue creates a new initialized pointer previously not captured inside shadow memory.
-                if place.ty(&&self.locals).unwrap().kind().is_raw_ptr() {
+                if place.ty(&self.locals).unwrap().kind().is_raw_ptr() {
                     if let Rvalue::AddressOf(..) = rvalue {
                         self.push_target(MemoryInitOp::Set {
                             operand: Operand::Copy(place.clone()),
@@ -121,6 +122,58 @@ impl MirVisitor for CheckUninitVisitor {
                         });
                     }
                 }
+
+                // TODO: add support for ADTs which could have unions as subfields. Currently,
+                // if a union as a subfield is detected, `assert!(false)` will be injected from
+                // the type layout code.
+                let is_inside_union = {
+                    let mut contains_union = false;
+                    let mut place_to_add_projections =
+                        Place { local: place.local, projection: vec![] };
+                    for projection_elem in place.projection.iter() {
+                        if place_to_add_projections.ty(&self.locals).unwrap().kind().is_union() {
+                            contains_union = true;
+                            break;
+                        }
+                        place_to_add_projections.projection.push(projection_elem.clone());
+                    }
+                    contains_union
+                };
+
+                // Need to copy some information about union initialization, since lvalue is
+                // either a union or a field inside a union.
+                if is_inside_union {
+                    if let Rvalue::Use(operand) = rvalue {
+                        // This is a union-to-union assignment, so we need to copy the
+                        // initialization state.
+                        if place.ty(&self.locals).unwrap().kind().is_union() {
+                            self.push_target(MemoryInitOp::AssignUnion {
+                                lvalue: place.clone(),
+                                rvalue: operand.clone(),
+                            });
+                        } else {
+                            // This is assignment to a field of a union.
+                            self.push_target(MemoryInitOp::SetRef {
+                                operand: Operand::Copy(place.clone()),
+                                value: true,
+                                position: InsertPosition::After,
+                            });
+                        }
+                    }
+                }
+
+                // Create a union from scratch as an aggregate. We handle it here because we
+                // need to know which field is getting assigned.
+                if let Rvalue::Aggregate(AggregateKind::Adt(adt_def, _, _, _, union_field), _) =
+                    rvalue
+                {
+                    if adt_def.kind() == AdtKind::Union {
+                        self.push_target(MemoryInitOp::CreateUnion {
+                            operand: Operand::Copy(place.clone()),
+                            field: union_field.unwrap(), // Safe to unwrap because we know this is a union.
+                        });
+                    }
+                }
             }
             StatementKind::Deinit(place) => {
                 self.super_statement(stmt, location);
@@ -350,12 +403,19 @@ impl MirVisitor for CheckUninitVisitor {
                         });
                     }
                 }
-                ProjectionElem::Field(idx, target_ty) => {
-                    if target_ty.kind().is_union()
-                        && (!ptx.is_mutating() || place.projection.len() > idx + 1)
+                ProjectionElem::Field(_, _) => {
+                    if intermediate_place.ty(&self.locals).unwrap().kind().is_union()
+                        && !ptx.is_mutating()
                     {
-                        self.push_target(MemoryInitOp::Unsupported {
-                            reason: "Kani does not support reasoning about memory initialization of unions.".to_string(),
+                        // Generate a place which includes all projections until (and including)
+                        // union field access.
+                        let union_field_access_place = Place {
+                            local: place.local,
+                            projection: place.projection[..idx + 1].to_vec(),
+                        };
+                        // Accessing a place inside the union, need to check if it is initialized.
+                        self.push_target(MemoryInitOp::CheckRef {
+                            operand: Operand::Copy(union_field_access_place.clone()),
                         });
                     }
                 }