1#![allow(non_snake_case)]
19
20use std::collections::HashMap;
54use std::sync::{Arc, RwLock};
55
56use super::routing::{RiRoute, RiRouteHandler};
57use crate::core::lock::RwLockExtensions;
58
59#[derive(Debug, Clone, PartialEq)]
61pub enum SegmentType {
62 Static,
64 Param,
66 Wildcard,
68}
69
70#[derive(Debug, Clone)]
72pub struct PathSegment {
73 pub value: String,
75 pub segment_type: SegmentType,
77}
78
79impl PathSegment {
80 pub fn new(value: &str) -> Self {
90 let segment_type = if value.starts_with('*') {
91 SegmentType::Wildcard
92 } else if value.starts_with(':') {
93 SegmentType::Param
94 } else {
95 SegmentType::Static
96 };
97
98 Self {
99 value: value.to_string(),
100 segment_type,
101 }
102 }
103
104 pub fn param_name(&self) -> Option<String> {
112 match self.segment_type {
113 SegmentType::Param => Some(self.value[1..].to_string()),
114 SegmentType::Wildcard => Some(self.value[1..].to_string()),
115 SegmentType::Static => None,
116 }
117 }
118}
119
120#[derive(Debug, Clone)]
122pub struct RouteMatch {
123 pub route: RiRoute,
125 pub params: HashMap<String, String>,
127}
128
129#[derive(Clone)]
136pub struct RadixNode {
137 pub segment: PathSegment,
139 pub handler: Option<RiRouteHandler>,
141 pub route_data: Option<RiRoute>,
143 pub children: HashMap<String, Arc<RwLock<RadixNode>>>,
145 pub param_child: Option<Arc<RwLock<RadixNode>>>,
147 pub wildcard_child: Option<Arc<RwLock<RadixNode>>>,
149}
150
151impl RadixNode {
152 pub fn new(segment: PathSegment) -> Self {
162 Self {
163 segment,
164 handler: None,
165 route_data: None,
166 children: HashMap::new(),
167 param_child: None,
168 wildcard_child: None,
169 }
170 }
171
172 pub fn root() -> Self {
178 Self::new(PathSegment {
179 value: String::new(),
180 segment_type: SegmentType::Static,
181 })
182 }
183
184 pub fn is_terminal(&self) -> bool {
190 self.handler.is_some()
191 }
192
193 pub fn add_static_child(&mut self, segment: PathSegment) -> Arc<RwLock<RadixNode>> {
203 let key = segment.value.clone();
204 let node = Arc::new(RwLock::new(RadixNode::new(segment)));
205 self.children.insert(key, node.clone());
206 node
207 }
208
209 pub fn get_or_create_child(&mut self, segment: PathSegment) -> Arc<RwLock<RadixNode>> {
219 match segment.segment_type {
220 SegmentType::Static => {
221 if let Some(child) = self.children.get(&segment.value) {
222 child.clone()
223 } else {
224 self.add_static_child(segment)
225 }
226 }
227 SegmentType::Param => {
228 if let Some(child) = &self.param_child {
229 child.clone()
230 } else {
231 let node = Arc::new(RwLock::new(RadixNode::new(segment)));
232 self.param_child = Some(node.clone());
233 node
234 }
235 }
236 SegmentType::Wildcard => {
237 if let Some(child) = &self.wildcard_child {
238 child.clone()
239 } else {
240 let node = Arc::new(RwLock::new(RadixNode::new(segment)));
241 self.wildcard_child = Some(node.clone());
242 node
243 }
244 }
245 }
246 }
247}
248
249pub struct RiRadixTree {
254 root: Arc<RwLock<RadixNode>>,
256 route_count: RwLock<usize>,
258}
259
260impl Default for RiRadixTree {
261 fn default() -> Self {
262 Self::new()
263 }
264}
265
266impl RiRadixTree {
267 pub fn new() -> Self {
273 Self {
274 root: Arc::new(RwLock::new(RadixNode::root())),
275 route_count: RwLock::new(0),
276 }
277 }
278
279 pub fn parse_path(path: &str) -> Vec<PathSegment> {
289 path.split('/')
290 .filter(|s| !s.is_empty())
291 .map(PathSegment::new)
292 .collect()
293 }
294
295 pub fn insert(&self, route: RiRoute) {
301 let segments = Self::parse_path(&route.path);
302 let handler = route.handler.clone();
303 let route_clone = route.clone();
304
305 let mut current = self.root.clone();
306
307 if segments.is_empty() {
308 let mut node = match current.write_safe("root for insert") {
309 Ok(n) => n,
310 Err(_) => return,
311 };
312 node.handler = Some(handler);
313 node.route_data = Some(route_clone);
314 let mut count = match self.route_count.write_safe("count for insert") {
315 Ok(c) => c,
316 Err(_) => return,
317 };
318 *count += 1;
319 return;
320 }
321
322 for segment in segments {
323 let segment_type = segment.segment_type.clone();
324 let segment_value = segment.value.clone();
325
326 let next = {
327 let node = match current.read_safe("node for insert read") {
328 Ok(n) => n,
329 Err(_) => return,
330 };
331
332 match segment_type {
333 SegmentType::Static => {
334 if let Some(child) = node.children.get(&segment_value) {
335 child.clone()
336 } else {
337 drop(node);
338 let mut node_mut = match current.write_safe("node for insert write") {
339 Ok(n) => n,
340 Err(_) => return,
341 };
342 node_mut.get_or_create_child(segment)
343 }
344 }
345 SegmentType::Param => {
346 if let Some(child) = &node.param_child {
347 child.clone()
348 } else {
349 drop(node);
350 let mut node_mut = match current.write_safe("node for insert write param") {
351 Ok(n) => n,
352 Err(_) => return,
353 };
354 node_mut.get_or_create_child(segment)
355 }
356 }
357 SegmentType::Wildcard => {
358 if let Some(child) = &node.wildcard_child {
359 child.clone()
360 } else {
361 drop(node);
362 let mut node_mut = match current.write_safe("node for insert write wildcard") {
363 Ok(n) => n,
364 Err(_) => return,
365 };
366 node_mut.get_or_create_child(segment)
367 }
368 }
369 }
370 };
371
372 current = next;
373 }
374
375 let mut node = match current.write_safe("final node for insert") {
376 Ok(n) => n,
377 Err(_) => return,
378 };
379 node.handler = Some(handler);
380 node.route_data = Some(route_clone);
381
382 let mut count = match self.route_count.write_safe("count for insert final") {
383 Ok(c) => c,
384 Err(_) => return,
385 };
386 *count += 1;
387 }
388
389 pub fn find(&self, path: &str) -> Option<RouteMatch> {
399 let segments: Vec<&str> = path.split('/').filter(|s| !s.is_empty()).collect();
400 let mut params = HashMap::new();
401
402 let root = match self.root.read_safe("root for find") {
403 Ok(r) => r,
404 Err(_) => return None,
405 };
406
407 if segments.is_empty() {
408 if root.is_terminal() {
409 return root.route_data.clone().map(|route| RouteMatch { route, params });
410 }
411 return None;
412 }
413
414 Self::find_recursive(Arc::new(RwLock::new((*root).clone())), &segments, 0, &mut params)
415 }
416
417 fn find_recursive(
430 node: Arc<RwLock<RadixNode>>,
431 segments: &[&str],
432 index: usize,
433 params: &mut HashMap<String, String>,
434 ) -> Option<RouteMatch> {
435 if index >= segments.len() {
436 let node_guard = match node.read_safe("node for terminal check") {
437 Ok(n) => n,
438 Err(_) => return None,
439 };
440 if node_guard.is_terminal() {
441 return node_guard.route_data.clone().map(|route| RouteMatch {
442 route,
443 params: params.clone()
444 });
445 }
446 return None;
447 }
448
449 let current_segment = segments[index];
450 let node_guard = match node.read_safe("node for find recursive") {
451 Ok(n) => n,
452 Err(_) => return None,
453 };
454
455 if let Some(child) = node_guard.children.get(current_segment) {
456 let result = Self::find_recursive(child.clone(), segments, index + 1, params);
457 if result.is_some() {
458 return result;
459 }
460 }
461
462 if let Some(param_child) = &node_guard.param_child {
463 let param_node = match param_child.read_safe("param child for find") {
464 Ok(n) => n,
465 Err(_) => return None,
466 };
467 if let Some(param_name) = param_node.segment.param_name() {
468 params.insert(param_name, current_segment.to_string());
469 }
470 drop(param_node);
471
472 let result = Self::find_recursive(param_child.clone(), segments, index + 1, params);
473 if result.is_some() {
474 return result;
475 }
476
477 if let Some(name) = params.keys().last().cloned() {
478 if name == current_segment {
479 params.remove(&name);
480 }
481 }
482 }
483
484 if let Some(wildcard_child) = &node_guard.wildcard_child {
485 let wildcard_node = match wildcard_child.read_safe("wildcard child for find") {
486 Ok(n) => n,
487 Err(_) => return None,
488 };
489 if let Some(wildcard_name) = wildcard_node.segment.param_name() {
490 let remaining_path = segments[index..].join("/");
491 params.insert(wildcard_name, remaining_path);
492 }
493
494 if wildcard_node.is_terminal() {
495 return wildcard_node.route_data.clone().map(|route| RouteMatch {
496 route,
497 params: params.clone()
498 });
499 }
500 }
501
502 None
503 }
504
505 pub fn remove(&self, path: &str) -> bool {
515 let segments = Self::parse_path(path);
516
517 if segments.is_empty() {
518 let mut root = match self.root.write_safe("root for remove") {
519 Ok(r) => r,
520 Err(_) => return false,
521 };
522 if root.handler.is_some() {
523 root.handler = None;
524 root.route_data = None;
525 let mut count = match self.route_count.write_safe("count for remove") {
526 Ok(c) => c,
527 Err(_) => return false,
528 };
529 *count = count.saturating_sub(1);
530 return true;
531 }
532 return false;
533 }
534
535 let result = Self::remove_recursive(self.root.clone(), &segments, 0);
536 if result {
537 let mut count = match self.route_count.write_safe("count for remove success") {
538 Ok(c) => c,
539 Err(_) => return result,
540 };
541 *count = count.saturating_sub(1);
542 }
543 result
544 }
545
546 fn remove_recursive(
558 node: Arc<RwLock<RadixNode>>,
559 segments: &[PathSegment],
560 index: usize,
561 ) -> bool {
562 if index >= segments.len() {
563 let mut node_guard = match node.write_safe("node for remove terminal") {
564 Ok(n) => n,
565 Err(_) => return false,
566 };
567 if node_guard.handler.is_some() {
568 node_guard.handler = None;
569 node_guard.route_data = None;
570 return true;
571 }
572 return false;
573 }
574
575 let segment = &segments[index];
576 let next_node = {
577 let node_guard = match node.read_safe("node for remove read") {
578 Ok(n) => n,
579 Err(_) => return false,
580 };
581
582 match segment.segment_type {
583 SegmentType::Static => node_guard.children.get(&segment.value).cloned(),
584 SegmentType::Param => node_guard.param_child.clone(),
585 SegmentType::Wildcard => node_guard.wildcard_child.clone(),
586 }
587 };
588
589 if let Some(next) = next_node {
590 let result = Self::remove_recursive(next.clone(), segments, index + 1);
591 if result {
592 let next_guard = match next.read_safe("next for remove cleanup") {
593 Ok(n) => n,
594 Err(_) => return result,
595 };
596 if next_guard.children.is_empty()
597 && next_guard.param_child.is_none()
598 && next_guard.wildcard_child.is_none()
599 && !next_guard.is_terminal() {
600 drop(next_guard);
601 let mut node_guard = match node.write_safe("node for remove cleanup") {
602 Ok(n) => n,
603 Err(_) => return result,
604 };
605 match segment.segment_type {
606 SegmentType::Static => {
607 node_guard.children.remove(&segment.value);
608 }
609 SegmentType::Param => {
610 node_guard.param_child = None;
611 }
612 SegmentType::Wildcard => {
613 node_guard.wildcard_child = None;
614 }
615 }
616 }
617 }
618 return result;
619 }
620
621 false
622 }
623
624 pub fn route_count(&self) -> usize {
630 match self.route_count.read_safe("route count") {
631 Ok(count) => *count,
632 Err(_) => 0,
633 }
634 }
635
636 pub fn clear(&self) {
638 let mut root = match self.root.write_safe("root for clear") {
639 Ok(r) => r,
640 Err(_) => return,
641 };
642 *root = RadixNode::root();
643
644 let mut count = match self.route_count.write_safe("count for clear") {
645 Ok(c) => c,
646 Err(_) => return,
647 };
648 *count = 0;
649 }
650}
651
652#[cfg(test)]
653mod tests {
654 use super::*;
655 use std::sync::Arc;
656 use std::pin::Pin;
657 use std::future::Future;
658
659 fn create_test_handler() -> RiRouteHandler {
660 Arc::new(|_req| {
661 Box::pin(async move {
662 Ok(crate::gateway::RiGatewayResponse::new(200, b"OK".to_vec(), String::new()))
663 }) as Pin<Box<dyn Future<Output = RiResult<crate::gateway::RiGatewayResponse>> + Send>>
664 })
665 }
666
667 #[test]
668 fn test_path_segment_static() {
669 let segment = PathSegment::new("users");
670 assert_eq!(segment.segment_type, SegmentType::Static);
671 assert_eq!(segment.value, "users");
672 assert!(segment.param_name().is_none());
673 }
674
675 #[test]
676 fn test_path_segment_param() {
677 let segment = PathSegment::new(":id");
678 assert_eq!(segment.segment_type, SegmentType::Param);
679 assert_eq!(segment.value, ":id");
680 assert_eq!(segment.param_name(), Some("id".to_string()));
681 }
682
683 #[test]
684 fn test_path_segment_wildcard() {
685 let segment = PathSegment::new("*path");
686 assert_eq!(segment.segment_type, SegmentType::Wildcard);
687 assert_eq!(segment.value, "*path");
688 assert_eq!(segment.param_name(), Some("path".to_string()));
689 }
690
691 #[test]
692 fn test_parse_path() {
693 let segments = RiRadixTree::parse_path("/api/v1/users");
694 assert_eq!(segments.len(), 3);
695 assert_eq!(segments[0].value, "api");
696 assert_eq!(segments[1].value, "v1");
697 assert_eq!(segments[2].value, "users");
698 }
699
700 #[test]
701 fn test_parse_path_with_param() {
702 let segments = RiRadixTree::parse_path("/users/:id");
703 assert_eq!(segments.len(), 2);
704 assert_eq!(segments[0].segment_type, SegmentType::Static);
705 assert_eq!(segments[1].segment_type, SegmentType::Param);
706 }
707
708 #[test]
709 fn test_parse_path_with_wildcard() {
710 let segments = RiRadixTree::parse_path("/files/*path");
711 assert_eq!(segments.len(), 2);
712 assert_eq!(segments[0].segment_type, SegmentType::Static);
713 assert_eq!(segments[1].segment_type, SegmentType::Wildcard);
714 }
715
716 #[test]
717 fn test_radix_tree_insert_and_find() {
718 let tree = RiRadixTree::new();
719 let route = RiRoute::new(
720 "GET".to_string(),
721 "/api/users".to_string(),
722 create_test_handler(),
723 );
724
725 tree.insert(route);
726 assert_eq!(tree.route_count(), 1);
727
728 let result = tree.find("/api/users");
729 assert!(result.is_some());
730 }
731
732 #[test]
733 fn test_radix_tree_find_with_param() {
734 let tree = RiRadixTree::new();
735 let route = RiRoute::new(
736 "GET".to_string(),
737 "/users/:id".to_string(),
738 create_test_handler(),
739 );
740
741 tree.insert(route);
742
743 let result = tree.find("/users/123");
744 assert!(result.is_some());
745 let match_result = result.unwrap();
746 assert_eq!(match_result.params.get("id"), Some(&"123".to_string()));
747 }
748
749 #[test]
750 fn test_radix_tree_find_with_wildcard() {
751 let tree = RiRadixTree::new();
752 let route = RiRoute::new(
753 "GET".to_string(),
754 "/files/*path".to_string(),
755 create_test_handler(),
756 );
757
758 tree.insert(route);
759
760 let result = tree.find("/files/docs/readme.txt");
761 assert!(result.is_some());
762 let match_result = result.unwrap();
763 assert_eq!(match_result.params.get("path"), Some(&"docs/readme.txt".to_string()));
764 }
765
766 #[test]
767 fn test_radix_tree_remove() {
768 let tree = RiRadixTree::new();
769 let route = RiRoute::new(
770 "GET".to_string(),
771 "/api/users".to_string(),
772 create_test_handler(),
773 );
774
775 tree.insert(route);
776 assert_eq!(tree.route_count(), 1);
777
778 let removed = tree.remove("/api/users");
779 assert!(removed);
780 assert_eq!(tree.route_count(), 0);
781
782 let result = tree.find("/api/users");
783 assert!(result.is_none());
784 }
785
786 #[test]
787 fn test_radix_tree_no_match() {
788 let tree = RiRadixTree::new();
789 let route = RiRoute::new(
790 "GET".to_string(),
791 "/api/users".to_string(),
792 create_test_handler(),
793 );
794
795 tree.insert(route);
796
797 let result = tree.find("/api/posts");
798 assert!(result.is_none());
799 }
800
801 #[test]
802 fn test_radix_tree_clear() {
803 let tree = RiRadixTree::new();
804
805 for i in 0..10 {
806 let route = RiRoute::new(
807 "GET".to_string(),
808 format!("/api/route_{}", i),
809 create_test_handler(),
810 );
811 tree.insert(route);
812 }
813
814 assert_eq!(tree.route_count(), 10);
815
816 tree.clear();
817 assert_eq!(tree.route_count(), 0);
818 }
819}