Skip to content

Commit

Permalink
evalengine: Simplify logic for collation handling
Browse files Browse the repository at this point in the history
Signed-off-by: Dirkjan Bussink <[email protected]>
  • Loading branch information
dbussink committed Feb 12, 2024
1 parent 3ae4423 commit 789bbad
Show file tree
Hide file tree
Showing 2 changed files with 29 additions and 112 deletions.
58 changes: 2 additions & 56 deletions go/vt/vtgate/evalengine/compiler_asm.go
Original file line number Diff line number Diff line change
Expand Up @@ -2986,61 +2986,7 @@ func (asm *assembler) Like_collate(expr *LikeExpr, collation colldata.Collation)
}, "LIKE VARCHAR(SP-2), VARCHAR(SP-1) COLLATE '%s'", collation.Name())
}

func (asm *assembler) Locate_coerce3(coercion *compiledCoercion) {
asm.adjustStack(-2)

asm.emit(func(env *ExpressionEnv) int {
substr := env.vm.stack[env.vm.sp-3].(*evalBytes)
str := env.vm.stack[env.vm.sp-2].(*evalBytes)
pos := env.vm.stack[env.vm.sp-1].(*evalInt64)
env.vm.sp -= 2

if pos.i < 1 {
env.vm.stack[env.vm.sp-1] = env.vm.arena.newEvalInt64(0)
return 1
}

var bsub, bstr []byte
bsub, env.vm.err = coercion.left(nil, substr.bytes)
if env.vm.err != nil {
return 0
}
bstr, env.vm.err = coercion.right(nil, str.bytes)
if env.vm.err != nil {
return 0
}

found := colldata.Index(coercion.col, bstr, bsub, int(pos.i)-1)
env.vm.stack[env.vm.sp-1] = env.vm.arena.newEvalInt64(int64(found) + 1)
return 1
}, "LOCATE VARCHAR(SP-3), VARCHAR(SP-2) INT64(SP-1) COERCE AND COLLATE '%s'", coercion.col.Name())
}

func (asm *assembler) Locate_coerce2(coercion *compiledCoercion) {
asm.adjustStack(-1)

asm.emit(func(env *ExpressionEnv) int {
substr := env.vm.stack[env.vm.sp-2].(*evalBytes)
str := env.vm.stack[env.vm.sp-1].(*evalBytes)
env.vm.sp--

var bsub, bstr []byte
bsub, env.vm.err = coercion.left(nil, substr.bytes)
if env.vm.err != nil {
return 0
}
bstr, env.vm.err = coercion.right(nil, str.bytes)
if env.vm.err != nil {
return 0
}

found := colldata.Index(coercion.col, bstr, bsub, 0)
env.vm.stack[env.vm.sp-1] = env.vm.arena.newEvalInt64(int64(found) + 1)
return 1
}, "LOCATE VARCHAR(SP-2), VARCHAR(SP-1) COERCE AND COLLATE '%s'", coercion.col.Name())
}

