builtins: product, sort, sum

Also: implementing internal.member_3 for compliance testing

Resolves #12

Signed-off-by: Dmitry Frenkel <[email protected]>

Author: DFrenkel

Sources/Rego/Builtins/Aggregates.swift

@@ -9,7 +9,7 @@ extension BuiltinFuncs {
 
         guard let len = args[0].count else {
             throw BuiltinError.argumentTypeMismatch(
-                arg: "collection", got: args[0].typeName, want: "string|array|object|set")
+                arg: "collection", got: args[0].typeName, want: "any<array, object, set, string>")
         }
 
         return .number(NSNumber(value: len))
@@ -27,7 +27,7 @@ extension BuiltinFuncs {
             return s.max() ?? .undefined
         default:
             throw BuiltinError.argumentTypeMismatch(
-                arg: "collection", got: args[0].typeName, want: "array|set")
+                arg: "collection", got: args[0].typeName, want: "any<array[any], set[any]>")
         }
     }
 
@@ -43,7 +43,74 @@ extension BuiltinFuncs {
             return s.min() ?? .undefined
         default:
             throw BuiltinError.argumentTypeMismatch(
-                arg: "collection", got: args[0].typeName, want: "array|set")
+                arg: "collection", got: args[0].typeName, want: "any<array[any], set[any]>")
+        }
+    }
+
+    static func sort(ctx: BuiltinContext, args: [AST.RegoValue]) async throws -> AST.RegoValue {
+        guard args.count == 1 else {
+            throw BuiltinError.argumentCountMismatch(got: args.count, want: 1)
+        }
+
+        switch args[0] {
+        case .array(let a):
+            return .array(a.sorted())
+        case .set(let s):
+            return .array(s.sorted())
+        default:
+            throw BuiltinError.argumentTypeMismatch(
+                arg: "collection", got: args[0].typeName, want: "any<array[any], set[any]>")
+        }
+    }
+
+    static func sum(ctx: BuiltinContext, args: [AST.RegoValue]) async throws -> AST.RegoValue {
+        return try doReduce(args: args, initialValue: .number(0), op: BuiltinFuncs.doPlus)
+    }
+
+    static func product(ctx: BuiltinContext, args: [AST.RegoValue]) async throws -> AST.RegoValue {
+        return try doReduce(args: args, initialValue: .number(1), op: BuiltinFuncs.doMul)
+    }
+
+    /// Returns reduction over an array or set of RegoValues with a given async Builtin being used as an reducer operation.
+    /// Returns the normalized metric unit symbol for a given symbol.
+    /// - Parameters:
+    ///   - ctx: The builtin context.
+    ///   - args: The arguments to reduce.
+    ///   - initialValue: The initial value to start with.
+    ///   - op: The Arithmetic builtin operation to be applied to the partial result an the next value in the sequence to produce the next result.
+    private static func doReduce(
+        args: [AST.RegoValue],
+        initialValue: AST.RegoValue,
+        op: ([AST.RegoValue]) throws -> AST.RegoValue
+    ) throws -> AST.RegoValue {
+        guard args.count == 1 else {
+            throw BuiltinError.argumentCountMismatch(got: args.count, want: 1)
+        }
+
+        // We will iterate over this sequence
+        var sequence: any Sequence<RegoValue>
+        switch args[0] {
+        case .array(let a):
+            sequence = a
+        case .set(let s):
+            sequence = s
+        default:
+            throw BuiltinError.argumentTypeMismatch(
+                arg: "collection", got: args[0].typeName, want: "any<array[number], set[number]>")
+        }
+
+        do {
+            let result = try sequence.reduce(
+                initialValue,
+                { x, y in
+                    try op([x, y])
+                })
+            return result
+        } catch is RegoError {
+            let receivedTypes = Set(sequence.map({ $0.typeName })).sorted().joined(separator: ", ")
+            throw BuiltinError.argumentTypeMismatch(
+                arg: "collection", got: "\(args[0].typeName)[any<\(receivedTypes)>]",
+                want: "any<array[number], set[number]>")
         }
     }
 }

Sources/Rego/Builtins/Arithmetic.swift

