Add ability to transmute with generic consts

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.

Author: JulianKnodt

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")
+                }
+            }
             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,35 @@ impl<'tcx> SizeSkeleton<'tcx> {
                     ),
                 }
             }
+            ty::Array(inner, len) if len.ty() == tcx.types.usize => {
+                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 +390,9 @@ impl<'tcx> SizeSkeleton<'tcx> {
                                 }
                                 ptr = Some(field);
                             }
+                            SizeSkeleton::Generic(_) => {
+                                return Err(err);
+                            }
                         }
                     }
                     Ok(ptr)
@@ -410,11 +448,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;
+    for _ in done.drain_filter(|c| {
+        let Some(c) = c.try_eval_target_usize(tcx, param_env) else {
+            return false;
+        };
+        let Some(next) = c.checked_mul(k) else {
+            overflow = true;
+            return true;
+        };
+        k *= next;
+        true
+    }) {}
+    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>;
 }

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

@@ -0,0 +1,34 @@
+#![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:6:5
+   |
+LL |     std::mem::transmute(v)
+   |     ^^^^^^^^^^^^^^^^^^^
+   |
+   = note: source type: `[[u32; H+1]; W]` (generic size)
+   = note: target type: `[[u32; W+1]; H]` (generic size)
+
+error[E0512]: cannot transmute between types of different sizes, or dependently-sized types
+  --> $DIR/transmute-fail.rs:15: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:11: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:11: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:22:5
+   |
+LL |     std::mem::transmute(v)
+   |     ^^^^^^^^^^^^^^^^^^^
+   |
+   = note: source type: `[[u32; H]; W]` (generic size)
+   = note: target type: `[u32; W * H * H]` (generic size)
+
+error[E0512]: cannot transmute between types of different sizes, or dependently-sized types
+  --> $DIR/transmute-fail.rs:29: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,82 @@
+// run-pass
+#![feature(generic_const_exprs)]
+#![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]);
+}