fromsoftware_shared/
subclass.rs

1use std::ptr::NonNull;
2
3use pelite::pe64::{Pe, Rva, Va};
4use thiserror::Error;
5
6use crate::Program;
7
8/// The error type returend when a superclass isn't an instance of a subclass.
9#[derive(Error, Debug)]
10#[error("superclass is not an instance of {0}")]
11pub struct TryFromSuperclassError(String);
12
13impl TryFromSuperclassError {
14    /// Returns a [TryFromSuperclassError] for the given subclass name.
15    pub fn new(subclass: String) -> Self {
16        TryFromSuperclassError(subclass)
17    }
18}
19
20/// A trait for C++ types that have multiple different subclasses. Implementing
21/// this for a superclass and [Subclass] for its subclasses makes it possible to
22/// safely check for the object's actual type based on its vtable.
23///
24/// ## Safety
25///
26/// In order to implement this for a struct, you must guarantee:
27///
28/// * The struct uses C-style layout.
29/// * The first element of the struct is a pointer to a C++-style vtable.
30/// * There's a 1-to-1 correspondence between vtable address and C++ class.
31pub unsafe trait Superclass: Sized {
32    /// The RVA of this class's virtual method table.
33    fn vmt_rva() -> Rva;
34
35    /// The VA of this class's virtual method table.
36    fn vmt_va() -> Va {
37        Program::current()
38            .rva_to_va(Self::vmt_rva())
39            .expect("VMT address not in executable!")
40    }
41
42    /// Returns the [Va] for the runtime virtual method table for this.
43    fn vmt(&self) -> Va {
44        *unsafe { NonNull::from_ref(self).cast::<Va>().as_ref() }
45    }
46
47    /// Returns whether this is an instance of `T`.
48    ///
49    /// **Note:** Because this just checks the address of the virtual method
50    /// table, it will return `false` for *subclasses* of `T` even though C++
51    /// considers them to be of type `T`.
52    fn is_subclass<T: Subclass<Self>>(&self) -> bool {
53        self.vmt() == Self::vmt_va()
54    }
55
56    /// Returns this as a `T` if it is one. Otherwise, returns `None`.
57    ///
58    /// **Note:** Because this just checks the address of the virtual method
59    /// table, it will return `None` for *subclasses* of `T` even though C++
60    /// considers them to be of type `T`.
61    fn as_subclass<T: Subclass<Self>>(&self) -> Option<&T> {
62        if self.is_subclass::<T>() {
63            // Safety: We require that VMTs indicate object type.
64            Some(unsafe { NonNull::from_ref(self).cast::<T>().as_ref() })
65        } else {
66            None
67        }
68    }
69
70    /// Returns this as a mutable `T` if it is one. Otherwise, returns `None`.
71    ///
72    /// **Note:** Because this just checks the address of the virtual method
73    /// table, it will return `None` for *subclasses* of `T` even though C++
74    /// considers them to be of type `T`.
75    fn as_subclass_mut<T: Subclass<Self>>(&mut self) -> Option<&mut T> {
76        if self.is_subclass::<T>() {
77            // Safety: We require that VMTs indicate object type.
78            Some(unsafe { NonNull::from_ref(self).cast::<T>().as_mut() })
79        } else {
80            None
81        }
82    }
83}
84
85/// A trait for C++ subclasses of the superclass `T`. Implementing this trait
86/// makes it possible for Rust code to be generic over all subclasses of a given
87/// C++ supertype.
88///
89/// ## Safety
90///
91/// In order to implement this for a struct, you must guarantee:
92///
93/// * The struct uses C-style layout.
94/// * An initial subsequence of the struct is a valid isntance of `T`.
95pub unsafe trait Subclass<T: Superclass> {
96    /// The RVA of this class's virtual method table.
97    fn vmt_rva() -> Rva;
98
99    /// The VA of this class's virtual method table.
100    fn vmt_va() -> Va {
101        Program::current()
102            .rva_to_va(Self::vmt_rva())
103            .expect("VMT address not in executable!")
104    }
105
106    /// Returns this as a `T`.
107    fn superclass(&self) -> &T {
108        // The implementer has vouched that this type's struct layout begins
109        // with its superclass.
110        unsafe { NonNull::from_ref(self).cast::<T>().as_ref() }
111    }
112
113    /// Returns this as a mutable `T`.
114    fn superclass_mut(&mut self) -> &mut T {
115        // The implementer has vouched that this type's struct layout begins
116        // with its superclass.
117        unsafe { NonNull::from_ref(self).cast::<T>().as_mut() }
118    }
119}
120
121// Safety: Superclass has the same safety requirements as subclass.
122unsafe impl<T> Subclass<T> for T
123where
124    T: Superclass,
125{
126    fn vmt_rva() -> Rva {
127        <T as Superclass>::vmt_rva()
128    }
129}