@@ -9,6 +9,10 @@ let _floor = floor
 
 extension BuiltinFuncs {
     static func plus(ctx: BuiltinContext, args: [AST.RegoValue]) async throws -> AST.RegoValue {
+        return try doPlus(args: args)
+    }
+
+    static func doPlus(args: [AST.RegoValue]) throws -> AST.RegoValue {
         guard args.count == 2 else {
             throw BuiltinError.argumentCountMismatch(got: args.count, want: 2)
         }
@@ -48,6 +52,10 @@ extension BuiltinFuncs {
     // Multiplies two numbers.
     // Returns: the product of `x` and `y`
     static func mul(ctx: BuiltinContext, args: [AST.RegoValue]) async throws -> AST.RegoValue {
+        return try doMul(args: args)
+    }
+
+    static func doMul(args: [AST.RegoValue]) throws -> AST.RegoValue {
         guard args.count == 2 else {
             throw BuiltinError.argumentCountMismatch(got: args.count, want: 2)
         }

Sources/Rego/Builtins/Casts.swift

@@ -34,7 +34,7 @@ extension BuiltinFuncs {
             return .number(NSNumber(value: x))
         default:
             throw BuiltinError.argumentTypeMismatch(
-                arg: "x", got: args[0].typeName, want: "boolean|null|number|string")
+                arg: "x", got: args[0].typeName, want: "any<boolean, null, number, string>")
         }
     }
 

Sources/Rego/Builtins/Collections.swift

@@ -21,4 +21,47 @@ extension BuiltinFuncs {
             return .boolean(false)
         }
     }
+
+    // internal.member_3
+    // memberOfWithKey is a membership check with key:
+    // memberOfWithKey(k: any, x: any, y: any) checks if y has property or index k and it is equal to x
+    // For objects, we are checking the keys AND the values.
+    static func isMemberOfWithKey(ctx: BuiltinContext, args: [AST.RegoValue]) async throws -> AST.RegoValue {
+        guard args.count == 3 else {
+            throw BuiltinError.argumentCountMismatch(got: args.count, want: 3)
+        }
+
+        let key = args[0]
+        let value = args[1]
+        // See https://github.com/open-policy-agent/opa/blob/b942136a4ad049262fd72026421dac6bdd705059/v1/topdown/aggregates.go#L247
+        let match = args[2][key]
+        if match != nil {
+            return .boolean(RegoValue.compare(value, match!) == .orderedSame)
+        }
+        return .boolean(false)
+    }
+}
+
+extension RegoValue {
+    /// Some RegoValues implement Get(key) interface. We will just implement it here as a subscript extension.
+    /// See https://github.com/open-policy-agent/opa/blob/7ddaff2cc3dd749af25bab7d6a1f5a9cdbfe9833/v1/ast/term.go#L380
+    fileprivate subscript(key: RegoValue) -> RegoValue? {
+        switch self {
+        case .object(let o):
+            return o[key]
+        case .array(let a):
+            // key must be an integer position
+            guard !key.isFloat, let index = key.integerValue else { return nil }
+            // check bounds
+            guard index >= 0 && index < a.count && index < Int32.max else { return nil }
+            return a[Int(index)]
+        case .set(let s):
+            if s.contains(key) {
+                return key
+            }
+            return nil
+        default:
+            return nil
+        }
+    }
 }

Sources/Rego/Builtins/Registry.swift

@@ -38,6 +38,9 @@ public struct BuiltinRegistry: Sendable {
             "count": BuiltinFuncs.count,
             "max": BuiltinFuncs.max,
             "min": BuiltinFuncs.min,
+            "product": BuiltinFuncs.product,
+            "sort": BuiltinFuncs.sort,
+            "sum": BuiltinFuncs.sum,
 
             // Arithmetic
             "plus": BuiltinFuncs.plus,
@@ -65,6 +68,7 @@ public struct BuiltinRegistry: Sendable {
 
             // Collections
             "internal.member_2": BuiltinFuncs.isMemberOf,
+            "internal.member_3": BuiltinFuncs.isMemberOfWithKey,
 
             // Comparison
             "gt": BuiltinFuncs.greaterThan,