Auto merge of #106281 - JulianKnodt:transmute_const_generics, r=b-naber

Add ability to transmute (somewhat) with generic consts in arrays

Previously if the expression contained generic consts and did not have a directly equivalent type, transmuting the type in this way was forbidden, despite the two sizes being identical. Instead, we should be able to lazily tell if the two consts are identical, and if so allow them to be transmuted.

This is done by normalizing the forms of expressions into sorted order of multiplied terms, which is not generic over all expressions, but should handle most cases.

This allows for some _basic_ transmutations between types that are equivalent in size without requiring additional stack space at runtime.

I only see one other location at which `SizeSkeleton` is being used, and it checks for equality so this shouldn't affect anywhere else that I can tell.

See [this Stackoverflow post](https://stackoverflow.com/questions/73085012/transmute-nested-const-generic-array-rust) for what was previously necessary to convert between types. This PR makes converting nested `T -> [T; 1]` transmutes possible, and `[uB*2; N] -> [uB; N * 2]` possible as well.

I'm not sure whether this is something that would be wanted, and if it is it definitely should not be insta-stable, so I'd add a feature gate.

Author: bors

compiler/rustc_feature/src/active.rs

@@ -518,6 +518,8 @@ declare_features! (
     /// Allows dyn upcasting trait objects via supertraits.
     /// Dyn upcasting is casting, e.g., `dyn Foo -> dyn Bar` where `Foo: Bar`.
     (active, trait_upcasting, "1.56.0", Some(65991), None),
+    /// Allows for transmuting between arrays with sizes that contain generic consts.
+    (active, transmute_generic_consts, "CURRENT_RUSTC_VERSION", Some(109929), None),
     /// Allows #[repr(transparent)] on unions (RFC 2645).
     (active, transparent_unions, "1.37.0", Some(60405), None),
     /// Allows inconsistent bounds in where clauses.

compiler/rustc_hir_typeck/src/intrinsicck.rs

@@ -84,6 +84,13 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
         let skeleton_string = |ty: Ty<'tcx>, sk| match sk {
             Ok(SizeSkeleton::Known(size)) => format!("{} bits", size.bits()),
             Ok(SizeSkeleton::Pointer { tail, .. }) => format!("pointer to `{tail}`"),
+            Ok(SizeSkeleton::Generic(size)) => {
+                if let Some(size) = size.try_eval_target_usize(tcx, self.param_env) {
+                    format!("{size} bytes")
+                } else {
+                    format!("generic size {size}")
+                }
+            }
             Err(LayoutError::Unknown(bad)) => {
                 if bad == ty {
                     "this type does not have a fixed size".to_owned()

compiler/rustc_middle/src/ty/layout.rs

@@ -281,6 +281,12 @@ pub enum SizeSkeleton<'tcx> {
     /// Any statically computable Layout.
     Known(Size),
 
+    /// This is a generic const expression (i.e. N * 2), which may contain some parameters.
+    /// It must be of type usize, and represents the size of a type in bytes.
+    /// It is not required to be evaluatable to a concrete value, but can be used to check
+    /// that another SizeSkeleton is of equal size.
+    Generic(ty::Const<'tcx>),
+
     /// A potentially-fat pointer.
     Pointer {
         /// If true, this pointer is never null.
@@ -326,6 +332,37 @@ impl<'tcx> SizeSkeleton<'tcx> {
                     ),
                 }
             }
+            ty::Array(inner, len)
+                if len.ty() == tcx.types.usize && tcx.features().transmute_generic_consts =>
+            {
+                match SizeSkeleton::compute(inner, tcx, param_env)? {
+                    // This may succeed because the multiplication of two types may overflow
+                    // but a single size of a nested array will not.
+                    SizeSkeleton::Known(s) => {
+                        if let Some(c) = len.try_eval_target_usize(tcx, param_env) {
+                            let size = s
+                                .bytes()
+                                .checked_mul(c)
+                                .ok_or_else(|| LayoutError::SizeOverflow(ty))?;
+                            return Ok(SizeSkeleton::Known(Size::from_bytes(size)));
+                        }
+                        let len = tcx.expand_abstract_consts(len);
+                        let prev = ty::Const::from_target_usize(tcx, s.bytes());
+                        let Some(gen_size) = mul_sorted_consts(tcx, param_env, len, prev) else {
+                            return Err(LayoutError::SizeOverflow(ty));
+                        };
+                        Ok(SizeSkeleton::Generic(gen_size))
+                    }
+                    SizeSkeleton::Pointer { .. } => Err(err),
+                    SizeSkeleton::Generic(g) => {
+                        let len = tcx.expand_abstract_consts(len);
+                        let Some(gen_size) = mul_sorted_consts(tcx, param_env, len, g) else {
+                            return Err(LayoutError::SizeOverflow(ty));
+                        };
+                        Ok(SizeSkeleton::Generic(gen_size))
+                    }
+                }
+            }
 
             ty::Adt(def, substs) => {
                 // Only newtypes and enums w/ nullable pointer optimization.
@@ -355,6 +392,9 @@ impl<'tcx> SizeSkeleton<'tcx> {
                                 }
                                 ptr = Some(field);
                             }
+                            SizeSkeleton::Generic(_) => {
+                                return Err(err);
+                            }
                         }
                     }
                     Ok(ptr)
@@ -410,11 +450,66 @@ impl<'tcx> SizeSkeleton<'tcx> {
             (SizeSkeleton::Pointer { tail: a, .. }, SizeSkeleton::Pointer { tail: b, .. }) => {
                 a == b
             }
+            // constants are always pre-normalized into a canonical form so this
+            // only needs to check if their pointers are identical.
+            (SizeSkeleton::Generic(a), SizeSkeleton::Generic(b)) => a == b,
             _ => false,
         }
     }
 }
 
