Rego - how to mimic set generation using functions

Question

I have a rule that I expect to be reused by a variety of modules. I figured, let's turn that into a function, have the modules pass their input into a function and use a set comprehension like approach, but I'm running into the "functions must not produce multiple outputs for same inputs" error.

Here's a contrived example of what I want to accomplish. I'm thinking I'm going about this the wrong way and there's another approach to this type of problem in Rego.

Classic generator:

arr = [1,2,3,4]

result[entry] {
    itm := arr[i]
    r := itm % 2
    r == 0
    entry := { "type": "even", "val": itm }
}
result[entry] {
    itm := arr[i]
    r := itm % 2
    r == 1
    entry := { "type": "odd", "val": itm }
}

This works as expected.

          "result": [
            {
              "type": "even",
              "val": 2
            },
            {
              "type": "even",
              "val": 4
            },
            {
              "type": "odd",
              "val": 1
            },
            {
              "type": "odd",
              "val": 3
            }
          ]

Here's the function approach that will trigger this error. I am passing that t_label variable to give the function some argument, but it's not really important.

f(t_label) := q {
    q := [ entry | itm := arr[i]
           r := itm % 2
           r == 0
           entry := { t_label: "even", "val": itm }
           ]
}
f(t_label) := q {
    q := [ entry | itm := arr[i]
           r := itm % 2
           r == 1
           entry := { t_label: "odd", "val": itm }
           ]
}

Is this a thing that is done? How is this problem generally approached using Rego?

Answer 1

You're right — unlike rules, functions can't be partial. If you really need something like that for a function you could either have a function to aggregate the result of two (or more) other function calls:

even(t_label) := {entry |
    itm := arr[_]
    itm % 2 == 0
    entry := {t_label: "even", "val": itm}
}

odd(t_label) := {entry |
    itm := arr[_]
    itm % 2 == 1
    entry := {t_label: "odd", "val": itm}
}

f(t_label) := even(t_label) | odd(t_label)

both := f("foo")

For your particlar example though, I think we'd be excused using a little "or hack" using map based branching:

f(t_label) := {entry |
    itm := arr[_]
    entry := {t_label: {0: "even", 1: "odd"}[itm % 2], "val": itm}
}