OPA Rego collection helpers

collection.rego

package collection

# Map

my_add(x, y) = x + y

map_add[x] = val {
	col := input.col
	delta := input.delta

	# rego doesn't seem to support higher-order function, thus, hard-coding
	val := my_add(col[x], delta)
}

test_map_add {
	3 == my_add(1, 2)

	some i
	mapped := map_add with input as {
		"col": [1, 2, 3, 4],
		"delta": 2,
	}

	[3, 4, 5, 6][i] == mapped[i]
	[8, 4, 7, 6][i] != mapped[i]
}

map_each_mul[x] = val {
	col := input.col
	delta := input.delta

	val := col[x] * delta[x]
}

test_map_each_mul {
	2 == 1 * 2

	some i
	mapped := map_each_mul with input as {
		"col": [1, 2, 3, 4],
		"delta": [1, 10, 100, 1000],
	}

	[1, 20, 300, 4000][i] == mapped[i]
	[8, 4, 7, 6][i] != mapped[i]
}



# Filter
# there is actually this way too, https://github.com/open-policy-agent/opa/issues/3155#issuecomment-780159138

filter[x] = val {
	col := input.col
	filter_with := input.filter_with
	col[x] != filter_with

	val := col[x]
}

test_filter {
	some i
	filtered := filter with input as {
		"col": [1, null, 2, 3, null, 4],
		"filter_with": null,
	}

	[1, 2, 3, 4][i] == filtered[i]
	[5, 2, 5, 4][i] != filtered[i]
}

test_both {
	some i

	filtered := filter with input as {
		"col": [1, null, 2, 3, null, 4],
		"filter_with": null,
	}

	[1, 2, 3, 4][i] == filtered[i]
	[5, 2, 5, 4][i] != filtered[i]

	mapped := map_add with input as {
		"col": filtered,
		"delta": 2,
	}

	[3, 4, 5, 6][i] == mapped[i]
	[8, 4, 7, 6][i] != mapped[i]
}

# Reduce

# Sadly, I don't think it's possible with Rego but if you know how, please let me know


# Helpers

# because many times your function or rule's result could be object rather than array
# using https://www.openpolicyagent.org/docs/latest/policy-language/#array-comprehensions
object_to_array(obj) = [l[_] | l := obj]

test_object_to_array {
	[0, 1] == object_to_array({
		"0": 0,
		"1": 1,
	})
	
	[0, 1] == object_to_array({
		"2": 0,
		"4": 1,
	})
}


array_reverse(arr) = [reversed |
	some i

	# this line is crucial to make the whole expression safe
	arr[i]
	revered_i := (count(arr) - i) - 1
	reversed := arr[revered_i]
]

test_array_reverse {
	array_reverse([1, 2]) == [2, 1]
	array_reverse([5, 4, 3, 2, 1]) == [1, 2, 3, 4, 5]
}

power.rego

package power

# start with a base document since rego can't handle unlimited set of data
array_length_11 = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]

# now this kind of function is possible
identity[n] = result {
	result := array_length_11[n] + n
}

test_identity {
	identity[3] == 3
	identity[10] == 10
  
	not identity[11] == 11
	not identity[99] == 99
}

# just a helper function to DRY the code below
slice(x) = array.slice(array_length_11, 0, x)

# fill_as_obj(x, n) = r {
# 	r := map_add with input as {
# 		"col": slice(n),
# 		"delta": x,
# 	}
# }

# test_fill_as_obj {
# 	slice(2) == [0, 0]
# 	fill_as_obj(5, 3)[_] == 5
# }

# should be same as `fill(x, n) = [r | r := fill_as_obj(x, n)[_]]`
fill(x, n) = [r[_] |
	r := map_add with input as {
		"col": slice(n),
		"delta": x,
	}
]

test_fill {
	some i
	fill(10, 3)[i] == 10
	product(fill(10, 3)) == 1000
}

pow(x, n) = 1 {
	n == 0
}

pow(x, n) = product(fill(x, n))

# yay, now we can do power!
test_power {
	pow(10, 2) == 100
	pow(10, 3) == 1000
	pow(5, 2) == 25

	# n being 11 is the limit because of array_length_11
	pow(2, 11) == 2048
	not pow(2, 12) == 2048 * 2
}

power_table[[x, n]] = p {
	# this line is crucial to make x and n safe
	array_length_11[x] == array_length_11[n]

	p := pow(x, n)
}

test_power_table {
	some i
	power_table[[2, i]] == [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024][i]
}

power_list(x) = [r | r := power_table[[x, _]]]

test_power_list {
	power_list(2) == [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024]
}

power_range(x, n) = array.slice(power_list(x), 0, n)

test_power_range {
	power_range(2, 5) == [1, 2, 4, 8, 16]
}