1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218

use std::marker::PhantomData;
use std::mem;
use std::ops::Add;
use std::fmt;

/// Represents a pointer to a field of type `U` within the type `T`
pub struct FieldOffset<T, U>(
    /// Offset in bytes of the field within the struct
    usize,
    /// A pointer-to-member can be thought of as a function from
    /// `&T` to `&U` with matching lifetimes
    PhantomData<for<'a> Fn(&'a T) -> &'a U>
);

impl<T, U> FieldOffset<T, U> {
    /// Construct a field offset via a lambda which returns a reference
    /// to the field in question.
    ///
    /// The lambda *must not* access the value passed in.
    pub unsafe fn new<F: for<'a> FnOnce(&'a T) -> &'a U>(f: F) -> Self {
        // Construct a "fake" T. It's not valid, but the lambda shouldn't
        // actually access it (which is why this is unsafe)
        let x = mem::zeroed();
        let offset = {
            let x = &x;
            // Pass a reference to the zeroed T to the lambda
            // The lambda gives us back a reference to (what we hope is)
            // a field of T, of type U
            let y = f(x);
            // Compute the offset of the field via the difference between the
            // references `x` and `y`. Overflow is an error: in debug builds it
            // will be caught here, in release it will wrap around and be caught
            // on the next line.
            (y as *const U as usize) - (x as *const T as usize)
        };
        // Don't run destructor on "fake" T
        mem::forget(x);
        // Sanity check: ensure that the field offset plus the field size
        // is no greater than the size of the containing struct. This is
        // not sufficient to make the function *safe*, but it does catch
        // obvious errors like returning a reference to a boxed value,
        // which is owned by `T` and so has the correct lifetime, but is not
        // actually a field.
        assert!(offset + mem::size_of::<U>() <= mem::size_of::<T>());
        // Construct an instance using the offset
        Self::new_from_offset(offset)
    }
    /// Construct a field offset directly from a byte offset.
    pub unsafe fn new_from_offset(offset: usize) -> Self {
        FieldOffset(offset, PhantomData)
    }
    // Methods for applying the pointer to member
    /// Apply the field offset to a native pointer.
    pub fn apply_ptr<'a>(&self, x: *const T) -> *const U {
        ((x as usize) + self.0) as *const U
    }
    /// Apply the field offset to a native mutable pointer.
    pub fn apply_ptr_mut<'a>(&self, x: *mut T) -> *mut U {
        ((x as usize) + self.0) as *mut U
    }
    /// Apply the field offset to a reference.
    pub fn apply<'a>(&self, x: &'a T) -> &'a U {
        unsafe { &*self.apply_ptr(x) }
    }
    /// Apply the field offset to a mutable reference.
    pub fn apply_mut<'a>(&self, x: &'a mut T) -> &'a mut U {
        unsafe { &mut *self.apply_ptr_mut(x) }
    }
    /// Get the raw byte offset for this field offset.
    pub fn get_byte_offset(&self) -> usize {
        self.0
    }
    // Methods for unapplying the pointer to member
    /// Unapply the field offset to a native pointer.
    ///
    /// *Warning: very unsafe!*
    pub unsafe fn unapply_ptr<'a>(&self, x: *const U) -> *const T {
        ((x as usize) - self.0) as *const T
    }
    /// Unapply the field offset to a native mutable pointer.
    ///
    /// *Warning: very unsafe!*
    pub unsafe fn unapply_ptr_mut<'a>(&self, x: *mut U) -> *mut T {
        ((x as usize) - self.0) as *mut T
    }
    /// Unapply the field offset to a reference.
    ///
    /// *Warning: very unsafe!*
    pub unsafe fn unapply<'a>(&self, x: &'a U) -> &'a T {
        &*self.unapply_ptr(x)
    }
    /// Unapply the field offset to a mutable reference.
    ///
    /// *Warning: very unsafe!*
    pub unsafe fn unapply_mut<'a>(&self, x: &'a mut U) -> &'a mut T {
        &mut *self.unapply_ptr_mut(x)
    }
}

/// Allow chaining pointer-to-members.
///
/// Applying the resulting field offset is equivalent to applying the first
/// field offset, then applying the second field offset.
///
/// The requirements on the generic type parameters ensure this is a safe operation.
impl<T, U, V> Add<FieldOffset<U, V>> for FieldOffset<T, U> {
    type Output = FieldOffset<T, V>;

    fn add(self, other: FieldOffset<U, V>) -> FieldOffset<T, V> {
        FieldOffset(self.0 + other.0, PhantomData)
    }
}

/// The debug implementation prints the byte offset of the field in hexadecimal.
impl<T, U> fmt::Debug for FieldOffset<T, U> {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        write!(f, "FieldOffset({:#x})", self.0)
    }
}

impl<T, U> Copy for FieldOffset<T, U> { }
impl<T, U> Clone for FieldOffset<T, U> {
    fn clone(&self) -> Self { *self }
}

/// This macro allows safe construction of a FieldOffset,
/// by generating a known to be valid lambda to pass to the
/// constructor. It takes a type and the identifier of a field
/// within that type as input.
///
/// Examples:
///
/// Offset of field `Foo().bar`
///
/// `offset_of!(Foo => bar)`
///
/// Offset of nested field `Foo().bar.x`
///
/// `offset_of!(Foo => bar: Bar => x)`
#[macro_export]
macro_rules! offset_of {
    ($t: path => $f: ident) => {
        unsafe { $crate::FieldOffset::<$t, _>::new(|x| {
            let $t { ref $f, .. } = *x;
            $f
        }) }
    };
    ($t: path => $f: ident: $($rest: tt)*) => {
        offset_of!($t => $f) + offset_of!($($rest)*)
    };
}

#[cfg(test)]
mod tests {
    // Example structs
    #[derive(Debug)]
    struct Foo {
        a: u32,
        b: f64,
        c: bool
    }

    #[derive(Debug)]
    struct Bar {
        x: u32,
        y: Foo,
    }

    #[test]
    fn test_simple() {
        // Get a pointer to `b` within `Foo`
        let foo_b = offset_of!(Foo => b);

        // Construct an example `Foo`
        let mut x = Foo {
            a: 1,
            b: 2.0,
            c: false
        };

        // Apply the pointer to get at `b` and read it
        {
            let y = foo_b.apply(&x);
            assert!(*y == 2.0);
        }

        // Apply the pointer to get at `b` and mutate it
        {
            let y = foo_b.apply_mut(&mut x);
            *y = 42.0;
        }
        assert!(x.b == 42.0);
    }

    #[test]
    fn test_nested() {
        // Construct an example `Foo`
        let mut x = Bar {
            x: 0,
            y: Foo {
                a: 1,
                b: 2.0,
                c: false
            }
        };

        // Combine the pointer-to-members
        let bar_y_b = offset_of!(Bar => y: Foo => b);

        // Apply the pointer to get at `b` and mutate it
        {
            let y = bar_y_b.apply_mut(&mut x);
            *y = 42.0;
        }
        assert!(x.y.b == 42.0);
    }
}