beego:router tree

+package beego
+
+import (
+	"path"
+	"regexp"
+	"strings"
+
+	"github.com/astaxie/beego/utils"
+)
+
+type Tree struct {
+	//search fix route first
+	fixrouters map[string]*Tree
+
+	//if set, failure to match fixrouters search then search wildcard
+	wildcard *Tree
+
+	//if set, failure to match wildcard search
+	leaf *leafInfo
+}
+
+func NewTree() *Tree {
+	return &Tree{
+		fixrouters: make(map[string]*Tree),
+	}
+}
+
+// call addseg function
+func (t *Tree) AddRouter(pattern string, runObject interface{}) {
+	t.addseg(splitPath(pattern), runObject, nil, "")
+}
+
+// "/"
+// "admin" ->
+func (t *Tree) addseg(segments []string, route interface{}, wildcards []string, reg string) {
+	if len(segments) == 0 {
+		if reg != "" {
+			filterCards := []string{}
+			for _, v := range wildcards {
+				if v == ":" || v == "." {
+					continue
+				}
+				filterCards = append(filterCards, v)
+			}
+			t.leaf = &leafInfo{runObject: route, wildcards: wildcards, regexps: regexp.MustCompile("^" + reg + "$")}
+		} else {
+			t.leaf = &leafInfo{runObject: route, wildcards: wildcards}
+		}
+
+	} else {
+		seg := segments[0]
+		iswild, params, regexpStr := splitSegment(seg)
+		if iswild {
+			if t.wildcard == nil {
+				t.wildcard = NewTree()
+			}
+			t.wildcard.addseg(segments[1:], route, append(wildcards, params...), reg+regexpStr)
+		} else {
+			subTree, ok := t.fixrouters[seg]
+			if !ok {
+				subTree = NewTree()
+				t.fixrouters[seg] = subTree
+			}
+			subTree.addseg(segments[1:], route, wildcards, reg)
+		}
+	}
+}
+
+// match router to runObject & params
+func (t *Tree) Match(pattern string) (runObject interface{}, params map[string]string) {
+	if len(pattern) == 0 || pattern[0] != '/' {
+		return nil, nil
+	}
+
+	return t.match(splitPath(pattern), nil)
+}
+
+func (t *Tree) match(segments []string, wildcardValues []string) (runObject interface{}, params map[string]string) {
+	// Handle leaf nodes:
+	if len(segments) == 0 {
+		if t.leaf != nil {
+			if ok, pa := t.leaf.match(wildcardValues); ok {
+				return t.leaf.runObject, pa
+			}
+		}
+		return nil, nil
+	}
+
+	var seg string
+	seg, segments = segments[0], segments[1:]
+
+	subTree, ok := t.fixrouters[seg]
+	if ok {
+		runObject, params = subTree.match(segments, wildcardValues)
+	}
+	if runObject == nil && t.wildcard != nil {
+		runObject, params = t.wildcard.match(segments, append(wildcardValues, seg))
+	}
+	if runObject == nil {
+		if t.leaf != nil {
+			if ok, pa := t.leaf.match(append(wildcardValues, seg)); ok {
+				return t.leaf.runObject, pa
+			}
+		}
+	}
+
+	return runObject, params
+}
+
+type leafInfo struct {
+	// names of wildcards that lead to this leaf. eg, ["id" "name"] for the wildcard ":id" and ":name"
+	wildcards []string
+
+	// if the leaf is regexp
+	regexps *regexp.Regexp
+
+	runObject interface{}
+}
+
+func (leaf *leafInfo) match(wildcardValues []string) (ok bool, params map[string]string) {
+	if leaf.regexps == nil {
+		// has error
+		if len(wildcardValues) == 0 && len(leaf.wildcards) > 0 {
+			if utils.InSlice(":", leaf.wildcards) {
+				return true, nil
+			}
+			Error("bug of router")
+			return false, nil
+		} else if len(wildcardValues) == 0 { // static path
+			return true, nil
+		}
+		// match *
+		if len(leaf.wildcards) == 1 && leaf.wildcards[0] == ":splat" {
+			params = make(map[string]string)
+			params[":splat"] = path.Join(wildcardValues...)
+			return true, params
+		}
+		// match *.*
+		if len(leaf.wildcards) == 3 && leaf.wildcards[0] == "." {
+			params = make(map[string]string)
+			lastone := wildcardValues[len(wildcardValues)-1]
+			strs := strings.SplitN(lastone, ".", 2)
+			if len(strs) == 2 {
+				params[":ext"] = strs[1]
+			} else {
+				params[":ext"] = ""
+			}
+			params[":path"] = path.Join(wildcardValues[:len(wildcardValues)-1]...) + "/" + strs[0]
+			return true, params
+		}
+		// match :id
+		params = make(map[string]string)
+		j := 0
+		for _, v := range leaf.wildcards {
+			if v == ":" {
+				continue
+			}
+			params[v] = wildcardValues[j]
+			j += 1
+		}
+		if len(params) != len(wildcardValues) {
+			Error("bug of router")
+			return false, nil
+		}
+		return true, params
+	}
+
+	if !leaf.regexps.MatchString(strings.Join(wildcardValues, "")) {
+		return false, nil
+	}
+	params = make(map[string]string)
+	matches := leaf.regexps.FindStringSubmatch(strings.Join(wildcardValues, ""))
+	for i, match := range matches[1:] {
+		params[leaf.wildcards[i]] = match
+	}
+	return true, params
+}
+
+// "/" -> []
+// "/admin" -> ["admin"]
+// "/admin/" -> ["admin"]
+// "/admin/users" -> ["admin", "users"]
+func splitPath(key string) []string {
+	elements := strings.Split(key, "/")
+	if elements[0] == "" {
+		elements = elements[1:]
+	}
+	if elements[len(elements)-1] == "" {
+		elements = elements[:len(elements)-1]
+	}
+	return elements
+}
+
+// "admin" -> false, nil, ""
+// ":id" -> true, [:id], ""
+// "?:id" -> true, [: id], ""        : meaning can empty
+// ":id:int" -> true, [:id], ([0-9]+)
+// ":name:string" -> true, [:name], ([\w]+)
+// ":id([0-9]+)" -> true, [:id], ([0-9]+)
+// ":id([0-9]+)_:name" -> true, [:id :name], ([0-9]+)_(.+)
+// "cms_:id_:page.html" -> true, [:id :page], cms_(.+)_(.+).html
+// "*" -> true, [:splat], ""
+// "*.*" -> true,[. :path :ext], ""      . meaning separator
+func splitSegment(key string) (bool, []string, string) {
+	if strings.HasPrefix(key, "*") {
+		if key == "*.*" {
+			return true, []string{".", ":path", ":ext"}, ""
+		} else {
+			return true, []string{":splat"}, ""
+		}
+	}
+	if strings.ContainsAny(key, ":") {
+		var paramsNum int
+		var out []rune
+		var start bool
+		var startexp bool
+		var param []rune
+		var expt []rune
+		var skipnum int
+		params := []string{}
+		reg := regexp.MustCompile(`[a-zA-Z0-9]+`)
+		for i, v := range key {
+			if skipnum > 0 {
+				skipnum -= 1
+				continue
+			}
+			if start {
+				//:id:int and :name:string
+				if v == ':' {
+					if len(key) >= i+4 {
+						if key[i+1:i+4] == "int" {
+							out = append(out, []rune("([0-9]+)")...)
+							params = append(params, ":"+string(param))
+							start = false
+							startexp = false
+							skipnum = 3
+							param = make([]rune, 0)
+							paramsNum += 1
+							continue
+						}
+					}
+					if len(key) >= i+7 {
+						if key[i+1:i+7] == "string" {
+							out = append(out, []rune(`([\w]+)`)...)
+							params = append(params, ":"+string(param))
+							paramsNum += 1
+							start = false
+							startexp = false
+							skipnum = 6
+							param = make([]rune, 0)
+							continue
+						}
+					}
+				}
+				// params only support a-zA-Z0-9
+				if reg.MatchString(string(v)) {
+					param = append(param, v)
+					continue
+				}
+				if v != '(' {
+					out = append(out, []rune(`(.+)`)...)
+					params = append(params, ":"+string(param))
+					param = make([]rune, 0)
+					paramsNum += 1
+					start = false
+					startexp = false
+				}
+			}
+			if startexp {
+				if v != ')' {
+					expt = append(expt, v)
+					continue
+				}
+			}
+			if v == ':' {
+				param = make([]rune, 0)
+				start = true
+			} else if v == '(' {
+				startexp = true
+				start = false
+				params = append(params, ":"+string(param))
+				paramsNum += 1
+				expt = make([]rune, 0)
+				expt = append(expt, '(')
+			} else if v == ')' {
+				startexp = false
+				expt = append(expt, ')')
+				out = append(out, expt...)
+				param = make([]rune, 0)
+			} else if v == '?' {
+				params = append(params, ":")
+			} else {
+				out = append(out, v)
+			}
+		}
+		if len(param) > 0 {
+			if paramsNum > 0 {
+				out = append(out, []rune(`(.+)`)...)
+			}
+			params = append(params, ":"+string(param))
+		}
+		return true, params, string(out)
+	} else {
+		return false, nil, ""
+	}
+}

