Propagate description of custom builtins

This allows processes, such as doc generation, to have additional
information about custom builtin functions.

Fixes: open-policy-agent#6449

Signed-off-by: Luiz Carvalho <lucarval@redhat.com>

rego/rego.go

@@ -599,6 +599,7 @@ type Rego struct {
 // Function represents a built-in function that is callable in Rego.
 type Function struct {
 	Name             string
+	Description      string
 	Decl             *types.Function
 	Memoize          bool
 	Nondeterministic bool
@@ -629,6 +630,7 @@ type (
 func RegisterBuiltin1(decl *Function, impl Builtin1) {
 	ast.RegisterBuiltin(&ast.Builtin{
 		Name:             decl.Name,
+		Description:      decl.Description,
 		Decl:             decl.Decl,
 		Nondeterministic: decl.Nondeterministic,
 	})
@@ -642,6 +644,7 @@ func RegisterBuiltin1(decl *Function, impl Builtin1) {
 func RegisterBuiltin2(decl *Function, impl Builtin2) {
 	ast.RegisterBuiltin(&ast.Builtin{
 		Name:             decl.Name,
+		Description:      decl.Description,
 		Decl:             decl.Decl,
 		Nondeterministic: decl.Nondeterministic,
 	})
@@ -655,6 +658,7 @@ func RegisterBuiltin2(decl *Function, impl Builtin2) {
 func RegisterBuiltin3(decl *Function, impl Builtin3) {
 	ast.RegisterBuiltin(&ast.Builtin{
 		Name:             decl.Name,
+		Description:      decl.Description,
 		Decl:             decl.Decl,
 		Nondeterministic: decl.Nondeterministic,
 	})
@@ -668,6 +672,7 @@ func RegisterBuiltin3(decl *Function, impl Builtin3) {
 func RegisterBuiltin4(decl *Function, impl Builtin4) {
 	ast.RegisterBuiltin(&ast.Builtin{
 		Name:             decl.Name,
+		Description:      decl.Description,
 		Decl:             decl.Decl,
 		Nondeterministic: decl.Nondeterministic,
 	})
@@ -681,6 +686,7 @@ func RegisterBuiltin4(decl *Function, impl Builtin4) {
 func RegisterBuiltinDyn(decl *Function, impl BuiltinDyn) {
 	ast.RegisterBuiltin(&ast.Builtin{
 		Name:             decl.Name,
+		Description:      decl.Description,
 		Decl:             decl.Decl,
 		Nondeterministic: decl.Nondeterministic,
 	})

rego/rego_test.go

@@ -2555,3 +2555,131 @@ func TestRegoLazyObjCopyMaps(t *testing.T) {
 		t.Errorf("expected no change in foo, found one: %v", foo)
 	}
 }
+
+func TestDescriptionRegisterBuiltin1(t *testing.T) {
+	description := "custom-arity-1"
+
+	decl := &Function{
+		Name:        "foo",
+		Description: description,
+		Decl: types.NewFunction(
+			types.Args(types.S),
+			types.S,
+		),
+	}
+
+	RegisterBuiltin1(decl, func(_ BuiltinContext, _ *ast.Term) (*ast.Term, error) {
+		return ast.StringTerm("bar"), nil
+	})
+	defer unregisterBuiltin("foo")
+
+	got := ast.Builtins[len(ast.Builtins)-1].Description
+	if got != description {
+		t.Fatalf("expected %q, got %q", description, got)
+	}
+}
+
+func TestDescriptionRegisterBuiltin2(t *testing.T) {
+	description := "custom-arity-2"
+
+	decl := &Function{
+		Name:        "foo",
+		Description: description,
+		Decl: types.NewFunction(
+			types.Args(types.S, types.S),
+			types.S,
+		),
+	}
+
+	RegisterBuiltin2(decl, func(_ BuiltinContext, _, _ *ast.Term) (*ast.Term, error) {
+		return ast.StringTerm("bar"), nil
+	})
+	defer unregisterBuiltin("foo")
+
+	got := ast.Builtins[len(ast.Builtins)-1].Description
+	if got != description {
+		t.Fatalf("expected %q, got %q", description, got)
+	}
+}
+
+func TestDescriptionRegisterBuiltin3(t *testing.T) {
+	description := "custom-arity-3"
+
+	decl := &Function{
+		Name:        "foo",
+		Description: description,
+		Decl: types.NewFunction(
+			types.Args(types.S, types.S, types.S),
+			types.S,
+		),
+	}
+
+	RegisterBuiltin3(decl, func(_ BuiltinContext, _, _, _ *ast.Term) (*ast.Term, error) {
+		return ast.StringTerm("bar"), nil
+	})
+	defer unregisterBuiltin("foo")
+
+	got := ast.Builtins[len(ast.Builtins)-1].Description
+	if got != description {
+		t.Fatalf("expected %q, got %q", description, got)
+	}
+}
+
+func TestDescriptionRegisterBuiltin4(t *testing.T) {
+	description := "custom-arity-4"
+
+	decl := &Function{
+		Name:        "foo",
+		Description: description,
+		Decl: types.NewFunction(
+			types.Args(types.S, types.S, types.S, types.S),
+			types.S,
+		),
+	}
+
+	RegisterBuiltin4(decl, func(_ BuiltinContext, _, _, _, _ *ast.Term) (*ast.Term, error) {
+		return ast.StringTerm("bar"), nil
+	})
+	defer unregisterBuiltin("foo")
+
+	got := ast.Builtins[len(ast.Builtins)-1].Description
+	if got != description {
+		t.Fatalf("expected %q, got %q", description, got)
+	}
+}
+
+func TestDescriptionRegisterBuiltinDyn(t *testing.T) {
+	description := "custom-arity-dyn"
+
+	decl := &Function{
+		Name:        "foo",
+		Description: description,
+		Decl: types.NewFunction(
+			types.Args(types.S),
+			types.S,
+		),
+	}
+
+	RegisterBuiltinDyn(decl, func(_ BuiltinContext, _ []*ast.Term) (*ast.Term, error) {
+		return ast.StringTerm("bar"), nil
+	})
+	defer unregisterBuiltin("foo")
+
+	got := ast.Builtins[len(ast.Builtins)-1].Description
+	if got != description {
+		t.Fatalf("expected %q, got %q", description, got)
+	}
+}
+
+// unregisterBuiltin removes the builtin of the given name from ast.Builtins. This assists in
+// cleaning up custom functions added as part of certain test cases.
+func unregisterBuiltin(name string) {
+	builtins := make([]*ast.Builtin, 0, len(ast.Builtins))
+	for _, builtin := range ast.Builtins {
+		if builtin.Name == name {
+			continue
+		}
+		builtins = append(builtins, builtin)
+	}
+	ast.Builtins = builtins
+}