Skip to content

Commit 855ab4b

Browse files
youknowoneyouknowone
committed
DictInner::indices: Vec<i64> => Vec<IndexEntry>
1 parent 0652485 commit 855ab4b

File tree

1 file changed

+48
-45
lines changed

1 file changed

+48
-45
lines changed

‎vm/src/dictdatatype.rs‎

Lines changed: 48 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -37,34 +37,24 @@ impl<T> fmt::Debug for Dict<T> {
3737
}
3838
}
3939

40-
#[derive(Debug, Copy, Clone)]
41-
enum IndexEntry {
42-
Dummy,
43-
Free,
44-
Index(usize),
45-
}
40+
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
41+
#[repr(transparent)]
42+
struct IndexEntry(i64);
4643

4744
impl IndexEntry {
48-
const FREE: i64 = -1;
49-
const DUMMY: i64 = -2;
50-
}
45+
const FREE: Self = Self(-1);
46+
const DUMMY: Self = Self(-2);
5147

52-
impl From<i64> for IndexEntry {
53-
fn from(idx: i64) -> Self {
54-
match idx {
55-
IndexEntry::FREE => IndexEntry::Free,
56-
IndexEntry::DUMMY => IndexEntry::Dummy,
57-
x => IndexEntry::Index(x as usize),
58-
}
48+
unsafe fn from_index_unchecked(idx: usize) -> Self {
49+
debug_assert!((idx as isize) >= 0);
50+
Self(idx as i64)
5951
}
60-
}
6152

62-
impl From<IndexEntry> for i64 {
63-
fn from(idx: IndexEntry) -> Self {
64-
match idx {
65-
IndexEntry::Free => IndexEntry::FREE,
66-
IndexEntry::Dummy => IndexEntry::DUMMY,
67-
IndexEntry::Index(i) => i as i64,
53+
fn index(self) -> Option<usize> {
54+
if self.0 >= 0 {
55+
Some(self.0 as usize)
56+
} else {
57+
None
6858
}
6959
}
7060
}
@@ -73,7 +63,7 @@ impl From<IndexEntry> for i64 {
7363
struct DictInner<T> {
7464
used: usize,
7565
filled: usize,
76-
indices: Vec<i64>,
66+
indices: Vec<IndexEntry>,
7767
entries: Vec<Option<DictEntry<T>>>,
7868
}
7969

@@ -175,7 +165,10 @@ impl<T> DictInner<T> {
175165
let index_index = idxs.next();
176166
let idx = &mut self.indices[index_index];
177167
if *idx == IndexEntry::FREE {
178-
*idx = entry_idx as i64;
168+
*idx = unsafe {
169+
// entry_idx never grow up to usize::MAX
170+
IndexEntry::from_index_unchecked(entry_idx)
171+
};
179172
entry.index = index_index;
180173
break;
181174
}
@@ -203,9 +196,9 @@ impl<T> DictInner<T> {
203196
};
204197
let entry_index = self.entries.len();
205198
self.entries.push(Some(entry));
206-
self.indices[index] = entry_index as i64;
199+
self.indices[index] = unsafe { IndexEntry::from_index_unchecked(entry_index) };
207200
self.used += 1;
208-
if let IndexEntry::Free = index_entry {
201+
if let IndexEntry::FREE = index_entry {
209202
self.filled += 1;
210203
if let Some(new_size) = self.should_resize() {
211204
self.resize(new_size)
@@ -252,7 +245,7 @@ impl<T: Clone> Dict<T> {
252245
let _removed = loop {
253246
let (entry_index, index_index) = self.lookup(vm, key, hash, None)?;
254247
let mut inner = self.write();
255-
if let IndexEntry::Index(index) = entry_index {
248+
if let Some(index) = entry_index.index() {
256249
// Update existing key
257250
if let Some(entry) = inner.entries.get_mut(index) {
258251
let entry = entry
@@ -279,7 +272,7 @@ impl<T: Clone> Dict<T> {
279272

280273
pub fn contains<K: DictKey + ?Sized>(&self, vm: &VirtualMachine, key: &K) -> PyResult<bool> {
281274
let (entry, _) = self.lookup(vm, key, key.key_hash(vm)?, None)?;
282-
Ok(matches!(entry, IndexEntry::Index(_)))
275+
Ok(entry.index().is_some())
283276
}
284277

285278
/// Retrieve a key
@@ -297,7 +290,7 @@ impl<T: Clone> Dict<T> {
297290
) -> PyResult<Option<T>> {
298291
let ret = loop {
299292
let (entry, index_index) = self.lookup(vm, key, hash, None)?;
300-
if let IndexEntry::Index(index) = entry {
293+
if let Some(index) = entry.index() {
301294
let inner = self.read();
302295
if let Some(entry) = inner.entries.get(index) {
303296
let entry = extract_dict_entry(entry);
@@ -385,7 +378,7 @@ impl<T: Clone> Dict<T> {
385378
let _removed = loop {
386379
let lookup = self.lookup(vm, key, hash, None)?;
387380
let (entry, index_index) = lookup;
388-
if let IndexEntry::Index(_) = entry {
381+
if entry.index().is_some() {
389382
match self.pop_inner(lookup) {
390383
ControlFlow::Break(Some(entry)) => break Some(entry),
391384
_ => continue,
@@ -407,8 +400,8 @@ impl<T: Clone> Dict<T> {
407400
let hash = key.key_hash(vm)?;
408401
let res = loop {
409402
let lookup = self.lookup(vm, key, hash, None)?;
410-
let (entry, index_index) = lookup;
411-
if let IndexEntry::Index(index) = entry {
403+
let (index_entry, index_index) = lookup;
404+
if let Some(index) = index_entry.index() {
412405
let inner = self.read();
413406
if let Some(entry) = inner.entries.get(index) {
414407
let entry = extract_dict_entry(entry);
@@ -424,7 +417,13 @@ impl<T: Clone> Dict<T> {
424417
} else {
425418
let value = default();
426419
let mut inner = self.write();
427-
inner.unchecked_push(index_index, hash, key.to_pyobject(vm), value.clone(), entry);
420+
inner.unchecked_push(
421+
index_index,
422+
hash,
423+
key.to_pyobject(vm),
424+
value.clone(),
425+
index_entry,
426+
);
428427
break value;
429428
}
430429
};
@@ -445,8 +444,8 @@ impl<T: Clone> Dict<T> {
445444
let hash = key.key_hash(vm)?;
446445
let res = loop {
447446
let lookup = self.lookup(vm, key, hash, None)?;
448-
let (entry, index_index) = lookup;
449-
if let IndexEntry::Index(index) = entry {
447+
let (index_entry, index_index) = lookup;
448+
if let Some(index) = index_entry.index() {
450449
let inner = self.read();
451450
if let Some(entry) = inner.entries.get(index) {
452451
let entry = extract_dict_entry(entry);
@@ -464,7 +463,7 @@ impl<T: Clone> Dict<T> {
464463
let key = key.to_pyobject(vm);
465464
let mut inner = self.write();
466465
let ret = (key.clone(), value.clone());
467-
inner.unchecked_push(index_index, hash, key, value, entry);
466+
inner.unchecked_push(index_index, hash, key, value, index_entry);
468467
break ret;
469468
}
470469
};
@@ -541,22 +540,26 @@ impl<T: Clone> Dict<T> {
541540
});
542541
loop {
543542
let index_index = idxs.next();
544-
match IndexEntry::from(inner.indices[index_index]) {
545-
IndexEntry::Dummy => {
543+
match inner.indices[index_index] {
544+
IndexEntry::DUMMY => {
546545
if freeslot.is_none() {
547546
freeslot = Some(index_index);
548547
}
549548
}
550-
IndexEntry::Free => {
549+
IndexEntry::FREE => {
551550
let idxs = match freeslot {
552-
Some(free) => (IndexEntry::Dummy, free),
553-
None => (IndexEntry::Free, index_index),
551+
Some(free) => (IndexEntry::DUMMY, free),
552+
None => (IndexEntry::FREE, index_index),
554553
};
555554
return Ok(idxs);
556555
}
557-
IndexEntry::Index(i) => {
556+
idx => {
557+
let i = idx.index().unwrap_or_else(|| unsafe {
558+
// DUMMY and FREE is already checked above.
559+
std::hint::unreachable_unchecked()
560+
});
558561
let entry = &inner.entries[i].as_ref().unwrap();
559-
let ret = (IndexEntry::Index(i), index_index);
562+
let ret = (idx, index_index);
560563
if key.key_is(&entry.key) {
561564
break 'outer ret;
562565
} else if entry.hash == hash_value {
@@ -593,7 +596,7 @@ impl<T: Clone> Dict<T> {
593596
pred: impl Fn(&T) -> Result<bool, E>,
594597
) -> Result<PopInnerResult<T>, E> {
595598
let (entry_index, index_index) = lookup;
596-
let entry_index = if let IndexEntry::Index(entry_index) = entry_index {
599+
let entry_index = if let Some(entry_index) = entry_index.index() {
597600
entry_index
598601
} else {
599602
return Ok(ControlFlow::Break(None));

0 commit comments

Comments
 (0)