1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
//! Deoptimization tracking for JIT-compiled functions.
//!
//! Tracks which JIT-compiled functions depend on specific shape IDs,
//! so that when a shape transitions (e.g., a HashMap gains a property),
//! all functions that guarded on that shape can be invalidated.
use std::collections::{HashMap, HashSet};
use shape_value::shape_graph::ShapeId;
/// Tracks shape dependencies for JIT-compiled functions.
///
/// When a function is compiled with shape guards (e.g., guarding that an
/// object has shape X for inline caching), the shape IDs it depends on
/// are registered here. When a shape transition occurs, all functions
/// that depend on the transitioning shape are invalidated.
pub struct DeoptTracker {
/// function_id → set of ShapeIds it depends on
dependencies: HashMap<u16, HashSet<ShapeId>>,
/// shape_id → set of function_ids that depend on it
shape_dependents: HashMap<ShapeId, HashSet<u16>>,
}
impl DeoptTracker {
/// Create an empty deopt tracker.
pub fn new() -> Self {
Self {
dependencies: HashMap::new(),
shape_dependents: HashMap::new(),
}
}
/// Register shape dependencies for a compiled function.
///
/// Called after successful JIT compilation when the compilation result
/// includes shape guard IDs.
pub fn register(&mut self, function_id: u16, shape_ids: &[ShapeId]) {
if shape_ids.is_empty() {
return;
}
let dep_set = self
.dependencies
.entry(function_id)
.or_insert_with(HashSet::new);
for &sid in shape_ids {
dep_set.insert(sid);
self.shape_dependents
.entry(sid)
.or_insert_with(HashSet::new)
.insert(function_id);
}
}
/// Invalidate all functions that depend on the given shape.
///
/// Returns the list of function IDs that were invalidated (need to
/// have their JIT code removed from the native_code_table).
pub fn invalidate_shape(&mut self, shape_id: ShapeId) -> Vec<u16> {
let dependents = match self.shape_dependents.remove(&shape_id) {
Some(set) => set,
None => return Vec::new(),
};
let mut invalidated = Vec::with_capacity(dependents.len());
for func_id in dependents {
// Remove all of this function's dependencies
if let Some(dep_shapes) = self.dependencies.remove(&func_id) {
// Clean up reverse mappings for other shapes this function depended on
for sid in &dep_shapes {
if *sid != shape_id {
if let Some(funcs) = self.shape_dependents.get_mut(sid) {
funcs.remove(&func_id);
if funcs.is_empty() {
self.shape_dependents.remove(sid);
}
}
}
}
}
invalidated.push(func_id);
}
invalidated
}
/// Clear all dependencies for a function (e.g., when it's recompiled).
pub fn clear_function(&mut self, function_id: u16) {
if let Some(dep_shapes) = self.dependencies.remove(&function_id) {
for sid in dep_shapes {
if let Some(funcs) = self.shape_dependents.get_mut(&sid) {
funcs.remove(&function_id);
if funcs.is_empty() {
self.shape_dependents.remove(&sid);
}
}
}
}
}
/// Number of functions being tracked.
pub fn tracked_function_count(&self) -> usize {
self.dependencies.len()
}
/// Number of shapes being watched.
pub fn watched_shape_count(&self) -> usize {
self.shape_dependents.len()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_register_and_invalidate() {
let mut tracker = DeoptTracker::new();
let s1 = ShapeId(1);
let s2 = ShapeId(2);
tracker.register(0, &[s1, s2]);
tracker.register(1, &[s1]);
assert_eq!(tracker.tracked_function_count(), 2);
assert_eq!(tracker.watched_shape_count(), 2);
// Invalidate shape 1 — both functions depend on it
let invalidated = tracker.invalidate_shape(s1);
assert_eq!(invalidated.len(), 2);
assert!(invalidated.contains(&0));
assert!(invalidated.contains(&1));
// Both functions fully removed
assert_eq!(tracker.tracked_function_count(), 0);
// Shape 2 no longer watched (function 0 was the only dependent)
assert_eq!(tracker.watched_shape_count(), 0);
}
#[test]
fn test_invalidate_no_dependents() {
let mut tracker = DeoptTracker::new();
let invalidated = tracker.invalidate_shape(ShapeId(99));
assert!(invalidated.is_empty());
}
#[test]
fn test_clear_function() {
let mut tracker = DeoptTracker::new();
let s1 = ShapeId(1);
tracker.register(0, &[s1]);
tracker.register(1, &[s1]);
tracker.clear_function(0);
assert_eq!(tracker.tracked_function_count(), 1);
// Shape 1 still watched by function 1
let invalidated = tracker.invalidate_shape(s1);
assert_eq!(invalidated, vec![1]);
}
#[test]
fn test_register_empty_shapes() {
let mut tracker = DeoptTracker::new();
tracker.register(0, &[]);
assert_eq!(tracker.tracked_function_count(), 0);
}
#[test]
fn test_duplicate_registration() {
let mut tracker = DeoptTracker::new();
let s1 = ShapeId(1);
tracker.register(0, &[s1]);
tracker.register(0, &[s1]); // duplicate
assert_eq!(tracker.tracked_function_count(), 1);
assert_eq!(tracker.watched_shape_count(), 1);
}
#[test]
fn test_invalidate_partial_overlap() {
let mut tracker = DeoptTracker::new();
let s1 = ShapeId(1);
let s2 = ShapeId(2);
let s3 = ShapeId(3);
tracker.register(0, &[s1, s2]); // depends on s1, s2
tracker.register(1, &[s2, s3]); // depends on s2, s3
// Invalidate s2 — both functions invalidated
let invalidated = tracker.invalidate_shape(s2);
assert_eq!(invalidated.len(), 2);
// All cleaned up
assert_eq!(tracker.tracked_function_count(), 0);
assert_eq!(tracker.watched_shape_count(), 0);
}
}