+package beego
+
+import "testing"
+
+type testinfo struct {
+	url        string
+	requesturl string
+	params     map[string]string
+}
+
+var routers []testinfo
+
+func init() {
+	routers = make([]testinfo, 0)
+	routers = append(routers, testinfo{"/:id", "/123", map[string]string{":id": "123"}})
+	routers = append(routers, testinfo{"/", "/", nil})
+	routers = append(routers, testinfo{"/customer/login", "/customer/login", nil})
+	routers = append(routers, testinfo{"/*", "/customer/123", map[string]string{":splat": "customer/123"}})
+	routers = append(routers, testinfo{"/*.*", "/nice/api.json", map[string]string{":path": "nice/api", ":ext": "json"}})
+	routers = append(routers, testinfo{"/v1/shop/:id:int", "/v1/shop/123", map[string]string{":id": "123"}})
+	routers = append(routers, testinfo{"/v1/shop/:id/:name", "/v1/shop/123/nike", map[string]string{":id": "123", ":name": "nike"}})
+	routers = append(routers, testinfo{"/v1/shop/:id/account", "/v1/shop/123/account", map[string]string{":id": "123"}})
+	routers = append(routers, testinfo{"/v1/shop/:name:string", "/v1/shop/nike", map[string]string{":name": "nike"}})
+	routers = append(routers, testinfo{"/v1/shop/:id([0-9]+)", "/v1/shop//123", map[string]string{":id": "123"}})
+	routers = append(routers, testinfo{"/v1/shop/:id([0-9]+)_:name", "/v1/shop/123_nike", map[string]string{":id": "123", ":name": "nike"}})
+	routers = append(routers, testinfo{"/v1/shop/:id_cms.html", "/v1/shop/123_cms.html", map[string]string{":id": "123"}})
+	routers = append(routers, testinfo{"/v1/shop/cms_:id_:page.html", "/v1/shop/cms_123_1.html", map[string]string{":id": "123", ":page": "1"}})
+}
+
+func TestTreeRouters(t *testing.T) {
+	for _, r := range routers {
+		tr := NewTree()
+		tr.AddRouter(r.url, "astaxie")
+		obj, param := tr.Match(r.requesturl)
+		if obj == nil || obj.(string) != "astaxie" {
+			t.Fatal(r.url + " can't get obj ")
+		}
+		if r.params != nil {
+			for k, v := range r.params {
+				if vv, ok := param[k]; !ok {
+					t.Fatal(r.url + r.requesturl + " get param empty:" + k)
+				} else if vv != v {
+					t.Fatal(r.url + " " + r.requesturl + " should be:" + v + " get param:" + vv)
+				}
+			}
+		}
+	}
+}
+
+func TestSplitPath(t *testing.T) {
+	a := splitPath("/")
+	if len(a) != 0 {
+		t.Fatal("/ should retrun []")
+	}
+	a = splitPath("/admin")
+	if len(a) != 1 || a[0] != "admin" {
+		t.Fatal("/admin should retrun [admin]")
+	}
+	a = splitPath("/admin/")
+	if len(a) != 1 || a[0] != "admin" {
+		t.Fatal("/admin/ should retrun [admin]")
+	}
+	a = splitPath("/admin/users")
+	if len(a) != 2 || a[0] != "admin" || a[1] != "users" {
+		t.Fatal("/admin should retrun [admin users]")
+	}
+	a = splitPath("/admin/:id:int")
+	if len(a) != 2 || a[0] != "admin" || a[1] != ":id:int" {
+		t.Fatal("/admin should retrun [admin :id:int]")
+	}
+}
+
+func TestSplitSegment(t *testing.T) {
+	b, w, r := splitSegment("admin")
+	if b || len(w) != 0 || r != "" {
+		t.Fatal("admin should return false, nil, ''")
+	}
+	b, w, r = splitSegment("*")
+	if !b || len(w) != 1 || w[0] != ":splat" || r != "" {
+		t.Fatal("* should return true, [:splat], ''")
+	}
+	b, w, r = splitSegment("*.*")
+	if !b || len(w) != 3 || w[1] != ":path" || w[2] != ":ext" || w[0] != "." || r != "" {
+		t.Fatal("admin should return true,[. :path :ext], ''")
+	}
+	b, w, r = splitSegment(":id")
+	if !b || len(w) != 1 || w[0] != ":id" || r != "" {
+		t.Fatal(":id should return true, [:id], ''")
+	}
+	b, w, r = splitSegment("?:id")
+	if !b || len(w) != 2 || w[0] != ":" || w[1] != ":id" || r != "" {
+		t.Fatal("?:id should return true, [: :id], ''")
+	}
+	b, w, r = splitSegment(":id:int")
+	if !b || len(w) != 1 || w[0] != ":id" || r != "([0-9]+)" {
+		t.Fatal(":id:int should return true, [:id], '([0-9]+)'")
+	}
+	b, w, r = splitSegment(":name:string")
+	if !b || len(w) != 1 || w[0] != ":name" || r != `([\w]+)` {
+		t.Fatal(`:name:string should return true, [:name], '([\w]+)'`)
+	}
+	b, w, r = splitSegment(":id([0-9]+)")
+	if !b || len(w) != 1 || w[0] != ":id" || r != `([0-9]+)` {
+		t.Fatal(`:id([0-9]+) should return true, [:id], '([0-9]+)'`)
+	}
+	b, w, r = splitSegment(":id([0-9]+)_:name")
+	if !b || len(w) != 2 || w[0] != ":id" || w[1] != ":name" || r != `([0-9]+)_(.+)` {
+		t.Fatal(`:id([0-9]+)_:name should return true, [:id :name], '([0-9]+)_(.+)'`)
+	}
+	b, w, r = splitSegment(":id_cms.html")
+	if !b || len(w) != 1 || w[0] != ":id" || r != `(.+)_cms.html` {
+		t.Fatal(":id_cms.html should return true, [:id], '(.+)_cms.html'")
+	}
+	b, w, r = splitSegment("cms_:id_:page.html")
+	if !b || len(w) != 2 || w[0] != ":id" || w[1] != ":page" || r != `cms_(.+)_(.+).html` {
+		t.Fatal(":id_cms.html should return true, [:id :page], cms_(.+)_(.+).html")
+	}
+}