fixed remaining rewriter bugs

Signed-off-by: Andres Taylor <[email protected]>

go/vt/sqlparser/ast_funcs.go

@@ -21,6 +21,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"io"
+	"slices"
 	"strconv"
 	"strings"
 
@@ -2864,3 +2865,72 @@ func ExtractAllTables(stmt Statement) []string {
 	}, stmt)
 	return tables
 }
+
+// ExtractINFromOR rewrites the OR expression into an IN clause.
+// Each side of each ORs has to be an equality comparison expression and the column names have to
+// match for all sides of each comparison.
+// This rewriter takes a query that looks like this WHERE a = 1 and b = 11 or a = 2 and b = 12 or a = 3 and b = 13
+// And rewrite that to WHERE (a, b) IN ((1,11), (2,12), (3,13))
+func ExtractINFromOR(expr *OrExpr) []Expr {
+	var varNames []*ColName
+	var values []Exprs
+	orSlice := orToSlice(expr)
+	for _, expr := range orSlice {
+		andSlice := andToSlice(expr)
+		if len(andSlice) == 0 {
+			return nil
+		}
+
+		var currentVarNames []*ColName
+		var currentValues []Expr
+		for _, comparisonExpr := range andSlice {
+			if comparisonExpr.Operator != EqualOp {
+				return nil
+			}
+
+			var colName *ColName
+			if left, ok := comparisonExpr.Left.(*ColName); ok {
+				colName = left
+				currentValues = append(currentValues, comparisonExpr.Right)
+			}
+
+			if right, ok := comparisonExpr.Right.(*ColName); ok {
+				if colName != nil {
+					return nil
+				}
+				colName = right
+				currentValues = append(currentValues, comparisonExpr.Left)
+			}
+
+			if colName == nil {
+				return nil
+			}
+
+			currentVarNames = append(currentVarNames, colName)
+		}
+
+		if len(varNames) == 0 {
+			varNames = currentVarNames
+		} else if !slices.EqualFunc(varNames, currentVarNames, func(col1, col2 *ColName) bool { return col1.Equal(col2) }) {
+			return nil
+		}
+
+		values = append(values, currentValues)
+	}
+
+	var nameTuple ValTuple
+	for _, name := range varNames {
+		nameTuple = append(nameTuple, name)
+	}
+
+	var valueTuple ValTuple
+	for _, value := range values {
+		valueTuple = append(valueTuple, ValTuple(value))
+	}
+
+	return []Expr{&ComparisonExpr{
+		Operator: InOp,
+		Left:     nameTuple,
+		Right:    valueTuple,
+	}}
+}