Fix how we handle niche optimization with ZST (rust-lang#1205)

* Fix how we handle niche optimization with ZST

The compiler uses niche optimization when there is only one variant that
contains nonzero sized fields. This variant is represented as the
dataful variant. However, other variants may have one or more ZST fields
that can be deconstructed and referenced in the code.

Before this change, we were not generating code for them and it generate
type mismatch issues when the user tried to access them.

This change ensures that we represent all variants that have any field,
so their access is valid.

* Enable debug assert inside projection mismatch

We should probably turn this into an assert and revert the logic added
by model-checking/kani#1057. But this is a much
larger change that I prefer creating a separate PR for.

Co-authored-by: Zyad Hassan <[email protected]>

Author: celinval

cprover_bindings/src/goto_program/stmt.rs

@@ -4,7 +4,6 @@ use self::StmtBody::*;
 use super::{BuiltinFn, Expr, Location};
 use crate::{InternString, InternedString};
 use std::fmt::Debug;
-use tracing::debug;
 
 ///////////////////////////////////////////////////////////////////////////////////////////////
 /// Datatypes
@@ -166,28 +165,13 @@ macro_rules! stmt {
 impl Stmt {
     /// `lhs = rhs;`
     pub fn assign(lhs: Expr, rhs: Expr, loc: Location) -> Self {
-        //Temporarily work around https://github.com/model-checking/kani/issues/95
-        //by disabling the assert and soundly assigning nondet
-        //assert_eq!(lhs.typ(), rhs.typ());
-        if lhs.typ() != rhs.typ() {
-            debug!(
-                "WARNING: assign statement with unequal types lhs {:?} rhs {:?}",
-                lhs.typ(),
-                rhs.typ()
-            );
-            let assert_stmt = Stmt::assert_false(
-                "sanity_check",
-                &format!(
-                    "Reached assignment statement with unequal types {:?} {:?}",
-                    lhs.typ(),
-                    rhs.typ()
-                ),
-                loc.clone(),
-            );
-            let nondet_value = lhs.typ().nondet();
-            let nondet_assign_stmt = stmt!(Assign { lhs, rhs: nondet_value }, loc.clone());
-            return Stmt::block(vec![assert_stmt, nondet_assign_stmt], loc);
-        }
+        assert_eq!(
+            lhs.typ(),
+            rhs.typ(),
+            "Error: assign statement with unequal types lhs {:?} rhs {:?}",
+            lhs.typ(),
+            rhs.typ()
+        );
         stmt!(Assign { lhs, rhs }, loc)
     }
 

kani-compiler/src/codegen_cprover_gotoc/codegen/place.rs

@@ -169,10 +169,12 @@ impl<'tcx> ProjectedPlace<'tcx> {
         if let Some((expr_ty, ty_from_mir)) =
             Self::check_expr_typ_mismatch(&goto_expr, &mir_typ_or_variant, ctx)
         {
-            warn!(
+            let msg = format!(
                 "Unexpected type mismatch in projection:\n{:?}\nExpr type\n{:?}\nType from MIR\n{:?}",
                 goto_expr, expr_ty, ty_from_mir
             );
+            warn!("{}", msg);
+            debug_assert!(false, "{}", msg);
             return Err(UnimplementedData::new(
                 "Projection mismatch",
                 "https://github.com/model-checking/kani/issues/277",
@@ -499,18 +501,18 @@ impl<'tcx> GotocCtx<'tcx> {
                 match t.kind() {
                     ty::Adt(def, _) => {
                         let variant = def.variants().get(idx).unwrap();
+                        let case_name = variant.name.to_string();
                         let typ = TypeOrVariant::Variant(variant);
                         let expr = match &self.layout_of(t).variants {
                             Variants::Single { .. } => before.goto_expr,
                             Variants::Multiple { tag_encoding, .. } => match tag_encoding {
-                                TagEncoding::Direct => {
-                                    let case_name = variant.name.to_string();
-                                    before
-                                        .goto_expr
-                                        .member("cases", &self.symbol_table)
-                                        .member(&case_name, &self.symbol_table)
+                                TagEncoding::Direct => before
+                                    .goto_expr
+                                    .member("cases", &self.symbol_table)
+                                    .member(&case_name, &self.symbol_table),
+                                TagEncoding::Niche { .. } => {
+                                    before.goto_expr.member(&case_name, &self.symbol_table)
                                 }
-                                TagEncoding::Niche { .. } => before.goto_expr,
                             },
                         };
                         ProjectedPlace::try_new(

kani-compiler/src/codegen_cprover_gotoc/codegen/typ.rs

@@ -1084,7 +1084,17 @@ impl<'tcx> GotocCtx<'tcx> {
                 Variants::Multiple { tag_encoding, variants, .. } => {
                     match tag_encoding {
                         TagEncoding::Direct => {
-                            // direct encoding of tags
+                            // For direct encoding of tags, we generate a type with two fields:
+                            // ```
+                            // struct tag-<> { // enum type
+                            //    case: <discriminant  type>,
+                            //    cases: tag-<>-union,
+                            // }
+                            // ```
+                            // The `case` field type determined by the enum representation
+                            // (`#[repr]`) and it represents which variant is being used.
+                            // The `cases` field is a union of all variant types where the name
+                            // of each union field is the name of the corresponding discriminant.
                             let discr_t = ctx.codegen_enum_discr_typ(ty);
                             let int = ctx.codegen_ty(discr_t);
                             let discr_offset = ctx.layout_of(discr_t).size.bits_usize();
@@ -1098,31 +1108,48 @@ impl<'tcx> GotocCtx<'tcx> {
                             }
                             fields.push(Type::datatype_component(
                                 "cases",
-                                ctx.codegen_enum_cases_union(
-                                    name,
-                                    adtdef,
-                                    subst,
-                                    variants,
-                                    initial_offset,
+                                ctx.ensure_union(
+                                    &format!("{}-union", name.to_string()),
+                                    NO_PRETTY_NAME,
+                                    |ctx, name| {
+                                        ctx.codegen_enum_cases(
+                                            name,
+                                            adtdef,
+                                            subst,
+                                            variants,
+                                            initial_offset,
+                                        )
+                                    },
                                 ),
                             ));
                             fields
                         }
                         TagEncoding::Niche { dataful_variant, .. } => {
-                            // niche encoding is an optimization, which uses invalid values for discriminant
-                            // for example, Option<&i32> becomes just a pointer to i32, and pattern
-                            // matching becomes checking whether the pointer is null or not. direct
-                            // encoding, on the other hand, would have been maintaining a field
-                            // storing the discriminant, which is a few bytes larger.
+                            // Enumerations with multiple variants and niche encoding have a
+                            // specific format that can be used to optimize its layout and reduce
+                            // memory consumption.
                             //
-                            // dataful_variant is pretty much the only variant which contains the valid data
-                            let variant = &adtdef.variants()[*dataful_variant];
-                            ctx.codegen_variant_struct_fields(
-                                variant,
-                                subst,
-                                &variants[*dataful_variant],
-                                0,
-                            )
+                            // These enumerations have one and only one variant with non-ZST
+                            // fields which is referred to by the `dataful_variant` index. Their
+                            // final size and alignment is equal to the one from the
+                            // `dataful_variant`. All other variants either don't have any field
+                            // or all fields types are ZST.
+                            //
+                            // Because of that, we can represent these enums as simple structures
+                            // where each field represent one variant. This allows them to be
+                            // referred to correctly.
+                            //
+                            // Note: I tried using a union instead but it had a significant runtime
+                            // penalty.
+                            tracing::trace!(
+                                ?name,
+                                ?variants,
+                                ?dataful_variant,
+                                ?tag_encoding,
+                                ?subst,
+                                "codegen_enum: Niche"
+                            );
+                            ctx.codegen_enum_cases(name, adtdef, subst, variants, 0)
                         }
                     }
                 }
@@ -1211,32 +1238,37 @@ impl<'tcx> GotocCtx<'tcx> {
         }
     }
 
-    fn codegen_enum_cases_union(
+    /// Codegen the type for each variant represented in this enum.
+    /// As an optimization, we ignore the ones that don't have any field, since they
+    /// are only manipulated via discriminant operations.
+    fn codegen_enum_cases(
         &mut self,
         name: InternedString,
         def: &'tcx AdtDef,
         subst: &'tcx InternalSubsts<'tcx>,
         layouts: &IndexVec<VariantIdx, Layout>,
         initial_offset: usize,
-    ) -> Type {
-        // TODO Should we have a pretty name here?
-        self.ensure_union(&format!("{}-union", name.to_string()), NO_PRETTY_NAME, |ctx, name| {
-            def.variants()
-                .iter_enumerated()
-                .map(|(i, case)| {
-                    Type::datatype_component(
+    ) -> Vec<DatatypeComponent> {
+        def.variants()
+            .iter_enumerated()
+            .filter_map(|(i, case)| {
+                if case.fields.is_empty() {
+                    // Skip variant types that cannot be referenced.
+                    None
+                } else {
+                    Some(Type::datatype_component(
                         &case.name.to_string(),
-                        ctx.codegen_enum_case_struct(
+                        self.codegen_enum_case_struct(
                             name,
                             case,
                             subst,
                             &layouts[i],
                             initial_offset,
                         ),
-                    )
-                })
-                .collect()
-        })
+                    ))
+                }
+            })
+            .collect()
     }
 
     fn codegen_enum_case_struct(

tests/kani/Enum/niche_many_variants.rs

@@ -0,0 +1,96 @@
+// Copyright Kani Contributors
+// SPDX-License-Identifier: Apache-2.0 OR MIT
+
+//! Testcase for niche encoding where there are multiple variants but only one contains
+//! non-zero sized data with niche, making it a great candidate for niche optimization.
+#[derive(PartialEq)]
+enum MyEnum {
+    NoFields,
+    DataFul(bool),
+    UnitFields((), ()),
+    ZSTField(ZeroSized),
+    ZSTStruct { field: ZeroSized, unit: () },
+}
+
+#[derive(PartialEq)]
+struct ZeroSized {}
+
+impl ZeroSized {
+    fn works(&self) -> bool {
+        true
+    }
+}
+
+impl MyEnum {
+    fn create_no_field() -> MyEnum {
+        MyEnum::NoFields
+    }
+
+    fn create_data_ful(data: bool) -> MyEnum {
+        MyEnum::DataFul(data)
+    }
+
+    fn create_unit() -> MyEnum {
+        MyEnum::UnitFields((), ())
+    }
+
+    fn create_zst_field() -> MyEnum {
+        MyEnum::ZSTField(ZeroSized {})
+    }
+
+    fn create_zst_struct() -> MyEnum {
+        MyEnum::ZSTStruct { field: ZeroSized {}, unit: () }
+    }
+}
+
+/// Ensure we are testing a case of niche optimization.
+#[kani::proof]
+fn check_is_niche() {
+    assert_eq!(std::mem::size_of::<MyEnum>(), 1);
+    assert_eq!(std::mem::size_of::<bool>(), 1);
+}
+
+/// Check the behavior for the variant without any field.
+#[kani::proof]
+fn check_niche_no_fields() {
+    let x = MyEnum::create_no_field();
+    assert!(matches!(x, MyEnum::NoFields));
+}
+
+/// Check the behavior for the dataful variant.
+#[kani::proof]
+fn check_niche_data_ful() {
+    let x = MyEnum::create_data_ful(true);
+    assert!(matches!(x, MyEnum::DataFul(true)));
+}
+
+/// Check the behavior for the variant with multiple unit fields.
+#[kani::proof]
+fn check_niche_unit_fields() {
+    let x = MyEnum::create_unit();
+    assert_eq!(x, MyEnum::UnitFields((), ()));
+    if let MyEnum::UnitFields(ref v, ..) = &x {
+        assert_eq!(std::mem::size_of_val(v), 0);
+    }
+}
+
+/// Check the behavior for the variant with one ZST field.
+#[kani::proof]
+fn check_niche_zst_field() {
+    let x = MyEnum::create_zst_field();
+    assert_eq!(x, MyEnum::ZSTField(ZeroSized {}));
+    if let MyEnum::ZSTField(ref field) = &x {
+        assert!(field.works());
+    }
+}
+
+/// Check the behavior for the variant representing a struct with one ZST field.
+#[kani::proof]
+fn check_niche_zst_struct() {
+    let x = MyEnum::create_zst_struct();
+    assert!(matches!(x, MyEnum::ZSTStruct { .. }));
+    if let MyEnum::ZSTStruct { ref field, ref unit } = &x {
+        assert_eq!(std::mem::size_of_val(unit), 0);
+        assert!(field.works());
+    }
+}