1use super::newton::newton_raphson_with_constants;
2use super::{Match, SearchConfig, SearchContext, SearchStats, SearchTimer};
3
4use crate::expr::EvaluatedExpr;
5
6use crate::pool::TopKPool;
7
8use crate::thresholds::{
9 ADAPTIVE_COMPLEXITY_SCALE, ADAPTIVE_EXACT_MATCH_FACTOR, ADAPTIVE_POOL_FULLNESS_SCALE,
10 BASE_SEARCH_RADIUS_FACTOR, DEGENERATE_RANGE_TOLERANCE, DEGENERATE_TEST_THRESHOLD,
11 EXACT_MATCH_TOLERANCE, MAX_SEARCH_RADIUS_FACTOR, NEWTON_FINAL_TOLERANCE,
12 STRICT_GATE_CAPACITY_FRACTION, STRICT_GATE_FACTOR, TIER_0_MAX, TIER_1_MAX, TIER_2_MAX,
13};
14
15pub struct ExprDatabase {
18 rhs_sorted: Vec<EvaluatedExpr>,
20}
21
22#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
31pub enum ComplexityTier {
32 Tier0,
34 Tier1,
36 Tier2,
38 Tier3,
40}
41
42impl ComplexityTier {
43 #[inline]
45 pub fn from_complexity(complexity: u32) -> Self {
46 if complexity <= TIER_0_MAX {
47 ComplexityTier::Tier0
48 } else if complexity <= TIER_1_MAX {
49 ComplexityTier::Tier1
50 } else if complexity <= TIER_2_MAX {
51 ComplexityTier::Tier2
52 } else {
53 ComplexityTier::Tier3
54 }
55 }
56}
57
58pub struct TieredExprDatabase {
64 tiers: [Vec<EvaluatedExpr>; 4],
66 total_count: usize,
68}
69
70impl TieredExprDatabase {
71 pub fn new() -> Self {
73 Self {
74 tiers: [Vec::new(), Vec::new(), Vec::new(), Vec::new()],
75 total_count: 0,
76 }
77 }
78
79 pub fn insert(&mut self, expr: EvaluatedExpr) {
81 let tier = ComplexityTier::from_complexity(expr.expr.complexity());
82 let tier_idx = tier as usize;
83 self.tiers[tier_idx].push(expr);
84 self.total_count += 1;
85 }
86
87 pub fn finalize(&mut self) {
89 for tier in &mut self.tiers {
90 tier.sort_by(|a, b| a.value.total_cmp(&b.value));
91 }
92 }
93
94 pub fn total_count(&self) -> usize {
96 self.total_count
97 }
98
99 #[allow(dead_code)]
101 pub fn tier_count(&self, tier: ComplexityTier) -> usize {
102 self.tiers[tier as usize].len()
103 }
104
105 #[cfg(test)]
106 pub(super) fn tier(&self, tier: ComplexityTier) -> &[EvaluatedExpr] {
107 &self.tiers[tier as usize]
108 }
109
110 #[allow(dead_code)]
112 pub fn range_in_tier(&self, tier: ComplexityTier, low: f64, high: f64) -> &[EvaluatedExpr] {
113 let tier_vec = &self.tiers[tier as usize];
114 let start = tier_vec.partition_point(|e| e.value < low);
115 let end = tier_vec.partition_point(|e| e.value <= high);
116 &tier_vec[start..end]
117 }
118
119 pub fn count_in_range(&self, low: f64, high: f64) -> usize {
121 self.tiers
122 .iter()
123 .map(|tier| {
124 let start = tier.partition_point(|e| e.value < low);
125 let end = tier.partition_point(|e| e.value <= high);
126 end.saturating_sub(start)
127 })
128 .sum()
129 }
130
131 pub fn iter_tiers_in_range(&self, low: f64, high: f64) -> TieredRangeIter<'_> {
134 TieredRangeIter::new(self, low, high)
135 }
136}
137
138impl Default for TieredExprDatabase {
139 fn default() -> Self {
140 Self::new()
141 }
142}
143
144pub struct TieredRangeIter<'a> {
146 db: &'a TieredExprDatabase,
147 low: f64,
148 high: f64,
149 current_tier: usize,
150 current_start: usize,
151 current_end: usize,
152}
153
154impl<'a> TieredRangeIter<'a> {
155 fn new(db: &'a TieredExprDatabase, low: f64, high: f64) -> Self {
156 let mut iter = Self {
157 db,
158 low,
159 high,
160 current_tier: 0,
161 current_start: 0,
162 current_end: 0,
163 };
164 iter.find_next_nonempty_tier();
165 iter
166 }
167
168 fn calculate_tier_range(&self, tier_idx: usize) -> (usize, usize) {
170 let tier_vec = &self.db.tiers[tier_idx];
171 let start = tier_vec.partition_point(|e| e.value < self.low);
172 let end = tier_vec.partition_point(|e| e.value <= self.high);
173 (start, end)
174 }
175
176 fn find_next_nonempty_tier(&mut self) {
178 while self.current_tier < 4 {
179 let (start, end) = self.calculate_tier_range(self.current_tier);
180 self.current_start = start;
181 self.current_end = end;
182
183 if self.current_start < self.current_end {
184 return;
186 }
187 self.current_tier += 1;
188 }
189 }
190}
191
192impl<'a> Iterator for TieredRangeIter<'a> {
193 type Item = &'a EvaluatedExpr;
194
195 fn next(&mut self) -> Option<Self::Item> {
196 while self.current_tier < 4 {
197 if self.current_start < self.current_end {
198 let expr = &self.db.tiers[self.current_tier][self.current_start];
199 self.current_start += 1;
200 return Some(expr);
201 }
202 self.current_tier += 1;
203 self.find_next_nonempty_tier();
204 }
205 None
206 }
207}
208
209#[inline]
230pub(super) fn calculate_adaptive_search_radius(
231 derivative: f64,
232 complexity: u32,
233 pool_size: usize,
234 pool_capacity: usize,
235 best_error: f64,
236 accept_error: f64,
237) -> f64 {
238 let deriv_abs = derivative.abs();
239
240 let base_radius = BASE_SEARCH_RADIUS_FACTOR * deriv_abs;
242
243 let normalized_complexity = (complexity as f64) / 50.0;
246 let complexity_factor = 1.0 / (1.0 + ADAPTIVE_COMPLEXITY_SCALE * normalized_complexity);
247
248 let pool_fraction = if pool_capacity > 0 {
250 pool_size as f64 / pool_capacity as f64
251 } else {
252 0.0
253 };
254 let pool_factor = (1.0 - ADAPTIVE_POOL_FULLNESS_SCALE * pool_fraction).max(0.1);
255
256 let exact_factor = if best_error < NEWTON_FINAL_TOLERANCE {
258 ADAPTIVE_EXACT_MATCH_FACTOR
259 } else {
260 1.0
261 };
262
263 let radius = base_radius * complexity_factor * pool_factor * exact_factor;
265
266 let gate_factor = if pool_capacity > 0
270 && pool_size as f64 >= pool_capacity as f64 * STRICT_GATE_CAPACITY_FRACTION
271 {
272 STRICT_GATE_FACTOR
273 } else {
274 1.0
275 };
276 let coarse_error_cap = accept_error.max(NEWTON_FINAL_TOLERANCE) * gate_factor;
277 let gated_radius_cap = deriv_abs * coarse_error_cap;
278
279 let min_radius = 0.1 * deriv_abs; radius
282 .max(min_radius)
283 .min(MAX_SEARCH_RADIUS_FACTOR * deriv_abs)
284 .min(gated_radius_cap)
285}
286
287impl ExprDatabase {
288 pub fn new() -> Self {
289 Self {
290 rhs_sorted: Vec::new(),
291 }
292 }
293
294 pub fn insert_rhs(&mut self, mut exprs: Vec<EvaluatedExpr>) {
297 exprs.sort_by(|a, b| a.value.total_cmp(&b.value));
300 self.rhs_sorted = exprs;
301 }
302
303 pub fn rhs_count(&self) -> usize {
305 self.rhs_sorted.len()
306 }
307
308 #[inline]
311 pub fn range(&self, low: f64, high: f64) -> &[EvaluatedExpr] {
312 let start = self.rhs_sorted.partition_point(|e| e.value < low);
314 let end = self.rhs_sorted.partition_point(|e| e.value <= high);
315 &self.rhs_sorted[start..end]
316 }
317
318 #[allow(dead_code)]
323 pub fn find_matches(&self, lhs_exprs: &[EvaluatedExpr], config: &SearchConfig) -> Vec<Match> {
324 let (matches, _stats) = self.find_matches_with_stats(lhs_exprs, config);
325 matches
326 }
327
328 pub fn find_matches_with_context(
330 &self,
331 lhs_exprs: &[EvaluatedExpr],
332 context: &SearchContext<'_>,
333 ) -> Vec<Match> {
334 let (matches, _stats) = self.find_matches_with_stats_and_context(lhs_exprs, context);
335 matches
336 }
337
338 pub fn find_matches_with_stats(
340 &self,
341 lhs_exprs: &[EvaluatedExpr],
342 config: &SearchConfig,
343 ) -> (Vec<Match>, SearchStats) {
344 let context = SearchContext::new(config);
345 self.find_matches_with_stats_and_context(lhs_exprs, &context)
346 }
347
348 pub fn find_matches_with_stats_and_context(
350 &self,
351 lhs_exprs: &[EvaluatedExpr],
352 context: &SearchContext<'_>,
353 ) -> (Vec<Match>, SearchStats) {
354 let config = context.config;
355 let mut stats = SearchStats::new();
356 let search_start = SearchTimer::start();
357
358 let initial_max_error = config.max_error.max(1e-12);
360
361 let mut pool = TopKPool::new_with_diagnostics(
363 config.max_matches,
364 initial_max_error,
365 config.show_db_adds,
366 config.ranking_mode,
367 );
368
369 let mut sorted_lhs: Vec<_> = lhs_exprs.iter().collect();
371 sorted_lhs.sort_by_key(|e| e.expr.complexity());
372
373 let mut early_exit = false;
375
376 for lhs in sorted_lhs {
377 if !match_one_lhs(self, lhs, context, &mut pool, &mut stats) {
378 early_exit = true;
379 break;
380 }
381 }
382
383 stats.pool_insertions = pool.stats.insertions;
385 stats.pool_rejections_error = pool.stats.rejections_error;
386 stats.pool_rejections_dedupe = pool.stats.rejections_dedupe;
387 stats.pool_evictions = pool.stats.evictions;
388 stats.pool_final_size = pool.len();
389 stats.pool_best_error = pool.best_error;
390 stats.search_time = search_start.elapsed();
391 stats.early_exit = early_exit;
392
393 (pool.into_sorted(), stats)
395 }
396
397 #[cfg(feature = "parallel")]
423 pub fn find_matches_turbo_with_stats_and_context(
424 &self,
425 lhs_exprs: &[EvaluatedExpr],
426 context: &SearchContext<'_>,
427 ) -> (Vec<Match>, SearchStats) {
428 use rayon::prelude::*;
429
430 let config = context.config;
431 let search_start = SearchTimer::start();
432 let initial_max_error = config.max_error.max(1e-12);
433
434 let mut sorted_lhs: Vec<_> = lhs_exprs.iter().collect();
435 sorted_lhs.sort_by_key(|e| e.expr.complexity());
436
437 let threads = rayon::current_num_threads().max(1);
440 let target_bands = (threads * 4).max(1);
441 let band_size = sorted_lhs.len().div_ceil(target_bands).max(1);
442
443 let (all_matches, mut stats) = sorted_lhs
444 .par_chunks(band_size)
445 .map(|band| {
446 let mut local_pool = TopKPool::new_with_diagnostics(
447 config.max_matches,
448 initial_max_error,
449 config.show_db_adds,
450 config.ranking_mode,
451 );
452 let mut local_stats = SearchStats::new();
453 for lhs in band {
454 if !match_one_lhs(self, lhs, context, &mut local_pool, &mut local_stats) {
455 local_stats.early_exit = true;
458 break;
459 }
460 }
461 (local_pool.into_sorted(), local_stats)
462 })
463 .reduce(
464 || (Vec::new(), SearchStats::new()),
465 |(mut acc_matches, mut acc_stats), (band_matches, band_stats)| {
466 acc_stats.merge_search_counters(&band_stats);
467 acc_matches.extend(band_matches);
468 (acc_matches, acc_stats)
469 },
470 );
471
472 let mut pool = TopKPool::new_with_diagnostics(
475 config.max_matches,
476 initial_max_error,
477 config.show_db_adds,
478 config.ranking_mode,
479 );
480 for m in all_matches {
481 pool.try_insert(m);
482 }
483
484 stats.pool_insertions = pool.stats.insertions;
485 stats.pool_rejections_error = pool.stats.rejections_error;
486 stats.pool_rejections_dedupe = pool.stats.rejections_dedupe;
487 stats.pool_evictions = pool.stats.evictions;
488 stats.pool_final_size = pool.len();
489 stats.pool_best_error = pool.best_error;
490 stats.search_time = search_start.elapsed();
491
492 (pool.into_sorted(), stats)
493 }
494}
495
496fn match_one_lhs(
503 db: &ExprDatabase,
504 lhs: &EvaluatedExpr,
505 context: &SearchContext<'_>,
506 pool: &mut TopKPool,
507 stats: &mut SearchStats,
508) -> bool {
509 let config = context.config;
510
511 let exact_ceiling = pool.exact_complexity_ceiling();
521 if let Some(ceiling) = exact_ceiling {
522 if lhs.expr.complexity() >= ceiling {
523 return false;
524 }
525 }
526
527 if lhs.value.abs() < config.zero_value_threshold {
531 if config.show_pruned_range {
532 eprintln!(
533 " [pruned range] value={:.6e} reason=\"near-zero\" expr=\"{}\"",
534 lhs.value,
535 lhs.expr.to_infix()
536 );
537 }
538 return true;
539 }
540
541 if super::should_skip_degenerate_lhs(lhs, config.target, &context.eval) {
544 return true;
545 }
546 if lhs.derivative.abs() < DEGENERATE_TEST_THRESHOLD {
547 let val_error = DEGENERATE_RANGE_TOLERANCE;
550 let low = lhs.value - val_error;
551 let high = lhs.value + val_error;
552
553 stats.lhs_tested += 1;
554 let rhs_slice = db.range(low, high);
555 stats.record_candidate_window(lhs, rhs_slice.len());
556 for rhs in rhs_slice {
557 if !config.rhs_symbol_allowed(&rhs.expr) {
558 continue;
559 }
560 stats.candidates_tested += 1;
561 if config.show_match_checks {
562 eprintln!(
563 " [match] checking lhs={:.6} rhs={:.6}",
564 lhs.value, rhs.value
565 );
566 }
567 let val_diff = (lhs.value - rhs.value).abs();
568 if val_diff < val_error && pool.would_accept(0.0, true) {
569 let m = Match {
570 lhs: lhs.clone(),
571 rhs: rhs.clone(),
572 x_value: config.target,
573 error: 0.0,
574 complexity: lhs.expr.complexity() + rhs.expr.complexity(),
575 };
576 pool.try_insert(m);
577 }
578 }
579 return true;
580 }
581
582 stats.lhs_tested += 1;
583
584 let search_radius = calculate_adaptive_search_radius(
586 lhs.derivative,
587 lhs.expr.complexity(),
588 pool.len(),
589 config.max_matches,
590 pool.best_error,
591 pool.accept_error,
592 );
593 let low = lhs.value - search_radius;
594 let high = lhs.value + search_radius;
595
596 let rhs_slice = db.range(low, high);
597 stats.record_candidate_window(lhs, rhs_slice.len());
598 for rhs in rhs_slice {
599 if !config.rhs_symbol_allowed(&rhs.expr) {
600 continue;
601 }
602
603 if let Some(ceiling) = exact_ceiling {
614 if lhs.expr.complexity() + rhs.expr.complexity() > ceiling {
615 continue;
616 }
617 }
618
619 stats.candidates_tested += 1;
620 if config.show_match_checks {
621 eprintln!(
622 " [match] checking lhs={:.6} rhs={:.6}",
623 lhs.value, rhs.value
624 );
625 }
626
627 let val_diff = lhs.value - rhs.value;
629 let x_delta = -val_diff / lhs.derivative;
630 let coarse_error = x_delta.abs();
631
632 let is_potentially_exact = coarse_error < NEWTON_FINAL_TOLERANCE;
635 if !pool.would_accept_strict(coarse_error, is_potentially_exact) {
636 stats.strict_gate_rejections += 1;
637 continue;
638 }
639
640 if !config.refine_with_newton {
641 let refined_x = config.target + x_delta;
642 let refined_error = x_delta;
643 let is_exact = refined_error.abs() < EXACT_MATCH_TOLERANCE;
644
645 if pool.would_accept(refined_error.abs(), is_exact) {
646 let m = Match {
647 lhs: lhs.clone(),
648 rhs: rhs.clone(),
649 x_value: refined_x,
650 error: refined_error,
651 complexity: lhs.expr.complexity() + rhs.expr.complexity(),
652 };
653
654 pool.try_insert(m);
655
656 if config.stop_at_exact && is_exact {
657 return false;
658 }
659 if let Some(threshold) = config.stop_below {
660 if refined_error.abs() < threshold {
661 return false;
662 }
663 }
664 }
665 continue;
666 }
667
668 stats.newton_calls += 1;
670 let initial_guess = config.target + x_delta;
671 if let Some(refined_x) = newton_raphson_with_constants(
672 &lhs.expr,
673 rhs.value,
674 initial_guess,
675 config.newton_iterations,
676 &context.eval,
677 config.show_newton,
678 config.derivative_margin,
679 ) {
680 stats.newton_success += 1;
681 let refined_error = refined_x - config.target;
682 let is_exact = refined_error.abs() < EXACT_MATCH_TOLERANCE;
683
684 if pool.would_accept(refined_error.abs(), is_exact) {
685 let m = Match {
686 lhs: lhs.clone(),
687 rhs: rhs.clone(),
688 x_value: refined_x,
689 error: refined_error,
690 complexity: lhs.expr.complexity() + rhs.expr.complexity(),
691 };
692
693 pool.try_insert(m);
694
695 if config.stop_at_exact && is_exact {
696 return false;
697 }
698 if let Some(threshold) = config.stop_below {
699 if refined_error.abs() < threshold {
700 return false;
701 }
702 }
703 }
704 }
705 }
706
707 true
708}
709
710impl Default for ExprDatabase {
711 fn default() -> Self {
712 Self::new()
713 }
714}
715
716#[cfg(test)]
717mod tests {
718 use super::*;
719
720 #[test]
721 fn test_adaptive_radius_is_capped_by_accept_error() {
722 let derivative = 10.0;
723 let radius = calculate_adaptive_search_radius(derivative, 10, 0, 16, 1.0, 1e-3);
724
725 assert!(
726 (radius - 0.01).abs() < 1e-12,
727 "radius should be capped by derivative * accept_error"
728 );
729 }
730
731 #[test]
732 fn test_adaptive_radius_uses_strict_gate_cap_when_pool_is_full() {
733 let derivative = 100.0;
734 let accept_error = 0.02;
735 let radius = calculate_adaptive_search_radius(derivative, 10, 16, 16, 1.0, accept_error);
736
737 let expected_cap = derivative * accept_error * STRICT_GATE_FACTOR;
738 assert!(
739 (radius - expected_cap).abs() < 1e-12,
740 "strict-gate fullness should tighten the scan radius cap"
741 );
742 }
743}