cql3: Add has_eq_restriction_on_column function

Adds a function that checks whether a given expression has eq restrction
on the specified column.

It finds restrictions like
col = ...
or
(col, col2) = ...

IN restrictions don't count, they aren't EQ restrictions

Signed-off-by: Jan Ciolek <jan.ciolek@scylladb.com>

cql3/expr/expression.cc

@@ -1071,6 +1071,80 @@ bool is_on_collection(const binary_operator& b) {
     return false;
 }
 
+bool contains_column(const column_definition& column, const expression& e) {
+    const column_value* find_res = find_in_expression<column_value>(e,
+        [&](const column_value& column_val) -> bool {
+            return (*column_val.col == column);
+        }
+    );
+    return find_res != nullptr;
+}
+
+bool contains_nonpure_function(const expression& e) {
+    const function_call* find_res = find_in_expression<function_call>(e,
+        [&](const function_call& fc) {
+            return std::visit(overloaded_functor {
+                [](const functions::function_name&) -> bool {
+                    on_internal_error(expr_logger,
+                    "contains_nonpure_function called on unprepared expression - expected function pointer, got name");
+                },
+                [](const ::shared_ptr<functions::function>& fun) -> bool {
+                    return !fun->is_pure();
+                }
+            }, fc.func);
+        }
+    );
+    return find_res != nullptr;
+}
+
+bool has_eq_restriction_on_column(const column_definition& column, const expression& e) {
+    std::function<bool(const expression&)> column_in_lhs = [&](const expression& e) -> bool {
+        return visit(overloaded_functor {
+            [&](const column_value& cv) {
+                // Use column_defintion::operator== for comparison,
+                // columns with the same name but different schema will not be equal.
+                return *cv.col == column;
+            },
+            [&](const tuple_constructor& tc) {
+                for (const expression& elem : tc.elements) {
+                    if (column_in_lhs(elem)) {
+                        return true;
+                    }
+                }
+
+                return false;
+            },
+            [&](const auto&) {return false;}
+        }, e);
+    };
+
+    // Look for binary operator describing eq relation with this column on lhs
+    const binary_operator* eq_restriction_search_res = find_binop(e, [&](const binary_operator& b) {
+        if (b.op != oper_t::EQ) {
+            return false;
+        }
+
+        if (!column_in_lhs(b.lhs)) {
+            return false;
+        }
+
+        // These conditions are not allowed to occur in the current code,
+        // but they might be allowed in the future.
+        // They are added now to avoid surprises later.
+        //
+        // These conditions detect cases like:
+        // WHERE column1 = column2
+        // WHERE column1 = row_number()
+        if (contains_column(column, b.rhs) || contains_nonpure_function(b.rhs)) {
+            return false;
+        }
+
+        return true;
+    });
+
+    return eq_restriction_search_res != nullptr;
+}
+
 expression replace_column_def(const expression& expr, const column_definition* new_cdef) {
     return search_and_replace(expr, [&] (const expression& expr) -> std::optional<expression> {
         if (expr::is<column_value>(expr)) {

cql3/expr/expression.hh

@@ -655,6 +655,20 @@ inline auto find_clustering_order(const expression& e) {
 /// True iff binary_operator involves a collection.
 extern bool is_on_collection(const binary_operator&);
 
+// Checks whether the given column occurs in the expression.
+// Uses column_defintion::operator== for comparison, columns with the same name but different schema will not be equal.
+bool contains_column(const column_definition& column, const expression& e);
+
+// Checks whether this expression contains a nonpure function.
+// The expression must be prepared, so that function names are converted to function pointers.
+bool contains_nonpure_function(const expression&);
+
+// Checks whether the given column has an EQ restriction in the expression.
+// EQ restriction is `col = ...` or `(col, col2) = ...`
+// IN restriction is NOT an EQ restriction, this function will not look for IN restrictions.
+// Uses column_defintion::operator== for comparison, columns with the same name but different schema will not be equal.
+bool has_eq_restriction_on_column(const column_definition& column, const expression& e);
+
 /// Replaces every column_definition in an expression with this one.  Throws if any LHS is not a single
 /// column_value.
 extern expression replace_column_def(const expression&, const column_definition*);

cql3/restrictions/statement_restrictions.cc

@@ -598,6 +598,14 @@ std::pair<std::optional<secondary_index::index>, ::shared_ptr<cql3::restrictions
     return {chosen_index, chosen_index_restrictions};
 }
 
+bool statement_restrictions::has_eq_restriction_on_column(const column_definition& column) const {
+    if (!_where.has_value()) {
+        return false;
+    }
+
+    return expr::has_eq_restriction_on_column(column, *_where);
+}
+
 std::vector<const column_definition*> statement_restrictions::get_column_defs_for_filtering(database& db) const {
     std::vector<const column_definition*> column_defs_for_filtering;
     if (need_filtering()) {

cql3/restrictions/statement_restrictions.hh

@@ -215,6 +215,12 @@ public:
         return has_token(_partition_key_restrictions->expression);
     }
 
+    // Checks whether the given column has an EQ restriction.
+    // EQ restriction is `col = ...` or `(col, col2) = ...`
+    // IN restriction is NOT an EQ restriction, this function will not look for IN restrictions.
+    // Uses column_defintion::operator== for comparison, columns with the same name but different schema will not be equal.
+    bool has_eq_restriction_on_column(const column_definition&) const;
+
     /**
      * Builds a possibly empty collection of column definitions that will be used for filtering
      * @param db - the database context