Merge pull request #7 from lxzan/dev

fix tree map

Author: lxzan

internal/treemap/treemap.go

@@ -5,64 +5,77 @@ import (
 )
 
 func New[T any](sep string) *TreeMap[T] {
-	return &TreeMap[T]{sep: sep, children: make(map[string]*TreeMap[T])}
+	return &TreeMap[T]{
+		sep:  sep,
+		root: &Element[T]{children: make(map[string]*Element[T])},
+	}
 }
 
-type TreeMap[T any] struct {
-	children map[string]*TreeMap[T]
-	sep      string
-	key      string
-	value    T
-}
+type (
+	TreeMap[T any] struct {
+		sep  string
+		root *Element[T]
+	}
 
-func (c *TreeMap[T]) Put(key string, val T) {
-	if !strings.Contains(key, c.sep) {
-		c.children[key] = &TreeMap[T]{key: key, value: val}
-	} else {
-		c.doPut(c, &TreeMap[T]{key: key, value: val}, key)
+	Element[T any] struct {
+		ok       bool
+		children map[string]*Element[T]
+		key      string
+		value    T
 	}
+)
+
+func (c *TreeMap[T]) Put(key string, val T) {
+	node := &Element[T]{key: key, value: val, ok: true}
+	c.doPut(c.root, node, key)
 }
 
-func (c *TreeMap[T]) doPut(far, son *TreeMap[T], key string) {
+func (c *TreeMap[T]) doPut(far, son *Element[T], key string) {
 	if far.children == nil {
-		far.children = make(map[string]*TreeMap[T])
+		far.children = make(map[string]*Element[T])
 	}
 
 	index := strings.Index(key, c.sep)
 	if index < 0 {
-		far.children[key] = son
+		if node := far.children[key]; node == nil {
+			far.children[key] = son
+		} else {
+			node.value, node.ok = son.value, son.ok
+		}
 		return
 	}
 
 	s0, s1 := key[:index], key[index+1:]
 	next := far.children[s0]
 	if far.children[s0] == nil {
-		next = &TreeMap[T]{}
+		next = &Element[T]{}
 		far.children[s0] = next
 	}
 	c.doPut(next, son, s1)
 }
 
 func (c *TreeMap[T]) Get(key string) (v T, exist bool) {
-	if result, ok := c.children[key]; ok {
+	if result, ok := c.root.children[key]; ok && result.ok {
 		v, exist = result.value, true
 	}
 	return
 }
 
 func (c *TreeMap[T]) Match(key string, cb func(k string, v T)) {
-	c.doMatch(c, key, cb)
+	c.doMatch(c.root, key, cb)
 }
 
-func (c *TreeMap[T]) doMatch(cur *TreeMap[T], key string, cb func(string, T)) {
+func (c *TreeMap[T]) doMatch(cur *Element[T], key string, cb func(string, T)) {
 	index := strings.Index(key, c.sep)
 	if index < 0 {
 		if key == "*" {
 			for _, node := range cur.children {
-				cb(node.key, node.value)
+				if node.ok {
+					cb(node.key, node.value)
+				}
 			}
 		} else {
-			if result, ok := cur.children[key]; ok {
+			if result, ok := cur.children[key]; ok && result.ok {
 				cb(result.key, result.value)
 			}
 		}

internal/treemap/treemap_test.go

@@ -5,6 +5,7 @@ import (
 	"github.com/lxzan/event_emitter/internal/helper"
 	"github.com/stretchr/testify/assert"
 	"strings"
+	"sync"
 	"testing"
 )
 
@@ -81,3 +82,26 @@ func TestTreeMap_Get(t *testing.T) {
 	_, ok := tm.Get("~aha")
 	assert.False(t, ok)
 }
+
+func TestTreeMap_Match(t *testing.T) {
+	var tm = New[uint8](".")
+	tm.Put("api.v1.actions.eat", 1)
+	tm.Put("api.v1.actions", 2)
+	var wg = new(sync.WaitGroup)
+	wg.Add(2)
+	var sum = uint8(0)
+	tm.Match("api.v1.*.eat", func(k string, v uint8) {
+		sum += v
+		wg.Done()
+	})
+	tm.Match("api.v1.*", func(k string, v uint8) {
+		sum += v
+		wg.Done()
+	})
+	tm.Match("api.v1.any", func(k string, v uint8) {
+		sum += v
+		wg.Done()
+	})
+	wg.Wait()
+	assert.Equal(t, sum, uint8(3))
+}