tf_idf_vectorizer/utils/datastruct/map/
index_map.rs

1use hashbrown::HashTable;
2use std::borrow::Borrow;
3use std::fmt::Debug;
4use std::hash::Hasher;
5use std::iter::FusedIterator;
6
7/// IndexMap
8/// 連続領域を保証するHashMap
9/// 
10/// # Safety
11/// table, hashes は table_* メソッドが責任をもつこと 更新とか
12/// 
13/// O(1) operations for:
14/// - key -> value
15/// - key -> index
16/// - index -> key
17/// - index -> value
18/// - insert
19/// - swap_remove
20#[derive(Clone, Debug)]
21pub struct IndexMap<K, V, S = ahash::RandomState> {
22    pub(crate) values: Vec<V>,
23    pub(crate) index_set: IndexSet<K, S>,
24}
25
26impl<K, V, S> IndexMap<K, V, S>
27where
28    K: Eq + std::hash::Hash + Clone,
29    S: std::hash::BuildHasher,
30{
31    pub fn with_hasher(hash_builder: S) -> Self {
32        IndexMap {
33            values: Vec::new(),
34            index_set: IndexSet::with_hasher(hash_builder),
35        }
36    }
37
38    pub fn with_capacity(capacity: usize) -> Self
39    where
40        S: Default,
41    {
42        IndexMap {
43            values: Vec::with_capacity(capacity),
44            index_set: IndexSet::with_capacity(capacity),
45        }
46    }
47
48    pub fn new() -> Self
49    where
50        S: Default,
51    {
52        IndexMap {
53            values: Vec::new(),
54            index_set: IndexSet::new(),
55        }
56    }
57
58    pub fn with_hasher_capacity(hash_builder: S, capacity: usize) -> Self {
59        IndexMap {
60            values: Vec::with_capacity(capacity),
61            index_set: IndexSet::with_hasher_capacity(hash_builder, capacity),
62        }
63    }
64
65    pub fn shrink_to_fit(&mut self) {
66        self.values.shrink_to_fit();
67        self.index_set.shrink_to_fit();
68    }
69
70    #[inline]
71    pub fn len(&self) -> usize {
72        self.values.len()
73    }
74
75    #[inline]
76    pub fn iter_values(&self) -> std::slice::Iter<'_, V> {
77        self.values.iter()
78    }
79
80    #[inline]
81    pub fn iter_keys(&self) -> std::slice::Iter<'_, K> {
82        self.index_set.keys.iter()
83    }
84
85    #[inline]
86    pub fn iter(&self) -> IndexMapIter<'_, K, V, S> {
87        IndexMapIter {
88            map: self,
89            index: 0,
90            end: self.len(),
91        }
92    }
93
94    #[inline]
95    pub fn values(&self) -> &Vec<V> {
96        &self.values
97    }
98
99    #[inline]
100    pub fn keys(&self) -> &Vec<K> {
101        &self.index_set.keys()
102    }
103
104    #[inline]
105    pub fn as_index_set(&self) -> &IndexSet<K, S> {
106        &self.index_set
107    }
108
109    #[inline]
110    pub fn get<Q: ?Sized>(&self, key: &Q) -> Option<&V> 
111    where 
112        K: Borrow<Q>,
113        Q: std::hash::Hash + Eq,
114    {
115        if let Some(idx) = self.index_set.table_get(key) {
116            unsafe {
117                Some(self.values.get_unchecked(idx))
118            }
119        } else {
120            None
121        }
122    }
123
124    #[inline]
125    pub fn get_mut<Q: ?Sized>(&mut self, key: &Q) -> Option<&mut V> 
126    where 
127        K: Borrow<Q>,
128        Q: std::hash::Hash + Eq,
129    {
130        if let Some(idx) = self.index_set.table_get(key) {
131            unsafe {
132                Some(self.values.get_unchecked_mut(idx))
133            }
134        } else {
135            None
136        }
137    }
138
139    #[inline]
140    pub fn get_with_index(&self, index: usize) -> Option<&V> {
141        self.values.get(index)
142    }
143
144    #[inline]
145    pub fn get_with_index_mut(&mut self, index: usize) -> Option<&mut V> {
146        self.values.get_mut(index)
147    }
148
149    #[inline]
150    pub fn get_key_with_index(&self, index: usize) -> Option<&K> {
151        self.index_set.get_with_index(index)
152    }
153
154    #[inline]
155    pub fn get_key_value_with_index(&self, index: usize) -> Option<(&K, &V)> {
156        if index < self.len() {
157            unsafe {
158                Some((
159                    self.index_set.keys.get_unchecked(index),
160                    self.values.get_unchecked(index),
161                ))
162            }
163        } else {
164            None
165        }
166    }
167
168    #[inline]
169    pub fn get_index<Q: ?Sized>(&self, key: &Q) -> Option<usize> 
170    where 
171        K: Borrow<Q>,
172        Q: std::hash::Hash + Eq,
173    {
174        self.index_set.table_get(key)
175    }
176
177    pub fn contains_key<Q: ?Sized>(&self, key: &Q) -> bool 
178    where 
179        K: Borrow<Q>,
180        Q: std::hash::Hash + Eq,
181    {
182        self.index_set.contains_key(key)
183    }
184
185    #[inline]
186    pub fn insert(&mut self, key: K, value: V) -> InsertResult<K, V> {
187        if let Some(idx) = self.index_set.table_get(&key) {
188            // K が Rc の場合を考慮して すべて差し替える
189            unsafe {
190                self.index_set.table_override(&key, &idx);
191            }
192            let old_value = Some(std::mem::replace(&mut self.values[idx], value));
193            let old_key = Some(std::mem::replace(&mut self.index_set.keys[idx], key.clone()));
194            InsertResult::Override {
195                old_value: old_value.unwrap(),
196                old_key:  old_key.unwrap(),
197                index: idx,
198            }
199        } else {
200            // New key, insert entry
201            let idx = self.values.len();
202            unsafe {
203                self.index_set.table_append(&key, &idx);
204            }
205            self.index_set.keys.push(key.clone());
206            self.values.push(value);
207            InsertResult::New {
208                index: idx,
209            }
210        }
211    }
212
213    #[inline]
214    pub fn entry_mut<'a>(&'a mut self, key: K) -> EntryMut<'a, K, V, S> {
215        if let Some(idx) = self.index_set.table_get(&key) {
216            unsafe {
217                EntryMut::Occupied {
218                    key: key,
219                    value: self.values.get_unchecked_mut(idx),
220                    index: idx,
221                }
222            }
223        } else {
224            EntryMut::Vacant { key , map: self }
225        }
226    }
227
228    #[inline]
229    pub fn swap_remove<Q: ?Sized>(&mut self, key: &Q) -> RemoveResult<K, V>
230    where 
231        K: Borrow<Q>,
232        Q: std::hash::Hash + Eq,
233    {
234        if let Some(idx) = self.index_set.table_get(key) {
235            let last_idx = self.values.len() - 1;
236            if idx == last_idx {
237                // 最後の要素を削除する場合
238                unsafe {
239                    self.index_set.table_swap_remove(key);
240                }
241                return RemoveResult::Removed {
242                    old_value: self.values.pop().unwrap(),
243                    old_key: self.index_set.keys.pop().unwrap(),
244                    index: idx,
245                };
246            } else {
247                let last_idx_key = self.index_set.keys[last_idx].clone();
248                unsafe {
249                    // keyとの整合性があるうちに削除予定のをtableから消す ここでhashesがswap_removeされる
250                    // last_idxの要素がswapで移動してくる
251                    self.index_set.table_swap_remove(key);
252                    // 移動させられた要素のtableを再登録
253                    // 登録されていた前のidxに対するkeyはまだ整合性が取れているので問題ない
254                    self.index_set.table_override(last_idx_key.borrow(), &idx);
255                }
256                // swap_remove ここで実際にtableのidxとvalues, keys, hashesの整合性が回復
257                let value = self.values.swap_remove(idx);
258                self.index_set.keys.swap_remove(idx);
259                RemoveResult::Removed { 
260                    old_value: value,
261                    old_key: last_idx_key,
262                    index: idx,
263                }
264            }
265        } else {
266            RemoveResult::None
267        }
268    }
269
270    #[inline]
271    pub fn from_kv_vec(k_vec: Vec<K>, v_vec: Vec<V>) -> Self
272    where
273        S: std::hash::BuildHasher + Default,
274    {
275        let hash_builder = S::default();
276        let mut map = IndexMap::with_hasher(hash_builder);
277        for (k, v) in k_vec.into_iter().zip(v_vec.into_iter()) {
278            let idx = map.values.len();
279            unsafe {
280                map.index_set.table_append(&k, &idx);
281            }
282            map.index_set.keys.push(k);
283            map.values.push(v);
284        }
285        map
286    }
287}
288
289pub struct IndexMapIter<'a, K, V, S> {
290    pub map: &'a IndexMap<K, V, S>,
291    pub index: usize,
292    pub end: usize, // exclusive
293}
294
295impl<'a, K, V, S> Iterator for IndexMapIter<'a, K, V, S>
296where
297    K: Eq + std::hash::Hash,
298    S: std::hash::BuildHasher,
299{
300    type Item = (&'a K, &'a V);
301
302    #[inline]
303    fn next(&mut self) -> Option<Self::Item> {
304        if self.index < self.end {
305            unsafe {
306                let i = self.index;
307                self.index += 1;
308                let k = self.map.index_set.keys.get_unchecked(i);
309                let v = self.map.values.get_unchecked(i);
310                Some((k, v))
311            }
312        } else {
313            None
314        }
315    }
316
317    #[inline]
318    fn size_hint(&self) -> (usize, Option<usize>) {
319        let remaining = self.end - self.index;
320        (remaining, Some(remaining))
321    }
322
323    #[inline]
324    fn nth(&mut self, n: usize) -> Option<Self::Item> {
325        // n 個スキップ
326        let new = self.index.saturating_add(n);
327        self.index = new.min(self.end);
328        self.next()
329    }
330}
331
332impl<'a, K, V, S> ExactSizeIterator for IndexMapIter<'a, K, V, S>
333where
334    K: Eq + std::hash::Hash,
335    S: std::hash::BuildHasher,
336{
337    #[inline]
338    fn len(&self) -> usize {
339        self.end - self.index
340    }
341}
342
343impl<'a, K, V, S> DoubleEndedIterator for IndexMapIter<'a, K, V, S>
344where
345    K: Eq + std::hash::Hash,
346    S: std::hash::BuildHasher,
347{
348    #[inline]
349    fn next_back(&mut self) -> Option<Self::Item> {
350        if self.index < self.end {
351            unsafe {
352                self.end -= 1;
353                let i = self.end;
354                let k = self.map.index_set.keys.get_unchecked(i);
355                let v = self.map.values.get_unchecked(i);
356                Some((k, v))
357            }
358        } else {
359            None
360        }
361    }
362}
363
364impl<'a, K, V, S> FusedIterator for IndexMapIter<'a, K, V, S>
365where
366    K: Eq + std::hash::Hash,
367    S: std::hash::BuildHasher,
368{}
369
370pub enum EntryMut<'a, K, V, S> {
371    Occupied { key: K, value: &'a mut V, index: usize },
372    Vacant { key: K , map: &'a mut IndexMap<K, V, S> },
373}
374
375impl<'a, K, V, S> EntryMut<'a, K, V, S>
376where 
377    K: Eq + std::hash::Hash + Clone,
378    S: std::hash::BuildHasher,
379{
380    #[inline]
381    pub fn is_occupied(&self) -> bool {
382        matches!(self, EntryMut::Occupied { .. })
383    }
384
385    #[inline]
386    pub fn or_insert_with<F>(self, value: F) -> &'a mut V
387    where
388        F: FnOnce() -> V,
389        K: Clone,
390    {
391        match self {
392            EntryMut::Occupied { value: v, .. } => v,
393            EntryMut::Vacant { key, map } => {
394                map.insert(key.clone(), value());
395                map.get_mut(&key).unwrap()
396            }
397        }
398    }
399}
400
401#[derive(Debug, PartialEq)]
402pub enum InsertResult<K, V> {
403    Override {
404        old_value: V,
405        old_key: K,
406        index: usize,
407    },
408    New {
409        index: usize,
410    }
411}
412
413#[derive(Debug, PartialEq)]
414pub enum RemoveResult<K, V> {
415    Removed {
416        old_value: V,
417        old_key: K,
418        index: usize,
419    },
420    None,
421}
422
423#[derive(Clone, Debug)]
424pub struct IndexSet<K, S = ahash::RandomState> {
425    pub(crate) keys: Vec<K>,
426    pub(crate) hashes: Vec<u64>,
427    pub(crate) table: HashTable<usize>,
428    pub(crate) hash_builder: S,
429}
430
431impl<K, S> IndexSet<K, S>
432where
433    K: Eq + std::hash::Hash,
434    S: std::hash::BuildHasher,
435{
436    pub fn with_hasher(hash_builder: S) -> Self {
437        IndexSet {
438            keys: Vec::new(),
439            hashes: Vec::new(),
440            table: HashTable::new(),
441            hash_builder,
442        }
443    }
444
445    pub fn new() -> Self
446    where
447        S: Default,
448    {
449        IndexSet {
450            keys: Vec::new(),
451            hashes: Vec::new(),
452            table: HashTable::new(),
453            hash_builder: S::default(),
454        }
455    }
456
457    pub fn with_capacity(capacity: usize) -> Self
458    where
459        S: Default,
460    {
461        IndexSet {
462            keys: Vec::with_capacity(capacity),
463            hashes: Vec::with_capacity(capacity),
464            table: HashTable::with_capacity(capacity),
465            hash_builder: S::default(),
466        }
467    }
468
469    pub fn with_hasher_capacity(hash_builder: S, capacity: usize) -> Self {
470        IndexSet {
471            keys: Vec::with_capacity(capacity),
472            hashes: Vec::with_capacity(capacity),
473            table: HashTable::with_capacity(capacity),
474            hash_builder,
475        }
476    }
477
478    pub fn capacity(&self) -> usize {
479        self.keys.capacity()
480    }
481
482    pub fn shrink_to_fit(&mut self) {
483        self.keys.shrink_to_fit();
484        self.hashes.shrink_to_fit();
485    }
486
487    /// hash util
488    #[inline]
489    fn hash_key<Q: ?Sized>(&self, key: &Q) -> u64 
490    where 
491        K: Borrow<Q>,
492        Q: std::hash::Hash + Eq,
493    {
494        let mut hasher = self.hash_builder.build_hasher();
495        key.hash(&mut hasher);
496        hasher.finish()
497    }
498
499    /// override
500    /// 完全な整合性が必要
501    /// keyに対するidxを更新し、更新したidxを返す
502    /// 存在しない場合はNone
503    #[inline]
504    unsafe fn table_override<Q: ?Sized>(&mut self, key: &Q, idx: &usize) -> Option<usize> 
505    where 
506        K: Borrow<Q>,
507        Q: std::hash::Hash + Eq,
508    {
509        let hash = self.hash_key(key);
510        match self.table.find_entry(hash, |&i| self.keys[i].borrow() == key) {
511            Ok(mut occ) => {
512                // idxの上書きだけ
513                *occ.get_mut() = *idx;
514                Some(*idx)
515            }
516            Err(_) => {
517                None
518            }
519        }
520    }
521
522    /// append
523    /// 完全な整合性が必要
524    /// hashesとtableを更新する
525    #[inline]
526    unsafe fn table_append<Q: ?Sized>(&mut self, key: &Q, idx: &usize) 
527    where 
528        K: Borrow<Q>,
529        Q: std::hash::Hash + Eq,
530    {
531        let hash = self.hash_key(key);
532        self.hashes.push(hash);
533        self.table.insert_unique(
534            hash,
535            *idx,
536            |&i| self.hashes[i]
537        );
538    }
539
540    /// get
541    /// とくに注意なし 不可変参照なので
542    #[inline]
543    fn table_get<Q: ?Sized>(&self, key: &Q) -> Option<usize> 
544    where 
545        K: Borrow<Q>,
546        Q: std::hash::Hash + Eq,
547    {
548        let hash = self.hash_key(key);
549        self.table.find(
550            hash, 
551            |&i| self.keys[i].borrow() == key
552        ).copied()
553    }
554
555    /// remove
556    /// 完全な整合性が必要
557    /// hashesはswap_removeされます
558    #[inline]
559    unsafe fn table_swap_remove<Q: ?Sized>(&mut self, key: &Q) -> Option<usize> 
560    where 
561        K: Borrow<Q>,
562        Q: std::hash::Hash + Eq,
563    {
564        let hash = self.hash_key(key);
565        if let Ok(entry) = self.table.find_entry(
566            hash,
567            |&i| self.keys[i].borrow() == key
568        ) {
569            let (odl_idx, _) = entry.remove();
570            self.hashes.swap_remove(odl_idx);
571            Some(odl_idx)
572        } else {
573            None
574        }
575    }
576
577    #[inline]
578    pub fn contains_key<Q: ?Sized>(&self, key: &Q) -> bool 
579    where 
580        K: Borrow<Q>,
581        Q: std::hash::Hash + Eq,
582    {
583        self.table_get(key).is_some()
584    }
585
586    #[inline]
587    pub fn len(&self) -> usize {
588        self.keys.len()
589    }
590
591    #[inline]
592    pub fn get_index<Q: ?Sized>(&self, key: &Q) -> Option<usize> 
593    where 
594        K: Borrow<Q>,
595        Q: std::hash::Hash + Eq,
596    {
597        self.table_get(key)
598    }
599
600    #[inline]
601    pub fn iter(&self) -> std::slice::Iter<'_, K> {
602        self.keys.iter()
603    }
604
605    #[inline]
606    pub fn keys(&self) -> &Vec<K> {
607        &self.keys
608    }
609
610    #[inline]
611    pub fn get_with_index(&self, index: usize) -> Option<&K> {
612        self.keys.get(index)
613    }
614}
615
616
617
618#[cfg(test)]
619mod tests {
620    use super::*;
621    use std::collections::HashMap;
622
623    // まずは比較しやすい型でテスト
624    type M = IndexMap<u64, i64>;
625
626    fn assert_internal_invariants(map: &M) {
627        // 長さが揃っていること
628        assert_eq!(map.values.len(), map.index_set.keys.len(), "values/keys len mismatch");
629        assert_eq!(map.values.len(), map.index_set.hashes.len(), "values/hashes len mismatch");
630
631        // table_get が返す idx が範囲内で、keys/values と一致すること
632        for (i, k) in map.index_set.keys.iter().enumerate() {
633            let idx = map.index_set.table_get(k).expect("table_get must find existing key");
634            assert_eq!(idx, i, "table idx mismatch for key");
635        }
636
637        // 逆方向も確認
638        // 重複キー禁止 + contains/get の整合
639        for i in 0..map.len() {
640            let k = &map.index_set.keys[i];
641            assert!(map.contains_key(k), "contains_key false for existing key");
642            let v = map.get(k).expect("get must return for existing key");
643            assert_eq!(*v, map.values[i], "get value mismatch");
644        }
645
646        // キー重複が無いこと
647        // O(n^2) だけどユニットテストならOK
648        for i in 0..map.index_set.keys.len() {
649            for j in (i + 1)..map.index_set.keys.len() {
650                assert!(map.index_set.keys[i] != map.index_set.keys[j], "duplicate keys detected");
651            }
652        }
653    }
654
655    fn assert_equals_oracle(map: &M, oracle: &HashMap<u64, i64>) {
656        assert_eq!(map.len(), oracle.len(), "len mismatch");
657
658        // 全キーが一致し、値も一致すること
659        for (k, v) in oracle.iter() {
660            let got = map.get(k).copied();
661            assert_eq!(got, Some(*v), "value mismatch for key={k}");
662        }
663
664        // map 側に余計なキーがないこと（oracle で確認）
665        for (k, v) in map.iter() {
666            assert_eq!(oracle.get(k).copied(), Some(*v), "extra/mismatch entry key={k}");
667        }
668    }
669
670    #[test]
671    fn swap_remove_last_and_middle() {
672        let mut m = M::new();
673        for i in 0..10 {
674            m.insert(i, (i as i64) * 10);
675        }
676
677        // last remove
678        let v = m.swap_remove(&9);
679        assert_eq!(v, RemoveResult::Removed {
680            old_value: 90,
681            old_key: 9,
682            index: 9,
683        });
684        assert!(m.get(&9).is_none());
685
686        // middle remove
687        let v = m.swap_remove(&3);
688        assert_eq!(v, RemoveResult::Removed {
689            old_value: 30,
690            old_key: 3,
691            index: 3,
692        });
693        assert!(m.get(&3).is_none());
694
695        assert_internal_invariants(&m);
696    }
697
698    #[test]
699    fn entry_or_insert_with_works() {
700        let mut m = M::new();
701
702        let v = m.entry_mut(7).or_insert_with(|| 123);
703        assert_eq!(*v, 123);
704
705        // 2回目は既存参照が返る
706        let v2 = m.entry_mut(7).or_insert_with(|| 999);
707        assert_eq!(*v2, 123);
708
709        assert_internal_invariants(&m);
710    }
711
712    #[test]
713    fn compare_with_std_hashmap_small_scripted() {
714        let mut m = M::new();
715        let mut o = HashMap::<u64, i64>::new();
716
717        // 混ぜた操作を固定シナリオで
718        for i in 0..50u64 {
719            m.insert(i, i as i64);
720            o.insert(i, i as i64);
721        }
722
723        for i in 0..50u64 {
724            if i % 3 == 0 {
725                let a = m.swap_remove(&i);
726                let b = o.remove(&i);
727                assert_eq!(a, RemoveResult::Removed {
728                    old_value: b.unwrap(),
729                    old_key: i,
730                    index: 0, // index はテストしない
731                });
732            }
733        }
734
735        for i in 0..50u64 {
736            if i % 5 == 0 {
737                m.insert(i, (i as i64) * 100);
738                o.insert(i, (i as i64) * 100);
739            }
740        }
741
742        assert_internal_invariants(&m);
743        assert_equals_oracle(&m, &o);
744    }
745
746    #[test]
747    fn randomized_ops_compare_with_oracle() {
748        use rand::{rngs::StdRng, Rng, SeedableRng};
749
750        let mut rng = StdRng::seed_from_u64(0xC0FFEE);
751        let mut m = M::new();
752        let mut o = HashMap::<u64, i64>::new();
753
754        // ある程度衝突や削除を踏む
755        const STEPS: usize = 30_000;
756        const KEY_SPACE: u64 = 2_000;
757
758        for _ in 0..STEPS {
759            let op = rng.gen_range(0..100);
760            let k = rng.gen_range(0..KEY_SPACE);
761            match op {
762                // insert (多め)
763                0..=59 => {
764                    let v = rng.gen_range(-1_000_000..=1_000_000);
765                    let a = m.insert(k, v);
766                    let b = o.insert(k, v);
767
768                    match (a, b) {
769                        (InsertResult::New { .. }, None) => {}
770                        (InsertResult::Override { old_key, old_value, .. }, Some(old)) => {
771                            assert_eq!(old_key, k);
772                            assert_eq!(old_value, old);
773                        }
774                        _ => panic!("insert mismatch"),
775                    }
776                }
777                // swap_remove
778                60..=79 => {
779                    let a = m.swap_remove(&k);
780                    let b = o.remove(&k);
781                    assert_eq!(a, RemoveResult::Removed {
782                        old_value: b.unwrap(),
783                        old_key: k,
784                        index: 0, // index はテストしない
785                    });
786                }
787                // get
788                80..=94 => {
789                    let a = m.get(&k).copied();
790                    let b = o.get(&k).copied();
791                    assert_eq!(a, b);
792                }
793                // contains
794                _ => {
795                    let a = m.contains_key(&k);
796                    let b = o.contains_key(&k);
797                    assert_eq!(a, b);
798                }
799            }
800
801            // たまに内部整合をチェック（重いので間引く）
802            if rng.gen_ratio(1, 200) {
803                assert_internal_invariants(&m);
804                assert_equals_oracle(&m, &o);
805            }
806        }
807
808        // 最後に必ず一致
809        assert_internal_invariants(&m);
810        assert_equals_oracle(&m, &o);
811    }
812
813    #[test]
814    fn empty_map_basics() {
815        let m = M::new();
816
817        assert_eq!(m.len(), 0);
818        assert!(m.get(&123).is_none());
819        assert!(!m.contains_key(&123));
820        // 空でも長さ整合は成立
821        assert_eq!(m.values.len(), 0);
822        assert_eq!(m.index_set.keys.len(), 0);
823        assert_eq!(m.index_set.hashes.len(), 0);
824    }
825
826    #[test]
827    fn swap_remove_single_element_roundtrip() {
828        let mut m = M::new();
829        m.insert(42, -7);
830        assert_internal_invariants(&m);
831
832        let v = m.swap_remove(&42);
833        assert_eq!(v, RemoveResult::Removed {
834            old_value: -7,
835            old_key: 42,
836            index: 0,
837        });
838        assert_eq!(m.len(), 0);
839        assert!(m.get(&42).is_none());
840        assert!(!m.contains_key(&42));
841
842        assert_internal_invariants(&m);
843    }
844
845    #[test]
846    fn remove_then_reinsert_same_key() {
847        let mut m = M::new();
848
849        m.insert(1, 10);
850        m.insert(2, 20);
851        m.insert(3, 30);
852        assert_internal_invariants(&m);
853
854        assert_eq!(m.swap_remove(&2), RemoveResult::Removed {
855            old_value: 20,
856            old_key: 2,
857            index: 1,
858        });
859        assert!(m.get(&2).is_none());
860        assert_internal_invariants(&m);
861
862        // 同じキーを再挿入しても table が壊れないこと
863        assert_eq!(m.insert(2, 200), InsertResult::New { index: 1 });
864        assert_eq!(m.get(&2).copied(), Some(200));
865        assert_internal_invariants(&m);
866    }
867
868    #[test]
869    fn from_kv_vec_builds_valid_map() {
870        let keys = vec![1u64, 2u64, 3u64, 10u64];
871        let values = vec![10i64, 20i64, 30i64, 100i64];
872
873        let m = M::from_kv_vec(keys.clone(), values.clone());
874        assert_eq!(m.len(), 4);
875
876        // 順序と内容が一致
877        assert_eq!(m.index_set.keys, keys);
878        assert_eq!(m.values, values);
879
880        assert_internal_invariants(&m);
881    }
882
883    #[test]
884    fn iter_order_matches_internal_storage_even_after_removes() {
885        let mut m = M::new();
886        for i in 0..8u64 {
887            m.insert(i, (i as i64) + 100);
888        }
889        assert_internal_invariants(&m);
890
891        // いくつか消して、内部順序が変わっても iter が keys/values と整合すること
892        assert_eq!(m.swap_remove(&0), RemoveResult::Removed {
893            old_value: 100,
894            old_key: 0,
895            index: 0,
896        });
897        assert_eq!(m.swap_remove(&5), RemoveResult::Removed {
898            old_value: 105,
899            old_key: 5,
900            index: 5,
901        });
902        assert_internal_invariants(&m);
903
904        let collected: Vec<(u64, i64)> = m.iter().map(|(k, v)| (*k, *v)).collect();
905        let expected: Vec<(u64, i64)> = m.index_set.keys.iter().copied().zip(m.values.iter().copied()).collect();
906        assert_eq!(collected, expected);
907    }
908
909    #[test]
910    fn num_23257_issue() {
911        const ITER_NUM: u64 = 223259;
912        let mut map: IndexMap<u64, Vec<u64>> = IndexMap::new();
913        for i in 0..ITER_NUM {
914            map.entry_mut(0).or_insert_with(Vec::new).push(i);
915        }
916        assert_eq!(map.len(), 1);
917        assert_eq!(map.get(&0).unwrap().len() as u64, ITER_NUM);
918    }
919}
tf_idf_vectorizer/utils/datastruct/map/index_map.rs

tf_idf_vectorizer/utils/datastruct/map/
index_map.rs
