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]
234 pub fn names(&self) -> Vec<String> {
235 let mut names = vec!["BINARY".to_owned(), "NOCASE".to_owned(), "RTRIM".to_owned()];
236 let mut custom: Vec<String> = self
237 .custom_collations
238 .keys()
239 .filter(|name| !matches!(name.as_str(), "BINARY" | "NOCASE" | "RTRIM"))
240 .cloned()
241 .collect();
242 custom.sort_unstable_by_key(|name| name.to_ascii_uppercase());
243 names.extend(custom);
244 names
245 }
246}
247
248#[derive(Debug, Clone, Copy, PartialEq, Eq)]
255pub enum CollationSource {
256 Explicit,
258 Schema,
260 Default,
262}
263
264#[derive(Debug, Clone)]
266pub struct CollationAnnotation {
267 pub name: String,
269 pub source: CollationSource,
271}
272
273#[must_use]
283pub fn resolve_collation(lhs: &CollationAnnotation, rhs: &CollationAnnotation) -> String {
284 let result = match (lhs.source, rhs.source) {
286 (_, CollationSource::Explicit) if lhs.source != CollationSource::Explicit => &rhs.name,
287 (CollationSource::Default, CollationSource::Schema) => &rhs.name,
288 _ => &lhs.name,
289 };
290 debug!(
291 collation = %result,
292 lhs_source = ?lhs.source,
293 rhs_source = ?rhs.source,
294 context = "COMPARE",
295 "collation selection"
296 );
297 result.clone()
298}
299
300#[cfg(test)]
301mod tests {
302 use super::*;
303
304 #[test]
307 fn test_collation_binary_memcmp() {
308 let coll = BinaryCollation;
309 assert_eq!(coll.compare(b"abc", b"abc"), Ordering::Equal);
310 assert_eq!(coll.compare(b"abc", b"abd"), Ordering::Less);
311 assert_eq!(coll.compare(b"abd", b"abc"), Ordering::Greater);
312 assert_eq!(coll.compare(b"ABC", b"abc"), Ordering::Less);
314 assert_eq!(
316 coll.compare("café".as_bytes(), "café".as_bytes()),
317 Ordering::Equal
318 );
319 assert_ne!(coll.compare("über".as_bytes(), b"uber"), Ordering::Equal);
320 }
321
322 #[test]
323 fn test_collation_binary_basic() {
324 let coll = BinaryCollation;
325 assert_eq!(coll.compare(b"ABC", b"abc"), Ordering::Less);
327 assert_eq!(coll.compare(b"\x00", b"\x01"), Ordering::Less);
329 assert_eq!(coll.compare(b"\xff", b"\x00"), Ordering::Greater);
330 }
331
332 #[test]
333 fn test_collation_nocase_ascii() {
334 let coll = NoCaseCollation;
335 assert_eq!(coll.compare(b"ABC", b"abc"), Ordering::Equal);
336 assert_eq!(coll.compare(b"Alice", b"alice"), Ordering::Equal);
337 assert_eq!(coll.compare(b"[", b"a"), Ordering::Greater);
339 }
340
341 #[test]
342 fn test_collation_nocase_ascii_only() {
343 let coll = NoCaseCollation;
344 assert_ne!(
346 coll.compare("Ä".as_bytes(), "ä".as_bytes()),
347 Ordering::Equal,
348 "NOCASE must NOT fold non-ASCII"
349 );
350 assert_eq!(coll.compare(b"Z", b"z"), Ordering::Equal);
352 assert_eq!(coll.compare(b"[", b"["), Ordering::Equal);
353 assert_ne!(coll.compare(b"[", b"{"), Ordering::Equal);
355 }
356
357 #[test]
358 fn test_collation_rtrim() {
359 let coll = RtrimCollation;
360 assert_eq!(coll.compare(b"hello ", b"hello"), Ordering::Equal);
362 assert_eq!(coll.compare(b"hello", b"hello "), Ordering::Equal);
363 assert_eq!(coll.compare(b"hello ", b"hello "), Ordering::Equal);
364 assert_ne!(coll.compare(b"hello!", b"hello"), Ordering::Equal);
366 assert_ne!(coll.compare(b"hello ", b"hello!"), Ordering::Equal);
368 }
369
370 #[test]
371 fn test_collation_rtrim_tabs_not_stripped() {
372 let coll = RtrimCollation;
373 assert_ne!(
375 coll.compare(b"hello\t", b"hello"),
376 Ordering::Equal,
377 "RTRIM must NOT strip tabs"
378 );
379 assert_ne!(
381 coll.compare(b"hello\xc2\xa0", b"hello"),
382 Ordering::Equal,
383 "RTRIM must NOT strip non-breaking spaces"
384 );
385 }
386
387 #[test]
388 fn test_collation_properties_antisymmetric() {
389 let collations: Vec<Box<dyn CollationFunction>> = vec![
390 Box::new(BinaryCollation),
391 Box::new(NoCaseCollation),
392 Box::new(RtrimCollation),
393 ];
394
395 let pairs: &[(&[u8], &[u8])] = &[
396 (b"abc", b"def"),
397 (b"hello", b"world"),
398 (b"ABC", b"abc"),
399 (b"hello ", b"hello"),
400 ];
401
402 for coll in &collations {
403 for &(a, b) in pairs {
404 let forward = coll.compare(a, b);
405 let reverse = coll.compare(b, a);
406 assert_eq!(
407 forward,
408 reverse.reverse(),
409 "{}: compare({:?}, {:?}) = {forward:?}, but reverse = {reverse:?}",
410 coll.name(),
411 std::str::from_utf8(a).unwrap_or("?"),
412 std::str::from_utf8(b).unwrap_or("?"),
413 );
414 }
415 }
416 }
417
418 #[test]
419 fn test_collation_properties_transitive() {
420 let coll = BinaryCollation;
421 let a = b"apple";
422 let b = b"banana";
423 let c = b"cherry";
424
425 assert_eq!(coll.compare(a, b), Ordering::Less);
427 assert_eq!(coll.compare(b, c), Ordering::Less);
428 assert_eq!(coll.compare(a, c), Ordering::Less);
429 }
430
431 #[test]
432 fn test_collation_send_sync() {
433 fn assert_send_sync<T: Send + Sync>() {}
434 assert_send_sync::<BinaryCollation>();
435 assert_send_sync::<NoCaseCollation>();
436 assert_send_sync::<RtrimCollation>();
437 }
438
439 #[test]
442 fn test_registry_preloaded_builtins() {
443 let reg = CollationRegistry::new();
444 assert!(reg.contains("BINARY"));
445 assert!(reg.contains("NOCASE"));
446 assert!(reg.contains("RTRIM"));
447
448 let binary = reg.find("BINARY").expect("BINARY must be pre-registered");
449 assert_eq!(binary.compare(b"a", b"b"), Ordering::Less);
450
451 let nocase = reg.find("NOCASE").expect("NOCASE must be pre-registered");
452 assert_eq!(nocase.compare(b"ABC", b"abc"), Ordering::Equal);
453
454 let rtrim = reg.find("RTRIM").expect("RTRIM must be pre-registered");
455 assert_eq!(rtrim.compare(b"x ", b"x"), Ordering::Equal);
456 }
457
458 struct ReverseCollation;
459
460 impl CollationFunction for ReverseCollation {
461 fn name(&self) -> &str {
462 "REVERSE"
463 }
464
465 fn compare(&self, left: &[u8], right: &[u8]) -> Ordering {
466 right.cmp(left)
467 }
468 }
469
470 #[test]
471 fn test_registry_custom_collation_registration() {
472 let mut reg = CollationRegistry::new();
473
474 let prev = reg.register(ReverseCollation);
475 assert!(prev.is_none(), "no prior REVERSE collation");
476 assert!(reg.contains("REVERSE"));
477
478 let coll = reg.find("reverse").expect("case-insensitive lookup");
479 assert_eq!(coll.compare(b"a", b"z"), Ordering::Greater);
480 }
481
482 struct AlwaysEqualCollation;
483
484 impl CollationFunction for AlwaysEqualCollation {
485 fn name(&self) -> &str {
486 "BINARY"
487 }
488
489 fn compare(&self, _left: &[u8], _right: &[u8]) -> Ordering {
490 Ordering::Equal
491 }
492 }
493
494 #[test]
495 fn test_registry_overwrite_builtin() {
496 let mut reg = CollationRegistry::new();
497 assert!(reg.uses_builtin_implementation("binary"));
498
499 let prev = reg.register(AlwaysEqualCollation);
500 assert!(prev.is_some(), "should return previous BINARY collation");
501 assert!(!reg.uses_builtin_implementation("BINARY"));
502
503 let coll = reg.find("BINARY").unwrap();
504 assert_eq!(
505 coll.compare(b"a", b"z"),
506 Ordering::Equal,
507 "custom overwrite must take effect"
508 );
509 }
510
511 #[test]
512 fn test_registry_unregistered_returns_none() {
513 let reg = CollationRegistry::new();
514 assert!(reg.find("NONEXISTENT").is_none());
515 assert!(!reg.contains("NONEXISTENT"));
516 }
517
518 #[test]
519 fn test_registry_name_case_insensitive() {
520 let reg = CollationRegistry::new();
521 assert!(reg.find("BINARY").is_some());
523 assert!(reg.find("binary").is_some());
524 assert!(reg.find("Binary").is_some());
525 assert!(reg.find("bInArY").is_some());
526
527 assert!(reg.contains("nocase"));
529 assert!(reg.contains("NOCASE"));
530 assert!(reg.contains("NoCase"));
531 }
532
533 fn ann(name: &str, source: CollationSource) -> CollationAnnotation {
536 CollationAnnotation {
537 name: name.to_owned(),
538 source,
539 }
540 }
541
542 #[test]
543 fn test_collation_selection_explicit_wins() {
544 let result = resolve_collation(
546 &ann("NOCASE", CollationSource::Explicit),
547 &ann("BINARY", CollationSource::Default),
548 );
549 assert_eq!(result, "NOCASE");
550 }
551
552 #[test]
553 fn test_collation_selection_explicit_rhs_wins_over_default() {
554 let result = resolve_collation(
555 &ann("BINARY", CollationSource::Default),
556 &ann("RTRIM", CollationSource::Explicit),
557 );
558 assert_eq!(result, "RTRIM");
559 }
560
561 #[test]
562 fn test_collation_selection_leftmost_explicit_wins() {
563 let result = resolve_collation(
565 &ann("NOCASE", CollationSource::Explicit),
566 &ann("RTRIM", CollationSource::Explicit),
567 );
568 assert_eq!(result, "NOCASE");
569 }
570
571 #[test]
572 fn test_collation_selection_schema_over_default() {
573 let result = resolve_collation(
574 &ann("NOCASE", CollationSource::Schema),
575 &ann("BINARY", CollationSource::Default),
576 );
577 assert_eq!(result, "NOCASE");
578 }
579
580 #[test]
581 fn test_collation_selection_schema_rhs_over_default() {
582 let result = resolve_collation(
583 &ann("BINARY", CollationSource::Default),
584 &ann("NOCASE", CollationSource::Schema),
585 );
586 assert_eq!(result, "NOCASE");
587 }
588
589 #[test]
590 fn test_collation_selection_explicit_over_schema() {
591 let result = resolve_collation(
592 &ann("RTRIM", CollationSource::Explicit),
593 &ann("NOCASE", CollationSource::Schema),
594 );
595 assert_eq!(result, "RTRIM");
596 }
597
598 #[test]
599 fn test_collation_selection_default_binary() {
600 let result = resolve_collation(
601 &ann("BINARY", CollationSource::Default),
602 &ann("BINARY", CollationSource::Default),
603 );
604 assert_eq!(result, "BINARY");
605 }
606
607 #[test]
610 fn test_min_respects_collation() {
611 let binary = BinaryCollation;
613 let binary_min = if binary.compare(b"ABC", b"abc") == Ordering::Less {
614 "ABC"
615 } else {
616 "abc"
617 };
618 assert_eq!(binary_min, "ABC");
619
620 let nocase = NoCaseCollation;
622 assert_eq!(nocase.compare(b"ABC", b"abc"), Ordering::Equal);
623 }
624
625 #[test]
626 fn test_max_respects_collation() {
627 let binary = BinaryCollation;
628 let binary_max = if binary.compare(b"abc", b"ABC") == Ordering::Greater {
630 "abc"
631 } else {
632 "ABC"
633 };
634 assert_eq!(binary_max, "abc");
635 }
636
637 #[test]
638 fn test_collation_aware_sort() {
639 let nocase = NoCaseCollation;
641 let mut data: Vec<&[u8]> = vec![b"Banana", b"apple", b"Cherry", b"date"];
642 data.sort_by(|a, b| nocase.compare(a, b));
643
644 assert_eq!(data[0], b"apple");
646 assert_eq!(data[1], b"Banana");
647 assert_eq!(data[2], b"Cherry");
648 assert_eq!(data[3], b"date");
649 }
650
651 #[test]
652 fn test_collation_aware_group_by() {
653 let nocase = NoCaseCollation;
655 let items: Vec<&[u8]> = vec![b"ABC", b"abc", b"Abc", b"def", b"DEF"];
656 let mut groups: Vec<Vec<&[u8]>> = Vec::new();
657
658 let mut sorted = items;
660 sorted.sort_by(|a, b| nocase.compare(a, b));
661
662 let mut current_group: Vec<&[u8]> = vec![sorted[0]];
663 for window in sorted.windows(2) {
664 if nocase.compare(window[0], window[1]) != Ordering::Equal {
665 groups.push(std::mem::take(&mut current_group));
666 }
667 current_group.push(window[1]);
668 }
669 groups.push(current_group);
670
671 assert_eq!(groups.len(), 2);
673 assert_eq!(groups[0].len(), 3);
674 assert_eq!(groups[1].len(), 2);
675 }
676
677 #[test]
678 fn test_collation_aware_distinct() {
679 let nocase = NoCaseCollation;
681 let items: Vec<&[u8]> = vec![b"ABC", b"abc", b"Abc", b"def", b"DEF"];
682
683 let mut distinct: Vec<&[u8]> = Vec::new();
684 for item in &items {
685 let already = distinct
686 .iter()
687 .any(|d| nocase.compare(d, item) == Ordering::Equal);
688 if !already {
689 distinct.push(item);
690 }
691 }
692
693 assert_eq!(distinct.len(), 2);
695 }
696
697 #[test]
698 fn test_registry_default_impl() {
699 let reg = CollationRegistry::default();
701 assert!(reg.contains("BINARY"));
702 assert!(reg.contains("NOCASE"));
703 assert!(reg.contains("RTRIM"));
704 }
705
706 #[test]
707 fn test_collation_annotation_debug() {
708 let ann = CollationAnnotation {
709 name: "NOCASE".to_owned(),
710 source: CollationSource::Explicit,
711 };
712 let debug_str = format!("{ann:?}");
713 assert!(debug_str.contains("NOCASE"));
714 assert!(debug_str.contains("Explicit"));
715 }
716
717 #[test]
718 fn test_collation_source_equality() {
719 assert_eq!(CollationSource::Explicit, CollationSource::Explicit);
720 assert_ne!(CollationSource::Explicit, CollationSource::Schema);
721 assert_ne!(CollationSource::Schema, CollationSource::Default);
722 }
723}