Skip to main content

ri/gateway/
radix_tree.rs

1//! Copyright © 2025-2026 Wenze Wei. All Rights Reserved.
2//!
3//! This file is part of Ri.
4//! The Ri project belongs to the Dunimd Team.
5//!
6//! Licensed under the Apache License, Version 2.0 (the "License");
7//! You may not use this file except in compliance with the License.
8//! You may obtain a copy of the License at
9//!
10//!     http://www.apache.org/licenses/LICENSE-2.0
11//!
12//! Unless required by applicable law or agreed to in writing, software
13//! distributed under the License is distributed on an "AS IS" BASIS,
14//! WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15//! See the License for the specific language governing permissions and
16//! limitations under the License.
17
18#![allow(non_snake_case)]
19
20//! # Radix Tree Router
21//! 
22//! This module implements a Radix Tree (also known as a compact prefix tree) for
23//! efficient route matching with O(k) time complexity, where k is the path length.
24//! 
25//! ## Key Features
26//! 
27//! - **O(k) Lookup**: Route matching is independent of the number of registered routes
28//! - **Dynamic Parameters**: Supports `:param` for named path segments
29//! - **Wildcard Routes**: Supports `*path` for catching all remaining path segments
30//! - **Method-based Routing**: Each HTTP method has its own radix tree
31//! - **Thread-safe**: Uses RwLock for safe concurrent access
32//! 
33//! ## Path Pattern Syntax
34//! 
35//! - `/users/:id` - Matches `/users/123`, extracts `id = "123"`
36//! - `/files/*path` - Matches `/files/docs/readme.txt`, extracts `path = "docs/readme.txt"`
37//! - `/api/v1/users` - Exact match for static paths
38//! 
39//! ## Algorithm
40//! 
41//! The radix tree stores path segments as nodes, sharing common prefixes:
42//! 
43//! ```text
44//! Root
45//! ├── api
46//! │   └── v1
47//! │       ├── users (handler)
48//! │       └── users/
49//! │           └── :id (handler with param extraction)
50//! └── health (handler)
51//! ```
52
53use std::collections::HashMap;
54use std::sync::{Arc, RwLock};
55
56use super::routing::{RiRoute, RiRouteHandler};
57use crate::core::lock::RwLockExtensions;
58
59/// Represents the type of a path segment in the radix tree.
60#[derive(Debug, Clone, PartialEq)]
61pub enum SegmentType {
62    /// Static segment that must match exactly (e.g., "users", "api")
63    Static,
64    /// Named parameter segment (e.g., ":id", ":name")
65    Param,
66    /// Wildcard segment that matches all remaining path (e.g., "*path")
67    Wildcard,
68}
69
70/// Represents a single segment in a path pattern.
71#[derive(Debug, Clone)]
72pub struct PathSegment {
73    /// The segment value (e.g., "users", ":id", "*path")
74    pub value: String,
75    /// The type of this segment
76    pub segment_type: SegmentType,
77}
78
79impl PathSegment {
80    /// Creates a new path segment from a string.
81    /// 
82    /// # Parameters
83    /// 
84    /// - `value`: The segment string
85    /// 
86    /// # Returns
87    /// 
88    /// A new `PathSegment` with the appropriate type
89    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    /// Returns the parameter name for param or wildcard segments.
105    /// 
106    /// # Returns
107    /// 
108    /// - For param segments: the name without the `:` prefix
109    /// - For wildcard segments: the name without the `*` prefix
110    /// - For static segments: `None`
111    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/// Result of a route match operation.
121#[derive(Debug, Clone)]
122pub struct RouteMatch {
123    /// The matched route
124    pub route: RiRoute,
125    /// Extracted path parameters
126    pub params: HashMap<String, String>,
127}
128
129/// A node in the radix tree.
130/// 
131/// Each node represents a path segment and may contain:
132/// - A route handler (if this node is a terminal)
133/// - Child nodes for further path segments
134/// - Parameter and wildcard children for dynamic matching
135#[derive(Clone)]
136pub struct RadixNode {
137    /// The path segment this node represents
138    pub segment: PathSegment,
139    /// Route handler if this node is a terminal
140    pub handler: Option<RiRouteHandler>,
141    /// Route data if this node is a terminal
142    pub route_data: Option<RiRoute>,
143    /// Static children (exact match segments)
144    pub children: HashMap<String, Arc<RwLock<RadixNode>>>,
145    /// Parameter child (for `:param` segments)
146    pub param_child: Option<Arc<RwLock<RadixNode>>>,
147    /// Wildcard child (for `*path` segments)
148    pub wildcard_child: Option<Arc<RwLock<RadixNode>>>,
149}
150
151impl RadixNode {
152    /// Creates a new radix tree node with the given segment.
153    /// 
154    /// # Parameters
155    /// 
156    /// - `segment`: The path segment for this node
157    /// 
158    /// # Returns
159    /// 
160    /// A new `RadixNode` instance
161    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    /// Creates a root node (empty segment).
173    /// 
174    /// # Returns
175    /// 
176    /// A new `RadixNode` representing the tree root
177    pub fn root() -> Self {
178        Self::new(PathSegment {
179            value: String::new(),
180            segment_type: SegmentType::Static,
181        })
182    }
183
184    /// Checks if this node is a terminal (has a handler).
185    /// 
186    /// # Returns
187    /// 
188    /// `true` if this node has a handler, `false` otherwise
189    pub fn is_terminal(&self) -> bool {
190        self.handler.is_some()
191    }
192
193    /// Adds a child node for a static segment.
194    /// 
195    /// # Parameters
196    /// 
197    /// - `segment`: The segment for the child
198    /// 
199    /// # Returns
200    /// 
201    /// An `Arc<RwLock<RadixNode>>` to the child node
202    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    /// Gets or creates a child node for the given segment.
210    /// 
211    /// # Parameters
212    /// 
213    /// - `segment`: The segment to find or create
214    /// 
215    /// # Returns
216    /// 
217    /// An `Arc<RwLock<RadixNode>>` to the child node
218    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
249/// A Radix Tree for efficient route matching.
250/// 
251/// This implementation provides O(k) lookup time where k is the path length,
252/// making it significantly faster than linear search for large numbers of routes.
253pub struct RiRadixTree {
254    /// Root node of the tree
255    root: Arc<RwLock<RadixNode>>,
256    /// Number of routes in the tree
257    route_count: RwLock<usize>,
258}
259
260impl Default for RiRadixTree {
261    fn default() -> Self {
262        Self::new()
263    }
264}
265
266impl RiRadixTree {
267    /// Creates a new empty radix tree.
268    /// 
269    /// # Returns
270    /// 
271    /// A new `RiRadixTree` instance
272    pub fn new() -> Self {
273        Self {
274            root: Arc::new(RwLock::new(RadixNode::root())),
275            route_count: RwLock::new(0),
276        }
277    }
278
279    /// Parses a path into segments.
280    /// 
281    /// # Parameters
282    /// 
283    /// - `path`: The path to parse
284    /// 
285    /// # Returns
286    /// 
287    /// A vector of `PathSegment` instances
288    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    /// Inserts a route into the tree.
296    /// 
297    /// # Parameters
298    /// 
299    /// - `route`: The route to insert
300    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    /// Finds a matching route for the given path.
390    /// 
391    /// # Parameters
392    /// 
393    /// - `path`: The path to match
394    /// 
395    /// # Returns
396    /// 
397    /// An `Option<RouteMatch>` containing the matched route and extracted parameters
398    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    /// Recursively searches for a matching route.
418    /// 
419    /// # Parameters
420    /// 
421    /// - `node`: Current node being examined
422    /// - `segments`: Path segments to match
423    /// - `index`: Current segment index
424    /// - `params`: Accumulated path parameters
425    /// 
426    /// # Returns
427    /// 
428    /// An `Option<RouteMatch>` if a match is found
429    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    /// Removes a route from the tree.
506    /// 
507    /// # Parameters
508    /// 
509    /// - `path`: The path of the route to remove
510    /// 
511    /// # Returns
512    /// 
513    /// `true` if the route was removed, `false` if it wasn't found
514    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    /// Recursively removes a route from the tree.
547    /// 
548    /// # Parameters
549    /// 
550    /// - `node`: Current node
551    /// - `segments`: Path segments
552    /// - `index`: Current segment index
553    /// 
554    /// # Returns
555    /// 
556    /// `true` if the route was removed
557    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    /// Returns the number of routes in the tree.
625    /// 
626    /// # Returns
627    /// 
628    /// The number of routes
629    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    /// Clears all routes from the tree.
637    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}