Support unions in offset_of(_union)

src/offset_of.rs

@@ -121,6 +121,39 @@ macro_rules! offset_of_tuple {
     }};
 }
 
+/// Calculates the offset of the specified union member from the start of the union.
+///
+/// ## Examples
+/// ```
+/// use memoffset::offset_of_union;
+///
+/// #[repr(C, packed)]
+/// union Foo {
+///     foo32: i32,
+///     foo64: i64,
+/// }
+///
+/// fn main() {
+///     assert!(offset_of_union!(Foo, foo64) == 0);
+/// }
+/// ```
+///
+/// ## Note
+/// Due to macro_rules limitations, this check will accept structs with a single field as well as unions.
+/// Using this macro for a single-field struct will still work, but might lead to a future
+/// compatibility problem if the struct gains more fields.
+#[macro_export(local_inner_macros)]
+macro_rules! offset_of_union {
+    ($parent:path, $field:tt) => {{
+        // Get a base pointer (non-dangling if rustc supports `MaybeUninit`).
+        _memoffset__let_base_ptr!(base_ptr, $parent);
+        // Get field pointer.
+        let field_ptr = raw_field_union!(base_ptr, $parent, $field);
+        // Compute offset.
+        _memoffset_offset_from_unsafe!(field_ptr, base_ptr)
+    }};
+}
+
 #[cfg(test)]
 mod tests {
     #[test]
@@ -161,6 +194,21 @@ mod tests {
         assert_eq!(offset_of!(Tup, 1), 4);
     }
 
+    #[test]
+    fn offset_union() {
+        // Since we're specifying repr(C), all fields are supposed to be at offset 0
+        #[repr(C)]
+        union Foo {
+            a: u32,
+            b: [u8; 2],
+            c: i64,
+        }
+
+        assert_eq!(offset_of_union!(Foo, a), 0);
+        assert_eq!(offset_of_union!(Foo, b), 0);
+        assert_eq!(offset_of_union!(Foo, c), 0);
+    }
+
     #[test]
     fn path() {
         mod sub {
@@ -238,6 +286,22 @@ mod tests {
         );
     }
 
+    #[test]
+    fn test_raw_field_union() {
+        #[repr(C)]
+        union Foo {
+            a: u32,
+            b: [u8; 2],
+            c: i64,
+        }
+
+        let f = Foo { a: 0 };
+        let f_ptr = &f as *const _;
+        assert_eq!(f_ptr as usize + 0, raw_field_union!(f_ptr, Foo, a) as usize);
+        assert_eq!(f_ptr as usize + 0, raw_field_union!(f_ptr, Foo, b) as usize);
+        assert_eq!(f_ptr as usize + 0, raw_field_union!(f_ptr, Foo, c) as usize);
+    }
+
     #[cfg(feature = "unstable_const")]
     #[test]
     fn const_offset() {

src/raw_field.rs

@@ -84,6 +84,41 @@ macro_rules! _memoffset__field_check_tuple {
     };
 }
 
+/// Deref-coercion protection macro for unions.
+/// Unfortunately accepts single-field structs as well, which is not ideal,
+/// but ultimately pretty harmless.
+#[cfg(allow_clippy)]
+#[macro_export]
+#[doc(hidden)]
+macro_rules! _memoffset__field_check_union {
+    ($type:path, $field:tt) => {
+        // Make sure the field actually exists. This line ensures that a
+        // compile-time error is generated if $field is accessed through a
+        // Deref impl.
+        #[allow(clippy::unneeded_wildcard_pattern)]
+        // rustc1.19 requires unsafe here for the pattern; not needed in newer versions
+        #[allow(unused_unsafe)]
+        unsafe {
+            let $type { $field: _ };
+        }
+    };
+}
+#[cfg(not(allow_clippy))]
+#[macro_export]
+#[doc(hidden)]
+macro_rules! _memoffset__field_check_union {
+    ($type:path, $field:tt) => {
+        // Make sure the field actually exists. This line ensures that a
+        // compile-time error is generated if $field is accessed through a
+        // Deref impl.
+        // rustc1.19 requires unsafe here for the pattern; not needed in newer versions
+        #[allow(unused_unsafe)]
+        unsafe {
+            let $type { $field: _ };
+        }
+    };
+}
+
 /// Computes a const raw pointer to the given field of the given base pointer
 /// to the given parent type.
 ///
@@ -126,3 +161,31 @@ macro_rules! raw_field_tuple {
         }
     }};
 }
+
+/// Computes a const raw pointer to the given field of the given base pointer
+/// to the given parent tuple typle.
+///
+/// The `base` pointer *must not* be dangling, but it *may* point to
+/// uninitialized memory.
+///
+/// ## Note
+/// This macro is the same as `raw_field`, except for a different Deref-coercion check that
+/// supports unions.
+/// Due to macro_rules limitations, this check will accept structs with a single field as well as unions.
+/// Using this macro for a single-field struct will still work, but might lead to a future
+/// compatibility problem if the struct gains more fields.
+#[macro_export(local_inner_macros)]
+macro_rules! raw_field_union {
+    ($base:expr, $parent:path, $field:tt) => {{
+        _memoffset__field_check_union!($parent, $field);
+        let base = $base; // evaluate $base outside the `unsafe` block
+
+        // Get the field address.
+        // Crucially, we know that this will not trigger a deref coercion because
+        // of the field check we did above.
+        #[allow(unused_unsafe)] // for when the macro is used in an unsafe block
+        unsafe {
+            _memoffset__addr_of!((*(base as *const $parent)).$field)
+        }
+    }};
+}