diff --git a/tokenizer/lattice/lattice.go b/tokenizer/lattice/lattice.go index a7fb52b..aab4eea 100644 --- a/tokenizer/lattice/lattice.go +++ b/tokenizer/lattice/lattice.go @@ -4,11 +4,11 @@ import ( "fmt" "io" "strings" - "sync" "unicode" "unicode/utf8" "github.com/ikawaha/kagome-dict/dict" + "github.com/ikawaha/kagome/v2/tokenizer/lattice/mem" ) const ( @@ -32,11 +32,9 @@ const ( Extended ) -var latticePool = sync.Pool{ - New: func() interface{} { - return new(Lattice) - }, -} +var latticePool = mem.NewPool[Lattice](func() *Lattice { + return new(Lattice) +}) // Lattice represents a grid of morph nodes. type Lattice struct { @@ -49,7 +47,7 @@ type Lattice struct { // New returns a new lattice. func New(d *dict.Dict, u *dict.UserDict) *Lattice { - la := latticePool.Get().(*Lattice) + la := latticePool.Get() la.dic = d la.udic = u return la @@ -86,7 +84,7 @@ func (la *Lattice) addNode(pos, id, position, start int, class NodeClass, surfac case USER: // use default cost } - n := nodePool.Get().(*Node) + n := nodePool.Get() n.ID = id n.Position = position n.Start = start @@ -309,6 +307,7 @@ func posFeature(d *dict.Dict, u *dict.UserDict, t *Node) string { } // Dot outputs a lattice in the graphviz dot format. +// //nolint:gocyclo func (la *Lattice) Dot(w io.Writer) { bests := make(map[*Node]struct{}) diff --git a/tokenizer/lattice/mem/doc.go b/tokenizer/lattice/mem/doc.go new file mode 100644 index 0000000..2213700 --- /dev/null +++ b/tokenizer/lattice/mem/doc.go @@ -0,0 +1,2 @@ +// Package mem implements the memory utility. +package mem diff --git a/tokenizer/lattice/mem/pool.go b/tokenizer/lattice/mem/pool.go new file mode 100644 index 0000000..1366ed0 --- /dev/null +++ b/tokenizer/lattice/mem/pool.go @@ -0,0 +1,31 @@ +package mem + +import ( + "sync" +) + +// Pool represents memory pool of T. +type Pool[T any] struct { + internal *sync.Pool +} + +// NewPool returns a memory pool of T. +func NewPool[T any](f func() *T) Pool[T] { + return Pool[T]{ + internal: &sync.Pool{ + New: func() any { + return f() + }, + }, + } +} + +// Get gets instance of T from the memory pool. +func (p Pool[T]) Get() *T { + return p.internal.Get().(*T) +} + +// Put puts the instance to the memory pool. +func (p Pool[T]) Put(x *T) { + p.internal.Put(x) +} diff --git a/tokenizer/lattice/node.go b/tokenizer/lattice/node.go index fbd5f29..00d5d07 100644 --- a/tokenizer/lattice/node.go +++ b/tokenizer/lattice/node.go @@ -1,6 +1,8 @@ package lattice -import "sync" +import ( + "github.com/ikawaha/kagome/v2/tokenizer/lattice/mem" +) // BosEosID represents Reserved identifier of node id. const BosEosID int = -1 @@ -45,8 +47,6 @@ type Node struct { prev *Node } -var nodePool = sync.Pool{ - New: func() interface{} { - return new(Node) - }, -} +var nodePool = mem.NewPool[Node](func() *Node { + return new(Node) +})