evalEngine: Implement string INSERT (#15201)

Signed-off-by: Noble Mittal <noblemittal@outlook.com>

go/vt/vtgate/evalengine/compiler.go

@@ -327,6 +327,15 @@ func (c *compiler) compileNullCheck3(arg1, arg2, arg3 ctype) *jump {
  return nil
 }
 
+func (c *compiler) compileNullCheck4(arg1, arg2, arg3, arg4 ctype) *jump {
+ if arg1.nullable() || arg2.nullable() || arg3.nullable() || arg4.nullable() {
+ j := c.asm.jumpFrom()
+ c.asm.NullCheck4(j)
+ return j
+ }
+ return nil
+}
+
 func (c *compiler) compileNullCheckArg(ct ctype, offset int) *jump {
  if ct.nullable() {
  j := c.asm.jumpFrom()

go/vt/vtgate/evalengine/compiler_asm.go

@@ -2345,6 +2345,28 @@ func (asm *assembler) Fn_BIT_LENGTH() {
  }, "FN BIT_LENGTH VARCHAR(SP-1)")
 }
 
+func (asm *assembler) Fn_INSERT(col collations.TypedCollation) {
+ asm.adjustStack(-3)
+
+ asm.emit(func(env *ExpressionEnv) int {
+ str := env.vm.stack[env.vm.sp-4].(*evalBytes)
+ pos := env.vm.stack[env.vm.sp-3].(*evalInt64).i
+ l := env.vm.stack[env.vm.sp-2].(*evalInt64).i
+ newstr := env.vm.stack[env.vm.sp-1].(*evalBytes)
+
+ res := insert(str, newstr, int(pos), int(l))
+ if !validMaxLength(int64(len(res)), 1) {
+ env.vm.stack[env.vm.sp-4] = nil
+ env.vm.sp -= 3
+ return 1
+ }
+
+ env.vm.stack[env.vm.sp-4] = env.vm.arena.newEvalText(res, col)
+ env.vm.sp -= 3
+ return 1
+ }, "FN INSERT VARCHAR(SP-4) INT64(SP-3) INT64(SP-2) VARCHAR(SP-1)")
+}
+
 func (asm *assembler) Fn_LUCASE(upcase bool) {
  if upcase {
  asm.emit(func(env *ExpressionEnv) int {
@@ -3147,6 +3169,17 @@ func (asm *assembler) NullCheck3(j *jump) {
  }, "NULLCHECK SP-1, SP-2, SP-3")
 }
 
+func (asm *assembler) NullCheck4(j *jump) {
+ asm.emit(func(env *ExpressionEnv) int {
+ if env.vm.stack[env.vm.sp-4] == nil || env.vm.stack[env.vm.sp-3] == nil || env.vm.stack[env.vm.sp-2] == nil || env.vm.stack[env.vm.sp-1] == nil {
+ env.vm.stack[env.vm.sp-4] = nil
+ env.vm.sp -= 3
+ return j.offset()
+ }
+ return 1
+ }, "NULLCHECK SP-1, SP-2, SP-3, SP-4")
+}
+
 func (asm *assembler) NullCheckArg(j *jump, offset int) {
  asm.emit(func(env *ExpressionEnv) int {
  if env.vm.stack[env.vm.sp-1] == nil {

go/vt/vtgate/evalengine/fn_string.go

@@ -18,6 +18,7 @@ package evalengine
 
 import (
  "bytes"
+ "math"
 
  "vitess.io/vitess/go/mysql/collations"
  "vitess.io/vitess/go/mysql/collations/charset"
@@ -29,6 +30,11 @@ import (
 )
 
 type (
+ builtinInsert struct {
+ CallExpr
+ collate collations.ID
+ }
+
  builtinChangeCase struct {
  CallExpr
  upcase bool
@@ -106,6 +112,7 @@ type (
  }
 )
 
+var _ IR = (*builtinInsert)(nil)
 var _ IR = (*builtinChangeCase)(nil)
 var _ IR = (*builtinCharLength)(nil)
 var _ IR = (*builtinLength)(nil)
@@ -120,6 +127,122 @@ var _ IR = (*builtinLeftRight)(nil)
 var _ IR = (*builtinPad)(nil)
 var _ IR = (*builtinTrim)(nil)
 
+func insert(str, newstr *evalBytes, pos, l int) []byte {
+ pos--
+
+ cs := colldata.Lookup(str.col.Collation).Charset()
+ strLen := charset.Length(cs, str.bytes)
+
+ if pos < 0 || strLen <= pos {
+ return str.bytes
+ }
+ if l < 0 {
+ l = strLen
+ }
+
+ front := charset.Slice(cs, str.bytes, 0, pos)
+ var back []byte
+ if pos <= math.MaxInt-l && pos+l < strLen {
+ back = charset.Slice(cs, str.bytes, pos+l, strLen)
+ }
+
+ res := make([]byte, len(front)+len(newstr.bytes)+len(back))
+
+ copy(res[:len(front)], front)
+ copy(res[len(front):], newstr.bytes)
+ copy(res[len(front)+len(newstr.bytes):], back)
+
+ return res
+}
+
+func (call *builtinInsert) eval(env *ExpressionEnv) (eval, error) {
+ args, err := call.args(env)
+ if err != nil {
+ return nil, err
+ }
+ if args[0] == nil || args[1] == nil || args[2] == nil || args[3] == nil {
+ return nil, nil
+ }
+
+ str, ok := args[0].(*evalBytes)
+ if !ok {
+ str, err = evalToVarchar(args[0], call.collate, true)
+ if err != nil {
+ return nil, err
+ }
+ }
+
+ pos := evalToInt64(args[1]).i
+ l := evalToInt64(args[2]).i
+
+ newstr, err := evalToVarchar(args[3], str.col.Collation, true)
+ if err != nil {
+ return nil, err
+ }
+
+ res := insert(str, newstr, int(pos), int(l))
+ if !validMaxLength(int64(len(res)), 1) {
+ return nil, nil
+ }
+ return newEvalText(res, str.col), nil
+}
+
+func (call *builtinInsert) compile(c *compiler) (ctype, error) {
+ str, err := call.Arguments[0].compile(c)
+ if err != nil {
+ return ctype{}, err
+ }
+
+ pos, err := call.Arguments[1].compile(c)
+ if err != nil {
+ return ctype{}, err
+ }
+
+ l, err := call.Arguments[2].compile(c)
+ if err != nil {
+ return ctype{}, err
+ }
+
+ newstr, err := call.Arguments[3].compile(c)
+ if err != nil {
+ return ctype{}, err
+ }
+
+ skip := c.compileNullCheck4(str, pos, l, newstr)
+
+ _ = c.compileToInt64(pos, 3)
+ _ = c.compileToInt64(l, 2)
+
+ if err != nil {
+ return ctype{}, nil
+ }
+
+ col := str.Col
+
+ switch {
+ case str.isTextual():
+ default:
+ c.asm.Convert_xce(4, sqltypes.VarChar, c.collation)
+ col = typedCoercionCollation(sqltypes.VarChar, c.collation)
+ }
+
+ switch {
+ case newstr.isTextual():
+ fromCharset := colldata.Lookup(newstr.Col.Collation).Charset()
+ toCharset := colldata.Lookup(col.Collation).Charset()
+ if fromCharset != toCharset && !toCharset.IsSuperset(fromCharset) {
+ c.asm.Convert_xce(1, sqltypes.VarChar, col.Collation)
+ }
+ default:
+ c.asm.Convert_xce(1, sqltypes.VarChar, col.Collation)
+ }
+
+ c.asm.Fn_INSERT(col)
+ c.asm.jumpDestination(skip)
+
+ return ctype{Type: sqltypes.VarChar, Col: col, Flag: flagNullable}, nil
+}
+
 func (call *builtinChangeCase) eval(env *ExpressionEnv) (eval, error) {
  arg, err := call.arg1(env)
  if err != nil {

go/vt/vtgate/evalengine/testcases/cases.go

@@ -63,6 +63,7 @@ var Cases = []TestCase{
  {Run: TupleComparisons},
  {Run: Comparisons},
  {Run: InStatement},
+ {Run: FnInsert},
  {Run: FnLower},
  {Run: FnUpper},
  {Run: FnCharLength},
@@ -1314,6 +1315,28 @@ var JSONExtract_Schema = []*querypb.Field{
  },
 }
 
+func FnInsert(yield Query) {
+ for _, s := range insertStrings {
+ for _, ns := range insertStrings {
+ for _, l := range inputBitwise {
+ for _, p := range inputBitwise {
+ yield(fmt.Sprintf("INSERT(%s, %s, %s, %s)", s, p, l, ns), nil)
+ }
+ }
+ }
+ }
+
+ mysqlDocSamples := []string{
+ "INSERT('Quadratic', 3, 4, 'What')",
+ "INSERT('Quadratic', -1, 4, 'What')",
+ "INSERT('Quadratic', 3, 100, 'What')",
+ }
+
+ for _, q := range mysqlDocSamples {
+ yield(q, nil)
+ }
+}
+
 func FnLower(yield Query) {
  for _, str := range inputStrings {
  yield(fmt.Sprintf("LOWER(%s)", str), nil)

go/vt/vtgate/evalengine/testcases/inputs.go

@@ -199,6 +199,41 @@ var inputStrings = []string{
  // "_ucs2 'AabcÅå'",
 }
 
+var insertStrings = []string{
+ "NULL",
+ "\"\"",
+ "\"a\"",
+ "\"abc\"",
+ "1",
+ "-1",
+ "0123",
+ "0xAACC",
+ "3.1415926",
+ // MySQL has broken behavior for these inputs,
+ // see https://github.com/mysql/mysql-server/pull/517
+ // "\"Å å\"",
+ // "\"中文测试\"",
+ // "\"日本語テスト\"",
+ // "\"한국어 시험\"",
+ // "\"😊😂🤢\"",
+ // "_utf8mb4 'abcABCÅå'",
+ "DATE '2022-10-11'",
+ "TIME '11:02:23'",
+ "'123'",
+ "9223372036854775807",
+ "-9223372036854775808",
+ "999999999999999999999999",
+ "-999999999999999999999999",
+ "_binary 'Müller' ",
+ "_latin1 0xFF",
+ // TODO: support other multibyte encodings
+ // "_dec8 'ÒòÅå'",
+ // "_utf8mb3 'abcABCÅå'",
+ // "_utf16 'AabcÅå'",
+ // "_utf32 'AabcÅå'",
+ // "_ucs2 'AabcÅå'",
+}
+
 var inputConversionTypes = []string{
  "BINARY", "BINARY(1)", "BINARY(0)", "BINARY(16)", "BINARY(-1)",
  "CHAR", "CHAR(1)", "CHAR(0)", "CHAR(16)", "CHAR(-1)",

go/vt/vtgate/evalengine/translate_builtin.go

@@ -984,6 +984,35 @@ func (ast *astCompiler) translateCallable(call sqlparser.Callable) (IR, error) {
  return &builtinRegexpReplace{
  CallExpr: CallExpr{Arguments: args, Method: "REGEXP_REPLACE"},
  }, nil
+
+ case *sqlparser.InsertExpr:
+ str, err := ast.translateExpr(call.Str)
+ if err != nil {
+ return nil, err
+ }
+
+ pos, err := ast.translateExpr(call.Pos)
+ if err != nil {
+ return nil, err
+ }
+
+ len, err := ast.translateExpr(call.Len)
+ if err != nil {
+ return nil, err
+ }
+
+ newstr, err := ast.translateExpr(call.NewStr)
+ if err != nil {
+ return nil, err
+ }
+
+ args := []IR{str, pos, len, newstr}
+
+ var cexpr = CallExpr{Arguments: args, Method: "INSERT"}
+ return &builtinInsert{
+ CallExpr: cexpr,
+ collate: ast.cfg.Collation,
+ }, nil
  default:
  return nil, translateExprNotSupported(call)
  }

go/vt/vtgate/planbuilder/testdata/select_cases.json

@@ -3048,15 +3048,15 @@
  "QueryType": "SELECT",
  "Original": "select insert('Quadratic', 3, 4, 'What')",
  "Instructions": {
- "OperatorType": "Route",
- "Variant": "Reference",
- "Keyspace": {
-  "Name": "main",
-  "Sharded": false
- },
- "FieldQuery": "select insert('Quadratic', 3, 4, 'What') from dual where 1 != 1",
- "Query": "select insert('Quadratic', 3, 4, 'What') from dual",
- "Table": "dual"
+ "OperatorType": "Projection",
+ "Expressions": [
+  "'QuWhattic' as insert('Quadratic', 3, 4, 'What')"
+ ],
+ "Inputs": [
+  {
+  "OperatorType": "SingleRow"
+  }
+ ]
  },
  "TablesUsed": [
  "main.dual"