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

compiler/rustc_feature/src/active.rs

@@ -516,6 +516,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

@@ -1493,6 +1493,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]);
+}