mlua/luau/
heap_dump.rs

1use std::collections::HashMap;
2use std::hash::Hash;
3use std::mem;
4use std::os::raw::c_char;
5
6use crate::state::ExtraData;
7
8use super::json::{self, Json};
9
10/// Represents a heap dump of a Luau memory state.
11#[cfg(any(feature = "luau", doc))]
12#[cfg_attr(docsrs, doc(cfg(feature = "luau")))]
13pub struct HeapDump {
14    data: Json<'static>, // refers to the contents of `buf`
15    buf: Box<str>,
16}
17
18impl HeapDump {
19    /// Dumps the current Lua heap state.
20    pub(crate) unsafe fn new(state: *mut ffi::lua_State) -> Option<Self> {
21        unsafe extern "C" fn category_name(state: *mut ffi::lua_State, cat: u8) -> *const c_char {
22            (&*ExtraData::get(state))
23                .mem_categories
24                .get(cat as usize)
25                .map(|s| s.as_ptr())
26                .unwrap_or(cstr!("unknown"))
27        }
28
29        let mut buf = Vec::new();
30        unsafe {
31            let file = libc::tmpfile();
32            if file.is_null() {
33                return None;
34            }
35            ffi::lua_gcdump(state, file as *mut _, Some(category_name));
36            libc::fseek(file, 0, libc::SEEK_END);
37            let len = libc::ftell(file) as usize;
38            libc::rewind(file);
39            if len > 0 {
40                buf.reserve(len);
41                libc::fread(buf.as_mut_ptr() as *mut _, 1, len, file);
42                buf.set_len(len);
43            }
44            libc::fclose(file);
45        }
46
47        let buf = String::from_utf8(buf).ok()?.into_boxed_str();
48        let data = json::parse(unsafe { mem::transmute::<&str, &'static str>(&buf) }).ok()?;
49        Some(HeapDump { data, buf })
50    }
51
52    /// Returns the raw JSON representation of the heap dump.
53    ///
54    /// The JSON structure is an internal detail and may change in future versions.
55    #[doc(hidden)]
56    pub fn to_json(&self) -> &str {
57        &self.buf
58    }
59
60    /// Returns the total size of the Lua heap in bytes.
61    pub fn size(&self) -> u64 {
62        self.data["stats"]["size"].as_u64().unwrap_or_default()
63    }
64
65    /// Returns a mapping from object type to (count, total size in bytes).
66    ///
67    /// If `category` is provided, only objects in that category are considered.
68    pub fn size_by_type<'a>(&'a self, category: Option<&str>) -> HashMap<&'a str, (usize, u64)> {
69        self.size_by_type_inner(category).unwrap_or_default()
70    }
71
72    fn size_by_type_inner<'a>(&'a self, category: Option<&str>) -> Option<HashMap<&'a str, (usize, u64)>> {
73        let category_id = match category {
74            // If we cannot find the category, return empty result
75            Some(cat) => Some(self.find_category_id(cat)?),
76            None => None,
77        };
78
79        let mut size_by_type = HashMap::new();
80        let objects = self.data["objects"].as_object()?;
81        for obj in objects.values() {
82            if let Some(cat_id) = category_id
83                && obj["cat"].as_i64()? != cat_id
84            {
85                continue;
86            }
87            update_size(&mut size_by_type, obj["type"].as_str()?, obj["size"].as_u64()?);
88        }
89        Some(size_by_type)
90    }
91
92    /// Returns a mapping from category name to total size in bytes.
93    pub fn size_by_category(&self) -> HashMap<&str, u64> {
94        let mut size_by_category = HashMap::new();
95        if let Some(categories) = self.data["stats"]["categories"].as_object() {
96            for cat in categories.values() {
97                if let Some(cat_name) = cat["name"].as_str() {
98                    size_by_category.insert(cat_name, cat["size"].as_u64().unwrap_or_default());
99                }
100            }
101        }
102        size_by_category
103    }
104
105    /// Returns a mapping from userdata type to (count, total size in bytes).
106    pub fn size_by_userdata<'a>(&'a self, category: Option<&str>) -> HashMap<&'a str, (usize, u64)> {
107        self.size_by_userdata_inner(category).unwrap_or_default()
108    }
109
110    fn size_by_userdata_inner<'a>(
111        &'a self,
112        category: Option<&str>,
113    ) -> Option<HashMap<&'a str, (usize, u64)>> {
114        let category_id = match category {
115            // If we cannot find the category, return empty result
116            Some(cat) => Some(self.find_category_id(cat)?),
117            None => None,
118        };
119
120        let mut size_by_userdata = HashMap::new();
121        let objects = self.data["objects"].as_object()?;
122        for obj in objects.values() {
123            if obj["type"] != "userdata" {
124                continue;
125            }
126            if let Some(cat_id) = category_id
127                && obj["cat"].as_i64()? != cat_id
128            {
129                continue;
130            }
131
132            // Determine userdata type from metatable
133            let mut ud_type = "unknown";
134            if let Some(metatable_addr) = obj["metatable"].as_str()
135                && let Some(t) = get_key(objects, &objects[metatable_addr], "__type")
136            {
137                ud_type = t;
138            }
139            update_size(&mut size_by_userdata, ud_type, obj["size"].as_u64()?);
140        }
141        Some(size_by_userdata)
142    }
143
144    /// Finds the category ID for a given category name.
145    fn find_category_id(&self, category: &str) -> Option<i64> {
146        let categories = self.data["stats"]["categories"].as_object()?;
147        for (cat_id, cat) in categories {
148            if cat["name"].as_str() == Some(category) {
149                return cat_id.parse().ok();
150            }
151        }
152        None
153    }
154}
155
156/// Updates the size mapping for a given key.
157fn update_size<K: Eq + Hash>(size_type: &mut HashMap<K, (usize, u64)>, key: K, size: u64) {
158    let (count, total_size) = size_type.entry(key).or_insert((0, 0));
159    *count += 1;
160    *total_size += size;
161}
162
163/// Retrieves the value associated with a given `key` from a Lua table `tbl`.
164fn get_key<'a>(objects: &'a HashMap<&'a str, Json>, tbl: &Json, key: &str) -> Option<&'a str> {
165    let pairs = tbl["pairs"].as_array()?;
166    for kv in pairs.chunks_exact(2) {
167        #[rustfmt::skip]
168        let (Some(key_addr), Some(val_addr)) = (kv[0].as_str(), kv[1].as_str()) else { continue; };
169        if objects[key_addr]["type"] == "string" && objects[key_addr]["data"].as_str() == Some(key) {
170            if objects[val_addr]["type"] == "string" {
171                return objects[val_addr]["data"].as_str();
172            } else {
173                break;
174            }
175        }
176    }
177    None
178}
mlua/luau/heap_dump.rs

mlua/luau/
heap_dump.rs
