perf: memoize IR plan calls to reduce reevaluation cost #36

shomron · 2025-10-20T20:57:49Z

Adds a per-evaluation call memoization stack to the IREvaluationContext. This currently supports rules, not arbitrary functions, based on arity. For a policy with many repeated calls, we see evaluation time drop from ~241ms to ~2ms. No regression in compliance tests.

anderseknert

Very nice! Can't wait to try this out.

anderseknert · 2025-10-20T22:00:21Z

Sources/Rego/IREvaluator.swift

    isDynamic: Bool = false
 ) async throws -> AST.RegoValue {
+    // Check memo cache if applicable to save repeated evaluation time for rules
+    let shouldMemoize = args.count == 2 // Currently support _rules_, not _functions_


// Currently support _rules_, not _functions_

Perhaps add a line or two here to explain why this is for anyone without that context. Which could very well be us a year or two from now :)

I'll leave it to @srenatus to provide additional context on this - I derived this from the WASM backend logic but do not understand the full nuance yet.

Would also be interested in the context here! :)

Taking a step back, in the IR, there are only functions, no (complete/partial or single-value/multi-value) rules as you know in Rego. The first two arguments are always present, it's data and input. Addition arguments come in through two means:

it's a user-defined (rego) function, so it has extra args

it's subject to a certain kind of with-replacement of (rego, builtin or user-defined) functions with a local variable. Those variables are turned into extra args, as otherwise the information wouldn't be available down the call stack where the with-replaced function might be used.

It's fair to say that (2.) is rather uncommon. You might do that in tests, but it's a concern for production policies, I would say.

Within a single eval of a policy (IR or Rego), data and input is static -- it won't change in between. So if an IR function only has two arguments, those will always be data and input, it's going to be a (Rego) rule. Caching these requires no further tracking of the values of extra arguments. If we, for example, wanted to cache this user-defined function,

allow(a, b) if { ... }

we'd need to put an element into the cache for any values of a and b we encounter. This can be undesirable, as you might call the function often with varying arguments -- that's why it's a function in the first place.

Caching only rules (IR functions with two arguments) is thus a middle ground between memory usage and performance improvements.

What about with you might ask, that can change input and data in arbitrary ways. It can, but when the planner in OPA encounters a with replacement, it'll increase a generation counter and start planning new functions. E.g.

package t allow if { # planned as `g0.data.t.allow` something # gets planned as `g0.data.t.something` something with input.foo as "bar" # this gets planned as `g1.data.t.something` } something if other_rule other_rule if input.foo == "bar" # both planned as `g0.data.t.other_rule` and as `g1.data.t.other_rule`, called accordingly in the function `g0.data.t.allow`

This also ensures that a simple cache mechanism, caching the values of the IR-names of the functions -- cache keys like g0.data.t.something and g1.data.t.something, not data.t.something -- will yield the correct results.

Does that help for context? :)

Great write down, thanks for making the effort!

Certainly! Linking to this comment would be great :)

This is the first time I've ever seen the little g0.data.x prefix explained, and that makes a lot of sense now that I know what it's for! 😄 I'd previously assumed it was an arbitrary compiler prefix, probably because I'd never plugged in a policy that would have generated a g1.data.x function variant.

philippzagar

Awesome to see this finally coming together, and I’m excited to see the performance gains in action! 🚀

philippzagar · 2025-10-21T06:50:11Z

Sources/Rego/Memo.swift

+/// Note we capture the arguments as unresolved operands and not resolved values,
+/// as hashing the values was proving extremely expensive. We instead rely on the
+/// invariant the the plan / evaluator will not modify a local after it has been initally set.
+internal struct InvocationKey: Hashable {


There’s no need to explicitly mark symbols as internal in Swift, it’s the default access level. Adding it for the sake of being explicit isn’t common practice (rarely done in Swift Server projects) and is generally discouraged.
This applies to all similar changes in this PR.

philippzagar · 2025-10-21T06:53:54Z

Sources/Rego/IREvaluator.swift

    isDynamic: Bool = false
 ) async throws -> AST.RegoValue {
+    // Check memo cache if applicable to save repeated evaluation time for rules
+    let shouldMemoize = args.count == 2 // Currently support _rules_, not _functions_


Would also be interested in the context here! :)

