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")]
426 pub fn find_matches_turbo_with_stats_and_context(
427 &self,
428 lhs_exprs: &[EvaluatedExpr],
429 context: &SearchContext<'_>,
430 ) -> (Vec<Match>, SearchStats) {
431 use rayon::prelude::*;
432
433 let config = context.config;
434 let search_start = SearchTimer::start();
435 let initial_max_error = config.max_error.max(1e-12);
436
437 let mut sorted_lhs: Vec<_> = lhs_exprs.iter().collect();
438 sorted_lhs.sort_by_key(|e| e.expr.complexity());
439
440 let threads = rayon::current_num_threads().max(1);
443 let target_bands = (threads * 4).max(1);
444 let band_size = sorted_lhs.len().div_ceil(target_bands).max(1);
445
446 let (all_matches, mut stats) = sorted_lhs
447 .par_chunks(band_size)
448 .map(|band| {
449 let mut local_pool = TopKPool::new_with_diagnostics(
450 config.max_matches,
451 initial_max_error,
452 config.show_db_adds,
453 config.ranking_mode,
454 );
455 let mut local_stats = SearchStats::new();
456 for lhs in band {
457 if !match_one_lhs(self, lhs, context, &mut local_pool, &mut local_stats) {
458 local_stats.early_exit = true;
461 break;
462 }
463 }
464 (local_pool.into_sorted(), local_stats)
465 })
466 .reduce(
467 || (Vec::new(), SearchStats::new()),
468 |(mut acc_matches, mut acc_stats), (band_matches, band_stats)| {
469 acc_stats.merge_search_counters(&band_stats);
470 acc_matches.extend(band_matches);
471 (acc_matches, acc_stats)
472 },
473 );
474
475 let mut pool = TopKPool::new_with_diagnostics(
478 config.max_matches,
479 initial_max_error,
480 config.show_db_adds,
481 config.ranking_mode,
482 );
483 for m in all_matches {
484 pool.try_insert(m);
485 }
486
487 stats.pool_insertions = pool.stats.insertions;
488 stats.pool_rejections_error = pool.stats.rejections_error;
489 stats.pool_rejections_dedupe = pool.stats.rejections_dedupe;
490 stats.pool_evictions = pool.stats.evictions;
491 stats.pool_final_size = pool.len();
492 stats.pool_best_error = pool.best_error;
493 stats.search_time = search_start.elapsed();
494
495 (pool.into_sorted(), stats)
496 }
497}
498
499fn match_one_lhs(
506 db: &ExprDatabase,
507 lhs: &EvaluatedExpr,
508 context: &SearchContext<'_>,
509 pool: &mut TopKPool,
510 stats: &mut SearchStats,
511) -> bool {
512 let config = context.config;
513
514 let exact_ceiling = pool.exact_complexity_ceiling();
524 if let Some(ceiling) = exact_ceiling {
525 if lhs.expr.complexity() >= ceiling {
526 return false;
527 }
528 }
529
530 if lhs.value.abs() < config.zero_value_threshold {
534 if config.show_pruned_range {
535 eprintln!(
536 " [pruned range] value={:.6e} reason=\"near-zero\" expr=\"{}\"",
537 lhs.value,
538 lhs.expr.to_infix()
539 );
540 }
541 return true;
542 }
543
544 if super::should_skip_degenerate_lhs(lhs, config.target, &context.eval) {
547 return true;
548 }
549 if lhs.derivative.abs() < DEGENERATE_TEST_THRESHOLD {
550 let val_error = DEGENERATE_RANGE_TOLERANCE;
553 let low = lhs.value - val_error;
554 let high = lhs.value + val_error;
555
556 stats.lhs_tested += 1;
557 let rhs_slice = db.range(low, high);
558 stats.record_candidate_window(lhs, rhs_slice.len());
559 for rhs in rhs_slice {
560 if !config.rhs_symbol_allowed(&rhs.expr) {
561 continue;
562 }
563 stats.candidates_tested += 1;
564 if config.show_match_checks {
565 eprintln!(
566 " [match] checking lhs={:.6} rhs={:.6}",
567 lhs.value, rhs.value
568 );
569 }
570 let val_diff = (lhs.value - rhs.value).abs();
571 if val_diff < val_error && pool.would_accept(0.0, true) {
572 let m = Match {
573 lhs: lhs.clone(),
574 rhs: rhs.clone(),
575 x_value: config.target,
576 error: 0.0,
577 complexity: lhs.expr.complexity() + rhs.expr.complexity(),
578 };
579 pool.try_insert(m);
580 }
581 }
582 return true;
583 }
584
585 stats.lhs_tested += 1;
586
587 let search_radius = calculate_adaptive_search_radius(
589 lhs.derivative,
590 lhs.expr.complexity(),
591 pool.len(),
592 config.max_matches,
593 pool.best_error,
594 pool.accept_error,
595 );
596 let low = lhs.value - search_radius;
597 let high = lhs.value + search_radius;
598
599 let rhs_slice = db.range(low, high);
600 stats.record_candidate_window(lhs, rhs_slice.len());
601 for rhs in rhs_slice {
602 if !config.rhs_symbol_allowed(&rhs.expr) {
603 continue;
604 }
605
606 if let Some(ceiling) = exact_ceiling {
617 if lhs.expr.complexity() + rhs.expr.complexity() > ceiling {
618 continue;
619 }
620 }
621
622 stats.candidates_tested += 1;
623 if config.show_match_checks {
624 eprintln!(
625 " [match] checking lhs={:.6} rhs={:.6}",
626 lhs.value, rhs.value
627 );
628 }
629
630 let val_diff = lhs.value - rhs.value;
632 let x_delta = -val_diff / lhs.derivative;
633 let coarse_error = x_delta.abs();
634
635 let is_potentially_exact = coarse_error < NEWTON_FINAL_TOLERANCE;
638 if !pool.would_accept_strict(coarse_error, is_potentially_exact) {
639 stats.strict_gate_rejections += 1;
640 continue;
641 }
642
643 if !config.refine_with_newton {
644 let refined_x = config.target + x_delta;
645 let refined_error = x_delta;
646 let is_exact = refined_error.abs() < EXACT_MATCH_TOLERANCE;
647
648 if pool.would_accept(refined_error.abs(), is_exact) {
649 let m = Match {
650 lhs: lhs.clone(),
651 rhs: rhs.clone(),
652 x_value: refined_x,
653 error: refined_error,
654 complexity: lhs.expr.complexity() + rhs.expr.complexity(),
655 };
656
657 pool.try_insert(m);
658
659 if config.stop_at_exact && is_exact {
660 return false;
661 }
662 if let Some(threshold) = config.stop_below {
663 if refined_error.abs() < threshold {
664 return false;
665 }
666 }
667 }
668 continue;
669 }
670
671 stats.newton_calls += 1;
673 let initial_guess = config.target + x_delta;
674 if let Some(refined_x) = newton_raphson_with_constants(
675 &lhs.expr,
676 rhs.value,
677 initial_guess,
678 config.newton_iterations,
679 &context.eval,
680 config.show_newton,
681 config.derivative_margin,
682 ) {
683 stats.newton_success += 1;
684 let refined_error = refined_x - config.target;
685 let is_exact = refined_error.abs() < EXACT_MATCH_TOLERANCE;
686
687 if pool.would_accept(refined_error.abs(), is_exact) {
688 let m = Match {
689 lhs: lhs.clone(),
690 rhs: rhs.clone(),
691 x_value: refined_x,
692 error: refined_error,
693 complexity: lhs.expr.complexity() + rhs.expr.complexity(),
694 };
695
696 pool.try_insert(m);
697
698 if config.stop_at_exact && is_exact {
699 return false;
700 }
701 if let Some(threshold) = config.stop_below {
702 if refined_error.abs() < threshold {
703 return false;
704 }
705 }
706 }
707 }
708 }
709
710 true
711}
712
713impl Default for ExprDatabase {
714 fn default() -> Self {
715 Self::new()
716 }
717}
718
719#[cfg(test)]
720mod tests {
721 use super::*;
722
723 #[test]
724 fn test_adaptive_radius_is_capped_by_accept_error() {
725 let derivative = 10.0;
726 let radius = calculate_adaptive_search_radius(derivative, 10, 0, 16, 1.0, 1e-3);
727
728 assert!(
729 (radius - 0.01).abs() < 1e-12,
730 "radius should be capped by derivative * accept_error"
731 );
732 }
733
734 #[test]
735 fn test_adaptive_radius_uses_strict_gate_cap_when_pool_is_full() {
736 let derivative = 100.0;
737 let accept_error = 0.02;
738 let radius = calculate_adaptive_search_radius(derivative, 10, 16, 16, 1.0, accept_error);
739
740 let expected_cap = derivative * accept_error * STRICT_GATE_FACTOR;
741 assert!(
742 (radius - expected_cap).abs() < 1e-12,
743 "strict-gate fullness should tighten the scan radius cap"
744 );
745 }
746}