From 2ebb5b6b585d6adc9d7d714a7d7592dbe950bdb6 Mon Sep 17 00:00:00 2001
From: kevin <wanjunfeng@gmail.com>
Date: Wed, 19 Aug 2020 14:54:59 +0800
Subject: [PATCH] support customized mask char on trie

---
 core/stringx/trie.go      | 22 ++++++++++++++++++++--
 core/stringx/trie_test.go | 12 ++++++------
 2 files changed, 26 insertions(+), 8 deletions(-)

diff --git a/core/stringx/trie.go b/core/stringx/trie.go
index 3c0abee8..b2a2115a 100644
--- a/core/stringx/trie.go
+++ b/core/stringx/trie.go
@@ -2,7 +2,11 @@ package stringx
 
 import "github.com/tal-tech/go-zero/core/lang"
 
+const defaultMask = '*'
+
 type (
+	TrieOption func(trie *trieNode)
+
 	Trie interface {
 		Filter(text string) (string, []string, bool)
 		FindKeywords(text string) []string
@@ -10,6 +14,7 @@ type (
 
 	trieNode struct {
 		node
+		mask rune
 	}
 
 	scope struct {
@@ -18,8 +23,15 @@ type (
 	}
 )
 
-func NewTrie(words []string) Trie {
+func NewTrie(words []string, opts ...TrieOption) Trie {
 	n := new(trieNode)
+
+	for _, opt := range opts {
+		opt(n)
+	}
+	if n.mask == 0 {
+		n.mask = defaultMask
+	}
 	for _, word := range words {
 		n.add(word)
 	}
@@ -114,6 +126,12 @@ func (n *trieNode) findKeywordScopes(chars []rune) []scope {
 
 func (n *trieNode) replaceWithAsterisk(chars []rune, start, stop int) {
 	for i := start; i < stop; i++ {
-		chars[i] = '*'
+		chars[i] = n.mask
+	}
+}
+
+func WithMask(mask rune) TrieOption {
+	return func(n *trieNode) {
+		n.mask = mask
 	}
 }
diff --git a/core/stringx/trie_test.go b/core/stringx/trie_test.go
index ca13e914..e74a822c 100644
--- a/core/stringx/trie_test.go
+++ b/core/stringx/trie_test.go
@@ -109,25 +109,25 @@ func TestTrie(t *testing.T) {
 func TestTrieSingleWord(t *testing.T) {
 	trie := NewTrie([]string{
 		"闹",
-	})
+	}, WithMask('#'))
 	output, keywords, ok := trie.Filter("今晚真热闹")
 	assert.ElementsMatch(t, []string{"闹"}, keywords)
 	assert.True(t, ok)
-	assert.Equal(t, "今晚真热*", output)
+	assert.Equal(t, "今晚真热#", output)
 }
 
 func TestTrieOverlap(t *testing.T) {
 	trie := NewTrie([]string{
 		"一二三四五",
 		"二三四五六七八",
-	})
+	}, WithMask('#'))
 	output, keywords, ok := trie.Filter("零一二三四五六七八九十")
 	assert.ElementsMatch(t, []string{
 		"一二三四五",
 		"二三四五六七八",
 	}, keywords)
 	assert.True(t, ok)
-	assert.Equal(t, "零********九十", output)
+	assert.Equal(t, "零########九十", output)
 }
 
 func TestTrieNested(t *testing.T) {
@@ -135,7 +135,7 @@ func TestTrieNested(t *testing.T) {
 		"一二三",
 		"一二三四五",
 		"一二三四五六七八",
-	})
+	}, WithMask('#'))
 	output, keywords, ok := trie.Filter("零一二三四五六七八九十")
 	assert.ElementsMatch(t, []string{
 		"一二三",
@@ -143,7 +143,7 @@ func TestTrieNested(t *testing.T) {
 		"一二三四五六七八",
 	}, keywords)
 	assert.True(t, ok)
-	assert.Equal(t, "零********九十", output)
+	assert.Equal(t, "零########九十", output)
 }
 
 func BenchmarkTrie(b *testing.B) {