Skip to content

Commit 92dd53e

Browse files
committed
fix weak table
1 parent f4fba79 commit 92dd53e

File tree

1 file changed

+210
-37
lines changed

1 file changed

+210
-37
lines changed

crates/luars/src/gc/mod.rs

Lines changed: 210 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -331,6 +331,8 @@ impl GC {
331331
while let Some(gc_id) = self.grayagain.pop() {
332332
self.mark_one(gc_id, pool);
333333
}
334+
// Clear weak table entries before sweep
335+
self.clear_weak_tables(pool);
334336
// Start sweep
335337
self.sweep_index = 0;
336338
self.state = GcState::Sweep;
@@ -505,28 +507,54 @@ impl GC {
505507

506508
match gc_id.gc_type() {
507509
GcObjectType::Table => {
508-
if let Some(table) = pool.tables.get_mut(gc_id.index()) {
509-
if table.header.is_gray() {
510+
// First pass: collect all needed info without mutating
511+
let (should_mark, entries, mt_value, weak_mode) = {
512+
if let Some(table) = pool.tables.get(gc_id.index()) {
513+
if table.header.is_gray() {
514+
let entries = table.data.iter_all();
515+
let mt = table.data.get_metatable();
516+
517+
// Check weak mode from metatable
518+
let weak = if let Some(mt_id) = mt.and_then(|v| v.as_table_id()) {
519+
if let Some(mt_table) = pool.tables.get(mt_id.0) {
520+
self.get_weak_mode(mt_table, pool)
521+
} else {
522+
None
523+
}
524+
} else {
525+
None
526+
};
527+
528+
(true, entries, mt, weak)
529+
} else {
530+
(false, vec![], None, None)
531+
}
532+
} else {
533+
(false, vec![], None, None)
534+
}
535+
};
536+
537+
if should_mark {
538+
// Now mark the table as black
539+
if let Some(table) = pool.tables.get_mut(gc_id.index()) {
510540
table.header.make_black();
511541
work += table.data.len();
512-
513-
// Collect references to mark
514-
let refs: Vec<LuaValue> = table
515-
.data
516-
.iter_all()
517-
.into_iter()
518-
.flat_map(|(k, v)| [k, v])
519-
.collect();
520-
let mt = table.data.get_metatable();
521-
522-
// Mark references
523-
for v in refs {
524-
self.mark_value(&v, pool);
542+
}
543+
544+
let (weak_keys, weak_values) = weak_mode.unwrap_or((false, false));
545+
546+
// Mark references (skip weak references)
547+
for (k, v) in entries {
548+
if !weak_keys {
549+
self.mark_value(&k, pool);
525550
}
526-
if let Some(mt) = mt {
527-
self.mark_value(&mt, pool);
551+
if !weak_values {
552+
self.mark_value(&v, pool);
528553
}
529554
}
555+
if let Some(mt) = mt_value {
556+
self.mark_value(&mt, pool);
557+
}
530558
}
531559
}
532560
GcObjectType::Function => {
@@ -1077,7 +1105,10 @@ impl GC {
10771105
// Phase 2: Mark from roots
10781106
self.mark_roots(roots, pool);
10791107

1080-
// Phase 3: Sweep (only traverse allgc, not entire pools!)
1108+
// Phase 3: Clear weak table entries (before sweep!)
1109+
self.clear_weak_tables(pool);
1110+
1111+
// Phase 4: Sweep (only traverse allgc, not entire pools!)
10811112
let collected = self.sweep(pool);
10821113

10831114
// Like Lua's setpause: set debt based on memory and pause factor
@@ -1141,32 +1172,64 @@ impl GC {
11411172
match value.kind() {
11421173
crate::lua_value::LuaValueKind::Table => {
11431174
if let Some(id) = value.as_table_id() {
1144-
if let Some(table) = pool.tables.get_mut(id.0) {
1145-
// For fixed tables: traverse if not yet traversed
1146-
// For non-fixed tables: traverse if white (not yet marked)
1147-
let is_fixed = table.header.is_fixed();
1148-
let should_traverse = if is_fixed {
1149-
traversed_fixed.insert(id.0) // Returns true if newly inserted
1175+
// First pass: collect info without mutating
1176+
let (should_traverse, is_fixed, entries, mt_value, weak_mode) = {
1177+
if let Some(table) = pool.tables.get(id.0) {
1178+
let is_fixed = table.header.is_fixed();
1179+
let should = if is_fixed {
1180+
!traversed_fixed.contains(&id.0)
1181+
} else {
1182+
table.header.is_white()
1183+
};
1184+
1185+
if should {
1186+
let entries = table.data.iter_all();
1187+
let mt = table.data.get_metatable();
1188+
1189+
// Check weak mode from metatable
1190+
let weak = if let Some(mt_id) = mt.and_then(|v| v.as_table_id()) {
1191+
if let Some(mt_table) = pool.tables.get(mt_id.0) {
1192+
self.get_weak_mode(mt_table, pool)
1193+
} else {
1194+
None
1195+
}
1196+
} else {
1197+
None
1198+
};
1199+
1200+
(true, is_fixed, entries, mt, weak)
1201+
} else {
1202+
(false, is_fixed, vec![], None, None)
1203+
}
11501204
} else {
1151-
table.header.is_white()
1152-
};
1153-
1154-
if should_traverse {
1155-
// Mark non-fixed tables as black
1156-
if !is_fixed {
1205+
(false, false, vec![], None, None)
1206+
}
1207+
};
1208+
1209+
if should_traverse {
1210+
// Mark the traversal
1211+
if is_fixed {
1212+
traversed_fixed.insert(id.0);
1213+
} else {
1214+
if let Some(table) = pool.tables.get_mut(id.0) {
11571215
table.header.make_black();
11581216
}
1159-
// Add table contents to worklist
1160-
let entries = table.data.iter_all();
1161-
let mt = table.data.get_metatable();
1162-
for (k, v) in entries {
1217+
}
1218+
1219+
let (weak_keys, weak_values) = weak_mode.unwrap_or((false, false));
1220+
1221+
// Add table contents to worklist (skip weak references)
1222+
for (k, v) in entries {
1223+
if !weak_keys {
11631224
worklist.push(k);
1164-
worklist.push(v);
11651225
}
1166-
if let Some(mt) = mt {
1167-
worklist.push(mt);
1226+
if !weak_values {
1227+
worklist.push(v);
11681228
}
11691229
}
1230+
if let Some(mt) = mt_value {
1231+
worklist.push(mt);
1232+
}
11701233
}
11711234
}
11721235
}
@@ -1252,6 +1315,116 @@ impl GC {
12521315
}
12531316
}
12541317

1318+
/// Check if a LuaValue points to a dead (white/unmarked) GC object
1319+
fn is_value_dead(&self, value: &LuaValue, pool: &ObjectPool) -> bool {
1320+
use crate::lua_value::LuaValueKind;
1321+
match value.kind() {
1322+
LuaValueKind::Table => {
1323+
if let Some(id) = value.as_table_id() {
1324+
if let Some(t) = pool.tables.get(id.0) {
1325+
return !t.header.is_fixed() && t.header.is_white();
1326+
}
1327+
}
1328+
false
1329+
}
1330+
LuaValueKind::Function => {
1331+
if let Some(id) = value.as_function_id() {
1332+
if let Some(f) = pool.functions.get(id.0) {
1333+
return !f.header.is_fixed() && f.header.is_white();
1334+
}
1335+
}
1336+
false
1337+
}
1338+
LuaValueKind::Thread => {
1339+
if let Some(id) = value.as_thread_id() {
1340+
if let Some(t) = pool.threads.get(id.0) {
1341+
return !t.header.is_fixed() && t.header.is_white();
1342+
}
1343+
}
1344+
false
1345+
}
1346+
LuaValueKind::String => {
1347+
if let Some(id) = value.as_string_id() {
1348+
if let Some(s) = pool.strings.get(id.0) {
1349+
return !s.header.is_fixed() && s.header.is_white();
1350+
}
1351+
}
1352+
false
1353+
}
1354+
// Numbers, booleans, nil, CFunction, Userdata - not GC managed or always live
1355+
_ => false,
1356+
}
1357+
}
1358+
1359+
/// Clear weak table entries that point to dead (white) objects
1360+
/// This must be called after marking but before sweeping
1361+
fn clear_weak_tables(&self, pool: &mut ObjectPool) {
1362+
// Collect tables with weak mode and their entries to remove
1363+
let mut tables_to_clear: Vec<(u32, Vec<LuaValue>)> = Vec::new();
1364+
1365+
for (id, table) in pool.tables.iter() {
1366+
// Check if this table has a metatable with __mode
1367+
if let Some(mt_id) = table.data.get_metatable().and_then(|v| v.as_table_id()) {
1368+
if let Some(mt) = pool.tables.get(mt_id.0) {
1369+
// Look for __mode key in metatable
1370+
// We need to find the string "__mode" in the metatable
1371+
let mode = self.get_weak_mode(mt, pool);
1372+
if let Some((weak_keys, weak_values)) = mode {
1373+
// Collect keys to remove
1374+
let mut keys_to_remove = Vec::new();
1375+
for (key, value) in table.data.iter_all() {
1376+
let key_dead = weak_keys && self.is_value_dead(&key, pool);
1377+
let value_dead = weak_values && self.is_value_dead(&value, pool);
1378+
if key_dead || value_dead {
1379+
keys_to_remove.push(key);
1380+
}
1381+
}
1382+
if !keys_to_remove.is_empty() {
1383+
tables_to_clear.push((id, keys_to_remove));
1384+
}
1385+
}
1386+
}
1387+
}
1388+
}
1389+
1390+
// Now actually remove the entries
1391+
for (table_id, keys) in tables_to_clear {
1392+
if let Some(table) = pool.tables.get_mut(table_id) {
1393+
for key in keys {
1394+
table.data.raw_set(key, LuaValue::nil());
1395+
}
1396+
}
1397+
}
1398+
}
1399+
1400+
/// Get weak mode from metatable's __mode field
1401+
/// Returns Some((weak_keys, weak_values)) if __mode is set
1402+
fn get_weak_mode(&self, metatable: &GcTable, pool: &ObjectPool) -> Option<(bool, bool)> {
1403+
// Find __mode key - it should be a string "k", "v", or "kv" (or "vk")
1404+
for (key, value) in metatable.data.iter_all() {
1405+
// Check if key is the string "__mode"
1406+
if let Some(key_str_id) = key.as_string_id() {
1407+
if let Some(key_str) = pool.strings.get(key_str_id.0) {
1408+
if key_str.data.as_str() == "__mode" {
1409+
// Found __mode, now check the value
1410+
if let Some(val_str_id) = value.as_string_id() {
1411+
if let Some(val_str) = pool.strings.get(val_str_id.0) {
1412+
let mode_str = val_str.data.as_str();
1413+
let weak_keys = mode_str.contains('k');
1414+
let weak_values = mode_str.contains('v');
1415+
if weak_keys || weak_values {
1416+
return Some((weak_keys, weak_values));
1417+
}
1418+
}
1419+
}
1420+
return None;
1421+
}
1422+
}
1423+
}
1424+
}
1425+
None
1426+
}
1427+
12551428
/// Sweep phase - iterate pools directly instead of allgc
12561429
/// This is much faster for allocation (no allgc.push) at cost of sweep traversal
12571430
fn sweep(&mut self, pool: &mut ObjectPool) -> usize {

0 commit comments

Comments
 (0)