func (asm *assembler) Locate_collate3(collation colldata.Collation) {
func (asm *assembler) Locate3(collation colldata.Collation) {
asm.adjustStack(-2)

asm.emit(func(env *ExpressionEnv) int {
Expand All @@ -3060,7 +3006,7 @@ func (asm *assembler) Locate_collate3(collation colldata.Collation) {
}, "LOCATE VARCHAR(SP-3), VARCHAR(SP-2) INT64(SP-1) COLLATE '%s'", collation.Name())
}

func (asm *assembler) Locate_collate2(collation colldata.Collation) {
func (asm *assembler) Locate2(collation colldata.Collation) {
asm.adjustStack(-1)

asm.emit(func(env *ExpressionEnv) int {
Expand Down
83 changes: 27 additions & 56 deletions go/vt/vtgate/evalengine/fn_string.go
Original file line number Diff line number Diff line change
Expand Up @@ -1281,6 +1281,19 @@ func (call *builtinLocate) eval(env *ExpressionEnv) (eval, error) {
return nil, err
}

if _, ok := str.(*evalBytes); !ok {
str, err = evalToVarchar(str, call.collate, true)
if err != nil {
return nil, err
}
}

col := str.(*evalBytes).col.Collation
substr, err = evalToVarchar(substr, col, true)
if err != nil {
return nil, err
}

pos := int64(1)
if len(call.Arguments) > 2 {
p, err := call.Arguments[2].eval(env)
Expand All @@ -1293,15 +1306,9 @@ func (call *builtinLocate) eval(env *ExpressionEnv) (eval, error) {
}
}

var col collations.TypedCollation
substr, str, col, err = mergeAndCoerceCollations(substr, str, env.collationEnv)
if err != nil {
return nil, err
}

var coll colldata.Collation
if typeIsTextual(substr.SQLType()) && typeIsTextual(str.SQLType()) {
coll = colldata.Lookup(col.Collation)
coll = colldata.Lookup(col)
} else {
coll = colldata.Lookup(collations.CollationBinaryID)
}
Expand Down Expand Up @@ -1331,73 +1338,37 @@ func (call *builtinLocate) compile(c *compiler) (ctype, error) {
_ = c.compileToInt64(l, 1)
}

if !substr.isTextual() {
c.asm.Convert_xc(len(call.Arguments), sqltypes.VarChar, c.collation, 0, false)
substr.Col = collations.TypedCollation{
if !str.isTextual() {
c.asm.Convert_xce(len(call.Arguments)-1, sqltypes.VarChar, c.collation)
str.Col = collations.TypedCollation{
Collation: c.collation,
Coercibility: collations.CoerceCoercible,
Repertoire: collations.RepertoireASCII,
}
}

if !str.isTextual() {
c.asm.Convert_xc(len(call.Arguments)-1, sqltypes.VarChar, c.collation, 0, false)
str.Col = collations.TypedCollation{
Collation: c.collation,
fromCharset := colldata.Lookup(substr.Col.Collation).Charset()
toCharset := colldata.Lookup(str.Col.Collation).Charset()
if !substr.isTextual() || (fromCharset != toCharset && !toCharset.IsSuperset(fromCharset)) {
c.asm.Convert_xce(len(call.Arguments), sqltypes.VarChar, str.Col.Collation)
substr.Col = collations.TypedCollation{
Collation: str.Col.Collation,
Coercibility: collations.CoerceCoercible,
Repertoire: collations.RepertoireASCII,
}
}

var merged collations.TypedCollation
var coerceLeft colldata.Coercion
var coerceRight colldata.Coercion

if substr.Col.Collation != str.Col.Collation {
merged, coerceLeft, coerceRight, err = colldata.Merge(c.env.CollationEnv(), substr.Col, str.Col, colldata.CoercionOptions{
ConvertToSuperset: true,
ConvertWithCoercion: true,
})
} else {
merged = substr.Col
}
if err != nil {
return ctype{}, err
}

var coll colldata.Collation
if typeIsTextual(substr.Type) && typeIsTextual(str.Type) {
coll = colldata.Lookup(merged.Collation)
coll = colldata.Lookup(str.Col.Collation)
} else {
coll = colldata.Lookup(collations.CollationBinaryID)
}

if coerceLeft == nil && coerceRight == nil {
if len(call.Arguments) > 2 {
c.asm.Locate_collate3(coll)
} else {
c.asm.Locate_collate2(coll)
}
if len(call.Arguments) > 2 {
c.asm.Locate3(coll)
} else {
if coerceLeft == nil {
coerceLeft = func(dst, in []byte) ([]byte, error) { return in, nil }
}
if coerceRight == nil {
coerceRight = func(dst, in []byte) ([]byte, error) { return in, nil }
}
if len(call.Arguments) > 2 {
c.asm.Locate_coerce3(&compiledCoercion{
col: colldata.Lookup(merged.Collation),
left: coerceLeft,
right: coerceRight,
})
} else {
c.asm.Locate_coerce2(&compiledCoercion{
col: colldata.Lookup(merged.Collation),
left: coerceLeft,
right: coerceRight,
})
}
c.asm.Locate2(coll)
}

c.asm.jumpDestination(skip1, skip2)
Expand Down

0 comments on commit 789bbad

Please sign in to comment.