router_trie.go

Documentation: goji.io

     1  // +build !goji_router_simple
     2  
     3  package goji
     4  
     5  import (
     6  	"net/http"
     7  	"sort"
     8  	"strings"
     9  
    10  	"goji.io/internal"
    11  )
    12  
    13  type router struct {
    14  	routes   []route
    15  	methods  map[string]*trieNode
    16  	wildcard trieNode
    17  }
    18  
    19  type route struct {
    20  	Pattern
    21  	http.Handler
    22  }
    23  
    24  type child struct {
    25  	prefix string
    26  	node   *trieNode
    27  }
    28  
    29  type trieNode struct {
    30  	routes   []int
    31  	children []child
    32  }
    33  
    34  func (rt *router) add(p Pattern, h http.Handler) {
    35  	i := len(rt.routes)
    36  	rt.routes = append(rt.routes, route{p, h})
    37  
    38  	var prefix string
    39  	if pp, ok := p.(pathPrefix); ok {
    40  		prefix = pp.PathPrefix()
    41  	}
    42  
    43  	var methods map[string]struct{}
    44  	if hm, ok := p.(httpMethods); ok {
    45  		methods = hm.HTTPMethods()
    46  	}
    47  	if methods == nil {
    48  		rt.wildcard.add(prefix, i)
    49  		for _, sub := range rt.methods {
    50  			sub.add(prefix, i)
    51  		}
    52  	} else {
    53  		if rt.methods == nil {
    54  			rt.methods = make(map[string]*trieNode)
    55  		}
    56  
    57  		for method := range methods {
    58  			if _, ok := rt.methods[method]; !ok {
    59  				rt.methods[method] = rt.wildcard.clone()
    60  			}
    61  			rt.methods[method].add(prefix, i)
    62  		}
    63  	}
    64  }
    65  
    66  func (rt *router) route(r *http.Request) *http.Request {
    67  	tn := &rt.wildcard
    68  	if tn2, ok := rt.methods[r.Method]; ok {
    69  		tn = tn2
    70  	}
    71  
    72  	ctx := r.Context()
    73  	path := ctx.Value(internal.Path).(string)
    74  	for path != "" {
    75  		i := sort.Search(len(tn.children), func(i int) bool {
    76  			return path[0] <= tn.children[i].prefix[0]
    77  		})
    78  		if i == len(tn.children) || !strings.HasPrefix(path, tn.children[i].prefix) {
    79  			break
    80  		}
    81  
    82  		path = path[len(tn.children[i].prefix):]
    83  		tn = tn.children[i].node
    84  	}
    85  	for _, i := range tn.routes {
    86  		if r2 := rt.routes[i].Match(r); r2 != nil {
    87  			return r2.WithContext(&match{
    88  				Context: r2.Context(),
    89  				p:       rt.routes[i].Pattern,
    90  				h:       rt.routes[i].Handler,
    91  			})
    92  		}
    93  	}
    94  	return r.WithContext(&match{Context: ctx})
    95  }
    96  
    97  // We can be a teensy bit more efficient here: we're maintaining a sorted list,
    98  // so we know exactly where to insert the new element. But since that involves
    99  // more bookkeeping and makes the code messier, let's cross that bridge when we
   100  // come to it.
   101  type byPrefix []child
   102  
   103  func (b byPrefix) Len() int {
   104  	return len(b)
   105  }
   106  func (b byPrefix) Less(i, j int) bool {
   107  	return b[i].prefix < b[j].prefix
   108  }
   109  func (b byPrefix) Swap(i, j int) {
   110  	b[i], b[j] = b[j], b[i]
   111  }
   112  
   113  func longestPrefix(a, b string) string {
   114  	mlen := len(a)
   115  	if len(b) < mlen {
   116  		mlen = len(b)
   117  	}
   118  	for i := 0; i < mlen; i++ {
   119  		if a[i] != b[i] {
   120  			return a[:i]
   121  		}
   122  	}
   123  	return a[:mlen]
   124  }
   125  
   126  func (tn *trieNode) add(prefix string, idx int) {
   127  	if len(prefix) == 0 {
   128  		tn.routes = append(tn.routes, idx)
   129  		for i := range tn.children {
   130  			tn.children[i].node.add(prefix, idx)
   131  		}
   132  		return
   133  	}
   134  
   135  	ch := prefix[0]
   136  	i := sort.Search(len(tn.children), func(i int) bool {
   137  		return ch <= tn.children[i].prefix[0]
   138  	})
   139  
   140  	if i == len(tn.children) || ch != tn.children[i].prefix[0] {
   141  		routes := append([]int(nil), tn.routes...)
   142  		tn.children = append(tn.children, child{
   143  			prefix: prefix,
   144  			node:   &trieNode{routes: append(routes, idx)},
   145  		})
   146  	} else {
   147  		lp := longestPrefix(prefix, tn.children[i].prefix)
   148  
   149  		if tn.children[i].prefix == lp {
   150  			tn.children[i].node.add(prefix[len(lp):], idx)
   151  			return
   152  		}
   153  
   154  		split := new(trieNode)
   155  		split.children = []child{
   156  			{tn.children[i].prefix[len(lp):], tn.children[i].node},
   157  		}
   158  		split.routes = append([]int(nil), tn.routes...)
   159  		split.add(prefix[len(lp):], idx)
   160  
   161  		tn.children[i].prefix = lp
   162  		tn.children[i].node = split
   163  	}
   164  
   165  	sort.Sort(byPrefix(tn.children))
   166  }
   167  
   168  func (tn *trieNode) clone() *trieNode {
   169  	clone := new(trieNode)
   170  	clone.routes = append(clone.routes, tn.routes...)
   171  	clone.children = append(clone.children, tn.children...)
   172  	for i := range clone.children {
   173  		clone.children[i].node = tn.children[i].node.clone()
   174  	}
   175  	return clone
   176  }
   177
View as plain text