Skip to main content

ipfrs_semantic/
query_rewriter.rs

1//! Semantic query rewriting and expansion system.
2//!
3//! Provides query transformation, synonym expansion, stemming normalization,
4//! stop word removal, phrase detection, and weighted boosting for semantic search.
5
6use std::collections::{HashMap, HashSet};
7use std::time::Instant;
8
9/// Rewrite rule types
10#[derive(Debug, Clone, PartialEq, Eq)]
11pub enum RewriteRuleType {
12    SynonymExpansion,
13    StemmingNormalization,
14    PhraseDetection,
15    StopWordRemoval,
16    SpellingCorrection,
17    ConceptExpansion,
18    WeightBoosting,
19}
20
21/// A single rewrite rule
22#[derive(Debug, Clone)]
23pub struct RewriteRule {
24    pub rule_type: RewriteRuleType,
25    pub pattern: String,
26    pub replacement: String,
27    pub boost: f64,
28    pub enabled: bool,
29    pub priority: u32,
30}
31
32/// Rewritten query term
33#[derive(Debug, Clone)]
34pub struct RewrittenTerm {
35    pub original: String,
36    pub rewritten: String,
37    pub boost: f64,
38    pub source_rule: RewriteRuleType,
39}
40
41/// Full rewrite result
42#[derive(Debug, Clone)]
43pub struct RewriteResult {
44    pub original_query: String,
45    pub rewritten_query: String,
46    pub terms: Vec<RewrittenTerm>,
47    pub expansions: Vec<String>,
48    pub rules_applied: usize,
49    pub processing_time_us: u64,
50}
51
52/// Configuration for the query rewriter
53#[derive(Debug, Clone)]
54pub struct QueryRewriterConfig {
55    pub max_expansions: usize,
56    pub min_term_length: usize,
57    pub enable_stemming: bool,
58    pub enable_stop_word_removal: bool,
59    pub default_boost: f64,
60    pub max_rules: usize,
61}
62
63impl Default for QueryRewriterConfig {
64    fn default() -> Self {
65        Self {
66            max_expansions: 10,
67            min_term_length: 2,
68            enable_stemming: true,
69            enable_stop_word_removal: true,
70            default_boost: 1.0,
71            max_rules: 1000,
72        }
73    }
74}
75
76/// Statistics for the query rewriter
77#[derive(Debug, Clone, Default)]
78pub struct QueryRewriterStats {
79    pub queries_processed: u64,
80    pub rules_applied: u64,
81    pub expansions_generated: u64,
82    pub avg_processing_time_us: f64,
83}
84
85/// Query rewriter engine
86pub struct QueryRewriter {
87    config: QueryRewriterConfig,
88    rules: Vec<RewriteRule>,
89    stop_words: HashSet<String>,
90    synonym_map: HashMap<String, Vec<(String, f64)>>,
91    stats: QueryRewriterStats,
92}
93
94impl QueryRewriter {
95    /// Create a new query rewriter with the given configuration.
96    pub fn new(config: QueryRewriterConfig) -> Self {
97        Self {
98            config,
99            rules: Vec::new(),
100            stop_words: HashSet::new(),
101            synonym_map: HashMap::new(),
102            stats: QueryRewriterStats::default(),
103        }
104    }
105
106    /// Add a rewrite rule. Returns an error if the max rules limit is reached.
107    pub fn add_rule(&mut self, rule: RewriteRule) -> Result<(), String> {
108        if self.rules.len() >= self.config.max_rules {
109            return Err(format!(
110                "Maximum number of rules ({}) reached",
111                self.config.max_rules
112            ));
113        }
114        self.rules.push(rule);
115        // Keep rules sorted by priority (lower number = higher priority)
116        self.rules.sort_by_key(|r| r.priority);
117        Ok(())
118    }
119
120    /// Add a synonym mapping with a boost factor.
121    pub fn add_synonym(&mut self, word: &str, synonym: &str, boost: f64) {
122        let key = word.to_lowercase();
123        let entry = self.synonym_map.entry(key).or_default();
124        entry.push((synonym.to_lowercase(), boost));
125    }
126
127    /// Add a stop word.
128    pub fn add_stop_word(&mut self, word: &str) {
129        self.stop_words.insert(word.to_lowercase());
130    }
131
132    /// Rewrite a query, applying all enabled rules, synonym expansion,
133    /// stemming, and stop word removal.
134    pub fn rewrite(&mut self, query: &str) -> RewriteResult {
135        let start = Instant::now();
136        let normalized = Self::normalize_query(query);
137        let original_query = normalized.clone();
138
139        if normalized.is_empty() {
140            let elapsed = start.elapsed().as_micros() as u64;
141            self.update_stats(0, 0, elapsed);
142            return RewriteResult {
143                original_query,
144                rewritten_query: String::new(),
145                terms: Vec::new(),
146                expansions: Vec::new(),
147                rules_applied: 0,
148                processing_time_us: elapsed,
149            };
150        }
151
152        let mut tokens = Self::tokenize(&normalized);
153        let mut all_terms: Vec<RewrittenTerm> = Vec::new();
154        let mut all_expansions: Vec<String> = Vec::new();
155        let mut total_rules_applied: usize = 0;
156
157        // Phase 1: Apply explicit rewrite rules (pattern-based)
158        let mut rewritten_tokens: Vec<String> = Vec::new();
159        for token in &tokens {
160            let mut current = token.clone();
161            for rule in &self.rules {
162                if !rule.enabled {
163                    continue;
164                }
165                if rule.pattern.to_lowercase() == current {
166                    all_terms.push(RewrittenTerm {
167                        original: token.clone(),
168                        rewritten: rule.replacement.clone(),
169                        boost: rule.boost,
170                        source_rule: rule.rule_type.clone(),
171                    });
172                    current = rule.replacement.clone();
173                    total_rules_applied += 1;
174                    break; // Only apply highest-priority matching rule per token
175                }
176            }
177            rewritten_tokens.push(current);
178        }
179        tokens = rewritten_tokens;
180
181        // Phase 2: Stop word removal
182        if self.config.enable_stop_word_removal {
183            let before_len = tokens.len();
184            tokens.retain(|t| !self.is_stop_word(t));
185            let removed = before_len - tokens.len();
186            if removed > 0 {
187                total_rules_applied += removed;
188            }
189        }
190
191        // Phase 3: Min term length filtering
192        tokens.retain(|t| t.len() >= self.config.min_term_length);
193
194        // Phase 4: Stemming
195        if self.config.enable_stemming {
196            let mut stemmed_tokens = Vec::new();
197            for token in &tokens {
198                let stemmed = Self::apply_stemming(token);
199                if stemmed != *token {
200                    all_terms.push(RewrittenTerm {
201                        original: token.clone(),
202                        rewritten: stemmed.clone(),
203                        boost: self.config.default_boost,
204                        source_rule: RewriteRuleType::StemmingNormalization,
205                    });
206                    total_rules_applied += 1;
207                }
208                stemmed_tokens.push(stemmed);
209            }
210            tokens = stemmed_tokens;
211        }
212
213        // Phase 5: Synonym expansion
214        let mut expansion_count: usize = 0;
215        for token in &tokens {
216            if expansion_count >= self.config.max_expansions {
217                break;
218            }
219            let synonyms = self.expand_synonyms(token);
220            for (syn, boost) in synonyms {
221                if expansion_count >= self.config.max_expansions {
222                    break;
223                }
224                all_terms.push(RewrittenTerm {
225                    original: token.clone(),
226                    rewritten: syn.clone(),
227                    boost,
228                    source_rule: RewriteRuleType::SynonymExpansion,
229                });
230                all_expansions.push(syn);
231                expansion_count += 1;
232                total_rules_applied += 1;
233            }
234        }
235
236        let rewritten_query = tokens.join(" ");
237        let elapsed = start.elapsed().as_micros() as u64;
238
239        self.update_stats(total_rules_applied, expansion_count, elapsed);
240
241        RewriteResult {
242            original_query,
243            rewritten_query,
244            terms: all_terms,
245            expansions: all_expansions,
246            rules_applied: total_rules_applied,
247            processing_time_us: elapsed,
248        }
249    }
250
251    /// Tokenize a query string: split on whitespace, lowercase, remove non-alphanumeric chars.
252    pub fn tokenize(query: &str) -> Vec<String> {
253        query
254            .split_whitespace()
255            .map(|w| {
256                w.chars()
257                    .filter(|c| c.is_alphanumeric())
258                    .collect::<String>()
259                    .to_lowercase()
260            })
261            .filter(|w| !w.is_empty())
262            .collect()
263    }
264
265    /// Simple suffix-stripping stemmer (Porter-like).
266    /// Handles: -ing, -ed, -ly, -tion, -ness, -ment
267    pub fn apply_stemming(word: &str) -> String {
268        if word.len() < 4 {
269            return word.to_string();
270        }
271
272        // Try suffixes from longest to shortest to avoid partial matches
273        let suffixes: &[&str] = &["tion", "ness", "ment", "ing", "ed", "ly"];
274        for suffix in suffixes {
275            if let Some(stem) = word.strip_suffix(suffix) {
276                // Only accept stems of reasonable length
277                if stem.len() >= 2 {
278                    return stem.to_string();
279                }
280            }
281        }
282
283        word.to_string()
284    }
285
286    /// Check if a word is a stop word.
287    pub fn is_stop_word(&self, word: &str) -> bool {
288        self.stop_words.contains(&word.to_lowercase())
289    }
290
291    /// Expand a term into its synonyms with associated boost values.
292    pub fn expand_synonyms(&self, term: &str) -> Vec<(String, f64)> {
293        self.synonym_map
294            .get(&term.to_lowercase())
295            .cloned()
296            .unwrap_or_default()
297    }
298
299    /// Return the number of rules.
300    pub fn rules_count(&self) -> usize {
301        self.rules.len()
302    }
303
304    /// Remove all rules matching the given pattern. Returns true if any were removed.
305    pub fn remove_rule(&mut self, pattern: &str) -> bool {
306        let before = self.rules.len();
307        self.rules
308            .retain(|r| r.pattern.to_lowercase() != pattern.to_lowercase());
309        self.rules.len() < before
310    }
311
312    /// Clear all rules.
313    pub fn clear_rules(&mut self) {
314        self.rules.clear();
315    }
316
317    /// Get current statistics.
318    pub fn stats(&self) -> &QueryRewriterStats {
319        &self.stats
320    }
321
322    /// Normalize a query: lowercase, trim, collapse multiple whitespace into single space.
323    pub fn normalize_query(query: &str) -> String {
324        let trimmed = query.trim().to_lowercase();
325        let mut result = String::with_capacity(trimmed.len());
326        let mut prev_was_space = false;
327        for ch in trimmed.chars() {
328            if ch.is_whitespace() {
329                if !prev_was_space {
330                    result.push(' ');
331                    prev_was_space = true;
332                }
333            } else {
334                result.push(ch);
335                prev_was_space = false;
336            }
337        }
338        result
339    }
340
341    /// Update internal statistics after processing a query.
342    fn update_stats(&mut self, rules_applied: usize, expansions: usize, elapsed_us: u64) {
343        let n = self.stats.queries_processed;
344        self.stats.queries_processed += 1;
345        self.stats.rules_applied += rules_applied as u64;
346        self.stats.expansions_generated += expansions as u64;
347
348        // Running average
349        let new_n = self.stats.queries_processed as f64;
350        self.stats.avg_processing_time_us =
351            (self.stats.avg_processing_time_us * n as f64 + elapsed_us as f64) / new_n;
352    }
353}
354
355#[cfg(test)]
356mod tests {
357    use super::*;
358
359    fn default_rewriter() -> QueryRewriter {
360        QueryRewriter::new(QueryRewriterConfig::default())
361    }
362
363    #[test]
364    fn test_basic_rewrite_single_rule() {
365        let mut rw = default_rewriter();
366        rw.add_rule(RewriteRule {
367            rule_type: RewriteRuleType::SpellingCorrection,
368            pattern: "teh".to_string(),
369            replacement: "the".to_string(),
370            boost: 1.0,
371            enabled: true,
372            priority: 1,
373        })
374        .expect("should add rule");
375
376        let result = rw.rewrite("teh cat");
377        assert_eq!(result.rewritten_query, "the cat");
378        assert!(result.rules_applied >= 1);
379    }
380
381    #[test]
382    fn test_multiple_rules_priority_order() {
383        let mut rw = default_rewriter();
384        // Lower priority number = applied first
385        rw.add_rule(RewriteRule {
386            rule_type: RewriteRuleType::SpellingCorrection,
387            pattern: "foo".to_string(),
388            replacement: "bar".to_string(),
389            boost: 1.0,
390            enabled: true,
391            priority: 10,
392        })
393        .expect("add rule");
394        rw.add_rule(RewriteRule {
395            rule_type: RewriteRuleType::SpellingCorrection,
396            pattern: "foo".to_string(),
397            replacement: "baz".to_string(),
398            boost: 1.0,
399            enabled: true,
400            priority: 1,
401        })
402        .expect("add rule");
403
404        // The priority=1 rule should be first and applied
405        let result = rw.rewrite("foo");
406        assert_eq!(result.rewritten_query, "baz");
407    }
408
409    #[test]
410    fn test_synonym_expansion() {
411        let mut rw = default_rewriter();
412        rw.add_synonym("fast", "quick", 0.9);
413        rw.add_synonym("fast", "rapid", 0.8);
414
415        let result = rw.rewrite("fast car");
416        assert!(result.expansions.contains(&"quick".to_string()));
417        assert!(result.expansions.contains(&"rapid".to_string()));
418    }
419
420    #[test]
421    fn test_stop_word_removal() {
422        let mut rw = default_rewriter();
423        rw.add_stop_word("the");
424        rw.add_stop_word("is");
425        rw.add_stop_word("a");
426
427        let result = rw.rewrite("the cat is a pet");
428        // "the", "is", "a" removed; "cat" and "pet" remain
429        assert!(!result.rewritten_query.contains("the"));
430        assert!(!result.rewritten_query.contains(" is "));
431        assert!(result.rewritten_query.contains("cat"));
432        assert!(result.rewritten_query.contains("pet"));
433    }
434
435    #[test]
436    fn test_stemming_ing() {
437        let stemmed = QueryRewriter::apply_stemming("running");
438        assert_eq!(stemmed, "runn");
439    }
440
441    #[test]
442    fn test_stemming_ed() {
443        let stemmed = QueryRewriter::apply_stemming("walked");
444        assert_eq!(stemmed, "walk");
445    }
446
447    #[test]
448    fn test_stemming_ly() {
449        let stemmed = QueryRewriter::apply_stemming("quickly");
450        assert_eq!(stemmed, "quick");
451    }
452
453    #[test]
454    fn test_stemming_tion() {
455        let stemmed = QueryRewriter::apply_stemming("creation");
456        assert_eq!(stemmed, "crea");
457    }
458
459    #[test]
460    fn test_stemming_ness() {
461        let stemmed = QueryRewriter::apply_stemming("darkness");
462        assert_eq!(stemmed, "dark");
463    }
464
465    #[test]
466    fn test_stemming_ment() {
467        let stemmed = QueryRewriter::apply_stemming("agreement");
468        assert_eq!(stemmed, "agree");
469    }
470
471    #[test]
472    fn test_stemming_short_word_unchanged() {
473        let stemmed = QueryRewriter::apply_stemming("go");
474        assert_eq!(stemmed, "go");
475    }
476
477    #[test]
478    fn test_phrase_detection_rule() {
479        let mut rw = default_rewriter();
480        rw.add_rule(RewriteRule {
481            rule_type: RewriteRuleType::PhraseDetection,
482            pattern: "machine".to_string(),
483            replacement: "machine_learning".to_string(),
484            boost: 1.5,
485            enabled: true,
486            priority: 1,
487        })
488        .expect("add rule");
489
490        let result = rw.rewrite("machine algorithms");
491        // The term "machine" should be rewritten to "machine_learning"
492        let has_phrase = result
493            .terms
494            .iter()
495            .any(|t| t.rewritten == "machine_learning");
496        assert!(has_phrase);
497    }
498
499    #[test]
500    fn test_boost_application() {
501        let mut rw = default_rewriter();
502        rw.add_synonym("search", "lookup", 2.5);
503
504        let result = rw.rewrite("search");
505        let syn_term = result
506            .terms
507            .iter()
508            .find(|t| t.rewritten == "lookup")
509            .expect("should have synonym term");
510        assert!((syn_term.boost - 2.5).abs() < f64::EPSILON);
511    }
512
513    #[test]
514    fn test_empty_query() {
515        let mut rw = default_rewriter();
516        let result = rw.rewrite("");
517        assert!(result.rewritten_query.is_empty());
518        assert!(result.terms.is_empty());
519        assert_eq!(result.rules_applied, 0);
520    }
521
522    #[test]
523    fn test_whitespace_only_query() {
524        let mut rw = default_rewriter();
525        let result = rw.rewrite("   \t\n  ");
526        assert!(result.rewritten_query.is_empty());
527    }
528
529    #[test]
530    fn test_very_long_query() {
531        let mut rw = default_rewriter();
532        let long_query = "word ".repeat(500);
533        let result = rw.rewrite(&long_query);
534        assert!(!result.rewritten_query.is_empty());
535    }
536
537    #[test]
538    fn test_rule_add_remove_clear() {
539        let mut rw = default_rewriter();
540        assert_eq!(rw.rules_count(), 0);
541
542        rw.add_rule(RewriteRule {
543            rule_type: RewriteRuleType::SpellingCorrection,
544            pattern: "abc".to_string(),
545            replacement: "def".to_string(),
546            boost: 1.0,
547            enabled: true,
548            priority: 1,
549        })
550        .expect("add");
551        assert_eq!(rw.rules_count(), 1);
552
553        let removed = rw.remove_rule("abc");
554        assert!(removed);
555        assert_eq!(rw.rules_count(), 0);
556
557        let not_removed = rw.remove_rule("nonexistent");
558        assert!(!not_removed);
559
560        rw.add_rule(RewriteRule {
561            rule_type: RewriteRuleType::SpellingCorrection,
562            pattern: "x".to_string(),
563            replacement: "y".to_string(),
564            boost: 1.0,
565            enabled: true,
566            priority: 1,
567        })
568        .expect("add");
569        rw.clear_rules();
570        assert_eq!(rw.rules_count(), 0);
571    }
572
573    #[test]
574    fn test_stats_tracking() {
575        let mut rw = default_rewriter();
576        assert_eq!(rw.stats().queries_processed, 0);
577
578        rw.rewrite("hello world");
579        assert_eq!(rw.stats().queries_processed, 1);
580
581        rw.rewrite("another query");
582        assert_eq!(rw.stats().queries_processed, 2);
583    }
584
585    #[test]
586    fn test_max_expansions_limit() {
587        let mut rw = QueryRewriter::new(QueryRewriterConfig {
588            max_expansions: 2,
589            ..Default::default()
590        });
591        rw.add_synonym("test", "exam", 1.0);
592        rw.add_synonym("test", "trial", 1.0);
593        rw.add_synonym("test", "assessment", 1.0);
594        rw.add_synonym("test", "evaluation", 1.0);
595
596        let result = rw.rewrite("test");
597        assert!(result.expansions.len() <= 2);
598    }
599
600    #[test]
601    fn test_min_term_length_filtering() {
602        let mut rw = QueryRewriter::new(QueryRewriterConfig {
603            min_term_length: 3,
604            enable_stop_word_removal: false,
605            enable_stemming: false,
606            ..Default::default()
607        });
608
609        let result = rw.rewrite("I am a big cat");
610        // "I", "am", "a" should be filtered (length < 3)
611        assert!(!result.rewritten_query.contains(" i "));
612        assert!(result.rewritten_query.contains("big"));
613        assert!(result.rewritten_query.contains("cat"));
614    }
615
616    #[test]
617    fn test_combined_stemming_synonyms_stopwords() {
618        let mut rw = default_rewriter();
619        rw.add_stop_word("the");
620        rw.add_synonym("run", "jog", 0.9);
621
622        // "the" removed, "running" stemmed to "runn" (simple stemmer)
623        let result = rw.rewrite("the running dog");
624        assert!(!result.rewritten_query.contains("the"));
625    }
626
627    #[test]
628    fn test_tokenization_correctness() {
629        let tokens = QueryRewriter::tokenize("Hello, World! Test-case 123");
630        assert_eq!(tokens, vec!["hello", "world", "testcase", "123"]);
631    }
632
633    #[test]
634    fn test_tokenization_empty() {
635        let tokens = QueryRewriter::tokenize("");
636        assert!(tokens.is_empty());
637    }
638
639    #[test]
640    fn test_query_normalization_whitespace() {
641        let normalized = QueryRewriter::normalize_query("  hello   world  ");
642        assert_eq!(normalized, "hello world");
643    }
644
645    #[test]
646    fn test_query_normalization_case() {
647        let normalized = QueryRewriter::normalize_query("Hello WORLD");
648        assert_eq!(normalized, "hello world");
649    }
650
651    #[test]
652    fn test_duplicate_rule_handling() {
653        let mut rw = default_rewriter();
654        // Adding two rules with same pattern but different replacements
655        rw.add_rule(RewriteRule {
656            rule_type: RewriteRuleType::SpellingCorrection,
657            pattern: "test".to_string(),
658            replacement: "exam".to_string(),
659            boost: 1.0,
660            enabled: true,
661            priority: 2,
662        })
663        .expect("add");
664        rw.add_rule(RewriteRule {
665            rule_type: RewriteRuleType::SpellingCorrection,
666            pattern: "test".to_string(),
667            replacement: "quiz".to_string(),
668            boost: 1.0,
669            enabled: true,
670            priority: 1,
671        })
672        .expect("add");
673
674        // Both stored, but only highest priority (priority=1) applied
675        assert_eq!(rw.rules_count(), 2);
676        let result = rw.rewrite("test");
677        assert_eq!(result.rewritten_query, "quiz");
678    }
679
680    #[test]
681    fn test_rule_priority_ordering() {
682        let mut rw = default_rewriter();
683        rw.add_rule(RewriteRule {
684            rule_type: RewriteRuleType::WeightBoosting,
685            pattern: "alpha".to_string(),
686            replacement: "beta".to_string(),
687            boost: 1.0,
688            enabled: true,
689            priority: 100,
690        })
691        .expect("add");
692        rw.add_rule(RewriteRule {
693            rule_type: RewriteRuleType::WeightBoosting,
694            pattern: "gamma".to_string(),
695            replacement: "delta".to_string(),
696            boost: 1.0,
697            enabled: true,
698            priority: 1,
699        })
700        .expect("add");
701
702        // Internal ordering: priority=1 first, then priority=100
703        assert_eq!(rw.rules[0].pattern, "gamma");
704        assert_eq!(rw.rules[1].pattern, "alpha");
705    }
706
707    #[test]
708    fn test_processing_time_tracked() {
709        let mut rw = default_rewriter();
710        let result = rw.rewrite("some query");
711        // processing_time_us should be set (at least 0)
712        assert!(result.processing_time_us < 1_000_000); // less than 1 second
713        assert!(rw.stats().avg_processing_time_us >= 0.0);
714    }
715
716    #[test]
717    fn test_config_defaults() {
718        let cfg = QueryRewriterConfig::default();
719        assert_eq!(cfg.max_expansions, 10);
720        assert_eq!(cfg.min_term_length, 2);
721        assert!(cfg.enable_stemming);
722        assert!(cfg.enable_stop_word_removal);
723        assert!((cfg.default_boost - 1.0).abs() < f64::EPSILON);
724        assert_eq!(cfg.max_rules, 1000);
725    }
726
727    #[test]
728    fn test_disabled_rule_not_applied() {
729        let mut rw = default_rewriter();
730        rw.add_rule(RewriteRule {
731            rule_type: RewriteRuleType::SpellingCorrection,
732            pattern: "hello".to_string(),
733            replacement: "hi".to_string(),
734            boost: 1.0,
735            enabled: false,
736            priority: 1,
737        })
738        .expect("add");
739
740        let result = rw.rewrite("hello");
741        // Rule is disabled, so "hello" should not become "hi"
742        // (stemming may still apply, but the rule replacement should not)
743        let has_rule_term = result
744            .terms
745            .iter()
746            .any(|t| t.source_rule == RewriteRuleType::SpellingCorrection);
747        assert!(!has_rule_term);
748    }
749
750    #[test]
751    fn test_max_rules_limit() {
752        let mut rw = QueryRewriter::new(QueryRewriterConfig {
753            max_rules: 2,
754            ..Default::default()
755        });
756
757        rw.add_rule(RewriteRule {
758            rule_type: RewriteRuleType::SpellingCorrection,
759            pattern: "a".to_string(),
760            replacement: "b".to_string(),
761            boost: 1.0,
762            enabled: true,
763            priority: 1,
764        })
765        .expect("first");
766
767        rw.add_rule(RewriteRule {
768            rule_type: RewriteRuleType::SpellingCorrection,
769            pattern: "c".to_string(),
770            replacement: "d".to_string(),
771            boost: 1.0,
772            enabled: true,
773            priority: 2,
774        })
775        .expect("second");
776
777        let res = rw.add_rule(RewriteRule {
778            rule_type: RewriteRuleType::SpellingCorrection,
779            pattern: "e".to_string(),
780            replacement: "f".to_string(),
781            boost: 1.0,
782            enabled: true,
783            priority: 3,
784        });
785        assert!(res.is_err());
786    }
787
788    #[test]
789    fn test_concept_expansion_rule_type() {
790        let mut rw = default_rewriter();
791        rw.add_rule(RewriteRule {
792            rule_type: RewriteRuleType::ConceptExpansion,
793            pattern: "ml".to_string(),
794            replacement: "machine_learning".to_string(),
795            boost: 1.2,
796            enabled: true,
797            priority: 1,
798        })
799        .expect("add");
800
801        let result = rw.rewrite("ml models");
802        let term = result
803            .terms
804            .iter()
805            .find(|t| t.source_rule == RewriteRuleType::ConceptExpansion);
806        assert!(term.is_some());
807    }
808
809    #[test]
810    fn test_stats_rules_applied_count() {
811        let mut rw = default_rewriter();
812        rw.add_synonym("dog", "canine", 1.0);
813        rw.add_stop_word("the");
814
815        rw.rewrite("the dog");
816        assert!(rw.stats().rules_applied > 0);
817        assert!(rw.stats().expansions_generated > 0);
818    }
819
820    #[test]
821    fn test_special_characters_in_query() {
822        let mut rw = default_rewriter();
823        let result = rw.rewrite("hello! @world #test");
824        // Special chars stripped by tokenizer
825        let tokens = QueryRewriter::tokenize("hello! @world #test");
826        assert_eq!(tokens, vec!["hello", "world", "test"]);
827        assert!(!result.rewritten_query.is_empty());
828    }
829
830    #[test]
831    fn test_synonym_case_insensitive() {
832        let mut rw = default_rewriter();
833        rw.add_synonym("Dog", "canine", 1.0);
834
835        let result = rw.rewrite("DOG");
836        // Synonym lookup is case-insensitive
837        assert!(result.expansions.contains(&"canine".to_string()));
838    }
839
840    #[test]
841    fn test_remove_rule_case_insensitive() {
842        let mut rw = default_rewriter();
843        rw.add_rule(RewriteRule {
844            rule_type: RewriteRuleType::SpellingCorrection,
845            pattern: "Hello".to_string(),
846            replacement: "hi".to_string(),
847            boost: 1.0,
848            enabled: true,
849            priority: 1,
850        })
851        .expect("add");
852
853        let removed = rw.remove_rule("hello");
854        assert!(removed);
855        assert_eq!(rw.rules_count(), 0);
856    }
857}