perf: reduce allocations in string.gsub and pattern matching

pm/cache_test.go

@@ -0,0 +1,21 @@
+package pm
+
+import (
+	"fmt"
+	"sync/atomic"
+	"testing"
+)
+
+func TestPatternCacheBounded(t *testing.T) {
+	var c Cache
+	for i := 0; i < maxCachedPatterns*3; i++ {
+		pat := fmt.Sprintf("x%d", i)
+		ms, err := c.Find(pat, []byte("a "+pat+" b"), 0, -1)
+		if err != nil || len(ms) != 1 {
+			t.Fatalf("pattern %q: matches=%d err=%v", pat, len(ms), err)
+		}
+	}
+	if n := atomic.LoadInt32(&c.n); n > maxCachedPatterns {
+		t.Fatalf("cache holds %d patterns, max is %d", n, maxCachedPatterns)
+	}
+}

pm/pm.go

@@ -3,6 +3,8 @@ package pm
 
 import (
 	"fmt"
+	"sync"
+	"sync/atomic"
 )
 
 const (
@@ -11,13 +13,21 @@ const (
 	maxRecursionLevel = 1000000
 )
 
+const maxCachedPatterns = 1024
+
 /* Error {{{ */
 
 type Error struct {
 	Pos     int
 	Message string
 }
 
+type compiled struct {
+	insts    []inst
+	mustHead bool
+	ncap     int
+}
+
 func newError(pos int, message string, args ...interface{}) *Error {
 	if len(args) == 0 {
 		return &Error{pos, message}
@@ -521,60 +531,89 @@ func compilePattern(p pattern, ps ...*iptr) []inst {
 	return ptr.insts
 }
 
+type Cache struct {
+	m sync.Map
+	n int32
+}
+
+func (c *Cache) getCompiled(p string) *compiled {
+	if c != nil {
+		if v, ok := c.m.Load(p); ok {
+			return v.(*compiled)
+		}
+	}
+	pat := parsePattern(newScanner([]byte(p)), true)
+	insts := compilePattern(pat)
+	ncap := 0
+	for _, in := range insts {
+		switch in.OpCode {
+		case opSave:
+			if in.Operand1+1 > ncap {
+				ncap = in.Operand1 + 1
+			}
+		case opPSave:
+			if in.Operand1+2 > ncap {
+				ncap = in.Operand1 + 2
+			}
+		}
+	}
+	cp := &compiled{insts, pat.MustHead, ncap}
+	if c != nil && atomic.LoadInt32(&c.n) < maxCachedPatterns {
+		if _, loaded := c.m.LoadOrStore(p, cp); !loaded {
+			atomic.AddInt32(&c.n, 1)
+		}
+	}
+	return cp
+}
+
 /* }}} parse */
 
 /* VM {{{ */
 
 // Simple recursive virtual machine based on the
 // "Regular Expression Matching: the Virtual Machine Approach" (https://swtch.com/~rsc/regexp/regexp2.html)
-func recursiveVM(src []byte, insts []inst, pc, sp, recLevel int, ms ...*MatchData) (bool, int, *MatchData) {
+func recursiveVM(src []byte, insts []inst, pc, sp, recLevel int, m *MatchData) (bool, int) {
 	recLevel++
 	if recLevel > maxRecursionLevel {
 		panic(newError(_UNKNOWN, "pattern/input too complex"))
 	}
-	var m *MatchData
-	if len(ms) == 0 {
-		m = newMatchState()
-	} else {
-		m = ms[0]
-	}
 redo:
 	inst := insts[pc]
 	switch inst.OpCode {
 	case opChar:
 		if sp >= len(src) || !inst.Class.Matches(int(src[sp])) {
-			return false, sp, m
+			return false, sp
 		}
 		pc++
 		sp++
 		goto redo
 	case opMatch:
-		return true, sp, m
+		return true, sp
 	case opTailMatch:
-		return sp >= len(src), sp, m
+		return sp >= len(src), sp
 	case opJmp:
 		pc = inst.Operand1
 		goto redo
 	case opSplit:
-		if ok, nsp, _ := recursiveVM(src, insts, inst.Operand1, sp, recLevel, m); ok {
-			return true, nsp, m
+		if ok, nsp := recursiveVM(src, insts, inst.Operand1, sp, recLevel, m); ok {
+			return true, nsp
 		}
 		pc = inst.Operand2
 		goto redo
 	case opSave:
 		s := m.setCapture(inst.Operand1, sp)
-		if ok, nsp, _ := recursiveVM(src, insts, pc+1, sp, recLevel, m); ok {
-			return true, nsp, m
+		if ok, nsp := recursiveVM(src, insts, pc+1, sp, recLevel, m); ok {
+			return true, nsp
 		}
 		m.restoreCapture(inst.Operand1, s)
-		return false, sp, m
+		return false, sp
 	case opPSave:
 		m.addPosCapture(inst.Operand1, sp+1)
 		pc++
 		goto redo
 	case opBrace:
 		if sp >= len(src) || int(src[sp]) != inst.Operand1 {
-			return false, sp, m
+			return false, sp
 		}
 		count := 1
 		for sp = sp + 1; sp < len(src); sp++ {
@@ -590,7 +629,7 @@ redo:
 				count++
 			}
 		}
-		return false, sp, m
+		return false, sp
 	case opNumber:
 		idx := inst.Operand1 * 2
 		if idx >= m.CaptureLength()-1 {
@@ -599,7 +638,7 @@ redo:
 		capture := src[m.Capture(idx):m.Capture(idx+1)]
 		for i := 0; i < len(capture); i++ {
 			if i+sp >= len(src) || capture[i] != src[i+sp] {
-				return false, sp, m
+				return false, sp
 			}
 		}
 		pc++
@@ -613,7 +652,11 @@ redo:
 
 /* API {{{ */
 
-func Find(p string, src []byte, offset, limit int) (matches []*MatchData, err error) {
+func Find(p string, src []byte, offset, limit int) ([]*MatchData, error) {
+	return (*Cache)(nil).Find(p, src, offset, limit)
+}
+
+func (c *Cache) Find(p string, src []byte, offset, limit int) (matches []*MatchData, err error) {
 	defer func() {
 		if v := recover(); v != nil {
 			if perr, ok := v.(*Error); ok {
@@ -623,19 +666,21 @@ func Find(p string, src []byte, offset, limit int) (matches []*MatchData, err er
 			}
 		}
 	}()
-	pat := parsePattern(newScanner([]byte(p)), true)
-	insts := compilePattern(pat)
+	cp := c.getCompiled(p)
 	matches = []*MatchData{}
+	m := &MatchData{make([]uint32, 0, cp.ncap)}
 	for sp := offset; sp <= len(src); {
-		ok, nsp, ms := recursiveVM(src, insts, 0, sp, 0)
+		m.captures = m.captures[:0]
+		ok, nsp := recursiveVM(src, cp.insts, 0, sp, 0, m)
 		sp++
 		if ok {
 			if sp < nsp {
 				sp = nsp
 			}
-			matches = append(matches, ms)
+			matches = append(matches, m)
+			m = &MatchData{make([]uint32, 0, cp.ncap)}
 		}
-		if len(matches) == limit || pat.MustHead {
+		if len(matches) == limit || cp.mustHead {
 			break
 		}
 	}

stringlib.go

@@ -112,7 +112,7 @@ func strFind(L *LState) int {
 		return 2
 	}
 
-	mds, err := pm.Find(pattern, unsafeFastStringToReadOnlyBytes(str), init, 1)
+	mds, err := L.G.patCache.Find(pattern, unsafeFastStringToReadOnlyBytes(str), init, 1)
 	if err != nil {
 		L.RaiseError(err.Error())
 	}
@@ -152,7 +152,7 @@ func strGsub(L *LState) int {
 	repl := L.CheckAny(3)
 	limit := L.OptInt(4, -1)
 
-	mds, err := pm.Find(pat, unsafeFastStringToReadOnlyBytes(str), 0, limit)
+	mds, err := L.G.patCache.Find(pat, unsafeFastStringToReadOnlyBytes(str), 0, limit)
 	if err != nil {
 		L.RaiseError(err.Error())
 	}
@@ -201,27 +201,27 @@ func capturedString(L *LState, m *pm.MatchData, str string, idx int) string {
 }
 
 func strGsubDoReplace(str string, info []replaceInfo) string {
-	offset := 0
-	buf := []byte(str)
-	for _, replace := range info {
-		oldlen := len(buf)
-		b1 := append([]byte(""), buf[0:offset+replace.Indicies[0]]...)
-		b2 := []byte("")
-		index2 := offset + replace.Indicies[1]
-		if index2 <= len(buf) {
-			b2 = append(b2, buf[index2:len(buf)]...)
-		}
-		buf = append(b1, replace.String...)
-		buf = append(buf, b2...)
-		offset += len(buf) - oldlen
-	}
-	return string(buf)
+	var sb strings.Builder
+	sb.Grow(len(str))
+	prev := 0
+	for _, r := range info {
+		sb.WriteString(str[prev:r.Indicies[0]])
+		sb.WriteString(r.String)
+		prev = r.Indicies[1]
+	}
+	sb.WriteString(str[prev:])
+	return sb.String()
 }
 
 func strGsubStr(L *LState, str string, repl string, matches []*pm.MatchData) string {
 	infoList := make([]replaceInfo, 0, len(matches))
+	plain := strings.IndexByte(repl, '%') < 0
 	for _, match := range matches {
 		start, end := match.Capture(0), match.Capture(1)
+		if plain {
+			infoList = append(infoList, replaceInfo{[]int{start, end}, repl})
+			continue
+		}
 		sc := newFlagScanner('%', "", "", repl)
 		for c, eos := sc.Next(); !eos; c, eos = sc.Next() {
 			if !sc.ChangeFlag {
@@ -327,7 +327,7 @@ func strGmatchIter(L *LState) int {
 func strGmatch(L *LState) int {
 	str := L.CheckString(1)
 	pattern := L.CheckString(2)
-	mds, err := pm.Find(pattern, []byte(str), 0, -1)
+	mds, err := L.G.patCache.Find(pattern, []byte(str), 0, -1)
 	if err != nil {
 		L.RaiseError(err.Error())
 	}
@@ -363,7 +363,7 @@ func strMatch(L *LState) int {
 		offset = 0
 	}
 
-	mds, err := pm.Find(pattern, unsafeFastStringToReadOnlyBytes(str), offset, 1)
+	mds, err := L.G.patCache.Find(pattern, unsafeFastStringToReadOnlyBytes(str), offset, 1)
 	if err != nil {
 		L.RaiseError(err.Error())
 	}

value.go

@@ -4,6 +4,8 @@ import (
 	"context"
 	"fmt"
 	"os"
+
+	"github.com/yuin/gopher-lua/pm"
 )
 
 type LValueType int
@@ -174,6 +176,7 @@ type Global struct {
 	builtinMts map[int]LValue
 	tempFiles  []*os.File
 	gccount    int32
+	patCache   pm.Cache
 }
 
 type LState struct {