1use std::collections::HashMap;
2use std::net::IpAddr;
3use std::num::NonZeroUsize;
4
5use lru::LruCache;
6use num_bigint::BigUint;
7use wp_model_core::model::{DataField, FValueStr, Value};
8
9const MAX_LOCAL_IDX: usize = 100_000;
11
12#[derive(Debug, Clone)]
13pub struct FieldQueryCache {
14 str_idx: HashMap<FValueStr, usize>,
15 i64_idx: HashMap<i64, usize>,
16 ip_idx: HashMap<IpAddr, usize>,
17 biguint_idx: HashMap<BigUint, usize>,
21 bool_idx: HashMap<bool, usize>,
23 float_idx: HashMap<u64, usize>,
25 text_idx: HashMap<String, usize>,
28 cache_data: LruCache<LocalCacheKey, Vec<DataField>>,
29 idx_num: usize,
30 generations: HashMap<u64, u64>,
33}
34
35pub type QueryLocalCache = FieldQueryCache;
36
37impl Default for FieldQueryCache {
38 fn default() -> Self {
39 Self::with_capacity(100)
40 }
41}
42
43impl FieldQueryCache {
44 pub fn with_capacity(size: usize) -> Self {
45 let size = size.max(1);
46 Self {
47 str_idx: HashMap::new(),
48 i64_idx: HashMap::new(),
49 ip_idx: HashMap::new(),
50 biguint_idx: HashMap::new(),
51 bool_idx: HashMap::new(),
52 float_idx: HashMap::new(),
53 text_idx: HashMap::new(),
54 cache_data: LruCache::new(NonZeroUsize::new(size).expect("non-zero cache size")),
55 idx_num: 0,
56 generations: HashMap::new(),
57 }
58 }
59
60 fn get_idx(&self, param: &DataField) -> Option<usize> {
61 match param.get_value() {
62 Value::Chars(v) => self.str_idx.get(v).copied(),
63 Value::Digit(v) => self.i64_idx.get(v).copied(),
64 Value::IpAddr(v) => self.ip_idx.get(v).copied(),
65 Value::BigUint(v) => self.biguint_idx.get(v).copied(),
66 Value::Bool(v) => self.bool_idx.get(v).copied(),
67 Value::Float(v) => self.float_idx.get(&v.to_bits()).copied(),
68 Value::Symbol(v) => self.text_idx.get(&v.to_string()).copied(),
70 Value::Time(v) => self.text_idx.get(&v.to_string()).copied(),
71 Value::Hex(v) => self.text_idx.get(&v.to_string()).copied(),
72 Value::IpNet(v) => self.text_idx.get(&v.to_string()).copied(),
73 Value::Domain(v) => self.text_idx.get(&v.to_string()).copied(),
74 Value::Url(v) => self.text_idx.get(&v.to_string()).copied(),
75 Value::Email(v) => self.text_idx.get(&v.to_string()).copied(),
76 Value::IdCard(v) => self.text_idx.get(&v.to_string()).copied(),
77 Value::MobilePhone(v) => self.text_idx.get(&v.to_string()).copied(),
78 _ => None,
79 }
80 }
81
82 fn up_text_idx(&mut self, key: String) -> Option<usize> {
84 if let Some(idx) = self.text_idx.get(&key) {
85 return Some(*idx);
86 }
87 self.idx_num += 1;
88 self.text_idx.insert(key, self.idx_num);
89 Some(self.idx_num)
90 }
91
92 fn try_up_idx(&mut self, param: &DataField) -> Option<usize> {
93 if self.idx_num >= MAX_LOCAL_IDX {
94 self.str_idx.clear();
96 self.i64_idx.clear();
97 self.ip_idx.clear();
98 self.biguint_idx.clear();
99 self.bool_idx.clear();
100 self.float_idx.clear();
101 self.text_idx.clear();
102 self.cache_data.clear();
103 self.generations.clear();
104 self.idx_num = 0;
105 }
106 match param.get_value() {
107 Value::Chars(v) => {
108 if let Some(idx) = self.str_idx.get(v) {
109 Some(*idx)
110 } else {
111 self.idx_num += 1;
112 self.str_idx.insert(v.clone(), self.idx_num);
113 Some(self.idx_num)
114 }
115 }
116 Value::Digit(v) => {
117 if let Some(idx) = self.i64_idx.get(v) {
118 Some(*idx)
119 } else {
120 self.idx_num += 1;
121 self.i64_idx.insert(*v, self.idx_num);
122 Some(self.idx_num)
123 }
124 }
125 Value::IpAddr(v) => {
126 if let Some(idx) = self.ip_idx.get(v) {
127 Some(*idx)
128 } else {
129 self.idx_num += 1;
130 self.ip_idx.insert(*v, self.idx_num);
131 Some(self.idx_num)
132 }
133 }
134 Value::BigUint(v) => {
135 if let Some(idx) = self.biguint_idx.get(v) {
136 Some(*idx)
137 } else {
138 self.idx_num += 1;
139 self.biguint_idx.insert(v.clone(), self.idx_num);
140 Some(self.idx_num)
141 }
142 }
143 Value::Bool(v) => {
144 if let Some(idx) = self.bool_idx.get(v) {
145 Some(*idx)
146 } else {
147 self.idx_num += 1;
148 self.bool_idx.insert(*v, self.idx_num);
149 Some(self.idx_num)
150 }
151 }
152 Value::Float(v) => {
153 let key = v.to_bits();
154 if let Some(idx) = self.float_idx.get(&key) {
155 Some(*idx)
156 } else {
157 self.idx_num += 1;
158 self.float_idx.insert(key, self.idx_num);
159 Some(self.idx_num)
160 }
161 }
162 Value::Symbol(v) => self.up_text_idx(v.to_string()),
164 Value::Time(v) => self.up_text_idx(v.to_string()),
165 Value::Hex(v) => self.up_text_idx(v.to_string()),
166 Value::IpNet(v) => self.up_text_idx(v.to_string()),
167 Value::Domain(v) => self.up_text_idx(v.to_string()),
168 Value::Url(v) => self.up_text_idx(v.to_string()),
169 Value::Email(v) => self.up_text_idx(v.to_string()),
170 Value::IdCard(v) => self.up_text_idx(v.to_string()),
171 Value::MobilePhone(v) => self.up_text_idx(v.to_string()),
172 _ => None,
173 }
174 }
175
176 fn prepare_generation_scoped(&mut self, scope: u64, generation: u64) {
179 if self.generations.get(&scope) == Some(&generation) {
180 return;
181 }
182 self.generations.insert(scope, generation);
183 let stale: Vec<LocalCacheKey> = self
184 .cache_data
185 .iter()
186 .filter(|(key, _)| key.scope_hash == scope)
187 .map(|(key, _)| key.clone())
188 .collect();
189 for key in stale {
190 self.cache_data.pop(&key);
191 }
192 }
193}
194
195#[derive(PartialEq, Eq, Hash, Debug, Clone)]
196pub enum EnumSizeIndex {
197 Idx1(usize),
198 Idx2(usize, usize),
199 Idx3(usize, usize, usize),
200 Idx4(usize, usize, usize, usize),
201 Idx5(usize, usize, usize, usize, usize),
202 Idx6(usize, usize, usize, usize, usize, usize),
203}
204
205#[derive(PartialEq, Eq, Hash, Debug, Clone)]
206struct LocalCacheKey {
207 scope_hash: u64,
208 idxs: EnumSizeIndex,
209}
210
211pub trait CacheAble<P, T, const N: usize> {
212 fn prepare_generation(&mut self, _scope: u64, _generation: u64) {}
214 fn save_scoped(&mut self, _scope_hash: u64, params: &[P; N], result: T) {
215 self.save(params, result);
216 }
217 fn fetch_scoped(&self, _scope_hash: u64, params: &[P; N]) -> Option<&T> {
218 self.fetch(params)
219 }
220 fn save(&mut self, params: &[P; N], result: T);
221 fn fetch(&self, params: &[P; N]) -> Option<&T>;
222}
223
224impl CacheAble<DataField, Vec<DataField>, 1> for FieldQueryCache {
225 fn prepare_generation(&mut self, scope: u64, generation: u64) {
226 self.prepare_generation_scoped(scope, generation);
227 }
228 fn save_scoped(&mut self, scope_hash: u64, params: &[DataField; 1], result: Vec<DataField>) {
229 if let Some(i0) = self.try_up_idx(¶ms[0]) {
230 self.cache_data.put(
231 LocalCacheKey {
232 scope_hash,
233 idxs: EnumSizeIndex::Idx1(i0),
234 },
235 result,
236 );
237 }
238 }
239
240 fn fetch_scoped(&self, scope_hash: u64, params: &[DataField; 1]) -> Option<&Vec<DataField>> {
241 if let Some(i0) = self.get_idx(¶ms[0]) {
242 return self.cache_data.peek(&LocalCacheKey {
243 scope_hash,
244 idxs: EnumSizeIndex::Idx1(i0),
245 });
246 }
247 None
248 }
249
250 fn save(&mut self, params: &[DataField; 1], result: Vec<DataField>) {
251 self.save_scoped(0, params, result);
252 }
253
254 fn fetch(&self, params: &[DataField; 1]) -> Option<&Vec<DataField>> {
255 self.fetch_scoped(0, params)
256 }
257}
258
259impl CacheAble<DataField, Vec<DataField>, 2> for FieldQueryCache {
260 fn prepare_generation(&mut self, scope: u64, generation: u64) {
261 self.prepare_generation_scoped(scope, generation);
262 }
263
264 fn save_scoped(&mut self, scope_hash: u64, params: &[DataField; 2], result: Vec<DataField>) {
265 if let (Some(i0), Some(i1)) = (self.try_up_idx(¶ms[0]), self.try_up_idx(¶ms[1])) {
266 self.cache_data.put(
267 LocalCacheKey {
268 scope_hash,
269 idxs: EnumSizeIndex::Idx2(i0, i1),
270 },
271 result,
272 );
273 }
274 }
275
276 fn fetch_scoped(&self, scope_hash: u64, params: &[DataField; 2]) -> Option<&Vec<DataField>> {
277 if let (Some(i0), Some(i1)) = (self.get_idx(¶ms[0]), self.get_idx(¶ms[1])) {
278 return self.cache_data.peek(&LocalCacheKey {
279 scope_hash,
280 idxs: EnumSizeIndex::Idx2(i0, i1),
281 });
282 }
283 None
284 }
285
286 fn save(&mut self, params: &[DataField; 2], result: Vec<DataField>) {
287 self.save_scoped(0, params, result);
288 }
289
290 fn fetch(&self, params: &[DataField; 2]) -> Option<&Vec<DataField>> {
291 self.fetch_scoped(0, params)
292 }
293}
294
295impl CacheAble<DataField, Vec<DataField>, 3> for FieldQueryCache {
296 fn prepare_generation(&mut self, scope: u64, generation: u64) {
297 self.prepare_generation_scoped(scope, generation);
298 }
299
300 fn save_scoped(&mut self, scope_hash: u64, params: &[DataField; 3], result: Vec<DataField>) {
301 if let (Some(i0), Some(i1), Some(i2)) = (
302 self.try_up_idx(¶ms[0]),
303 self.try_up_idx(¶ms[1]),
304 self.try_up_idx(¶ms[2]),
305 ) {
306 self.cache_data.put(
307 LocalCacheKey {
308 scope_hash,
309 idxs: EnumSizeIndex::Idx3(i0, i1, i2),
310 },
311 result,
312 );
313 }
314 }
315
316 fn fetch_scoped(&self, scope_hash: u64, params: &[DataField; 3]) -> Option<&Vec<DataField>> {
317 if let (Some(i0), Some(i1), Some(i2)) = (
318 self.get_idx(¶ms[0]),
319 self.get_idx(¶ms[1]),
320 self.get_idx(¶ms[2]),
321 ) {
322 return self.cache_data.peek(&LocalCacheKey {
323 scope_hash,
324 idxs: EnumSizeIndex::Idx3(i0, i1, i2),
325 });
326 }
327 None
328 }
329
330 fn save(&mut self, params: &[DataField; 3], result: Vec<DataField>) {
331 self.save_scoped(0, params, result);
332 }
333
334 fn fetch(&self, params: &[DataField; 3]) -> Option<&Vec<DataField>> {
335 self.fetch_scoped(0, params)
336 }
337}
338
339impl CacheAble<DataField, Vec<DataField>, 4> for FieldQueryCache {
340 fn prepare_generation(&mut self, scope: u64, generation: u64) {
341 self.prepare_generation_scoped(scope, generation);
342 }
343
344 fn save_scoped(&mut self, scope_hash: u64, params: &[DataField; 4], result: Vec<DataField>) {
345 if let (Some(i0), Some(i1), Some(i2), Some(i3)) = (
346 self.try_up_idx(¶ms[0]),
347 self.try_up_idx(¶ms[1]),
348 self.try_up_idx(¶ms[2]),
349 self.try_up_idx(¶ms[3]),
350 ) {
351 self.cache_data.put(
352 LocalCacheKey {
353 scope_hash,
354 idxs: EnumSizeIndex::Idx4(i0, i1, i2, i3),
355 },
356 result,
357 );
358 }
359 }
360
361 fn fetch_scoped(&self, scope_hash: u64, params: &[DataField; 4]) -> Option<&Vec<DataField>> {
362 if let (Some(i0), Some(i1), Some(i2), Some(i3)) = (
363 self.get_idx(¶ms[0]),
364 self.get_idx(¶ms[1]),
365 self.get_idx(¶ms[2]),
366 self.get_idx(¶ms[3]),
367 ) {
368 return self.cache_data.peek(&LocalCacheKey {
369 scope_hash,
370 idxs: EnumSizeIndex::Idx4(i0, i1, i2, i3),
371 });
372 }
373 None
374 }
375
376 fn save(&mut self, params: &[DataField; 4], result: Vec<DataField>) {
377 self.save_scoped(0, params, result);
378 }
379
380 fn fetch(&self, params: &[DataField; 4]) -> Option<&Vec<DataField>> {
381 self.fetch_scoped(0, params)
382 }
383}
384
385impl CacheAble<DataField, Vec<DataField>, 5> for FieldQueryCache {
386 fn prepare_generation(&mut self, scope: u64, generation: u64) {
387 self.prepare_generation_scoped(scope, generation);
388 }
389
390 fn save_scoped(&mut self, scope_hash: u64, params: &[DataField; 5], result: Vec<DataField>) {
391 if let (Some(i0), Some(i1), Some(i2), Some(i3), Some(i4)) = (
392 self.try_up_idx(¶ms[0]),
393 self.try_up_idx(¶ms[1]),
394 self.try_up_idx(¶ms[2]),
395 self.try_up_idx(¶ms[3]),
396 self.try_up_idx(¶ms[4]),
397 ) {
398 self.cache_data.put(
399 LocalCacheKey {
400 scope_hash,
401 idxs: EnumSizeIndex::Idx5(i0, i1, i2, i3, i4),
402 },
403 result,
404 );
405 }
406 }
407
408 fn fetch_scoped(&self, scope_hash: u64, params: &[DataField; 5]) -> Option<&Vec<DataField>> {
409 if let (Some(i0), Some(i1), Some(i2), Some(i3), Some(i4)) = (
410 self.get_idx(¶ms[0]),
411 self.get_idx(¶ms[1]),
412 self.get_idx(¶ms[2]),
413 self.get_idx(¶ms[3]),
414 self.get_idx(¶ms[4]),
415 ) {
416 return self.cache_data.peek(&LocalCacheKey {
417 scope_hash,
418 idxs: EnumSizeIndex::Idx5(i0, i1, i2, i3, i4),
419 });
420 }
421 None
422 }
423
424 fn save(&mut self, params: &[DataField; 5], result: Vec<DataField>) {
425 self.save_scoped(0, params, result);
426 }
427
428 fn fetch(&self, params: &[DataField; 5]) -> Option<&Vec<DataField>> {
429 self.fetch_scoped(0, params)
430 }
431}
432
433impl CacheAble<DataField, Vec<DataField>, 6> for FieldQueryCache {
434 fn prepare_generation(&mut self, scope: u64, generation: u64) {
435 self.prepare_generation_scoped(scope, generation);
436 }
437
438 fn save_scoped(&mut self, scope_hash: u64, params: &[DataField; 6], result: Vec<DataField>) {
439 if let (Some(i0), Some(i1), Some(i2), Some(i3), Some(i4), Some(i5)) = (
440 self.try_up_idx(¶ms[0]),
441 self.try_up_idx(¶ms[1]),
442 self.try_up_idx(¶ms[2]),
443 self.try_up_idx(¶ms[3]),
444 self.try_up_idx(¶ms[4]),
445 self.try_up_idx(¶ms[5]),
446 ) {
447 self.cache_data.put(
448 LocalCacheKey {
449 scope_hash,
450 idxs: EnumSizeIndex::Idx6(i0, i1, i2, i3, i4, i5),
451 },
452 result,
453 );
454 }
455 }
456
457 fn fetch_scoped(&self, scope_hash: u64, params: &[DataField; 6]) -> Option<&Vec<DataField>> {
458 if let (Some(i0), Some(i1), Some(i2), Some(i3), Some(i4), Some(i5)) = (
459 self.get_idx(¶ms[0]),
460 self.get_idx(¶ms[1]),
461 self.get_idx(¶ms[2]),
462 self.get_idx(¶ms[3]),
463 self.get_idx(¶ms[4]),
464 self.get_idx(¶ms[5]),
465 ) {
466 return self.cache_data.peek(&LocalCacheKey {
467 scope_hash,
468 idxs: EnumSizeIndex::Idx6(i0, i1, i2, i3, i4, i5),
469 });
470 }
471 None
472 }
473
474 fn save(&mut self, params: &[DataField; 6], result: Vec<DataField>) {
475 self.save_scoped(0, params, result);
476 }
477
478 fn fetch(&self, params: &[DataField; 6]) -> Option<&Vec<DataField>> {
479 self.fetch_scoped(0, params)
480 }
481}
482
483#[cfg(test)]
484#[allow(clippy::cloned_ref_to_slice_refs)]
487mod tests {
488 use super::*;
489 use num_bigint::BigUint;
490 use wp_model_core::model::DataType;
491
492 fn biguint_field(v: u64) -> DataField {
494 DataField::new(
495 DataType::BigInt,
496 "k".to_string(),
497 Value::BigUint(BigUint::from(v)),
498 )
499 }
500
501 fn result_row(v: u64) -> Vec<DataField> {
502 vec![DataField::from_digit("country_iso_code", v as i64)]
503 }
504
505 #[test]
506 fn biguint_single_param_hit_and_miss() {
507 let mut cache = FieldQueryCache::default();
508 let p = biguint_field(42);
509 let saved = result_row(1);
510
511 assert!(cache.fetch(&[p.clone()]).is_none());
513 cache.save(&[p.clone()], saved.clone());
514 assert_eq!(cache.fetch(&[p]), Some(&saved));
516 assert!(cache.fetch(&[biguint_field(43)]).is_none());
518 }
519
520 #[test]
521 fn biguint_multi_param_hit_and_partial_miss() {
522 let mut cache = FieldQueryCache::default();
523 let a = biguint_field(10);
524 let b = biguint_field(20);
525 let saved = result_row(7);
526 cache.save(&[a.clone(), b.clone()], saved.clone());
527
528 assert_eq!(cache.fetch(&[a.clone(), b.clone()]), Some(&saved));
530 assert!(cache.fetch(&[a, biguint_field(21)]).is_none());
532 }
533
534 #[test]
535 fn biguint_isolated_from_chars_digit_ip() {
536 let mut cache = FieldQueryCache::default();
537 let big = biguint_field(1);
538 let chars = DataField::from_chars("k", "1");
539 let digit = DataField::from_digit("k", 1);
540 let ip = DataField::from_ip("k", IpAddr::from([1, 1, 1, 1]));
541
542 cache.save(&[big.clone()], result_row(1));
543 cache.save(&[chars.clone()], result_row(2));
544 cache.save(&[digit.clone()], result_row(3));
545 cache.save(&[ip.clone()], result_row(4));
546
547 assert_eq!(cache.fetch(&[big]), Some(&result_row(1)));
549 assert_eq!(cache.fetch(&[chars]), Some(&result_row(2)));
550 assert_eq!(cache.fetch(&[digit]), Some(&result_row(3)));
551 assert_eq!(cache.fetch(&[ip]), Some(&result_row(4)));
552 }
553
554 #[test]
555 fn bool_and_float_param_cache() {
556 let mut cache = FieldQueryCache::default();
557 let b_true = DataField::from_bool("k", true);
558 let f_15 = DataField::from_float("k", 1.5);
559
560 assert!(cache.fetch(&[b_true.clone()]).is_none());
562 cache.save(&[b_true.clone()], result_row(1));
563 assert_eq!(cache.fetch(&[b_true.clone()]), Some(&result_row(1)));
564 assert!(cache.fetch(&[DataField::from_bool("k", false)]).is_none());
565
566 assert!(cache.fetch(&[f_15.clone()]).is_none());
568 cache.save(&[f_15.clone()], result_row(2));
569 assert_eq!(cache.fetch(&[f_15.clone()]), Some(&result_row(2)));
570 assert!(cache.fetch(&[DataField::from_float("k", 2.5)]).is_none());
571 }
572
573 #[test]
574 fn text_like_param_cache_and_isolation_from_chars() {
575 let mut cache = FieldQueryCache::default();
576 let domain = DataField::from_domain("k", "example.com");
577 let chars = DataField::from_chars("k", "example.com");
578
579 cache.save(&[domain.clone()], result_row(1));
580 cache.save(&[chars.clone()], result_row(2));
581
582 assert_eq!(
584 cache.fetch(&[DataField::from_domain("k", "example.com")]),
585 Some(&result_row(1))
586 );
587 assert!(
589 cache
590 .fetch(&[DataField::from_domain("k", "other.com")])
591 .is_none()
592 );
593 assert_eq!(cache.fetch(&[chars]), Some(&result_row(2)));
595 }
596
597 #[test]
598 fn mixed_type_multi_param_and_save_idempotent() {
599 let mut cache = FieldQueryCache::default();
600 let big = biguint_field(7);
601 let code = DataField::from_chars("code", "CN");
602 let flag = DataField::from_bool("ok", true);
603
604 cache.save(&[big.clone(), code.clone(), flag.clone()], result_row(1));
606 assert_eq!(
607 cache.fetch(&[big.clone(), code.clone(), flag.clone()]),
608 Some(&result_row(1))
609 );
610 assert!(
612 cache
613 .fetch(&[biguint_field(8), code.clone(), flag.clone()])
614 .is_none()
615 );
616
617 assert_eq!(cache.idx_num, 3);
619 cache.save(&[big.clone()], result_row(2));
620 cache.save(&[big.clone()], result_row(3));
621 cache.save(&[code.clone()], result_row(4));
622 cache.save(&[code.clone()], result_row(5));
623 assert_eq!(cache.idx_num, 3, "重复 save 应复用既有索引号");
624 assert_eq!(cache.fetch(&[big]), Some(&result_row(3)));
626 assert_eq!(cache.fetch(&[code]), Some(&result_row(5)));
627 }
628
629 #[test]
630 fn biguint_index_reset_when_cap_reached() {
631 let mut cache = FieldQueryCache::default();
632 let first = biguint_field(0);
633 cache.save(&[first.clone()], result_row(0));
634 assert_eq!(cache.fetch(&[first.clone()]), Some(&result_row(0)));
635
636 for i in 1..=MAX_LOCAL_IDX {
638 cache.save(&[biguint_field(i as u64)], result_row(i as u64));
639 }
640 assert!(
642 cache.fetch(&[first.clone()]).is_none(),
643 "触顶重置后旧缓存应被清空"
644 );
645 assert!(
646 !cache.biguint_idx.contains_key(&BigUint::from(0u64)),
647 "触顶重置后 BigUint 索引应被清理"
648 );
649
650 cache.save(&[first.clone()], result_row(0));
652 assert_eq!(cache.fetch(&[first]), Some(&result_row(0)));
653 }
654}