[DRAFT] [WORK IN PROGRESS] Support NOT and Binary Operators in Match Language Rewriter

mongosql/src/air/definitions.rs

@@ -720,6 +720,7 @@ pub struct Reduce {
 pub enum MatchQuery {
     Or(Vec<MatchQuery>),
     And(Vec<MatchQuery>),
+    Not(Box<MatchQuery>),
     Type(MatchLanguageType),
     Regex(MatchLanguageRegex),
     ElemMatch(ElemMatch),

mongosql/src/codegen/match_query.rs

@@ -26,6 +26,7 @@ impl MqlCodeGenerator {
         match q {
             Or(v) => self.codegen_match_logical_operator("$or", v),
             And(v) => self.codegen_match_logical_operator("$and", v),
+            Not(inner) => self.codegen_match_logical_operator("$not", vec![*inner]),
             Type(t) => self.codegen_match_type(t),
             Regex(r) => self.codegen_match_regex(r),
             ElemMatch(em) => self.codegen_match_elem_match(em),

mongosql/src/codegen/test/match_query.rs

@@ -214,6 +214,51 @@ mod match_constant_false {
     );
 }
 
+mod not {
+    test_codegen_match_query!(
+        single_comparison,
+        expected = Ok(bson!({"$not": [{"age": {"$gt": 10}}]})),
+        input = air::MatchQuery::Not(Box::new(air::MatchQuery::Comparison(
+            air::MatchLanguageComparison {
+                function: air::MatchLanguageComparisonOp::Gt,
+                input: Some("age".to_string().into()),
+                arg: air::LiteralValue::Integer(10),
+            }
+        )))
+    );
+
+    test_codegen_match_query!(
+        not_and,
+        expected =
+            Ok(bson!({"$not": [{"$and": [{"age": {"$gt": 10}}, {"name": {"$eq": "Alice"}}]}]})),
+        input = air::MatchQuery::Not(Box::new(air::MatchQuery::And(vec![
+            air::MatchQuery::Comparison(air::MatchLanguageComparison {
+                function: air::MatchLanguageComparisonOp::Gt,
+                input: Some("age".to_string().into()),
+                arg: air::LiteralValue::Integer(10),
+            }),
+            air::MatchQuery::Comparison(air::MatchLanguageComparison {
+                function: air::MatchLanguageComparisonOp::Eq,
+                input: Some("name".to_string().into()),
+                arg: air::LiteralValue::String("Alice".to_string()),
+            }),
+        ])))
+    );
+
+    // Double negation is structurally valid; simplification is the optimizer's responsibility.
+    test_codegen_match_query!(
+        nested_not,
+        expected = Ok(bson!({"$not": [{"$not": [{"age": {"$gt": 10}}]}]})),
+        input = air::MatchQuery::Not(Box::new(air::MatchQuery::Not(Box::new(
+            air::MatchQuery::Comparison(air::MatchLanguageComparison {
+                function: air::MatchLanguageComparisonOp::Gt,
+                input: Some("age".to_string().into()),
+                arg: air::LiteralValue::Integer(10),
+            })
+        ))))
+    );
+}
+
 mod in_op {
     test_codegen_match_query!(
         in_single,

mongosql/src/mir/definitions.rs

@@ -1047,6 +1047,7 @@ pub struct MatchLanguageLogical {
 pub enum MatchLanguageLogicalOp {
     Or,
     And,
+    Not,
 }
 
 #[derive(PartialEq, Debug, Clone)]

mongosql/src/mir/optimizer/mod.rs

@@ -36,9 +36,11 @@ static OPTIMIZERS: fn() -> Vec<Box<dyn Optimizer>> = || {
         Box::new(match_splitting::MatchSplittingOptimizer {}),
         Box::new(rewrite_to_match_language::MatchLanguageRewriter {}),
         Box::new(match_null_filtering::MatchNullFilteringOptimizer {}),
+        // NOTE: We probably need to update how we handle stage movement for the new match language changes. Namely that the match stages are too late in the pipeline. We want to push them up.
         Box::new(stage_movement::StageMovementOptimizer {}),
         Box::new(determine_join_semantics::JoinSemanticsOptimizer {}),
         Box::new(lower_joins::LowerJoinsOptimizer {}),
+        // NOTE: We may need to review this in context of our changes
         Box::new(prefilter_unwinds::PrefilterUnwindsOptimizer {}),
         Box::new(dead_code_elimination::DeadCodeEliminator {}),
     ]

mongosql/src/mir/optimizer/rewrite_to_match_language/mod.rs

@@ -1,9 +1,8 @@
-/// Optimizes IS,LIKE, IN, and NOT IN expressions such that they can be translated and
-/// codegened using match language. If a Filter stage's condition contains
-/// only IS and/or LIKE expressions, or disjunctions or conjunctions with
-/// only IS and/or LIKE expressions, then the condition can be rewritten to
-/// match language (mir::MatchQuery) and the Filter stage replaced with a
-/// MatchFilter stage.
+/// Optimizes IS, LIKE, IN, NOT IN, comparison, NOT, and bare field-access expressions
+/// such that they can be translated and codegened using match language. If a Filter
+/// stage's condition consists entirely of rewritable expressions — including logical
+/// conjunctions and disjunctions of them — the Filter stage is replaced with a
+/// MatchFilter stage that uses `mir::MatchQuery`.
 ///
 /// Also optimizes constant false conditions to a MatchFalse filter. This filter
 /// can be then treated specially in codegen because certain versions of mongodb
@@ -186,16 +185,100 @@ impl MatchLanguageRewriterVisitor {
         }))
     }
 
-    // Only rewrite a condition that consists of Is, Like, or a logical operation
-    // that contains only other rewritable expressions.
+    /// Rewrites a comparison scalar function to a native [`MatchQuery::Comparison`].
+    ///
+    /// Succeeds when one argument is a [`FieldAccess`] convertible to a [`FieldPath`] and
+    /// the other is a [`LiteralValue`]. When the literal is on the left, the operator is
+    /// commuted so that the field is always on the left in the output.
+    ///
+    /// # Returns
+    ///
+    /// - `Some(MatchQuery::Comparison { … })` when the expression matches the required shape.
+    /// - `None` when neither argument is a plain field access, or the field access cannot be
+    ///   converted to a `FieldPath`, or `sf.function` is not a comparison operator.
+    fn rewrite_comparison(sf: &ScalarFunctionApplication) -> Option<MatchQuery> {
+        let scalar_function_as_match_language_op = match sf.function {
+            ScalarFunction::Lt => MatchLanguageComparisonOp::Lt,
+            ScalarFunction::Lte => MatchLanguageComparisonOp::Lte,
+            ScalarFunction::Gt => MatchLanguageComparisonOp::Gt,
+            ScalarFunction::Gte => MatchLanguageComparisonOp::Gte,
+            ScalarFunction::Eq => MatchLanguageComparisonOp::Eq,
+            ScalarFunction::Neq => MatchLanguageComparisonOp::Ne,
+            _ => return None,
+        };
+
+        let (lhs, rhs) = (sf.args.first()?, sf.args.get(1)?);
+
+        let (field_path, literal, op): (FieldPath, &LiteralValue, MatchLanguageComparisonOp) =
+            match (lhs, rhs) {
+                (Expression::FieldAccess(field_access), Expression::Literal(literal)) => (
+                    field_access.try_into().ok()?,
+                    literal,
+                    scalar_function_as_match_language_op,
+                ),
+
+                (Expression::Literal(literal), Expression::FieldAccess(field_access)) => (
+                    field_access.try_into().ok()?,
+                    literal,
+                    scalar_function_as_match_language_op,
+                ),
+                _ => return None,
+            };
+
+        if field_path.is_nullable {
+            Some(MatchQuery::Logical(MatchLanguageLogical {
+                op: MatchLanguageLogicalOp::And,
+                args: vec![
+                    MatchQuery::Comparison(MatchLanguageComparison {
+                        function: MatchLanguageComparisonOp::Gt,
+                        input: Some(field_path.clone()),
+                        arg: LiteralValue::Null,
+                        cache: SchemaCache::new(),
+                    }),
+                    MatchQuery::Comparison(MatchLanguageComparison {
+                        function: op,
+                        input: Some(field_path),
+                        arg: literal.clone(),
+                        cache: SchemaCache::new(),
+                    }),
+                ],
+                cache: SchemaCache::new(),
+            }))
+        } else {
+            Some(MatchQuery::Comparison(MatchLanguageComparison {
+                function: op,
+                input: Some(field_path),
+                arg: literal.clone(),
+                cache: SchemaCache::new(),
+            }))
+        }
+    }
+
+    // Only rewrite a condition that consists of rewritable expressions or logical operations
+    // that contain only other rewritable expressions.
     fn rewrite_condition(condition: Expression) -> Option<MatchQuery> {
         match condition {
             Expression::Is(is) => Self::rewrite_is(is),
             Expression::Like(like) => Self::rewrite_like(like),
+            Expression::FieldAccess(fa) => fa.try_into().ok().map(|fp: FieldPath| {
+                MatchQuery::Comparison(MatchLanguageComparison {
+                    function: MatchLanguageComparisonOp::Eq,
+                    input: Some(fp),
+                    arg: LiteralValue::Boolean(true),
+                    cache: SchemaCache::new(),
+                })
+            }),
             Expression::ScalarFunction(sf) => match sf.function {
                 ScalarFunction::And => Self::rewrite_logical(MatchLanguageLogicalOp::And, sf.args),
                 ScalarFunction::Or => Self::rewrite_logical(MatchLanguageLogicalOp::Or, sf.args),
+                ScalarFunction::Not => Self::rewrite_logical(MatchLanguageLogicalOp::Not, sf.args),
                 ScalarFunction::In | ScalarFunction::NotIn => Self::rewrite_in(&sf),
+                ScalarFunction::Lt
+                | ScalarFunction::Lte
+                | ScalarFunction::Eq
+                | ScalarFunction::Gt
+                | ScalarFunction::Gte
+                | ScalarFunction::Neq => Self::rewrite_comparison(&sf),
                 _ => None,
             },
             // Note this relies on ConstantFolding to ensure that the a constant expression becomes

mongosql/src/mir/optimizer/rewrite_to_match_language/test.rs

@@ -167,12 +167,13 @@ fn invalid_like_pat() -> Expression {
     })
 }
 
-fn comp_expr() -> Expression {
+fn invalid_expr() -> Expression {
+    // ScalarFunction::Add is not match-language rewritable
     Expression::ScalarFunction(ScalarFunctionApplication::new(
-        ScalarFunction::Lt,
+        ScalarFunction::Add,
         vec![
             *mir_field_access("foo", "int", true),
-            Expression::Literal(LiteralValue::Integer(10)),
+            Expression::Literal(LiteralValue::Integer(1)),
         ],
     ))
 }
@@ -212,9 +213,9 @@ test_rewrite_to_match_language_no_op!(
     filter_stage(Expression::ScalarFunction(ScalarFunctionApplication {
         function: ScalarFunction::And,
         args: vec![
-            valid_is(),   // rewritable
-            valid_like(), // rewritable
-            comp_expr(),  // not rewritable - invalid expression
+            valid_is(),     // rewritable
+            valid_like(),   // rewritable
+            invalid_expr(), // not rewritable - Add is not a match-language operator
         ],
         is_nullable: true,
     }))
@@ -236,9 +237,9 @@ test_rewrite_to_match_language_no_op!(
     filter_stage(Expression::ScalarFunction(ScalarFunctionApplication {
         function: ScalarFunction::Or,
         args: vec![
-            valid_is(),   // rewritable
-            valid_like(), // rewritable
-            comp_expr(),  // not rewritable - invalid expression
+            valid_is(),     // rewritable
+            valid_like(),   // rewritable
+            invalid_expr(), // not rewritable - Add is not a match-language operator
         ],
         is_nullable: true,
     }))
@@ -498,3 +499,114 @@ test_rewrite_to_match_language_no_op!(
         is_nullable: false,
     }))
 );
+
+// --- FieldAccess-as-boolean ---
+
+test_rewrite_to_match_language!(
+    rewrite_valid_field_access,
+    expected = match_filter_stage(MatchQuery::Comparison(MatchLanguageComparison {
+        function: MatchLanguageComparisonOp::Eq,
+        input: Some(mir_field_path("foo", vec!["str"])),
+        arg: LiteralValue::Boolean(true),
+        cache: SchemaCache::new(),
+    })),
+    expected_changed = true,
+    input = filter_stage(*mir_field_access("foo", "str", true))
+);
+
+// --- Comparison rewrites ---
+
+test_rewrite_to_match_language!(
+    rewrite_valid_comparison_lt,
+    expected = match_filter_stage(MatchQuery::Comparison(MatchLanguageComparison {
+        function: MatchLanguageComparisonOp::Lt,
+        input: Some(mir_field_path("foo", vec!["int"])),
+        arg: LiteralValue::Integer(10),
+        cache: SchemaCache::new(),
+    })),
+    expected_changed = true,
+    input = filter_stage(Expression::ScalarFunction(ScalarFunctionApplication::new(
+        ScalarFunction::Lt,
+        vec![
+            *mir_field_access("foo", "int", true),
+            Expression::Literal(LiteralValue::Integer(10)),
+        ],
+    )))
+);
+
+test_rewrite_to_match_language!(
+    rewrite_valid_comparison_eq,
+    expected = match_filter_stage(MatchQuery::Comparison(MatchLanguageComparison {
+        function: MatchLanguageComparisonOp::Eq,
+        input: Some(mir_field_path("foo", vec!["int"])),
+        arg: LiteralValue::Integer(5),
+        cache: SchemaCache::new(),
+    })),
+    expected_changed = true,
+    input = filter_stage(Expression::ScalarFunction(ScalarFunctionApplication::new(
+        ScalarFunction::Eq,
+        vec![
+            *mir_field_access("foo", "int", true),
+            Expression::Literal(LiteralValue::Integer(5)),
+        ],
+    )))
+);
+
+// --- NOT rewrites ---
+
+test_rewrite_to_match_language!(
+    rewrite_valid_not_is,
+    expected = match_filter_stage(MatchQuery::Logical(MatchLanguageLogical {
+        op: MatchLanguageLogicalOp::Not,
+        args: vec![valid_match_is()],
+        cache: SchemaCache::new(),
+    })),
+    expected_changed = true,
+    input = filter_stage(Expression::ScalarFunction(ScalarFunctionApplication::new(
+        ScalarFunction::Not,
+        vec![valid_is()],
+    )))
+);
+
+test_rewrite_to_match_language!(
+    rewrite_valid_not_like,
+    expected = match_filter_stage(MatchQuery::Logical(MatchLanguageLogical {
+        op: MatchLanguageLogicalOp::Not,
+        args: vec![valid_match_like()],
+        cache: SchemaCache::new(),
+    })),
+    expected_changed = true,
+    input = filter_stage(Expression::ScalarFunction(ScalarFunctionApplication::new(
+        ScalarFunction::Not,
+        vec![valid_like()],
+    )))
+);
+
+// NOT foo.str — inner FieldAccess rewrites to Eq(str, true), then wrapped in Not.
+test_rewrite_to_match_language!(
+    rewrite_valid_not_field,
+    expected = match_filter_stage(MatchQuery::Logical(MatchLanguageLogical {
+        op: MatchLanguageLogicalOp::Not,
+        args: vec![MatchQuery::Comparison(MatchLanguageComparison {
+            function: MatchLanguageComparisonOp::Eq,
+            input: Some(mir_field_path("foo", vec!["str"])),
+            arg: LiteralValue::Boolean(true),
+            cache: SchemaCache::new(),
+        })],
+        cache: SchemaCache::new(),
+    })),
+    expected_changed = true,
+    input = filter_stage(Expression::ScalarFunction(ScalarFunctionApplication::new(
+        ScalarFunction::Not,
+        vec![*mir_field_access("foo", "str", true)],
+    )))
+);
+
+// NOT with a non-rewritable inner expression stays as a Filter.
+test_rewrite_to_match_language_no_op!(
+    cannot_rewrite_not_if_inner_not_rewritable,
+    filter_stage(Expression::ScalarFunction(ScalarFunctionApplication::new(
+        ScalarFunction::Not,
+        vec![invalid_expr()],
+    )))
+);

mongosql/src/translator/match_query.rs

@@ -62,6 +62,13 @@ impl MqlTranslator {
         Ok(match l.op {
             mir::MatchLanguageLogicalOp::Or => air::MatchQuery::Or(args),
             mir::MatchLanguageLogicalOp::And => air::MatchQuery::And(args),
+            mir::MatchLanguageLogicalOp::Not => {
+                let inner = args
+                    .into_iter()
+                    .next()
+                    .expect("MatchLanguageLogicalOp::Not requires exactly one argument");
+                air::MatchQuery::Not(Box::new(inner))
+            }
         })
     }