+/// When creating the layout for types with abstract conts in their size (i.e. [usize; 4 * N]),
+/// to ensure that they have a canonical order and can be compared directly we combine all
+/// constants, and sort the other terms. This allows comparison of expressions of sizes,
+/// allowing for things like transmutating between types that depend on generic consts.
+/// This returns `None` if multiplication of constants overflows.
+fn mul_sorted_consts<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    param_env: ty::ParamEnv<'tcx>,
+    a: ty::Const<'tcx>,
+    b: ty::Const<'tcx>,
+) -> Option<ty::Const<'tcx>> {
+    use crate::mir::BinOp::Mul;
+    use ty::ConstKind::Expr;
+    use ty::Expr::Binop;
+
+    let mut work = vec![a, b];
+    let mut done = vec![];
+    while let Some(n) = work.pop() {
+        if let Expr(Binop(Mul, l, r)) = n.kind() {
+            work.push(l);
+            work.push(r)
+        } else {
+            done.push(n);
+        }
+    }
+    let mut k = 1;
+    let mut overflow = false;
+    done.retain(|c| {
+        let Some(c) = c.try_eval_target_usize(tcx, param_env) else {
+            return true;
+        };
+        let Some(next) = c.checked_mul(k) else {
+            overflow = true;
+            return false;
+        };
+        k = next;
+        false
+    });
+    if overflow {
+        return None;
+    }
+    if k != 1 {
+        done.push(ty::Const::from_target_usize(tcx, k));
+    } else if k == 0 {
+        return Some(ty::Const::from_target_usize(tcx, 0));
+    }
+    done.sort_unstable();
+
+    // create a single tree from the buffer
+    done.into_iter().reduce(|acc, n| tcx.mk_const(Expr(Binop(Mul, n, acc)), n.ty()))
+}
+
 pub trait HasTyCtxt<'tcx>: HasDataLayout {
     fn tcx(&self) -> TyCtxt<'tcx>;
 }

compiler/rustc_span/src/symbol.rs

