@@ -1392,6 +1392,134 @@ impl<T: Clone> Dict<T> {
13921392 Ok ( removed)
13931393 }
13941394
1395+ /// Move an existing entry to the end and return its current value without
1396+ /// making the key temporarily absent. Hashing and equality run unlocked.
1397+ pub ( crate ) fn move_to_end < K : DictKey + ?Sized > (
1398+ & self ,
1399+ vm : & VirtualMachine ,
1400+ key : & K ,
1401+ ) -> PyResult < Option < T > > {
1402+ struct ProbeWitness {
1403+ index_index : IndexIndex ,
1404+ entry_index : IndexEntry ,
1405+ key : Option < ( HashValue , PyObjectRef ) > ,
1406+ }
1407+
1408+ let hash_value = key. key_hash ( vm) ?;
1409+ let mut idxs = None ;
1410+ let mut prefix = Vec :: < ProbeWitness > :: new ( ) ;
1411+ let mut compared: Option < ( EntryIndex , IndexIndex , PyObjectRef , bool ) > = None ;
1412+ ' lookup: loop {
1413+ // Keep the compared key alive until after the guard is released,
1414+ // including when equality cleared the dictionary and recycled its
1415+ // entry and bucket for a different key.
1416+ let comparison = compared. take ( ) ;
1417+ let mut inner = self . write ( ) ;
1418+ let mask = ( inner. indices . len ( ) - 1 ) as i64 ;
1419+ let probes = idxs. get_or_insert_with ( || GenIndexes :: new ( hash_value, mask) ) ;
1420+ if probes. mask != mask {
1421+ drop ( inner) ;
1422+ prefix. clear ( ) ;
1423+ idxs = None ;
1424+ continue ;
1425+ }
1426+ let matched = if let Some ( ( entry_index, index_index, candidate, equal) ) = & comparison {
1427+ let valid = inner. indices . get ( * index_index) . and_then ( |i| i. index ( ) )
1428+ == Some ( * entry_index)
1429+ && inner
1430+ . get_entry_checked ( * entry_index, * index_index)
1431+ . is_some_and ( |entry| entry. hash == hash_value && entry. key . is ( candidate) ) ;
1432+ // A false comparison can survive while another previously
1433+ // probed bucket changes (including after nested compaction).
1434+ // Revalidate the whole prefix before continuing past it.
1435+ let prefix_valid = * equal
1436+ || prefix. iter ( ) . all ( |witness| {
1437+ inner. indices . get ( witness. index_index ) == Some ( & witness. entry_index )
1438+ && witness. key . as_ref ( ) . is_none_or ( |( hash, key) | {
1439+ inner
1440+ . get_entry_checked (
1441+ witness. entry_index . index ( ) . unwrap ( ) ,
1442+ witness. index_index ,
1443+ )
1444+ . is_some_and ( |entry| entry. hash == * hash && entry. key . is ( key) )
1445+ } )
1446+ } ) ;
1447+ if !valid || !prefix_valid {
1448+ drop ( inner) ;
1449+ prefix. clear ( ) ;
1450+ idxs = None ;
1451+ continue ;
1452+ }
1453+ equal. then_some ( * entry_index)
1454+ } else {
1455+ None
1456+ } ;
1457+ let entry_index = match matched {
1458+ Some ( index) => index,
1459+ None => loop {
1460+ let index_index = probes. next ( ) ;
1461+ let index_entry = inner. indices [ index_index] ;
1462+ match index_entry {
1463+ IndexEntry :: FREE => return Ok ( None ) ,
1464+ IndexEntry :: DUMMY => {
1465+ prefix. push ( ProbeWitness {
1466+ index_index,
1467+ entry_index : index_entry,
1468+ key : None ,
1469+ } ) ;
1470+ continue ;
1471+ }
1472+ _ => { }
1473+ }
1474+ let entry_index = index_entry. index ( ) . unwrap ( ) ;
1475+ let entry = inner. entries [ entry_index] . as_ref ( ) . unwrap ( ) ;
1476+ if key. key_is ( & entry. key ) {
1477+ break entry_index;
1478+ }
1479+ prefix. push ( ProbeWitness {
1480+ index_index,
1481+ entry_index : index_entry,
1482+ key : Some ( ( entry. hash , entry. key . clone ( ) ) ) ,
1483+ } ) ;
1484+ if entry. hash == hash_value {
1485+ let candidate = entry. key . clone ( ) ;
1486+ drop ( inner) ;
1487+ drop ( comparison) ;
1488+ let equal = key. key_eq ( vm, & candidate) ?;
1489+ compared = Some ( ( entry_index, index_index, candidate, equal) ) ;
1490+ continue ' lookup;
1491+ }
1492+ } ,
1493+ } ;
1494+ let value = inner. entries [ entry_index] . as_ref ( ) . unwrap ( ) . value . clone ( ) ;
1495+ if inner. entries [ entry_index + 1 ..] . iter ( ) . any ( Option :: is_some) {
1496+ // Reserve before taking the entry. No reader can observe an
1497+ // absent entry, and no hashing, equality or finalizer runs in
1498+ // this critical section.
1499+ inner. entries . reserve ( 1 ) ;
1500+ self . invalidate_keys_version ( ) ;
1501+ let entry = inner. entries [ entry_index] . take ( ) . unwrap ( ) ;
1502+ let new_index = inner. entries . len ( ) ;
1503+ inner. indices [ entry. index ] = unsafe {
1504+ // SAFETY: new_index is a valid nonnegative entry index.
1505+ IndexEntry :: from_index_unchecked ( new_index)
1506+ } ;
1507+ inner. entries . push ( Some ( entry) ) ;
1508+ let holes = inner. entries . len ( ) - inner. used ;
1509+ if holes >= 2 && holes >= inner. used {
1510+ // Moves do not increase `filled`, so ordinary insertion's
1511+ // resize threshold cannot bound their accumulated holes.
1512+ // Two holes also ensure this operation changes DictSize,
1513+ // allowing existing iterators to detect the relocation.
1514+ let indices_size = inner. indices . len ( ) ;
1515+ inner. resize ( indices_size) ;
1516+ }
1517+ }
1518+ drop ( inner) ;
1519+ return Ok ( Some ( value) ) ;
1520+ }
1521+ }
1522+
13951523 pub ( crate ) fn pop_back ( & self ) -> Option < ( PyObjectRef , T ) > {
13961524 let inner = & mut * self . write ( ) ;
13971525 let entry = loop {
@@ -1838,6 +1966,114 @@ mod tests {
18381966 use super :: * ;
18391967 use crate :: { Interpreter , common:: ascii} ;
18401968
1969+ #[ test]
1970+ fn move_to_end_preserves_values_and_invalidates_layout ( ) {
1971+ Interpreter :: without_stdlib ( Default :: default ( ) ) . enter ( |vm| {
1972+ let dict = Dict :: default ( ) ;
1973+ let keys = [
1974+ vm. ctx . intern_str ( "first" ) ,
1975+ vm. ctx . intern_str ( "middle" ) ,
1976+ vm. ctx . intern_str ( "last" ) ,
1977+ ] ;
1978+ for ( value, key) in keys. iter ( ) . enumerate ( ) {
1979+ dict. insert ( vm, * key, value) . unwrap ( ) ;
1980+ }
1981+ let hint = dict. hint_for_key ( vm, keys[ 1 ] ) . unwrap ( ) . unwrap ( ) as usize ;
1982+ let version = dict. assign_keys_version ( ) ;
1983+ let size = dict. size ( ) ;
1984+ assert_ne ! ( version, 0 ) ;
1985+ assert_eq ! ( dict. move_to_end( vm, keys[ 1 ] ) . unwrap( ) , Some ( 1 ) ) ;
1986+ assert_eq ! ( dict. values( ) , vec![ 0 , 2 , 1 ] ) ;
1987+ assert_eq ! ( dict. keys_version( ) , 0 ) ;
1988+ assert_eq ! ( dict. get_index_if_keys_version( version, hint) , None ) ;
1989+ assert_eq ! ( dict. get_hint( vm, keys[ 1 ] , hint) . unwrap( ) , None ) ;
1990+ assert ! ( dict. next_entry_checked( 0 , & size, |_, v| * v) . is_err( ) ) ;
1991+ assert ! (
1992+ dict. prev_entry_checked( size. entries_size - 1 , & size, |_, v| * v)
1993+ . is_err( )
1994+ ) ;
1995+
1996+ let version = dict. assign_keys_version ( ) ;
1997+ let size = dict. size ( ) ;
1998+ assert_eq ! ( dict. move_to_end( vm, keys[ 1 ] ) . unwrap( ) , Some ( 1 ) ) ;
1999+ assert_eq ! ( dict. keys_version( ) , version) ;
2000+ assert_eq ! ( dict. size( ) , size) ;
2001+ assert_eq ! ( dict. move_to_end( vm, "absent" ) . unwrap( ) , None ) ;
2002+ assert_eq ! ( dict. keys_version( ) , version) ;
2003+ assert_eq ! ( dict. size( ) , size) ;
2004+ assert_eq ! ( dict. values( ) , vec![ 0 , 2 , 1 ] ) ;
2005+ } ) ;
2006+ }
2007+
2008+ #[ test]
2009+ fn move_to_end_compacts_and_keeps_indices_consistent ( ) {
2010+ Interpreter :: without_stdlib ( Default :: default ( ) ) . enter ( |vm| {
2011+ for count in [ 2 , 3 , 32 ] {
2012+ let dict = Dict :: default ( ) ;
2013+ for key in 0 ..count {
2014+ dict. insert ( vm, & key, key) . unwrap ( ) ;
2015+ }
2016+ for step in 0 ..1024 {
2017+ let key = step % count;
2018+ let size = dict. size ( ) ;
2019+ let hint = dict. hint_for_key ( vm, & key) . unwrap ( ) . unwrap ( ) as usize ;
2020+ let version = dict. assign_keys_version ( ) ;
2021+ assert_eq ! ( dict. move_to_end( vm, & key) . unwrap( ) , Some ( key) ) ;
2022+ assert_eq ! ( dict. get_index_if_keys_version( version, hint) , None ) ;
2023+ assert ! ( dict. next_entry_checked( 0 , & size, |_, v| * v) . is_err( ) ) ;
2024+ assert ! (
2025+ dict. prev_entry_checked( size. entries_size - 1 , & size, |_, v| * v)
2026+ . is_err( )
2027+ ) ;
2028+ let inner = dict. read ( ) ;
2029+ assert_eq ! ( inner. used, count) ;
2030+ assert_eq ! ( inner. filled, count) ;
2031+ assert ! ( inner. entries. len( ) < 2 * count) ;
2032+ assert ! ( inner. entries. capacity( ) <= 4 * count) ;
2033+ let values: Vec < _ > = inner. entries . iter ( ) . flatten ( ) . map ( |e| e. value ) . collect ( ) ;
2034+ assert_eq ! (
2035+ values,
2036+ ( 1 ..=count) . map( |i| ( key + i) % count) . collect:: <Vec <_>>( )
2037+ ) ;
2038+ for ( index, entry) in inner. entries . iter ( ) . enumerate ( ) {
2039+ if let Some ( entry) = entry {
2040+ assert_eq ! ( inner. indices[ entry. index] . index( ) , Some ( index) ) ;
2041+ assert_eq ! ( entry. hash, entry. value. key_hash( vm) . unwrap( ) ) ;
2042+ }
2043+ }
2044+ for ( bucket, index) in inner. indices . iter ( ) . enumerate ( ) {
2045+ if let Some ( index) = index. index ( ) {
2046+ assert_eq ! ( inner. entries[ index] . as_ref( ) . unwrap( ) . index, bucket) ;
2047+ }
2048+ }
2049+ }
2050+ }
2051+ } ) ;
2052+ }
2053+
2054+ #[ test]
2055+ fn move_to_end_restores_only_matching_shared_shapes ( ) {
2056+ Interpreter :: without_stdlib ( Default :: default ( ) ) . enter ( |vm| {
2057+ let first = vm. ctx . intern_str ( "move_first" ) ;
2058+ let last = vm. ctx . intern_str ( "move_last" ) ;
2059+ let dict = Dict :: default ( ) ;
2060+ let peer = Dict :: default ( ) ;
2061+ for table in [ & dict, & peer] {
2062+ table. insert ( vm, first, 1 ) . unwrap ( ) ;
2063+ table. insert ( vm, last, 2 ) . unwrap ( ) ;
2064+ }
2065+ let version = dict. assign_keys_version ( ) ;
2066+ assert_eq ! ( peer. assign_keys_version( ) , version) ;
2067+ dict. move_to_end ( vm, first) . unwrap ( ) ;
2068+ assert_ne ! ( dict. assign_keys_version( ) , version) ;
2069+ dict. move_to_end ( vm, last) . unwrap ( ) ;
2070+ // The second move compacts back to the original hole-free layout.
2071+ assert_eq ! ( dict. assign_keys_version( ) , version) ;
2072+ assert_eq ! ( dict. get_index_if_keys_version( version, 0 ) , Some ( 1 ) ) ;
2073+ assert_eq ! ( dict. get_index_if_keys_version( version, 1 ) , Some ( 2 ) ) ;
2074+ } ) ;
2075+ }
2076+
18412077 #[ test]
18422078 fn clone_compacts_deleted_entries ( ) {
18432079 Interpreter :: without_stdlib ( Default :: default ( ) ) . enter ( |vm| {
0 commit comments