Arithmetic expression support

Sources/MongoDBModel/FunctionEvaluation.swift

@@ -19,6 +19,13 @@ internal extension FetchRequest {
         if predicate?.containsFunction == true {
             return true
         }
+        // `$divide` aborts the whole query on division by zero and always produces
+        // floating point (no truncating integer division), and `$mod` diverges on
+        // floats — so division and remainder run in memory, where the semantics
+        // exactly match CoreModel's evaluation engine.
+        if predicate?.containsMemoryOnlyArithmetic == true {
+            return true
+        }
         return sortDescriptors.contains { descriptor in
             if case .function = descriptor.term { return true } else { return false }
         }
@@ -47,8 +54,9 @@ internal extension FetchRequest.Predicate {
         case .value:
             return self
         case let .comparison(comparison):
-            let usesFunction = comparison.left.containsFunction || comparison.right.containsFunction
-            return usesFunction ? .value(true) : self
+            let requiresMemory = comparison.left.containsFunction || comparison.right.containsFunction
+                || comparison.left.containsMemoryOnlyArithmetic || comparison.right.containsMemoryOnlyArithmetic
+            return requiresMemory ? .value(true) : self
         case let .compound(compound):
             switch compound {
             case let .and(subpredicates):
@@ -62,8 +70,39 @@ internal extension FetchRequest.Predicate {
     }
 }
 
+internal extension FetchRequest.Predicate {
+
+    /// Whether this predicate contains an arithmetic operation that must run in memory.
+    var containsMemoryOnlyArithmetic: Bool {
+        switch self {
+        case .value:
+            return false
+        case let .comparison(comparison):
+            return comparison.left.containsMemoryOnlyArithmetic || comparison.right.containsMemoryOnlyArithmetic
+        case let .compound(compound):
+            return compound.subpredicates.contains { $0.containsMemoryOnlyArithmetic }
+        }
+    }
+}
+
 internal extension FetchRequest.Predicate.Expression {
 
+    /// Whether this expression contains a division or remainder, which the server
+    /// cannot evaluate with CoreModel's semantics — see `requiresInMemoryEvaluation`.
+    var containsMemoryOnlyArithmetic: Bool {
+        switch self {
+        case let .arithmetic(expression):
+            if expression.function == .divide || expression.function == .modulus {
+                return true
+            }
+            return expression.left.containsMemoryOnlyArithmetic || expression.right.containsMemoryOnlyArithmetic
+        case let .function(function):
+            return function.arguments.contains { $0.containsMemoryOnlyArithmetic }
+        case .attribute, .relationship, .keyPath:
+            return false
+        }
+    }
+
     var containsFunction: Bool {
         switch self {
         case .function:
@@ -139,9 +178,10 @@ internal extension FetchRequest.Predicate.Expression {
         case .relationship:
             // relationships aren't compared by the in-memory function path
             return nil
-        case .arithmetic:
-            // - TODO: Evaluate arithmetic expressions in the in-memory fallback.
-            return nil
+        case let .arithmetic(arithmetic):
+            let lhs = arithmetic.left.evaluate(with: data, functions: functions)
+            let rhs = arithmetic.right.evaluate(with: data, functions: functions)
+            return AttributeValue.arithmetic(arithmetic.function, lhs, rhs)
         }
     }
 }
@@ -220,6 +260,56 @@ private extension AttributeValue {
         return nil
     }
 
+    /// An integer representation for integer value types, for integer arithmetic.
+    var integerValue: Int64? {
+        switch self {
+        case let .int16(value):     return Int64(value)
+        case let .int32(value):     return Int64(value)
+        case let .int64(value):     return value
+        default:                    return nil
+        }
+    }
+
+    /// Apply an arithmetic function to two values.
+    ///
+    /// Mirrors CoreModel's in-memory engine (which is internal to that module):
+    /// integer operands stay in integer arithmetic, so division truncates
+    /// (`7 / 2` is `3`); mixed or floating-point operands compute as `Double`.
+    /// Returns `nil` for non-numeric operands, division or remainder by zero,
+    /// an overflowing division, or a floating-point remainder.
+    static func arithmetic(
+        _ function: FetchRequest.Predicate.ArithmeticExpression.Function,
+        _ lhs: AttributeValue?,
+        _ rhs: AttributeValue?
+    ) -> AttributeValue? {
+        guard let lhs, let rhs else { return nil }
+        if let leftInteger = lhs.integerValue, let rightInteger = rhs.integerValue {
+            switch function {
+            case .add:      return .int64(leftInteger &+ rightInteger)
+            case .subtract: return .int64(leftInteger &- rightInteger)
+            case .multiply: return .int64(leftInteger &* rightInteger)
+            case .divide:
+                guard rightInteger != 0 else { return nil }
+                let (quotient, overflow) = leftInteger.dividedReportingOverflow(by: rightInteger)
+                return overflow ? nil : .int64(quotient)
+            case .modulus:
+                guard rightInteger != 0 else { return nil }
+                let (remainder, overflow) = leftInteger.remainderReportingOverflow(dividingBy: rightInteger)
+                return overflow ? nil : .int64(remainder)
+            }
+        }
+        guard let leftNumber = lhs.comparableDouble, let rightNumber = rhs.comparableDouble else {
+            return nil
+        }
+        switch function {
+        case .add:      return .double(leftNumber + rightNumber)
+        case .subtract: return .double(leftNumber - rightNumber)
+        case .multiply: return .double(leftNumber * rightNumber)
+        case .divide:   return rightNumber == 0 ? nil : .double(leftNumber / rightNumber)
+        case .modulus:  return nil // integers only
+        }
+    }
+
     static func areEqual(_ lhs: AttributeValue?, _ rhs: AttributeValue?, caseInsensitive: Bool) -> Bool {
         switch (lhs, rhs) {
         case (.none, .none), (.some(.null), .none), (.none, .some(.null)), (.some(.null), .some(.null)):

Sources/MongoDBModel/Predicate.swift

@@ -47,6 +47,14 @@ public extension BSONDocument {
 public extension BSONDocument {
 
     init?(predicate: FetchRequest.Predicate.Comparison) {
+        // `(a <op> b) <operator> constant` compiles to an aggregation `$expr` filter.
+        if case let .arithmetic(arithmetic) = predicate.left {
+            guard let document = BSONDocument(arithmetic: arithmetic, comparison: predicate) else {
+                return nil
+            }
+            self = document
+            return
+        }
         // { <field>: { $eq: <value> } }
         guard case let .keyPath(keyPath) = predicate.left,
               let comparisonOperator = ComparisonQueryOperator(predicate: predicate.type),
@@ -71,3 +79,85 @@ public extension BSONDocument {
             keyPath.rawValue: .document([comparisonOperator.rawValue: valueBSON])]
     }
 }
+
+// MARK: - Arithmetic
+
+internal extension BSONDocument {
+
+    /// The `$expr` filter for an `arithmetic <operator> constant` comparison.
+    ///
+    /// ```
+    /// { $expr: { $and: [ { $isNumber: A }, { $gt: [ A, constant ] } ] } }
+    /// ```
+    ///
+    /// The `$isNumber` guard makes a missing or null operand fail the comparison —
+    /// `$add` of a missing field yields `null`, and without the guard aggregation
+    /// comparisons would rank `null` *below* every number, so `arithmetic < constant`
+    /// would spuriously match. CoreModel's engine yields `nil` for such rows, which
+    /// matches no comparison.
+    ///
+    /// Only `.add`, `.subtract` and `.multiply` are translated; `.divide` and
+    /// `.modulus` are routed to in-memory evaluation before this is reached — see
+    /// `FetchRequest.requiresInMemoryEvaluation`.
+    init?(
+        arithmetic: FetchRequest.Predicate.ArithmeticExpression,
+        comparison: FetchRequest.Predicate.Comparison
+    ) {
+        guard comparison.options.isEmpty,
+              comparison.modifier == nil,
+              let comparisonOperator = ComparisonQueryOperator(predicate: comparison.type),
+              // only aggregation comparison operators are valid inside `$expr`
+              [.equalTo, .notEqualTo, .greaterThan, .greaterThanOrEqualTo, .lessThan, .lessThanOrEqualTo].contains(comparisonOperator),
+              case let .attribute(value) = comparison.right,
+              let constant = try? BSON(attributeValue: value),
+              let expression = BSON(aggregation: .arithmetic(arithmetic)) else {
+            return nil
+        }
+        self = [
+            "$expr": .document([
+                "$and": .array([
+                    .document(["$isNumber": expression]),
+                    .document([comparisonOperator.rawValue: .array([expression, constant])])
+                ])
+            ])
+        ]
+    }
+}
+
+internal extension BSON {
+
+    /// The aggregation expression for a predicate expression, for use inside `$expr`.
+    ///
+    /// A key path becomes a `$`-prefixed field path — dotted paths address composite
+    /// attribute elements natively. Returns `nil` for expressions the server cannot
+    /// evaluate (custom functions, relationships, division and remainder).
+    init?(aggregation expression: FetchRequest.Predicate.Expression) {
+        switch expression {
+        case let .attribute(value):
+            guard let bson = try? BSON(attributeValue: value) else {
+                return nil
+            }
+            self = bson
+        case let .keyPath(keyPath):
+            self = .string("$" + keyPath.rawValue)
+        case let .arithmetic(arithmetic):
+            let aggregationOperator: String
+            switch arithmetic.function {
+            case .add:      aggregationOperator = "$add"
+            case .subtract: aggregationOperator = "$subtract"
+            case .multiply: aggregationOperator = "$multiply"
+            case .divide, .modulus:
+                // `$divide` aborts the query on division by zero and has no truncating
+                // integer form; `$mod` diverges on floats. Evaluated in memory instead.
+                return nil
+            }
+            guard let left = BSON(aggregation: arithmetic.left),
+                  let right = BSON(aggregation: arithmetic.right) else {
+                return nil
+            }
+            self = .document([aggregationOperator: .array([left, right])])
+        case .function, .relationship:
+            return nil
+        }
+    }
+}