evalEngine: Implement ELT and FIELD

go/vt/vtgate/evalengine/compiler_asm.go

@@ -2345,6 +2345,127 @@ func (asm *assembler) Fn_BIT_LENGTH() {
 	}, "FN BIT_LENGTH VARCHAR(SP-1)")
 }
 
+func (asm *assembler) Fn_FIELD_i(args int) {
+	asm.adjustStack(-args + 1)
+	asm.emit(func(env *ExpressionEnv) int {
+		if env.vm.stack[env.vm.sp-args] == nil {
+			env.vm.stack[env.vm.sp-args] = env.vm.arena.newEvalInt64(0)
+			env.vm.sp -= args - 1
+			return 1
+		}
+
+		tar := env.vm.stack[env.vm.sp-args].(*evalInt64)
+
+		for i := range args - 1 {
+			if env.vm.stack[env.vm.sp-args+i+1] == nil {
+				continue
+			}
+
+			arg := env.vm.stack[env.vm.sp-args+i+1].(*evalInt64)
+
+			if tar.i == arg.i {
+				env.vm.stack[env.vm.sp-args] = env.vm.arena.newEvalInt64(int64(i + 1))
+				env.vm.sp -= args - 1
+				return 1
+			}
+		}
+
+		env.vm.stack[env.vm.sp-args] = env.vm.arena.newEvalInt64(0)
+		env.vm.sp -= args - 1
+		return 1
+	}, "FN FIELD INT64(SP-%d)...INT64(SP-1)", args)
+}
+
+func (asm *assembler) Fn_FIELD_b(args int) {
+	asm.adjustStack(-args + 1)
+	asm.emit(func(env *ExpressionEnv) int {
+		if env.vm.stack[env.vm.sp-args] == nil {
+			env.vm.stack[env.vm.sp-args] = env.vm.arena.newEvalInt64(0)
+			env.vm.sp -= args - 1
+			return 1
+		}
+
+		tar := env.vm.stack[env.vm.sp-args].(*evalBytes)
+
+		for i := range args - 1 {
+			if env.vm.stack[env.vm.sp-args+i+1] == nil {
+				continue
+			}
+
+			str := env.vm.stack[env.vm.sp-args+i+1].(*evalBytes)
+
+			// Compare target and current string
+			if len(tar.bytes) == len(str.bytes) {
+				eq := true
+				for i, b := range tar.bytes {
+					if str.bytes[i] != b {
+						eq = false
+						break
+					}
+				}
+
+				if eq {
+					env.vm.stack[env.vm.sp-args] = env.vm.arena.newEvalInt64(int64(i + 1))
+					env.vm.sp -= args - 1
+					return 1
+				}
+			}
+		}
+
+		env.vm.stack[env.vm.sp-args] = env.vm.arena.newEvalInt64(0)
+		env.vm.sp -= args - 1
+		return 1
+	}, "FN FIELD VARCHAR(SP-%d)...VARCHAR(SP-1)", args)
+}
+
+func (asm *assembler) Fn_FIELD_f(args int) {
+	asm.adjustStack(-args + 1)
+	asm.emit(func(env *ExpressionEnv) int {
+		if env.vm.stack[env.vm.sp-args] == nil {
+			env.vm.stack[env.vm.sp-args] = env.vm.arena.newEvalInt64(0)
+			env.vm.sp -= args - 1
+			return 1
+		}
+
+		tar := env.vm.stack[env.vm.sp-args].(*evalFloat)
+
+		for i := range args - 1 {
+			if env.vm.stack[env.vm.sp-args+i+1] == nil {
+				continue
+			}
+
+			arg := env.vm.stack[env.vm.sp-args+i+1].(*evalFloat)
+
+			if tar.f == arg.f {
+				env.vm.stack[env.vm.sp-args] = env.vm.arena.newEvalInt64(int64(i + 1))
+				env.vm.sp -= args - 1
+				return 1
+			}
+		}
+
+		env.vm.stack[env.vm.sp-args] = env.vm.arena.newEvalInt64(0)
+		env.vm.sp -= args - 1
+		return 1
+	}, "FN FIELD VARCHAR(SP-%d)...VARCHAR(SP-1)", args)
+}
+
+func (asm *assembler) Fn_ELT(args int, tt sqltypes.Type, tc collations.TypedCollation) {
+	asm.adjustStack(-args + 1)
+	asm.emit(func(env *ExpressionEnv) int {
+		i := env.vm.stack[env.vm.sp-args].(*evalInt64)
+
+		if i.i < 1 || int(i.i) >= args || env.vm.stack[env.vm.sp-args+int(i.i)] == nil {
+			env.vm.stack[env.vm.sp-args] = nil
+		} else {
+			b := env.vm.stack[env.vm.sp-args+int(i.i)].(*evalBytes)
+			env.vm.stack[env.vm.sp-args] = env.vm.arena.newEvalRaw(b.bytes, tt, tc)
+		}
+
+		env.vm.sp -= args - 1
+		return 1
+	}, "FN ELT INT64(SP-%d) VARCHAR(SP-%d)...VARCHAR(SP-1)", args, args-1)
+}
+
 func (asm *assembler) Fn_INSERT(col collations.TypedCollation) {
 	asm.adjustStack(-3)