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.