From 9f7e5e6be6841516d022490de02af821592bc73b Mon Sep 17 00:00:00 2001
From: Johan Fylling <johan.dev@fylling.se>
Date: Mon, 23 Jan 2023 21:47:24 +0100
Subject: [PATCH] compile: Allow object generating rules to be annotated as
 entrypoint (#5578)

Using ground ref of rule when generating entrypoint from annotation
as variable in head suffix cannot be included in generated entrypoint path.

Fixes: #5577
Signed-off-by: Johan Fylling <johan.dev@fylling.se>
---
 cmd/build_test.go  | 94 ++++++++++++++++++++++++++++++++++++++++++++++
 compile/compile.go |  2 +-
 2 files changed, 95 insertions(+), 1 deletion(-)

diff --git a/cmd/build_test.go b/cmd/build_test.go
index ee9a596e74..394b2a0b21 100644
--- a/cmd/build_test.go
+++ b/cmd/build_test.go
@@ -477,6 +477,100 @@ package test
 p[1]
 
 f(x) { p[x] }
+		`,
+			},
+			err: nil,
+		},
+		{
+			note: "set generation with annotated entrypoint",
+			files: map[string]string{
+				"test.rego": `
+package test
+
+# METADATA
+# entrypoint: true
+p[x] {
+	{"a", "b"}[x]
+}
+		`,
+			},
+			err: nil,
+		},
+		{
+			note: "set generation with annotated entrypoint (contains if)",
+			files: map[string]string{
+				"test.rego": `
+package test
+
+import future.keywords
+
+# METADATA
+# entrypoint: true
+p contains x if {
+	{"a", "b"}[x]
+}
+		`,
+			},
+			err: nil,
+		},
+		{
+			note: "object generation with annotated entrypoint",
+			files: map[string]string{
+				"test.rego": `
+package test
+
+# METADATA
+# entrypoint: true
+p[i] := x {
+	x := ["a", "b"][i]
+}
+		`,
+			},
+			err: nil,
+		},
+		{
+			note: "object generation with annotated entrypoint (if)",
+			files: map[string]string{
+				"test.rego": `
+package test
+
+import future.keywords
+
+# METADATA
+# entrypoint: true
+p[i] if {
+	{"a", "b"}[i]
+}
+		`,
+			},
+			err: nil,
+		},
+		{
+			note: "dots in head with annotated entrypoint",
+			files: map[string]string{
+				"test.rego": `
+package test
+
+# METADATA
+# entrypoint: true
+p.a.b if {
+	true
+}
+		`,
+			},
+			err: nil,
+		},
+		{
+			note: "dots in head object generation with annotated entrypoint",
+			files: map[string]string{
+				"test.rego": `
+package test
+
+# METADATA
+# entrypoint: true
+p.a.b[i] := x {
+	x := ["a", "b"][i]
+}
 		`,
 			},
 			err: nil,
diff --git a/compile/compile.go b/compile/compile.go
index ff71d3f77f..03d915adab 100644
--- a/compile/compile.go
+++ b/compile/compile.go
@@ -234,7 +234,7 @@ func addEntrypointsFromAnnotations(c *Compiler, ar []*ast.AnnotationsRef) error
 				}
 			case "rule":
 				if r := ref.GetRule(); r != nil {
-					entrypoint = r.Ref()
+					entrypoint = r.Ref().GroundPrefix()
 				}
 			default:
 				continue // Wrong scope type. Bail out early.