fixup replace

Author: James Cor

sql/expression/function/regexp_replace.go

@@ -17,29 +17,80 @@ package function
 import (
 	"fmt"
 	"strings"
+	"sync"
 
 	"gopkg.in/src-d/go-errors.v1"
 
+	regex "github.com/dolthub/go-icu-regex"
+
 	"github.com/dolthub/go-mysql-server/sql"
+	"github.com/dolthub/go-mysql-server/sql/expression"
 	"github.com/dolthub/go-mysql-server/sql/types"
 )
 
 // RegexpReplace implements the REGEXP_REPLACE function.
 // https://dev.mysql.com/doc/refman/8.0/en/regexp.html#function_regexp-replace
 type RegexpReplace struct {
-	args []sql.Expression
+	Text       sql.Expression
+	Pattern    sql.Expression
+	RText      sql.Expression
+	Position   sql.Expression
+	Occurrence sql.Expression
+	Flags      sql.Expression
+
+	cacheVal    bool
+	cachedVal   any
+	cacheRegex  bool
+	re          regex.Regex
+	compileOnce sync.Once
+	compileErr  error
 }
 
 var _ sql.FunctionExpression = (*RegexpReplace)(nil)
 var _ sql.CollationCoercible = (*RegexpReplace)(nil)
+var _ sql.Disposable = (*RegexpReplace)(nil)
 
 // NewRegexpReplace creates a new RegexpReplace expression.
 func NewRegexpReplace(args ...sql.Expression) (sql.Expression, error) {
-	if len(args) < 3 || len(args) > 6 {
+	var r *RegexpReplace
+	switch len(args) {
+	case 6:
+		r = &RegexpReplace{
+			Text:       args[0],
+			Pattern:    args[1],
+			RText:      args[2],
+			Position:   args[3],
+			Occurrence: args[4],
+			Flags:      args[5],
+		}
+	case 5:
+		r = &RegexpReplace{
+			Text:       args[0],
+			Pattern:    args[1],
+			RText:      args[2],
+			Position:   args[3],
+			Occurrence: args[4],
+		}
+	case 4:
+		r = &RegexpReplace{
+			Text:       args[0],
+			Pattern:    args[1],
+			RText:      args[2],
+			Position:   args[3],
+			Occurrence: expression.NewLiteral(0, types.Int32),
+		}
+	case 3:
+		r = &RegexpReplace{
+			Text:       args[0],
+			Pattern:    args[1],
+			RText:      args[2],
+			Position:   expression.NewLiteral(1, types.Int32),
+			Occurrence: expression.NewLiteral(0, types.Int32),
+		}
+	default:
 		return nil, sql.ErrInvalidArgumentNumber.New("regexp_replace", "3,4,5 or 6", len(args))
 	}
-
-	return &RegexpReplace{args: args}, nil
+	return r, nil
 }
 
 // FunctionName implements sql.FunctionExpression
@@ -57,14 +108,11 @@ func (r *RegexpReplace) Type() sql.Type { return types.LongText }
 
 // CollationCoercibility implements the interface sql.CollationCoercible.
 func (r *RegexpReplace) CollationCoercibility(ctx *sql.Context) (collation sql.CollationID, coercibility byte) {
-	if len(r.args) == 0 {
-		return sql.Collation_binary, 6
-	}
-	collation, coercibility = sql.GetCoercibility(ctx, r.args[0])
-	for i := 1; i < len(r.args) && i < 3; i++ {
-		nextCollation, nextCoercibility := sql.GetCoercibility(ctx, r.args[i])
-		collation, coercibility = sql.ResolveCoercibility(collation, coercibility, nextCollation, nextCoercibility)
-	}
+	collation, coercibility = sql.GetCoercibility(ctx, r.Text)
+	nextCollation, nextCoercibility := sql.GetCoercibility(ctx, r.Pattern)
+	collation, coercibility = sql.ResolveCoercibility(collation, coercibility, nextCollation, nextCoercibility)
+	nextCollation, nextCoercibility = sql.GetCoercibility(ctx, r.RText)
+	collation, coercibility = sql.ResolveCoercibility(collation, coercibility, nextCollation, nextCoercibility)
 	return collation, coercibility
 }
 
@@ -73,152 +121,163 @@ func (r *RegexpReplace) IsNullable() bool { return true }
 
 // Children implements the sql.Expression interface.
 func (r *RegexpReplace) Children() []sql.Expression {
-	return r.args
+	var children = []sql.Expression{r.Text, r.Pattern, r.RText, r.Position, r.Occurrence}
+	if r.Flags != nil {
+		children = append(children, r.Flags)
+	}
+	return children
 }
 
 // Resolved implements the sql.Expression interface.
 func (r *RegexpReplace) Resolved() bool {
-	for _, arg := range r.args {
-		if !arg.Resolved() {
-			return false
-		}
-	}
-	return true
+	return r.Text.Resolved() &&
+		r.Pattern.Resolved() &&
+		r.RText.Resolved() &&
+		r.Position.Resolved() &&
+		r.Occurrence.Resolved() &&
+		(r.Flags == nil || r.Flags.Resolved())
 }
 
 // WithChildren implements the sql.Expression interface.
 func (r *RegexpReplace) WithChildren(children ...sql.Expression) (sql.Expression, error) {
-	if len(children) != len(r.args) {
-		return nil, sql.ErrInvalidChildrenNumber.New(r, len(children), len(r.args))
+	required := 3
+	if r.Flags != nil {
+		required = 4
+	}
+	if len(children) != required {
+		return nil, sql.ErrInvalidChildrenNumber.New(r, len(children), required)
+	}
+
+	// Copy over the regex instance, in case it has already been set to avoid leaking it.
+	replace, err := NewRegexpReplace(children...)
+	if err != nil {
+		if r.re != nil {
+			if err = r.re.Close(); err != nil {
+				return nil, err
+			}
+		}
+		return nil, err
+	}
+	if r.re != nil {
+		replace.(*RegexpReplace).re = r.re
 	}
-	return NewRegexpReplace(children...)
+	return replace, nil
 }
 
 func (r *RegexpReplace) String() string {
 	var args []string
-	for _, e := range r.args {
+	for _, e := range r.Children() {
 		args = append(args, e.String())
 	}
 	return fmt.Sprintf("%s(%s)", r.FunctionName(), strings.Join(args, ","))
 }
 
+func (r *RegexpReplace) compile(ctx *sql.Context, row sql.Row) {
+	r.compileOnce.Do(func() {
+		r.cacheRegex = canBeCached(r.Pattern, r.Flags)
+		r.cacheVal = r.cacheRegex && canBeCached(r.Text, r.RText, r.Position, r.Occurrence)
+		if r.cacheRegex {
+			r.re, r.compileErr = compileRegex(ctx, r.Pattern, r.Text, r.Flags, r.FunctionName(), row)
+		}
+	})
+	if !r.cacheRegex {
+		if r.re != nil {
+			if r.compileErr = r.re.Close(); r.compileErr != nil {
+				return
+			}
+		}
+		r.re, r.compileErr = compileRegex(ctx, r.Pattern, r.Text, r.Flags, r.FunctionName(), row)
+	}
+}
+
 // Eval implements the sql.Expression interface.
 func (r *RegexpReplace) Eval(ctx *sql.Context, row sql.Row) (val interface{}, err error) {
-	// Evaluate string value
-	str, err := r.args[0].Eval(ctx, row)
+	span, ctx := ctx.Span("function.RegexpReplace")
+	defer span.End()
+
+	if r.cachedVal != nil {
+		return r.cachedVal, nil
+	}
+
+	r.compile(ctx, row)
+	if r.compileErr != nil {
+		return nil, r.compileErr
+	}
+	if r.re == nil {
+		return nil, nil
+	}
+
+	text, err := r.Text.Eval(ctx, row)
 	if err != nil {
 		return nil, err
 	}
-	if str == nil {
+	if text == nil {
 		return nil, nil
 	}
-	str, _, err = types.LongText.Convert(ctx, str)
+	text, _, err = types.LongText.Convert(ctx, text)
 	if err != nil {
 		return nil, err
 	}
 
-	// Convert to string
-	_str := str.(string)
-
-	// Handle flags
-	var flags sql.Expression = nil
-	if len(r.args) == 6 {
-		flags = r.args[5]
-	}
-
-	// Create regex, should handle null pattern and null flags
-	re, compileErr := compileRegex(ctx, r.args[1], r.args[0], flags, r.FunctionName(), row)
-	if compileErr != nil {
-		return nil, compileErr
+	rText, err := r.RText.Eval(ctx, row)
+	if err != nil {
+		return nil, err
 	}
-	if re == nil {
+	if rText == nil {
 		return nil, nil
 	}
-	defer func() {
-		if nErr := re.Close(); err == nil {
-			err = nErr
-		}
-	}()
-	if err = re.SetMatchString(ctx, _str); err != nil {
+	rText, _, err = types.LongText.Convert(ctx, rText)
+	if err != nil {
 		return nil, err
 	}
 
-	// Evaluate ReplaceStr
-	replaceStr, err := r.args[2].Eval(ctx, row)
+	pos, err := r.Position.Eval(ctx, row)
 	if err != nil {
 		return nil, err
 	}
-	if replaceStr == nil {
+	if pos == nil {
 		return nil, nil
 	}
-	replaceStr, _, err = types.LongText.Convert(ctx, replaceStr)
+	pos, _, err = types.Int32.Convert(ctx, pos)
 	if err != nil {
 		return nil, err
 	}
-
-	// Convert to string
-	_replaceStr := replaceStr.(string)
-
-	// Do nothing if str is empty
-	if len(_str) == 0 {
-		return _str, nil
+	if pos.(int32) <= 0 {
+		return nil, sql.ErrInvalidArgumentDetails.New(r.FunctionName(), fmt.Sprintf("%d", pos.(int32)))
 	}
 
-	// Default position is 1
-	_pos := 1
-
-	// Check if position argument was provided
-	if len(r.args) >= 4 {
-		// Evaluate position argument
-		pos, err := r.args[3].Eval(ctx, row)
-		if err != nil {
-			return nil, err
-		}
-		if pos == nil {
-			return nil, nil
-		}
-
-		// Convert to int32
-		pos, _, err = types.Int32.Convert(ctx, pos)
-		if err != nil {
-			return nil, err
-		}
-		// Convert to int
-		_pos = int(pos.(int32))
+	if len(text.(string)) != 0 && int(pos.(int32)) > len(text.(string)) {
+		return nil, errors.NewKind("Index out of bounds for regular expression search.").New()
 	}
 
-	// Non-positive position throws incorrect parameter
-	if _pos <= 0 {
-		return nil, sql.ErrInvalidArgumentDetails.New(r.FunctionName(), fmt.Sprintf("%d", _pos))
+	occurrence, err := r.Occurrence.Eval(ctx, row)
+	if err != nil {
+		return nil, err
 	}
-
-	// Handle out of bounds
-	if _pos > len(_str) {
-		return nil, errors.NewKind("Index out of bounds for regular expression search.").New()
+	if occurrence == nil {
+		return nil, nil
+	}
+	occurrence, _, err = types.Int32.Convert(ctx, occurrence)
+	if err != nil {
+		return nil, err
 	}
 
-	// Default occurrence is 0 (replace all occurrences)
-	_occ := 0
+	err = r.re.SetMatchString(ctx, text.(string))
+	if err != nil {
+		return nil, err
+	}
 
-	// Check if Occurrence argument was provided
-	if len(r.args) >= 5 {
-		occ, err := r.args[4].Eval(ctx, row)
-		if err != nil {
-			return nil, err
-		}
-		if occ == nil {
-			return nil, nil
-		}
+	result, err := r.re.Replace(ctx, rText.(string), int(pos.(int32)), int(occurrence.(int32)))
+	if err != nil {
+		return nil, err
+	}
 
-		// Convert occurrence to int32
-		occ, _, err = types.Int32.Convert(ctx, occ)
-		if err != nil {
-			return nil, err
-		}
+	return result, nil
+}
 
-		// Convert to int
-		_occ = int(occ.(int32))
+// Dispose implements the sql.Disposable interface.
+func (r *RegexpReplace) Dispose() {
+	if r.re != nil {
+		_ = r.re.Close()
 	}
-
-	return re.Replace(ctx, _replaceStr, _pos, _occ)
 }