Improve PE to execute against comprehensions dependent on unknowns

[tsandall]

In some cases, it would be beneficial to run PE on comprehension bodies, however, today, OPA will refuse to PE comprehension bodies if they depend on unknowns because the result may be not easily generated. For instance, if PE were to generate multiple output queries for the comprehension body, those would need to be lifted into rules, which may not even be representable in Rego today (e.g,. as partial sets requiring arguments.) That said, it would be nice if PE could be applied to comprehension bodies for simple cases like inlining references to constants.

Take the following file as an example (x.rego):

package x

p1 {
    v = input[i]
    v[q] = r
}

p2[1] {
    v = input[i]
    v[q] = r
}

p3 := {1 |
    v = input[i]
    v[q] = r
}

q = 1
r = 2

In this case, p1, p2, and p3 are almost identical rules. However, because p3 uses a comprehension, PE gives up and the result contains the unmodified body:

$ opa eval -d x.rego 'data.x.p1' -p -f source
# Query 1
2 = input[i1][1]

$ opa eval -d x.rego 'data.x.p2' -p -f source
# Query 1
data.partial.x.p2

# Module 1
package partial.x

p2[1] {
        2 = input[i1][1]
}

$ opa eval -d x.rego 'data.x.p3' -p -f source
# Query 1
_term_0_0 = {1 |
        v1 = input[i1]
        __local3__1 = data.x.q
        v1[__local3__1] = data.x.r
}

_term_0_0

In this case, there's no reason why PE could not recurse into the comprehension body and produce an output that can be easily represented:

_term_0_0 = {1 |
        2 = input[i1][1]
}

_term_0_0

One approach could be to run PE on the comprehension body and simply throw away the result, saving the original, if there were multiple outputs. Otherwise, the output of PE on the comprehension body could replace the original body. The implementation would need to ensure that the head of the comprehension is still safe.

[stale[bot]]

This issue has been automatically marked as inactive because it has not had any activity in the last 30 days.