1#![allow(clippy::unnecessary_literal_bound)]
18
19use std::cmp::Ordering;
20use std::collections::HashMap;
21use std::sync::{Arc, OnceLock};
22
23use tracing::{debug, info};
24
25pub trait CollationFunction: Send + Sync {
32 fn name(&self) -> &str;
34
35 fn compare(&self, left: &[u8], right: &[u8]) -> Ordering;
39}
40
41pub struct BinaryCollation;
48
49impl CollationFunction for BinaryCollation {
50 fn name(&self) -> &str {
51 "BINARY"
52 }
53
54 fn compare(&self, left: &[u8], right: &[u8]) -> Ordering {
55 left.cmp(right)
56 }
57}
58
59pub struct NoCaseCollation;
64
65impl CollationFunction for NoCaseCollation {
66 fn name(&self) -> &str {
67 "NOCASE"
68 }
69
70 fn compare(&self, left: &[u8], right: &[u8]) -> Ordering {
71 let l = left.iter().map(u8::to_ascii_uppercase);
72 let r = right.iter().map(u8::to_ascii_uppercase);
73 l.cmp(r)
74 }
75}
76
77pub struct RtrimCollation;
82
83impl CollationFunction for RtrimCollation {
84 fn name(&self) -> &str {
85 "RTRIM"
86 }
87
88 fn compare(&self, left: &[u8], right: &[u8]) -> Ordering {
89 let l = strip_trailing_spaces(left);
90 let r = strip_trailing_spaces(right);
91 l.cmp(r)
92 }
93}
94
95fn strip_trailing_spaces(s: &[u8]) -> &[u8] {
96 let mut end = s.len();
97 while end > 0 && s[end - 1] == b' ' {
98 end -= 1;
99 }
100 &s[..end]
101}
102
103fn builtin_collation(name: &str) -> Option<Arc<dyn CollationFunction>> {
104 type BuiltinCollations = (
105 Arc<dyn CollationFunction>,
106 Arc<dyn CollationFunction>,
107 Arc<dyn CollationFunction>,
108 );
109
110 static BUILTINS: OnceLock<BuiltinCollations> = OnceLock::new();
111 let (binary, nocase, rtrim) = BUILTINS.get_or_init(|| {
112 (
113 Arc::new(BinaryCollation) as Arc<dyn CollationFunction>,
114 Arc::new(NoCaseCollation) as Arc<dyn CollationFunction>,
115 Arc::new(RtrimCollation) as Arc<dyn CollationFunction>,
116 )
117 });
118 match name {
119 "BINARY" => Some(Arc::clone(binary)),
120 "NOCASE" => Some(Arc::clone(nocase)),
121 "RTRIM" => Some(Arc::clone(rtrim)),
122 _ => None,
123 }
124}
125
126#[derive(Clone)]
133pub struct CollationRegistry {
134 custom_collations: HashMap<String, Arc<dyn CollationFunction>>,
135}
136
137impl std::fmt::Debug for CollationRegistry {
138 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
139 f.debug_struct("CollationRegistry")
140 .field("collations", &self.names())
141 .finish()
142 }
143}
144
145impl Default for CollationRegistry {
146 fn default() -> Self {
147 Self::new()
148 }
149}
150
151impl CollationRegistry {
152 #[must_use]
154 pub fn new() -> Self {
155 Self {
156 custom_collations: HashMap::new(),
157 }
158 }
159
160 pub fn register<C: CollationFunction + 'static>(
165 &mut self,
166 collation: C,
167 ) -> Option<Arc<dyn CollationFunction>> {
168 let name = collation.name().to_owned();
169 self.register_captured(&name, collation)
170 }
171
172 pub fn register_captured<C: CollationFunction + 'static>(
178 &mut self,
179 name: &str,
180 collation: C,
181 ) -> Option<Arc<dyn CollationFunction>> {
182 let name = name.to_ascii_uppercase();
183 info!(collation_name = %name, deterministic = true, "custom collation registration");
184 self.custom_collations
185 .insert(name.clone(), Arc::new(collation))
186 .or_else(|| builtin_collation(&name))
187 }
188
189 #[must_use]
193 pub fn find(&self, name: &str) -> Option<Arc<dyn CollationFunction>> {
194 let canon = name.to_ascii_uppercase();
195 let result = self
196 .custom_collations
197 .get(&canon)
198 .cloned()
199 .or_else(|| builtin_collation(&canon));
200 debug!(
201 collation = %canon,
202 hit = result.is_some(),
203 "collation registry lookup"
204 );
205 result
206 }
207
208 #[must_use]
210 pub fn contains(&self, name: &str) -> bool {
211 let canon = name.to_ascii_uppercase();
212 self.custom_collations.contains_key(&canon) || builtin_collation(&canon).is_some()
213 }
214
215 #[must_use]
222 pub fn uses_builtin_implementation(&self, name: &str) -> bool {
223 let canon = name.to_ascii_uppercase();
224 matches!(canon.as_str(), "BINARY" | "NOCASE" | "RTRIM")
225 && !self.custom_collations.contains_key(&canon)
226 }
227
228 #[must_use]
237 pub fn any_builtin_overridden(&self) -> bool {
238 ["BINARY", "NOCASE", "RTRIM"]
239 .iter()
240 .any(|name| self.custom_collations.contains_key(*name))
241 }
242
243 #[must_use]
249 pub fn names(&self) -> Vec<String> {
250 let mut names = vec!["BINARY".to_owned(), "NOCASE".to_owned(), "RTRIM".to_owned()];
251 let mut custom: Vec<String> = self
252 .custom_collations
253 .keys()
254 .filter(|name| !matches!(name.as_str(), "BINARY" | "NOCASE" | "RTRIM"))
255 .cloned()
256 .collect();
257 custom.sort_unstable_by_key(|name| name.to_ascii_uppercase());
258 names.extend(custom);
259 names
260 }
261}
262
263#[derive(Debug, Clone, Copy, PartialEq, Eq)]
270pub enum CollationSource {
271 Explicit,
273 Schema,
275 Default,
277}
278
279#[derive(Debug, Clone)]
281pub struct CollationAnnotation {
282 pub name: String,
284 pub source: CollationSource,
286}
287
288#[must_use]
298pub fn resolve_collation(lhs: &CollationAnnotation, rhs: &CollationAnnotation) -> String {
299 let result = match (lhs.source, rhs.source) {
301 (_, CollationSource::Explicit) if lhs.source != CollationSource::Explicit => &rhs.name,
302 (CollationSource::Default, CollationSource::Schema) => &rhs.name,
303 _ => &lhs.name,
304 };
305 debug!(
306 collation = %result,
307 lhs_source = ?lhs.source,
308 rhs_source = ?rhs.source,
309 context = "COMPARE",
310 "collation selection"
311 );
312 result.clone()
313}
314
315#[cfg(test)]
316mod tests {
317 use super::*;
318
319 #[test]
322 fn test_collation_binary_memcmp() {
323 let coll = BinaryCollation;
324 assert_eq!(coll.compare(b"abc", b"abc"), Ordering::Equal);
325 assert_eq!(coll.compare(b"abc", b"abd"), Ordering::Less);
326 assert_eq!(coll.compare(b"abd", b"abc"), Ordering::Greater);
327 assert_eq!(coll.compare(b"ABC", b"abc"), Ordering::Less);
329 assert_eq!(
331 coll.compare("café".as_bytes(), "café".as_bytes()),
332 Ordering::Equal
333 );
334 assert_ne!(coll.compare("über".as_bytes(), b"uber"), Ordering::Equal);
335 }
336
337 #[test]
338 fn test_collation_binary_basic() {
339 let coll = BinaryCollation;
340 assert_eq!(coll.compare(b"ABC", b"abc"), Ordering::Less);
342 assert_eq!(coll.compare(b"\x00", b"\x01"), Ordering::Less);
344 assert_eq!(coll.compare(b"\xff", b"\x00"), Ordering::Greater);
345 }
346
347 #[test]
348 fn test_collation_nocase_ascii() {
349 let coll = NoCaseCollation;
350 assert_eq!(coll.compare(b"ABC", b"abc"), Ordering::Equal);
351 assert_eq!(coll.compare(b"Alice", b"alice"), Ordering::Equal);
352 assert_eq!(coll.compare(b"[", b"a"), Ordering::Greater);
354 }
355
356 #[test]
357 fn test_collation_nocase_ascii_only() {
358 let coll = NoCaseCollation;
359 assert_ne!(
361 coll.compare("Ä".as_bytes(), "ä".as_bytes()),
362 Ordering::Equal,
363 "NOCASE must NOT fold non-ASCII"
364 );
365 assert_eq!(coll.compare(b"Z", b"z"), Ordering::Equal);
367 assert_eq!(coll.compare(b"[", b"["), Ordering::Equal);
368 assert_ne!(coll.compare(b"[", b"{"), Ordering::Equal);
370 }
371
372 #[test]
373 fn test_collation_rtrim() {
374 let coll = RtrimCollation;
375 assert_eq!(coll.compare(b"hello ", b"hello"), Ordering::Equal);
377 assert_eq!(coll.compare(b"hello", b"hello "), Ordering::Equal);
378 assert_eq!(coll.compare(b"hello ", b"hello "), Ordering::Equal);
379 assert_ne!(coll.compare(b"hello!", b"hello"), Ordering::Equal);
381 assert_ne!(coll.compare(b"hello ", b"hello!"), Ordering::Equal);
383 }
384
385 #[test]
386 fn test_collation_rtrim_tabs_not_stripped() {
387 let coll = RtrimCollation;
388 assert_ne!(
390 coll.compare(b"hello\t", b"hello"),
391 Ordering::Equal,
392 "RTRIM must NOT strip tabs"
393 );
394 assert_ne!(
396 coll.compare(b"hello\xc2\xa0", b"hello"),
397 Ordering::Equal,
398 "RTRIM must NOT strip non-breaking spaces"
399 );
400 }
401
402 #[test]
403 fn test_collation_properties_antisymmetric() {
404 let collations: Vec<Box<dyn CollationFunction>> = vec![
405 Box::new(BinaryCollation),
406 Box::new(NoCaseCollation),
407 Box::new(RtrimCollation),
408 ];
409
410 let pairs: &[(&[u8], &[u8])] = &[
411 (b"abc", b"def"),
412 (b"hello", b"world"),
413 (b"ABC", b"abc"),
414 (b"hello ", b"hello"),
415 ];
416
417 for coll in &collations {
418 for &(a, b) in pairs {
419 let forward = coll.compare(a, b);
420 let reverse = coll.compare(b, a);
421 assert_eq!(
422 forward,
423 reverse.reverse(),
424 "{}: compare({:?}, {:?}) = {forward:?}, but reverse = {reverse:?}",
425 coll.name(),
426 std::str::from_utf8(a).unwrap_or("?"),
427 std::str::from_utf8(b).unwrap_or("?"),
428 );
429 }
430 }
431 }
432
433 #[test]
434 fn test_collation_properties_transitive() {
435 let coll = BinaryCollation;
436 let a = b"apple";
437 let b = b"banana";
438 let c = b"cherry";
439
440 assert_eq!(coll.compare(a, b), Ordering::Less);
442 assert_eq!(coll.compare(b, c), Ordering::Less);
443 assert_eq!(coll.compare(a, c), Ordering::Less);
444 }
445
446 #[test]
447 fn test_collation_send_sync() {
448 fn assert_send_sync<T: Send + Sync>() {}
449 assert_send_sync::<BinaryCollation>();
450 assert_send_sync::<NoCaseCollation>();
451 assert_send_sync::<RtrimCollation>();
452 }
453
454 #[test]
457 fn test_registry_preloaded_builtins() {
458 let reg = CollationRegistry::new();
459 assert!(reg.contains("BINARY"));
460 assert!(reg.contains("NOCASE"));
461 assert!(reg.contains("RTRIM"));
462
463 let binary = reg.find("BINARY").expect("BINARY must be pre-registered");
464 assert_eq!(binary.compare(b"a", b"b"), Ordering::Less);
465
466 let nocase = reg.find("NOCASE").expect("NOCASE must be pre-registered");
467 assert_eq!(nocase.compare(b"ABC", b"abc"), Ordering::Equal);
468
469 let rtrim = reg.find("RTRIM").expect("RTRIM must be pre-registered");
470 assert_eq!(rtrim.compare(b"x ", b"x"), Ordering::Equal);
471 }
472
473 struct ReverseCollation;
474
475 impl CollationFunction for ReverseCollation {
476 fn name(&self) -> &str {
477 "REVERSE"
478 }
479
480 fn compare(&self, left: &[u8], right: &[u8]) -> Ordering {
481 right.cmp(left)
482 }
483 }
484
485 #[test]
486 fn test_registry_custom_collation_registration() {
487 let mut reg = CollationRegistry::new();
488
489 let prev = reg.register(ReverseCollation);
490 assert!(prev.is_none(), "no prior REVERSE collation");
491 assert!(reg.contains("REVERSE"));
492
493 let coll = reg.find("reverse").expect("case-insensitive lookup");
494 assert_eq!(coll.compare(b"a", b"z"), Ordering::Greater);
495 }
496
497 struct AlwaysEqualCollation;
498
499 impl CollationFunction for AlwaysEqualCollation {
500 fn name(&self) -> &str {
501 "BINARY"
502 }
503
504 fn compare(&self, _left: &[u8], _right: &[u8]) -> Ordering {
505 Ordering::Equal
506 }
507 }
508
509 #[test]
510 fn test_registry_overwrite_builtin() {
511 let mut reg = CollationRegistry::new();
512 assert!(reg.uses_builtin_implementation("binary"));
513
514 let prev = reg.register(AlwaysEqualCollation);
515 assert!(prev.is_some(), "should return previous BINARY collation");
516 assert!(!reg.uses_builtin_implementation("BINARY"));
517
518 let coll = reg.find("BINARY").unwrap();
519 assert_eq!(
520 coll.compare(b"a", b"z"),
521 Ordering::Equal,
522 "custom overwrite must take effect"
523 );
524 }
525
526 #[test]
527 fn test_registry_unregistered_returns_none() {
528 let reg = CollationRegistry::new();
529 assert!(reg.find("NONEXISTENT").is_none());
530 assert!(!reg.contains("NONEXISTENT"));
531 }
532
533 #[test]
534 fn test_registry_name_case_insensitive() {
535 let reg = CollationRegistry::new();
536 assert!(reg.find("BINARY").is_some());
538 assert!(reg.find("binary").is_some());
539 assert!(reg.find("Binary").is_some());
540 assert!(reg.find("bInArY").is_some());
541
542 assert!(reg.contains("nocase"));
544 assert!(reg.contains("NOCASE"));
545 assert!(reg.contains("NoCase"));
546 }
547
548 fn ann(name: &str, source: CollationSource) -> CollationAnnotation {
551 CollationAnnotation {
552 name: name.to_owned(),
553 source,
554 }
555 }
556
557 #[test]
558 fn test_collation_selection_explicit_wins() {
559 let result = resolve_collation(
561 &ann("NOCASE", CollationSource::Explicit),
562 &ann("BINARY", CollationSource::Default),
563 );
564 assert_eq!(result, "NOCASE");
565 }
566
567 #[test]
568 fn test_collation_selection_explicit_rhs_wins_over_default() {
569 let result = resolve_collation(
570 &ann("BINARY", CollationSource::Default),
571 &ann("RTRIM", CollationSource::Explicit),
572 );
573 assert_eq!(result, "RTRIM");
574 }
575
576 #[test]
577 fn test_collation_selection_leftmost_explicit_wins() {
578 let result = resolve_collation(
580 &ann("NOCASE", CollationSource::Explicit),
581 &ann("RTRIM", CollationSource::Explicit),
582 );
583 assert_eq!(result, "NOCASE");
584 }
585
586 #[test]
587 fn test_collation_selection_schema_over_default() {
588 let result = resolve_collation(
589 &ann("NOCASE", CollationSource::Schema),
590 &ann("BINARY", CollationSource::Default),
591 );
592 assert_eq!(result, "NOCASE");
593 }
594
595 #[test]
596 fn test_collation_selection_schema_rhs_over_default() {
597 let result = resolve_collation(
598 &ann("BINARY", CollationSource::Default),
599 &ann("NOCASE", CollationSource::Schema),
600 );
601 assert_eq!(result, "NOCASE");
602 }
603
604 #[test]
605 fn test_collation_selection_explicit_over_schema() {
606 let result = resolve_collation(
607 &ann("RTRIM", CollationSource::Explicit),
608 &ann("NOCASE", CollationSource::Schema),
609 );
610 assert_eq!(result, "RTRIM");
611 }
612
613 #[test]
614 fn test_collation_selection_default_binary() {
615 let result = resolve_collation(
616 &ann("BINARY", CollationSource::Default),
617 &ann("BINARY", CollationSource::Default),
618 );
619 assert_eq!(result, "BINARY");
620 }
621
622 #[test]
625 fn test_min_respects_collation() {
626 let binary = BinaryCollation;
628 let binary_min = if binary.compare(b"ABC", b"abc") == Ordering::Less {
629 "ABC"
630 } else {
631 "abc"
632 };
633 assert_eq!(binary_min, "ABC");
634
635 let nocase = NoCaseCollation;
637 assert_eq!(nocase.compare(b"ABC", b"abc"), Ordering::Equal);
638 }
639
640 #[test]
641 fn test_max_respects_collation() {
642 let binary = BinaryCollation;
643 let binary_max = if binary.compare(b"abc", b"ABC") == Ordering::Greater {
645 "abc"
646 } else {
647 "ABC"
648 };
649 assert_eq!(binary_max, "abc");
650 }
651
652 #[test]
653 fn test_collation_aware_sort() {
654 let nocase = NoCaseCollation;
656 let mut data: Vec<&[u8]> = vec![b"Banana", b"apple", b"Cherry", b"date"];
657 data.sort_by(|a, b| nocase.compare(a, b));
658
659 assert_eq!(data[0], b"apple");
661 assert_eq!(data[1], b"Banana");
662 assert_eq!(data[2], b"Cherry");
663 assert_eq!(data[3], b"date");
664 }
665
666 #[test]
667 fn test_collation_aware_group_by() {
668 let nocase = NoCaseCollation;
670 let items: Vec<&[u8]> = vec![b"ABC", b"abc", b"Abc", b"def", b"DEF"];
671 let mut groups: Vec<Vec<&[u8]>> = Vec::new();
672
673 let mut sorted = items;
675 sorted.sort_by(|a, b| nocase.compare(a, b));
676
677 let mut current_group: Vec<&[u8]> = vec![sorted[0]];
678 for window in sorted.windows(2) {
679 if nocase.compare(window[0], window[1]) != Ordering::Equal {
680 groups.push(std::mem::take(&mut current_group));
681 }
682 current_group.push(window[1]);
683 }
684 groups.push(current_group);
685
686 assert_eq!(groups.len(), 2);
688 assert_eq!(groups[0].len(), 3);
689 assert_eq!(groups[1].len(), 2);
690 }
691
692 #[test]
693 fn test_collation_aware_distinct() {
694 let nocase = NoCaseCollation;
696 let items: Vec<&[u8]> = vec![b"ABC", b"abc", b"Abc", b"def", b"DEF"];
697
698 let mut distinct: Vec<&[u8]> = Vec::new();
699 for item in &items {
700 let already = distinct
701 .iter()
702 .any(|d| nocase.compare(d, item) == Ordering::Equal);
703 if !already {
704 distinct.push(item);
705 }
706 }
707
708 assert_eq!(distinct.len(), 2);
710 }
711
712 #[test]
713 fn test_registry_default_impl() {
714 let reg = CollationRegistry::default();
716 assert!(reg.contains("BINARY"));
717 assert!(reg.contains("NOCASE"));
718 assert!(reg.contains("RTRIM"));
719 }
720
721 #[test]
722 fn test_collation_annotation_debug() {
723 let ann = CollationAnnotation {
724 name: "NOCASE".to_owned(),
725 source: CollationSource::Explicit,
726 };
727 let debug_str = format!("{ann:?}");
728 assert!(debug_str.contains("NOCASE"));
729 assert!(debug_str.contains("Explicit"));
730 }
731
732 #[test]
733 fn test_collation_source_equality() {
734 assert_eq!(CollationSource::Explicit, CollationSource::Explicit);
735 assert_ne!(CollationSource::Explicit, CollationSource::Schema);
736 assert_ne!(CollationSource::Schema, CollationSource::Default);
737 }
738}