@@ -1496,6 +1496,7 @@ symbols! {
         trait_alias,
         trait_upcasting,
         transmute,
+        transmute_generic_consts,
         transmute_opts,
         transmute_trait,
         transparent,

tests/ui/const-generics/transmute-fail.rs

@@ -0,0 +1,35 @@
+#![feature(transmute_generic_consts)]
+#![feature(generic_const_exprs)]
+#![allow(incomplete_features)]
+
+fn foo<const W: usize, const H: usize>(v: [[u32;H+1]; W]) -> [[u32; W+1]; H] {
+  unsafe {
+    std::mem::transmute(v)
+    //~^ ERROR cannot transmute
+  }
+}
+
+fn bar<const W: bool, const H: usize>(v: [[u32; H]; W]) -> [[u32; W]; H] {
+  //~^ ERROR mismatched types
+  //~| ERROR mismatched types
+  unsafe {
+    std::mem::transmute(v)
+    //~^ ERROR cannot transmute between types
+  }
+}
+
+fn baz<const W: usize, const H: usize>(v: [[u32; H]; W]) -> [u32; W * H * H] {
+  unsafe {
+    std::mem::transmute(v)
+    //~^ ERROR cannot transmute
+  }
+}
+
+fn overflow(v: [[[u32; 8888888]; 9999999]; 777777777]) -> [[[u32; 9999999]; 777777777]; 8888888] {
+  unsafe {
+    std::mem::transmute(v)
+    //~^ ERROR cannot transmute
+  }
+}
+
+fn main() {}

tests/ui/const-generics/transmute-fail.stderr

@@ -0,0 +1,52 @@
+error[E0512]: cannot transmute between types of different sizes, or dependently-sized types
+  --> $DIR/transmute-fail.rs:7:5
+   |
+LL |     std::mem::transmute(v)
+   |     ^^^^^^^^^^^^^^^^^^^
+   |
+   = note: source type: `[[u32; H+1]; W]` (generic size [const expr])
+   = note: target type: `[[u32; W+1]; H]` (generic size [const expr])
+
+error[E0512]: cannot transmute between types of different sizes, or dependently-sized types
+  --> $DIR/transmute-fail.rs:16:5
+   |
+LL |     std::mem::transmute(v)
+   |     ^^^^^^^^^^^^^^^^^^^
+   |
+   = note: source type: `[[u32; H]; W]` (this type does not have a fixed size)
+   = note: target type: `[[u32; W]; H]` (size can vary because of [u32; W])
+
+error[E0308]: mismatched types
+  --> $DIR/transmute-fail.rs:12:53
+   |
+LL | fn bar<const W: bool, const H: usize>(v: [[u32; H]; W]) -> [[u32; W]; H] {
+   |                                                     ^ expected `usize`, found `bool`
+
+error[E0308]: mismatched types
+  --> $DIR/transmute-fail.rs:12:67
+   |
+LL | fn bar<const W: bool, const H: usize>(v: [[u32; H]; W]) -> [[u32; W]; H] {
+   |                                                                   ^ expected `usize`, found `bool`
+
+error[E0512]: cannot transmute between types of different sizes, or dependently-sized types
+  --> $DIR/transmute-fail.rs:23:5
+   |
+LL |     std::mem::transmute(v)
+   |     ^^^^^^^^^^^^^^^^^^^
+   |
+   = note: source type: `[[u32; H]; W]` (generic size [const expr])
+   = note: target type: `[u32; W * H * H]` (generic size [const expr])
+
+error[E0512]: cannot transmute between types of different sizes, or dependently-sized types
+  --> $DIR/transmute-fail.rs:30:5
+   |
+LL |     std::mem::transmute(v)
+   |     ^^^^^^^^^^^^^^^^^^^
+   |
+   = note: source type: `[[[u32; 8888888]; 9999999]; 777777777]` (values of the type `[[[u32; 8888888]; 9999999]; 777777777]` are too big for the current architecture)
+   = note: target type: `[[[u32; 9999999]; 777777777]; 8888888]` (values of the type `[[[u32; 9999999]; 777777777]; 8888888]` are too big for the current architecture)
+
+error: aborting due to 6 previous errors
+
+Some errors have detailed explanations: E0308, E0512.
+For more information about an error, try `rustc --explain E0308`.

tests/ui/const-generics/transmute.rs

@@ -0,0 +1,83 @@
+// run-pass
+#![feature(generic_const_exprs)]
+#![feature(transmute_generic_consts)]
+#![allow(incomplete_features)]
+
+fn transpose<const W: usize, const H: usize>(v: [[u32;H]; W]) -> [[u32; W]; H] {
+  unsafe {
+    std::mem::transmute(v)
+  }
+}
+
+fn ident<const W: usize, const H: usize>(v: [[u32; H]; W]) -> [[u32; H]; W] {
+  unsafe {
+    std::mem::transmute(v)
+  }
+}
+
+fn flatten<const W: usize, const H: usize>(v: [[u32; H]; W]) -> [u32; W * H] {
+  unsafe {
+    std::mem::transmute(v)
+  }
+}
+
+fn coagulate<const W: usize, const H: usize>(v: [u32; H*W]) -> [[u32; W];H] {
+  unsafe {
+    std::mem::transmute(v)
+  }
+}
+
+fn flatten_3d<const W: usize, const H: usize, const D: usize>(
+  v: [[[u32; D]; H]; W]
+) -> [u32; D * W * H] {
+  unsafe {
+    std::mem::transmute(v)
+  }
+}
+
+fn flatten_somewhat<const W: usize, const H: usize, const D: usize>(
+  v: [[[u32; D]; H]; W]
+) -> [[u32; D * W]; H] {
+  unsafe {
+    std::mem::transmute(v)
+  }
+}
+
+fn known_size<const L: usize>(v: [u16; L]) -> [u8; L * 2] {
+  unsafe {
+    std::mem::transmute(v)
+  }
+}
+
+fn condense_bytes<const L: usize>(v: [u8; L * 2]) -> [u16; L] {
+  unsafe {
+    std::mem::transmute(v)
+  }
+}
+
+fn singleton_each<const L: usize>(v: [u8; L]) -> [[u8;1]; L] {
+  unsafe {
+    std::mem::transmute(v)
+  }
+}
+
+fn transpose_with_const<const W: usize, const H: usize>(
+  v: [[u32; 2 * H]; W + W]
+) -> [[u32; W + W]; 2 * H] {
+  unsafe {
+    std::mem::transmute(v)
+  }
+}
+
+fn main() {
+  let _ = transpose([[0; 8]; 16]);
+  let _ = transpose_with_const::<8,4>([[0; 8]; 16]);
+  let _ = ident([[0; 8]; 16]);
+  let _ = flatten([[0; 13]; 5]);
+  let _: [[_; 5]; 13] = coagulate([0; 65]);
+  let _ = flatten_3d([[[0; 3]; 13]; 5]);
+  let _ = flatten_somewhat([[[0; 3]; 13]; 5]);
+  let _ = known_size([16; 13]);
+  let _: [u16; 5] = condense_bytes([16u8; 10]);
+  let _ = singleton_each([16; 10]);
+}