Last active
May 13, 2025 15:23
-
-
Save hsqStephenZhang/82d975dda35de22a83a4d6ac5973e42c to your computer and use it in GitHub Desktop.
hacks of rust vtable layout(multi traits)
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| use std::collections::{HashMap, HashSet}; | |
| #[derive(Debug, Clone, PartialEq, Eq, Hash)] | |
| struct TraitId(&'static str); // e.g., "A", "B", ... | |
| #[derive(Debug)] | |
| struct TraitDef { | |
| id: TraitId, | |
| supertraits: Vec<TraitId>, | |
| methods: Vec<&'static str>, | |
| } | |
| #[allow(unused)] | |
| #[derive(Debug)] | |
| struct TypeImpl { | |
| type_name: &'static str, | |
| implemented_traits: HashSet<TraitId>, | |
| } | |
| fn compute_prefix_set(traits: &HashMap<TraitId, TraitDef>, root: &TraitId) -> Vec<TraitId> { | |
| let mut result = Vec::new(); | |
| let mut current = root.clone(); | |
| loop { | |
| result.push(current.clone()); | |
| if let Some(trait_def) = traits.get(¤t) { | |
| if let Some(first_super) = trait_def.supertraits.get(0) { | |
| current = first_super.clone(); | |
| } else { | |
| break; | |
| } | |
| } else { | |
| break; | |
| } | |
| } | |
| result | |
| } | |
| fn postorder( | |
| traits: &HashMap<TraitId, TraitDef>, | |
| root: &TraitId, | |
| visited: &mut HashSet<TraitId>, | |
| out: &mut Vec<TraitId>, | |
| ) { | |
| if visited.contains(root) { | |
| return; | |
| } | |
| visited.insert(root.clone()); | |
| if let Some(def) = traits.get(root) { | |
| for sup in &def.supertraits { | |
| postorder(traits, sup, visited, out); | |
| } | |
| } | |
| out.push(root.clone()); | |
| } | |
| #[allow(unused)] | |
| #[derive(Debug)] | |
| enum VTableEntry { | |
| DropInPlace, | |
| Size, | |
| Align, | |
| Method { | |
| trait_id: TraitId, | |
| name: &'static str, | |
| }, | |
| TraitVPtr { | |
| trait_id: TraitId, | |
| }, // used for upcasting | |
| } | |
| #[allow(unused)] | |
| fn generate_vtable( | |
| traits: &HashMap<TraitId, TraitDef>, | |
| type_impl: &TypeImpl, | |
| root_trait: &TraitId, | |
| ) -> Vec<VTableEntry> { | |
| let mut layout = vec![ | |
| VTableEntry::DropInPlace, | |
| VTableEntry::Size, | |
| VTableEntry::Align, | |
| ]; | |
| let prefix_set: HashSet<TraitId> = compute_prefix_set(traits, root_trait).into_iter().collect(); | |
| let mut postorder_list = Vec::new(); | |
| let mut visited = HashSet::new(); | |
| postorder(traits, root_trait, &mut visited, &mut postorder_list); | |
| for tid in postorder_list { | |
| if let Some(def) = traits.get(&tid) { | |
| for method in &def.methods { | |
| layout.push(VTableEntry::Method { | |
| trait_id: tid.clone(), | |
| name: *method, | |
| }); | |
| } | |
| if !prefix_set.contains(&tid) { | |
| layout.push(VTableEntry::TraitVPtr { trait_id: tid }); | |
| } | |
| } | |
| } | |
| layout | |
| } | |
| #[test] | |
| fn hack() { | |
| let traits: HashMap<_, _> = vec![ | |
| TraitDef { | |
| id: TraitId("A"), | |
| supertraits: vec![], | |
| methods: vec!["foo_a"], | |
| }, | |
| TraitDef { | |
| id: TraitId("B"), | |
| supertraits: vec![TraitId("A")], | |
| methods: vec!["foo_b"], | |
| }, | |
| TraitDef { | |
| id: TraitId("C"), | |
| supertraits: vec![TraitId("A")], | |
| methods: vec!["foo_c"], | |
| }, | |
| TraitDef { | |
| id: TraitId("D"), | |
| supertraits: vec![TraitId("B"), TraitId("C")], | |
| methods: vec!["foo_d"], | |
| }, | |
| TraitDef { | |
| id: TraitId("E"), | |
| supertraits: vec![TraitId("D")], | |
| methods: vec!["foo_e"], | |
| }, | |
| ] | |
| .into_iter() | |
| .map(|t| (t.id.clone(), t)) | |
| .collect(); | |
| let type_impl = TypeImpl { | |
| type_name: "S", | |
| implemented_traits: traits.keys().cloned().collect(), | |
| }; | |
| let vtable = generate_vtable(&traits, &type_impl, &TraitId("E")); | |
| for (i, entry) in vtable.iter().enumerate() { | |
| println!("{:>2}: {:?}", i, entry); | |
| } | |
| } | |
| #[test] | |
| fn actual() { | |
| // E | |
| // | | |
| // D | |
| // / \ | |
| // B C | |
| // \ / | |
| // A | |
| trait A { | |
| fn foo_a(&self); | |
| } | |
| trait B: A { | |
| fn foo_b(&self); | |
| } | |
| trait C: A { | |
| fn foo_c(&self); | |
| } | |
| trait D: B + C { | |
| fn foo_d(&self); | |
| } | |
| trait E: D { | |
| fn foo_e(&self); | |
| } | |
| struct S; | |
| impl A for S { | |
| fn foo_a(&self) { | |
| println!("original foo_a"); | |
| } | |
| } | |
| impl B for S { | |
| fn foo_b(&self) { | |
| println!("original foo_b"); | |
| } | |
| } | |
| impl C for S { | |
| fn foo_c(&self) { | |
| println!("original foo_c"); | |
| } | |
| } | |
| impl D for S { | |
| fn foo_d(&self) { | |
| println!("original foo_d"); | |
| } | |
| } | |
| impl E for S { | |
| fn foo_e(&self) { | |
| println!("original foo_e"); | |
| } | |
| } | |
| use std::mem; | |
| use std::collections::HashMap; | |
| fn collect_impl_methods() -> HashMap<*const (), &'static str> { | |
| let mut map = HashMap::new(); | |
| map.insert(S::foo_a as *const (), "S::foo_a"); | |
| map.insert(S::foo_b as *const (), "S::foo_b"); | |
| map.insert(S::foo_c as *const (), "S::foo_c"); | |
| map.insert(S::foo_d as *const (), "S::foo_d"); | |
| map.insert(S::foo_e as *const (), "S::foo_e"); | |
| map | |
| } | |
| fn collect_dyn_vtable() -> HashMap<*const *const (), &'static str> { | |
| let mut table = HashMap::new(); | |
| unsafe { | |
| let obj = Box::new(S) as Box<dyn A>; | |
| let obj: &dyn A = &*obj; | |
| let (_data_ptr, vtable_ptr): (*const (), *const *const ()) = mem::transmute(obj); | |
| let vtable = vtable_ptr; | |
| table.insert(vtable, "traitVptr A"); | |
| } | |
| unsafe { | |
| let obj = Box::new(S) as Box<dyn B>; | |
| let obj: &dyn B = &*obj; | |
| let (_data_ptr, vtable_ptr): (*const (), *const *const ()) = mem::transmute(obj); | |
| let vtable = vtable_ptr; | |
| table.insert(vtable, "traitVptr B"); | |
| } | |
| unsafe { | |
| let obj = Box::new(S) as Box<dyn C>; | |
| let obj: &dyn C = &*obj; | |
| let (_data_ptr, vtable_ptr): (*const (), *const *const ()) = mem::transmute(obj); | |
| let vtable = vtable_ptr; | |
| table.insert(vtable, "traitVptr C"); | |
| } | |
| unsafe { | |
| let obj = Box::new(S) as Box<dyn D>; | |
| let obj: &dyn D = &*obj; | |
| let (_data_ptr, vtable_ptr): (*const (), *const *const ()) = mem::transmute(obj); | |
| let vtable = vtable_ptr; | |
| table.insert(vtable, "traitVptr D"); | |
| } | |
| unsafe { | |
| let obj = Box::new(S) as Box<dyn E>; | |
| let obj: &dyn E = &*obj; | |
| let (_data_ptr, vtable_ptr): (*const (), *const *const ()) = mem::transmute(obj); | |
| let vtable = vtable_ptr; | |
| table.insert(vtable, "traitVptr E"); | |
| } | |
| table | |
| } | |
| fn print_vtable( | |
| obj: &dyn E, | |
| methods_info: &HashMap<*const (), &'static str>, | |
| vtable_info: &HashMap<*const *const (), &'static str>, | |
| ) { | |
| unsafe { | |
| let obj = obj as *const dyn E; | |
| let (data_ptr, vtable_ptr): (*const (), *const *const ()) = mem::transmute(obj); | |
| println!("Data ptr: {:p}", data_ptr); | |
| println!("VTable ptr: {:p}", vtable_ptr); | |
| // 打印 vtable 前若干项 | |
| println!("--- VTable Layout ---"); | |
| for i in 0..9 { | |
| let ptr = *vtable_ptr.add(i); | |
| println!("[{}] {:p}, method:{:?}", i, ptr, methods_info.get(&ptr).or(vtable_info.get(&(ptr as *const *const ())))); | |
| } | |
| } | |
| } | |
| let s = Box::new(S) as Box<dyn E>; | |
| let obj: &dyn E = &*s; | |
| let methods_info = collect_impl_methods(); | |
| let vtable_info = collect_dyn_vtable(); | |
| print_vtable(obj, &methods_info, &vtable_info); | |
| } |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment