Expose builtin cache API

Sources/Rego/Builtins/BuiltinsCache.swift

@@ -1,22 +1,29 @@
 import AST
 import Foundation
 
-/// BuiltinsCache defines the caching strategy used by the top-down evaluation.
-/// Note that Golang implementation uses `type FooCachingKey string` approach to redefine string keys as distinct types.
-/// For Swift implementation, type aliasing does not create a new type, so we will have to use namespaces to separate keys.
-/// This implementation wraps a simple dictionary with composite keys made of namespace + key and values being RegoValues.
-/// Since Dictionary is not concurrency-safe, neither is BuiltinsCache, but since its intent is to be used within a single evaluation,
-/// we do not concurrency support.
-internal struct BuiltinsCache {
-    struct Namespace: Hashable, Sendable {
+/// Per-evaluation cache for OPA builtins.
+///
+/// Lightweight cache for (custom) OPA builtins.
+/// Stores `RegoValue` entries under namespaced string keys to avoid
+/// collisions between different builtins or key spaces.
+///
+/// It is designed for use within a single top-down policy evaluation, but you can also
+/// provide your own cache instance to reuse results between evaluations
+/// (for example, to avoid redundant costly network-bound builtin calls).
+public final class BuiltinsCache: Sendable {
+    /// Key namespace for isolating builtin domains.
+    public struct Namespace: Hashable, Sendable {
         let ns: String
         init(_ ns: String) {
             self.ns = ns
         }
 
+        /// Default namespace.
         public static let global: Namespace = Namespace("__global__")
 
-        // Helper for .namespace syntax within BuiltinsCache.subscript.
+        /// Create a custom namespace.
+        ///
+        /// - Parameter ns: The namespace name.
         public static func namespace(_ ns: String) -> Namespace {
             return Namespace(ns)
         }
@@ -33,32 +40,63 @@ internal struct BuiltinsCache {
         }
     }
 
-    private var cache: [CompositeKey: RegoValue]
+    private let lock = NSLock()
+    private nonisolated(unsafe) var _cache: [CompositeKey: RegoValue]
 
-    internal init() {
-        self.cache = [CompositeKey: RegoValue]()
+    /// Creates a new, empty builtin cache.
+    ///
+    /// The cache stores `RegoValue` entries keyed by a combination of a string and a namespace.
+    /// It is designed for use within a single top-down policy evaluation, but you can also
+    /// provide your own cache instance to reuse results between evaluations
+    /// (for example, to avoid redundant costly network-bound builtin calls).
+    public init() {
+        self._cache = [CompositeKey: RegoValue]()
     }
 
-    internal subscript(key: String, namespace: Namespace = .global) -> RegoValue? {
+    /// Accesses or updates a cached value scoped by a key and optional namespace.
+    ///
+    /// Use this subscript to read or write entries in the builtin cache.
+    /// Setting a value to `nil` removes the corresponding cache entry.
+    ///
+    /// - Parameters:
+    ///   - key: The cache key used to identify the entry within the given namespace.
+    ///   - namespace: The namespace to scope the key under to avoid collisions between different builtins.. Defaults to ``Namespace/global``.
+    ///
+    /// - Note: Writing to the same key multiple times simply overwrites the previous value,
+    ///         the cache does not enforce immutability or write-once semantics.
+    public subscript(key: String, namespace: Namespace = .global) -> RegoValue? {
         get {
-            return self.cache[CompositeKey(key, namespace: namespace)]
+            self.lock.withLock {
+                self._cache[CompositeKey(key, namespace: namespace)]
+            }
         }
         set(newValue) {
             let k = CompositeKey(key, namespace: namespace)
 
             guard let newValue else {
-                self.cache[k] = nil
+                self.lock.withLock {
+                    self._cache[k] = nil
+                }
                 return
             }
-            self.cache[k] = newValue
+            self.lock.withLock {
+                self._cache[k] = newValue
+            }
         }
     }
 
-    internal mutating func removeAll() {
-        self.cache.removeAll()
+    /// Remove all cached entries.
+    public func removeAll() {
+        self.lock.withLock {
+            self._cache.removeAll()
+        }
     }
 
-    internal var count: Int {
-        return self.cache.count
+    /// Current number of entries in cache.
+    public var count: Int {
+        self.lock.withLock {
+            self._cache.count
+        }
     }
 }
+

Sources/Rego/Builtins/Registry.swift

@@ -1,24 +1,38 @@
 import AST
 import Foundation
 
-public typealias Builtin = @Sendable (BuiltinContext, [AST.RegoValue]) async throws -> AST.RegoValue
-
+/// The OPA builtin signature.
+///
+/// Called by the evaluator with the per-evaluation `BuiltinContext` and the
+/// positional, already-resolved Rego arguments.
+///
+/// - Parameters:
+///   - context: Evaluation context (e.g., cache, tracer).
+///   - args: Positional arguments of the builtin as `AST.RegoValue`.
+/// - Returns: The returned value of the builtin as `AST.RegoValue`.
+/// - Throws: If the builtin logic throws.
+public typealias Builtin = @Sendable (_ context: BuiltinContext, _ args: [AST.RegoValue]) async throws -> AST.RegoValue
+
+/// Context available within builtin evaluation.
 public struct BuiltinContext {
+    /// Source location at which the builtin invocation occurred.
     public let location: OPA.Trace.Location
-    public var tracer: OPA.Trace.QueryTracer?
-    internal let cache: Ptr<BuiltinsCache>
+    /// Tracer that records events during evaluation.
+    public let tracer: OPA.Trace.QueryTracer?
+    /// Per-evaluation builtin cache.
+    ///
+    /// Shared by all builtin invocations within a single top-down policy evaluation.
+    public let cache: BuiltinsCache
 
     init(
         location: OPA.Trace.Location = .init(),
         tracer: OPA.Trace.QueryTracer? = nil,
-        cache: Ptr<BuiltinsCache>? = nil
+        cache: BuiltinsCache = .init()
     ) {
         self.location = location
         self.tracer = tracer
-        // Note that we will create a new cache if one hasn't been provided -
-        // some builtin evaluations (e.g. UUID) expect the cache.
-        // In most/all? cases, we expect a shared cache to be passed in here from EvaluationContext
-        self.cache = cache ?? Ptr<BuiltinsCache>(toCopyOf: BuiltinsCache())
+        // Shared cache to be passed from `EvaluationContext` or created on the fly
+        self.cache = cache
     }
 }
 

Sources/Rego/Builtins/Uuid.swift

@@ -17,10 +17,10 @@ extension BuiltinFuncs {
             throw BuiltinError.argumentTypeMismatch(arg: "k", got: args[0].typeName, want: "string")
         }
 
-        let existing: RegoValue? = ctx.cache.v[key, .namespace(uuidNamespace)]
+        let existing: RegoValue? = ctx.cache[key, .namespace(uuidNamespace)]
         guard let existing else {
             let newUUID = RegoValue.string(UUID().uuidString)
-            ctx.cache.v[key, .namespace(uuidNamespace)] = newUUID
+            ctx.cache[key, .namespace(uuidNamespace)] = newUUID
 
             return newUUID
         }

Sources/Rego/Engine.swift

@@ -101,11 +101,15 @@ extension OPA.Engine {
         /// - Parameters:
         ///   - input: The input data to evaluate the query against.
         ///   - tracer: (optional) The tracer to use for this evaluation.
+        ///   - cache: (optional) Per-evaluation builtin cache shared across all builtin invocations within a single top-down policy evaluation.
+        ///            You can provide a custom cache instance to reuse results between evaluations.
+        ///            By default, a new, empty cache is created for each evaluation.
         ///   - strictBuiltins: (optional) Whether to run in strict builtin evaluation mode.
         ///                     In strict mode, builtin errors abort evaluation, rather than returning undefined.
         public func evaluate(
             input: AST.RegoValue = .undefined,
             tracer: OPA.Trace.QueryTracer? = nil,
+            cache: BuiltinsCache = .init(),
             strictBuiltins: Bool = false
         ) async throws -> ResultSet {
             let ctx = EvaluationContext(
@@ -114,6 +118,7 @@ extension OPA.Engine {
                 store: self.store,
                 builtins: self.builtinRegistry,
                 tracer: tracer,
+                cache: cache,
                 strictBuiltins: strictBuiltins
             )
 

Sources/Rego/Evaluator.swift

@@ -13,6 +13,10 @@ internal struct EvaluationContext {
     public var store: OPA.Store
     public var builtins: BuiltinRegistry
     public var tracer: OPA.Trace.QueryTracer?
+    /// Per-evaluation builtin cache.
+    ///
+    /// Shared by all builtin invocations within a single top-down policy evaluation.
+    public let cache: BuiltinsCache
     public var strictBuiltins: Bool = false
 
     init(
@@ -21,13 +25,15 @@ internal struct EvaluationContext {
         store: OPA.Store = NullStore(),
         builtins: BuiltinRegistry = .defaultRegistry,
         tracer: OPA.Trace.QueryTracer? = nil,
+        cache: BuiltinsCache = .init(),
         strictBuiltins: Bool = false
     ) {
         self.query = query
         self.input = input
         self.store = store
         self.builtins = builtins
         self.tracer = tracer
+        self.cache = cache
         self.strictBuiltins = strictBuiltins
     }
 }

Sources/Rego/IREvaluator.swift

@@ -1032,7 +1032,8 @@ private func evalCall(
 
     let bctx = BuiltinContext(
         location: try frame.v.currentLocation(withContext: ctx, stmt: caller),
-        tracer: ctx.ctx.tracer
+        tracer: ctx.ctx.tracer,
+        cache: ctx.ctx.cache
     )
 
     return try await ctx.ctx.builtins.invoke(

Sources/Rego/Tracer.swift

@@ -16,7 +16,7 @@ extension OPA.Trace {
         case note
     }
 
-    /// Describes the type of traceable operation that occured
+    /// Describes the type of traceable operation that occurred
     public enum Operation: String, Codable, Hashable, Sendable {
         // Subset of the Go OPA topdown trace op's, add more as needed
         case enter
@@ -26,7 +26,7 @@ extension OPA.Trace {
         case note
     }
 
-    /// Describes the source location at which a trace event occured
+    /// Describes the source location at which a trace event occurred
     public struct Location: Codable, Hashable, Sendable {
         public var row: Int = 0
         public var col: Int = 0

Tests/RegoTests/BuiltinTests/BuiltinsCacheTests.swift

@@ -4,11 +4,15 @@ import Testing
 
 @testable import Rego
 
-@Suite("BuiltinsCacheTests")
-struct BuiltinsCacheTests {
+extension BuiltinTests {
+    @Suite("BuiltinsCacheTests", .tags(.builtins))
+    struct BuiltinsCacheTests {}
+}
+
+extension BuiltinTests.BuiltinsCacheTests {
     @Test("caching items with the same keys but different namespaces")
     func cachesDifferentNamespacesAsDifferentEntries() {
-        var cache = BuiltinsCache()
+        let cache = BuiltinsCache()
 
         // Set a value in the implied global namespace
         cache["3!"] = 6
@@ -35,7 +39,7 @@ struct BuiltinsCacheTests {
 
     @Test("cache removal and counting")
     func testRemoval() {
-        var cache = BuiltinsCache()
+        let cache = BuiltinsCache()
         cache["3!"] = 6
         #expect(cache.count == 1)
 
@@ -51,4 +55,62 @@ struct BuiltinsCacheTests {
         #expect(cache["3!"] == nil)
         #expect(cache.count == 0)
     }
+
+    static let cacheKey = "test_key"
+    static let builtinRegistry: BuiltinRegistry =
+        .init(
+            builtins: [
+                "custom_cache_builtin_1": { ctx, args in
+                    guard args.count == 1 else {
+                        throw BuiltinError.argumentCountMismatch(got: args.count, want: 1)
+                    }
+
+                    // Write to builtin cache
+                    ctx.cache[Self.cacheKey] = args[0]
+                    // Read from builtin cache
+                    let cacheEntry = try #require(ctx.cache[Self.cacheKey], "Expected entry is not present in cache")
+
+                    return cacheEntry
+                },
+                "custom_cache_builtin_2": { ctx, args in
+                    // Read from builtin cache
+                    let value = try #require(ctx.cache[Self.cacheKey], "Expected entry is not present in cache")
+                    // Reset entry
+                    ctx.cache[Self.cacheKey] = nil
+                    return value
+                }
+            ]
+        )
+
+    @Test("Shared cache across builtin evals")
+    func test() async throws {
+        // One builtin cache that is handed to both builtins
+        let cache = BuiltinsCache()
+        // Dummy value in cache
+        let testCacheValue: RegoValue = .string("test_value")
+
+        let result1 = await Result {
+            try await Self.builtinRegistry.invoke(
+                withContext: BuiltinContext(cache: cache),
+                name: "custom_cache_builtin_1",
+                args: [testCacheValue],
+                strict: true
+            )
+        }
+        try #expect(result1.get() == testCacheValue)
+        #expect(cache.count == 1)
+        #expect(cache[Self.cacheKey] == testCacheValue)
+
+        let result2 = await Result {
+            try await Self.builtinRegistry.invoke(
+                withContext: BuiltinContext(cache: cache),
+                name: "custom_cache_builtin_2",
+                args: [testCacheValue],
+                strict: true
+            )
+        }
+        try #expect(result2.get() == testCacheValue)
+        #expect(cache.count == 0)
+        #expect(cache[Self.cacheKey] == nil)
+    }
 }