From 819fad0d9f6bcded122b7c086a3c4a551bc6a824 Mon Sep 17 00:00:00 2001
From: "Adam C. Foltzer" <acfoltzer@fastly.com>
Date: Tue, 22 Aug 2023 13:13:40 -0700
Subject: [PATCH] wiggle: generate offset accessors for structs (#6884)

Adds methods of the shape `GeneratedStruct::offset_of_<field_name>() -> u32` to allow clients to get the
same offsets used in the generated `read` and `write` methods. I don't believe there is currently a
way to get these otherwise for applications that wish to selectively read and write to fields.

I considered a couple alternatives to this:

1. Add methods like `read_<field_name>` and `write_<field_name>`. The interface to these ends up
being awkward because the `GuestType` impls are currently for owned types, so these methods would be
forced to either take ownership of the entire struct to write one field, or else silently clone the
field.

2. Add `#[repr(C)]` to the generated struct definition so that clients could use tools like
`memoffset` to calculate the offset themselves. I didn't want to depend on the witx offsets
coinciding with `repr(C)`, though.

I extended the Wiggle `records` to exercise these new methods.
---
 crates/wiggle/generate/src/types/record.rs | 17 +++++++++++-
 crates/wiggle/tests/records.rs             | 31 ++++++++++++++++++++++
 crates/wiggle/tests/records.witx           |  7 +++++
 3 files changed, 54 insertions(+), 1 deletion(-)

diff --git a/crates/wiggle/generate/src/types/record.rs b/crates/wiggle/generate/src/types/record.rs
index c08090fd020d..167f317caa6b 100644
--- a/crates/wiggle/generate/src/types/record.rs
+++ b/crates/wiggle/generate/src/types/record.rs
@@ -2,7 +2,7 @@ use crate::lifetimes::{anon_lifetime, LifetimeExt};
 use crate::names;
 
 use proc_macro2::TokenStream;
-use quote::quote;
+use quote::{format_ident, quote};
 use witx::Layout;
 
 pub(super) fn define_struct(name: &witx::Id, s: &witx::RecordDatatype) -> TokenStream {
@@ -34,6 +34,17 @@ pub(super) fn define_struct(name: &witx::Id, s: &witx::RecordDatatype) -> TokenS
         quote!(pub #name: #type_)
     });
 
+    let member_offsets = s.member_layout().into_iter().map(|ml| {
+        let name = names::struct_member(&ml.member.name);
+        let offset = ml.offset as u32;
+        let method_name = format_ident!("offset_of_{}", name);
+        quote! {
+            pub const fn #method_name () -> u32 {
+                #offset
+            }
+        }
+    });
+
     let member_reads = s.member_layout().into_iter().map(|ml| {
         let name = names::struct_member(&ml.member.name);
         let offset = ml.offset as u32;
@@ -86,6 +97,10 @@ pub(super) fn define_struct(name: &witx::Id, s: &witx::RecordDatatype) -> TokenS
             #(#member_decls),*
         }
 
+        impl #struct_lifetime #ident #struct_lifetime {
+            #(#member_offsets)*
+        }
+
         impl<'a> wiggle::GuestType<'a> for #ident #struct_lifetime {
             #[inline]
             fn guest_size() -> u32 {
diff --git a/crates/wiggle/tests/records.rs b/crates/wiggle/tests/records.rs
index b33c50cd85fc..713015d27d6c 100644
--- a/crates/wiggle/tests/records.rs
+++ b/crates/wiggle/tests/records.rs
@@ -551,3 +551,34 @@ proptest! {
         e.test()
     }
 }
+
+#[test]
+fn pair_ints_offsets() {
+    assert_eq!(types::PairInts::offset_of_first(), 0);
+    assert_eq!(types::PairInts::offset_of_second(), 4);
+}
+
+#[test]
+fn pair_different_ints_offsets() {
+    assert_eq!(types::PairDifferentInts::offset_of_first(), 0);
+    assert_eq!(types::PairDifferentInts::offset_of_second(), 8);
+    assert_eq!(types::PairDifferentInts::offset_of_third(), 10);
+    assert_eq!(types::PairDifferentInts::offset_of_fourth(), 12);
+}
+
+#[test]
+fn pair_int_ptrs_offsets() {
+    assert_eq!(types::PairIntPtrs::offset_of_first(), 0);
+    assert_eq!(types::PairIntPtrs::offset_of_second(), 4);
+}
+
+#[test]
+fn pair_int_and_ptr_offsets() {
+    assert_eq!(types::PairIntAndPtr::offset_of_first(), 0);
+    assert_eq!(types::PairIntAndPtr::offset_of_second(), 4);
+}
+
+#[test]
+fn pair_record_of_list_offset() {
+    assert_eq!(types::RecordOfList::offset_of_arr(), 0);
+}
diff --git a/crates/wiggle/tests/records.witx b/crates/wiggle/tests/records.witx
index d09e37a38a51..aabc3dee6328 100644
--- a/crates/wiggle/tests/records.witx
+++ b/crates/wiggle/tests/records.witx
@@ -6,6 +6,13 @@
     (field $first s32)
     (field $second s32)))
 
+(typename $pair_different_ints
+  (record
+    (field $first s64)
+    (field $second s16)
+    (field $third s16)
+    (field $fourth s32)))
+
 (typename $pair_int_ptrs
   (record
     (field $first (@witx const_pointer s32))