A ternary search tree that handles wildcards

wildcard_tst.rs

#![allow(unused_imports)]
#![allow(unused_variables)]
#![allow(dead_code)]
#![allow(unused_mut)]

use std::cmp::{Ord, Ordering};
use std::collections::HashMap;
use std::convert::TryFrom;
use std::fmt;

fn main() {}

#[cfg(test)]
mod tests {
    use crate::{insert, longest_prefix, TstNode};

    fn new_node() -> TstNode<&'static str> {
        let mut root = TstNode {
            el: "zxc".into(),
            handler: Some("zxc"),
            children: Default::default(),
        };
        insert(
            &mut root,
            vec!["zxc".into(), "qwe".into(), "qwe".into()].iter(),
            "qwe2",
        );

        insert(&mut root, vec!["zxc".into(), "asd".into()].iter(), "asd");

        insert(&mut root, vec!["zxc".into(), "*".into()].iter(), "*");

        insert(
            &mut root,
            vec!["zxc".into(), ":asd".into(), "123".into()].iter(),
            ":",
        );

        root
    }

    #[test]
    fn test_1() {
        let mut root = new_node();

        println!("root: {:?}\n", root);
        println!("root.c[0]: {:?}\n", root.children[0]);
        println!("root.c[1]: {:?}\n", root.children[1]);

        let l = longest_prefix(&root, vec!["zxc".into(), "qwe".into(), "qwe".into()].iter());

        println!("got: {:?}", l);
        assert_eq!(l, Some("qwe2"));
    }

    #[test]
    fn test_2() {
        let mut root = new_node();

        let l = longest_prefix(&root, vec!["zxc".into(), "qwe".into()].iter());

        println!("got: {:?}", l);
        assert_eq!(l, Some("*"));
    }

    #[test]
    fn test_3() {
        let mut root = new_node();

        let l = longest_prefix(&root, vec!["zxc".into(), "asd".into()].iter());

        println!("got: {:?}", l);
        assert_eq!(l, Some("asd"));
    }

    #[test]
    fn test_4() {
        let mut root = new_node();

        let l = longest_prefix(&root, vec!["zxc".into()].iter());

        println!("got: {:?}", l);
        assert_eq!(l, Some("zxc"));
    }

    #[test]
    fn test_5() {
        let mut root = new_node();

        let l = longest_prefix(&root, vec!["zxc".into(), "123".into()].iter());

        println!("got: {:?}", l);
        assert_eq!(l, Some("*"));
    }

    #[test]
    fn test_51() {
        let mut root = new_node();

        let l = longest_prefix(&root, vec!["zxc".into(), "123".into(), "123".into()].iter());

        println!("got: {:?}", l);
        assert_eq!(l, Some(":"));
    }

    #[test]
    fn test_6() {
        let mut root = new_node();

        let l = longest_prefix(&root, vec!["098".into(), "asd".into(), "123".into()].iter());

        println!("got: {:?}", l);
        assert_eq!(l, None);
    }
}

#[derive(Eq, PartialEq, PartialOrd, Ord, Clone, Debug)]
pub enum PathElement {
    Dir(String),
    Param(String),
    Wildcard,
}

impl From<&str> for PathElement {
    fn from(s: &str) -> PathElement {
        assert!(!s.is_empty() && !s.contains("/"));

        if s.starts_with(":") {
            if s.len() < 2 {
                panic!("param must have name")
            }

            return PathElement::Param(s.get(1..).unwrap().to_string());
        }

        if s.starts_with("*") {
            if s.len() != 1 {
                panic!("wildcard must be by itself")
            }
            return PathElement::Wildcard;
        }

        PathElement::Dir(s.to_string())
    }
}

type Req = (String,);
type Res = ();

trait Handler<X> {
    fn handle(&self, ctx: X, req: Req, res: Res);
}

#[derive(Default)]
struct TstCtx {
    params: HashMap<String, String>,
    path: Vec<String>,
}

#[derive(Debug)]
struct Tst<V> {
    root: TstNode<V>,
}

impl<V> Tst<V> {
    fn new() -> Self {
        Tst {
            root: TstNode {
                el: PathElement::Wildcard,
                handler: None,
                children: Default::default(),
            },
        }
    }
}

struct TstNode<V> {
    el: PathElement,
    handler: Option<V>,
    children: [Option<Box<TstNode<V>>>; 3],
}

impl<V> fmt::Debug for TstNode<V> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.write_fmt(format_args!(
            "{:?}\n{:?}\t{:?}\t{:?}",
            self.el,
            self.lt().as_ref().map(|n| &n.el),
            self.eq().as_ref().map(|n| &n.el),
            self.gt().as_ref().map(|n| &n.el),
        ))
    }
}