philippzagar · 2025-10-21T06:54:43Z

Sources/Rego/IREvaluator.swift

-                block: stmt.block
-            )

+            let blockResult = try await ctx.callMemo.withPush {


Sweet, I like the with-style syntax in this use case, great idea that prohibits forgetting to pop.

philippzagar · 2025-10-21T07:01:38Z

Sources/Rego/IREvaluator.swift

+    init(ctx: EvaluationContext, policy: IndexedIRPolicy) {
+        self.ctx = ctx
+        self.policy = policy
+        self.callMemo = MemoStack.init()


I’d suggest moving this initialization directly into the variable declaration for consistency with the callDepth and maxCallDepth initializations. That would also make the explicit initializer unnecessary, but I don’t feel strongly about it.

philippzagar · 2025-10-21T18:58:20Z

Sources/Rego/Memo.swift

+    /// withPush returns the result of calling the provided closure with
+    /// a fresh memoCache pushed on the stack. The memoCache will only be
+    /// active during that call, and discarded when it completes.
+    func withPush<T>(_ body: () async throws -> T) async rethrows -> T {


You could use sending T here, but it doesn't make a big difference.

Also, probably a good idea to explicitly pass down the isolation domain via:
isolation: isolated (any Actor)? = #isolation,
so that we avoid an actor hop.

It's not clear to me that those are needed here? We're not doing an actor hop that I can see, and this stack is only used within a single evaluation, not across evaluations. Please let me know if I'm misunderstanding 🤔

Any method without isolation inheritance is going to hop defensively. The recommendation with Swift 6.2 is to adopt the upcoming language feature NonIsolatedNonSendingByDefault which makes all async methods nonisolated(nonsending)

@shomron That shouldn't be a huge performance loss in this case, however this applies to the entirety of swift-opa and should be kept in mind

philippzagar · 2025-10-21T19:00:05Z

Sources/Rego/Memo.swift

+        }
+    }
+
+    func push() {


Suggestion: Mark these funcs as @inlineable, should enable the compiler to inline these calls.

Thanks! I suggest we revisit this optimization if we see any measurable impact - my hunch is that we will not see excessive with statements in the hot path to the extent that a function call will be an issue.

Adds a per-evaluation call memoization stack to the IREvaluationContext. This currently supports rules, not arbitrary functions, based on arity. For a policy with many repeated calls, we see evaluation time drop from ~241ms to ~2ms. No regression in compliance tests. Signed-off-by: Oren Shomron <[email protected]>

srenatus

Approval with limited Swift knowledge ⚠️ Code change looks good in isolation.

shomron force-pushed the perf-memoize-ir-calls branch from 4a15eb5 to 7d55262 Compare October 20, 2025 20:58

shomron requested review from patrick-east and srenatus October 20, 2025 21:07

anderseknert previously approved these changes Oct 20, 2025

View reviewed changes

philippzagar approved these changes Oct 21, 2025

View reviewed changes

shomron dismissed anderseknert’s stale review via 4de8294 October 21, 2025 17:18

shomron force-pushed the perf-memoize-ir-calls branch from 7d55262 to 4de8294 Compare October 21, 2025 17:18

philippzagar approved these changes Oct 21, 2025

View reviewed changes

shomron force-pushed the perf-memoize-ir-calls branch from 4de8294 to 3811c03 Compare October 29, 2025 17:01

srenatus approved these changes Oct 29, 2025

View reviewed changes

shomron merged commit 29d8199 into open-policy-agent:main Oct 29, 2025
3 checks passed

shomron deleted the perf-memoize-ir-calls branch October 29, 2025 18:43

perf: memoize IR plan calls to reduce reevaluation cost #36

perf: memoize IR plan calls to reduce reevaluation cost #36

Uh oh!

Conversation

shomron commented Oct 20, 2025

Uh oh!

anderseknert left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

philippzagar Oct 21, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

philippzagar left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

philippzagar Oct 21, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

philippzagar Oct 21, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

srenatus left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

philippzagar Oct 21, 2025 •

edited

Loading

philippzagar Oct 21, 2025 •

edited

Loading

philippzagar Oct 21, 2025 •

edited

Loading

srenatus left a comment •

edited

Loading