Rust - Find the address of a field implementing a trait and call it

dynfields.rs

#![allow(dead_code)]

use std::ops::Add;

fn main() {
    check_if_a::<_A>();
    check_if_a::<_B>();
    check_if_a::<_C>();

    type T1<'a> = Cons<'a, _A, Cons<'a, _B, Cons<'a, _C, ()>>>;
    type T2<'a> = Cons<'a, _B, Cons<'a, _A, Cons<'a, _C, ()>>>;
    type T3<'a> = Cons<'a, _C, Cons<'a, _B, Cons<'a, _A, ()>>>;

    let a = _A { inner: 1 };
    let b = _B { inner: 2 };
    let c = _C { inner: 3 };

    let t1: T1<'_> = empty_driver_stack()
        .push_driver(&c)
        .push_driver(&b)
        .push_driver(&a);
    let t2: T2<'_> = empty_driver_stack()
        .push_driver(&c)
        .push_driver(&a)
        .push_driver(&b);
    let t3: T3<'_> = empty_driver_stack()
        .push_driver(&a)
        .push_driver(&b)
        .push_driver(&c);

    println!();
    println!("T1 size = {}", std::mem::size_of::<T1<'_>>());
    println!("T2 size = {}", std::mem::size_of::<T2<'_>>());
    println!("T3 size = {}", std::mem::size_of::<T3<'_>>());

    println!();
    println!("Offset of _A in T1 => {}", <T1 as FindOffset<_A>>::OFFSET);
    println!("Offset of _B in T1 => {}", <T2 as FindOffset<_A>>::OFFSET);
    println!("Offset of _C in T1 => {}", <T3 as FindOffset<_A>>::OFFSET);

    println!();
    t1.head.a();
    t1.find::<_A>().a();

    println!();
    t2.tail.head.a();
    t2.find::<_A>().a();

    println!();
    t3.tail.tail.head.a();
    t3.find::<_A>().a();
}

fn check_if_a<T: WhichAmI>() {
    println!(
        "Is {} the type _A? {}",
        std::any::type_name::<T>(),
        <T as WhichAmI>::WHICH.same_as(Which::A)
    );
}

#[repr(C)]
struct Cons<'a, H, T> {
    head: &'a H,
    tail: T,
}

impl<H, T: Copy> Copy for Cons<'_, H, T> {}

impl<H, T: Clone> Clone for Cons<'_, H, T> {
    fn clone(&self) -> Self {
        Self {
            head: self.head,
            tail: self.tail.clone(),
        }
    }
}

trait FindOffset<WHICH> {
    const OFFSET: isize;

    fn find_offset(&self) -> &WHICH {
        let this: *const Self = self as *const _;
        let this: *const *const WHICH = unsafe { this.byte_offset(Self::OFFSET).cast() };
        let this: &&WHICH = unsafe { std::mem::transmute(this) };
        *this
    }
}

impl<WHICH, T> FindOffset<WHICH> for T
where
    WHICH: WhichAmI,
    T: AuxFindOffset<WHICH, typenum::Z0>,
{
    const OFFSET: isize = <T as AuxFindOffset<WHICH, typenum::Z0>>::OFFSET;
}

trait AuxFindOffset<WHICH, INDEX> {
    const OFFSET: isize;
}

impl<WHICH, INDEX> AuxFindOffset<WHICH, INDEX> for () {
    const OFFSET: isize = 0;
}

const USIZE_LEN: isize = std::mem::size_of::<usize>() as isize;

impl<WHICH, INDEX, H, T> AuxFindOffset<WHICH, INDEX> for Cons<'_, H, T>
where
    WHICH: WhichAmI,
    INDEX: typenum::Integer + Add<typenum::P1>,
    H: WhichAmI,
    T: AuxFindOffset<WHICH, <INDEX as Add<typenum::P1>>::Output>,
{
    const OFFSET: isize = (USIZE_LEN
        * <INDEX as typenum::Integer>::ISIZE
        * <H as WhichAmI>::WHICH.same_as(<WHICH as WhichAmI>::WHICH) as isize)
        + <T as AuxFindOffset<WHICH, <INDEX as Add<typenum::P1>>::Output>>::OFFSET;
}

trait WhichAmI {
    const WHICH: Which;
}

impl Which {
    const fn same_as(self, other: Self) -> bool {
        self as u8 == other as u8
    }
}

macro_rules! impl_all {
    ($($trait:ident($method:ident) for $type:ident($inner:ty))*) => {
        $( impl_!($trait($method) for $type($inner)); )*

        #[repr(u8)]
        enum Which {
            $( $trait ),*
        }
    };
}

macro_rules! impl_ {
    ($trait:ident($method:ident) for $type:ident($inner:ty)) => {
        struct $type {
            inner: $inner,
        }

        trait $trait {
            fn $method(&self);
        }

        impl WhichAmI for $type {
            const WHICH: Which = Which::$trait;
        }

        impl $trait for $type {
            fn $method(&self) {
                println!(
                    "{}:{} = {}",
                    stringify!($type),
                    stringify!($method),
                    self.inner
                );
            }
        }
    };
}

impl_all! {
    A(a) for _A(u8)
    B(b) for _B(i32)
    C(c) for _C(u64)
    D(d) for _D(i8)
    E(e) for _E(i128)
    F(f) for _F(i64)
}

const fn empty_driver_stack() {}

trait DriverStack: Sized {
    fn push_driver<H>(self, head: &H) -> Cons<'_, H, Self>;
}

impl DriverStack for () {
    fn push_driver<H>(self, head: &H) -> Cons<'_, H, Self> {
        Cons { head, tail: self }
    }
}

impl<H, T> DriverStack for Cons<'_, H, T> {
    fn push_driver<H0>(self, head: &H0) -> Cons<'_, H0, Self> {
        Cons { head, tail: self }
    }
}

impl<H, T> Cons<'_, H, T> {
    fn find<WHICH>(&self) -> &WHICH
    where
        Self: FindOffset<WHICH>,
    {
        self.find_offset()
    }
}