impl<V> TstNode<V> {
    fn lt(&self) -> &Option<Box<TstNode<V>>> {
        &self.children[0]
    }
    fn eq(&self) -> &Option<Box<TstNode<V>>> {
        &self.children[1]
    }
    fn gt(&self) -> &Option<Box<TstNode<V>>> {
        &self.children[2]
    }
    fn lt_mut(&mut self) -> &mut Option<Box<TstNode<V>>> {
        &mut self.children[0]
    }
    fn eq_mut(&mut self) -> &mut Option<Box<TstNode<V>>> {
        &mut self.children[1]
    }
    fn gt_mut(&mut self) -> &mut Option<Box<TstNode<V>>> {
        &mut self.children[2]
    }
}

fn insert<'a, V, I>(mut node: &mut TstNode<V>, mut it: I, val: V)
where
    I: Iterator<Item = &'a PathElement>,
{
    // rust is a pita don't mess with this function

    let mut el = it.next().unwrap();

    let mut next: &mut Option<Box<TstNode<V>>>;
    loop {
        println!("path el {:?}", el);
        let cmp = el.cmp(&node.el);

        next = match cmp {
            Ordering::Less => node.lt_mut(),
            Ordering::Greater => node.gt_mut(),
            Ordering::Equal => {
                if let Some(ell) = it.next() {
                    el = ell;
                    node.eq_mut()
                } else {
                    node.handler = Some(val);
                    return;
                }
            }
        };

        if next.is_none() {
            break;
        }

        node = next.as_mut().unwrap().as_mut();
    }

    // only get here if we got to an empty spot
    node = next.get_or_insert(Box::new(TstNode {
        el: el.clone(),
        handler: None,
        children: Default::default(),
    }));

    while let Some(el) = it.next() {
        next = node.eq_mut();

        node = next.get_or_insert(Box::new(TstNode {
            el: el.clone(),
            handler: None,
            children: Default::default(),
        }));
    }

    node.handler = Some(val);
}

fn longest_prefix_rec<'a, V: Copy, I>(
    mut node: &TstNode<V>,
    mut it: I,
    cur: PathElement,
) -> Option<V>
where
    I: Iterator<Item = &'a String> + Clone,
{
    println!("looking for {:?}", cur);
    println!("in {:?}", node);

    let mut cmp = cur.cmp(&node.el);
    match cmp {
        Ordering::Less | Ordering::Greater => {
            println!("{:?}", cmp);
            
            match cmp {
                Ordering::Less => node.lt(),
                Ordering::Greater => node.gt(),
                _ => panic!("unreachable"),
            }.as_ref()
                .and_then(|next| longest_prefix_rec(next, it.clone(), cur.clone()))
                .or_else(|| {
                    // am i param?
                    if let PathElement::Param(_) = node.el {
                        Some(node)
                    } else {
                        None
                    }
                    .and_then(|node| {
                        let cur_next = it.next()?;
                        longest_prefix_rec(node.eq().as_ref()?, it.clone(),
                            PathElement::Dir(cur_next.to_string()))
                    })
                })
                .or_else(|| {
                    if let PathElement::Wildcard = node.el {
                        node.handler
                    } else {
                        None
                    }
                })
        }
        Ordering::Equal => {
            println!("equal");

            node.eq()
                .as_ref()
                .and_then(|next| {
                    let curv = PathElement::Dir(it.next()?.to_string());
                    
                    let ret =
                        longest_prefix_rec(
                            next,
                            it.clone(),
                            curv,
                        ).or_else(|| {
                            // am i param?
                            if let PathElement::Param(_) = next.el {
                                Some(next)
                            } else {
                                None
                            }.as_ref()
                            .and_then(|next| {
                                let cur_next = it.next()?;
                                longest_prefix_rec(node.eq().as_ref()?, it.clone(),
                                    PathElement::Dir(cur_next.to_string()))
                            })
                        
                        })
                        .or_else(|| {
                    
                            if let PathElement::Wildcard = node.el {
                                node.handler
                            } else {
                                None
                            }
                        });

                    println!("recursive res {:?}", ret.is_some());
                    ret
                })
                .or_else(|| {
                    println!("returning {:?}", node.el);
                    node.handler
                })
        }
    }.or_else(|| {
        match node.el {
            PathElement::Dir(_) => {
                // traverse tree looking for wildcard elements
                println!("looking for wildcards in {:?}", node);
                
                let mut next = node.gt();
                while next.is_some() &&
                    PathElement::Param("".to_string()) > next.as_ref().unwrap().el
                {
                    next = next.as_ref().unwrap().gt()
                }
                
                if let Some(nextv) = next {
                    println!("found wildcard {:?}", nextv);
                    longest_prefix_rec(nextv,
                        it.clone(),
                        cur)
                        
                } else {
                    None
                }
            },
            _ => None,
        }
    })    
}

fn longest_prefix<'a, V: Copy, I>(mut node: &TstNode<V>, mut it: I) -> Option<V>
where
    I: Iterator<Item = &'a String> + Clone,
{
    let c = PathElement::Dir(it.next()?.to_string());
    longest_prefix_rec(node, it, c)
}