bundle: Retain metadata annotations for Wasm entrypoints during inspection (#5603)

* Pruning METADATA blocks associated with Wasm compiled entrypoints from Rego source in bundle
* Adding metadata annotations to wasm entrypoint declarations in bundle .manifest file
* Reading metadata annotations from both Rego source and .manifest file in bundle during `inspect`

Fixes: #5588
Signed-off-by: Johan Fylling <johan.dev@fylling.se>

ast/annotations.go

@@ -37,6 +37,7 @@ type (
 		Schemas          []*SchemaAnnotation          `json:"schemas,omitempty"`
 		Custom           map[string]interface{}       `json:"custom,omitempty"`
 		node             Node
+		comments         []*Comment
 	}
 
 	// SchemaAnnotation contains a schema declaration for the document identified by the path.
@@ -74,6 +75,10 @@ type (
 		Annotations *Annotations `json:"annotations,omitempty"`
 		node        Node         // The node the annotations are applied to
 	}
+
+	AnnotationsRefSet []*AnnotationsRef
+
+	FlatAnnotationsRefSet AnnotationsRefSet
 )
 
 func (a *Annotations) String() string {
@@ -91,6 +96,15 @@ func (a *Annotations) SetLoc(l *Location) {
 	a.Location = l
 }
 
+// EndLoc returns the location of this annotation's last comment line.
+func (a *Annotations) EndLoc() *Location {
+	count := len(a.comments)
+	if count == 0 {
+		return a.Location
+	}
+	return a.comments[count-1].Location
+}
+
 // Compare returns an integer indicating if a is less than, equal to, or greater
 // than other.
 func (a *Annotations) Compare(other *Annotations) int {
@@ -162,8 +176,13 @@ func (a *Annotations) GetTargetPath() Ref {
 }
 
 func NewAnnotationsRef(a *Annotations) *AnnotationsRef {
+	var loc *Location
+	if a.node != nil {
+		loc = a.node.Loc()
+	}
+
 	return &AnnotationsRef{
-		Location:    a.node.Loc(),
+		Location:    loc,
 		Path:        a.GetTargetPath(),
 		Annotations: a,
 		node:        a.node,
@@ -668,7 +687,7 @@ func (as *AnnotationSet) GetPackageScope(pkg *Package) *Annotations {
 
 // Flatten returns a flattened list view of this AnnotationSet.
 // The returned slice is sorted, first by the annotations' target path, then by their target location
-func (as *AnnotationSet) Flatten() []*AnnotationsRef {
+func (as *AnnotationSet) Flatten() FlatAnnotationsRefSet {
 	// This preallocation often won't be optimal, but it's superior to starting with a nil slice.
 	refs := make([]*AnnotationsRef, 0, len(as.byPath.Children)+len(as.byRule)+len(as.byPackage))
 
@@ -686,13 +705,7 @@ func (as *AnnotationSet) Flatten() []*AnnotationsRef {
 
 	// Sort by path, then annotation location, for stable output
 	sort.SliceStable(refs, func(i, j int) bool {
-		if refs[i].Path.Compare(refs[j].Path) < 0 {
-			return true
-		}
-		if refs[i].Annotations.Location.Compare(refs[j].Annotations.Location) < 0 {
-			return true
-		}
-		return false
+		return refs[i].Compare(refs[j]) < 0
 	})
 
 	return refs
@@ -705,7 +718,7 @@ func (as *AnnotationSet) Flatten() []*AnnotationsRef {
 // 2. The 'package' scope entry, if any
 // 3. Entries for the 'subpackages' scope, if any; ordered from the closest package path to the fartest. E.g.: 'do.re.mi', 'do.re', 'do'
 // The returned slice is guaranteed to always contain at least one entry, corresponding to the given rule.
-func (as *AnnotationSet) Chain(rule *Rule) []*AnnotationsRef {
+func (as *AnnotationSet) Chain(rule *Rule) AnnotationsRefSet {
 	var refs []*AnnotationsRef
 
 	ruleAnnots := as.GetRuleScope(rule)
@@ -751,6 +764,26 @@ func (as *AnnotationSet) Chain(rule *Rule) []*AnnotationsRef {
 	return refs
 }
 
+func (ars FlatAnnotationsRefSet) Insert(ar *AnnotationsRef) FlatAnnotationsRefSet {
+	result := make(FlatAnnotationsRefSet, 0, len(ars)+1)
+
+	// insertion sort, first by path, then location
+	for i, current := range ars {
+		if ar.Compare(current) < 0 {
+			result = append(result, ar)
+			result = append(result, ars[i:]...)
+			break
+		}
+		result = append(result, current)
+	}
+
+	if len(result) < len(ars)+1 {
+		result = append(result, ar)
+	}
+
+	return result
+}
+
 func newAnnotationTree() *annotationTreeNode {
 	return &annotationTreeNode{
 		Value:    nil,
@@ -814,3 +847,15 @@ func (t *annotationTreeNode) flatten(refs []*AnnotationsRef) []*AnnotationsRef {
 	}
 	return refs
 }
+
+func (ar *AnnotationsRef) Compare(other *AnnotationsRef) int {
+	if c := ar.Path.Compare(other.Path); c != 0 {
+		return c
+	}
+
+	if c := ar.Annotations.Location.Compare(other.Annotations.Location); c != 0 {
+		return c
+	}
+
+	return ar.Annotations.Compare(other.Annotations)
+}

ast/parser.go

@@ -2188,6 +2188,7 @@ func (b *metadataParser) Parse() (*Annotations, error) {
 	}
 
 	var result Annotations
+	result.comments = b.comments
 	result.Scope = raw.Scope
 	result.Entrypoint = raw.Entrypoint
 	result.Title = raw.Title

bundle/bundle.go

@@ -127,8 +127,9 @@ type Manifest struct {
 
 // WasmResolver maps a wasm module to an entrypoint ref.
 type WasmResolver struct {
-	Entrypoint string `json:"entrypoint,omitempty"`
-	Module     string `json:"module,omitempty"`
+	Entrypoint  string             `json:"entrypoint,omitempty"`
+	Module      string             `json:"module,omitempty"`
+	Annotations []*ast.Annotations `json:"annotations,omitempty"`
 }
 
 // Init initializes the manifest. If you instantiate a manifest
@@ -166,6 +167,10 @@ func (m Manifest) Equal(other Manifest) bool {
 	return m.equalWasmResolversAndRoots(other)
 }
 
+func (m Manifest) Empty() bool {
+	return m.Equal(Manifest{})
+}
+
 // Copy returns a deep copy of the manifest.
 func (m Manifest) Copy() Manifest {
 	m.Init()
@@ -210,7 +215,7 @@ func (m Manifest) equalWasmResolversAndRoots(other Manifest) bool {
 	}
 
 	for i := 0; i < len(m.WasmResolvers); i++ {
-		if m.WasmResolvers[i] != other.WasmResolvers[i] {
+		if !m.WasmResolvers[i].Equal(&other.WasmResolvers[i]) {
 			return false
 		}
 	}
@@ -218,6 +223,37 @@ func (m Manifest) equalWasmResolversAndRoots(other Manifest) bool {
 	return m.rootSet().Equal(other.rootSet())
 }
 
+func (wr *WasmResolver) Equal(other *WasmResolver) bool {
+	if wr == nil && other == nil {
+		return true
+	}
+
+	if wr == nil || other == nil {
+		return false
+	}
+
+	if wr.Module != other.Module {
+		return false
+	}
+
+	if wr.Entrypoint != other.Entrypoint {
+		return false
+	}
+
+	annotLen := len(wr.Annotations)
+	if annotLen != len(other.Annotations) {
+		return false
+	}
+
+	for i := 0; i < annotLen; i++ {
+		if wr.Annotations[i].Compare(other.Annotations[i]) != 0 {
+			return false
+		}
+	}
+
+	return true
+}
+
 type stringSet map[string]struct{}
 
 func (ss stringSet) Equal(other stringSet) bool {
@@ -849,7 +885,7 @@ func (w *Writer) writePlan(tw *tar.Writer, bundle Bundle) error {
 
 func writeManifest(tw *tar.Writer, bundle Bundle) error {
 
-	if bundle.Manifest.Equal(Manifest{}) {
+	if bundle.Manifest.Empty() {
 		return nil
 	}
 
@@ -926,7 +962,7 @@ func hashBundleFiles(hash SignatureHasher, b *Bundle) ([]FileInfo, error) {
 	// parse the manifest into a JSON structure;
 	// then recursively order the fields of all objects alphabetically and then apply
 	// the hash function to result to compute the hash.
-	if !b.Manifest.Equal(Manifest{}) {
+	if !b.Manifest.Empty() {
 		mbs, err := json.Marshal(b.Manifest)
 		if err != nil {
 			return files, err

cmd/build_test.go

@@ -9,12 +9,13 @@ import (
 	"os"
 	"path"
 	"path/filepath"
+	"reflect"
 	"strings"
 	"testing"
 
 	"github.com/open-policy-agent/opa/ast"
 	"github.com/open-policy-agent/opa/loader"
-
+	"github.com/open-policy-agent/opa/util"
 	"github.com/open-policy-agent/opa/util/test"
 )
 
@@ -585,7 +586,7 @@ package test
 p[1]
 
 f(x) { p[x] }
-		`,
+`,
 			},
 			err: fmt.Errorf("plan compilation requires at least one entrypoint"),
 		},
@@ -619,6 +620,301 @@ f(x) { p[x] }
 	}
 }
 
+func TestBuildWasmWithAnnotations(t *testing.T) {
+	tests := []struct {
+		note        string
+		files       map[string]string
+		entrypoints []string
+		manifest    string
+	}{
+		{
+			note: "last rule is annotated entrypoint",
+			files: map[string]string{
+				"test.rego": `
+package test
+
+# METADATA
+# title: P1
+p1 := 1
+
+# METADATA
+# title: P2
+# entrypoint: true
+p2 := 2
+`,
+			},
+			manifest: `
+{
+	"revision":"",
+	"roots":[""],
+	"wasm":[{
+		"entrypoint":"test/p2",
+		"module":"/policy.wasm",
+		"annotations":[{
+			"scope":"rule",
+			"title":"P2",
+			"entrypoint":true
+		}]
+	}]
+}
+`,
+		},
+		{
+			note: "last rule is (not annotated) entrypoint",
+			files: map[string]string{
+				"test.rego": `
+package test
+
+# METADATA
+# title: P1
+p1 := 1
+
+# METADATA
+# title: P2
+p2 := 2
+`,
+			},
+			entrypoints: []string{"test/p2"},
+			manifest: `
+{
+	"revision":"",
+	"roots":[""],
+	"wasm":[{
+		"entrypoint":"test/p2",
+		"module":"/policy.wasm",
+		"annotations":[{
+			"scope":"rule",
+			"title":"P2"
+		}]
+	}]
+}
+`,
+		},
+		{
+			note: "rules in multiple files are entrypoints",
+			files: map[string]string{
+				"test1.rego": `
+package test
+
+# METADATA
+# title: P1
+p1 := 1
+
+# METADATA
+# title: P2
+# entrypoint: true
+p2 := 2
+`,
+				"test2.rego": `
+package test
+
+# METADATA
+# title: P3
+p3 := 3
+
+# METADATA
+# title: P4
+p4 := 4
+`,
+				"test3.rego": `
+package test.foo
+
+# METADATA
+# title: BAR
+# entrypoint: true
+bar := "baz"
+`,
+			},
+			entrypoints: []string{"test/p3"},
+			manifest: `
+{
+	"revision":"",
+	"roots":[""],
+	"wasm":[{
+		"entrypoint":"test/p3",
+		"annotations":[{"scope":"rule","title":"P3"}],
+		"module":"/policy.wasm"
+	},{
+		"entrypoint":"test/foo/bar",
+		"module":"/policy.wasm",
+		"annotations":[{
+			"scope":"rule",
+			"title":"BAR",
+			"entrypoint":true
+		}]
+	},{
+		"entrypoint":"test/p2",
+		"module":"/policy.wasm",
+		"annotations":[{
+			"scope":"rule",
+			"title":"P2",
+			"entrypoint":true
+		}]
+	}]
+}
+`,
+		},
+		{
+			note: "rule with multiple metadata blocks",
+			files: map[string]string{
+				"test.rego": `
+package test
+
+# METADATA
+# title: P doc
+# scope: document
+
+# METADATA
+# title: P
+# entrypoint: true
+p := 1
+`,
+			},
+			manifest: `
+{
+	"revision":"",
+	"roots":[""],
+	"wasm":[{
+		"entrypoint":"test/p",
+		"module":"/policy.wasm",
+		"annotations":[{
+			"scope":"document",
+			"title":"P doc"
+		},{
+			"scope":"rule",
+			"title":"P",
+			"entrypoint":true
+		}]
+	}]
+}
+`,
+		},
+
+		// Package annotations are not injected into manifest, as package definition is always retained in Rego source.
+		{
+			note: "package is annotated entrypoint",
+			files: map[string]string{
+				"test.rego": `
+# METADATA
+# title: PKG
+# entrypoint: true
+package test
+
+# METADATA
+# title: P1
+p1 := 1
+
+# METADATA
+# title: P2
+p2 := 2
+`,
+			},
+			manifest: `
+{
+	"revision":"",
+	"roots":[""],
+	"wasm":[{
+		"entrypoint":"test",
+		"module":"/policy.wasm"
+	}]
+}
+`,
+		},
+		{
+			note: "package is (not annotated) entrypoint",
+			files: map[string]string{
+				"test.rego": `
+package test
+
+# METADATA
+# title: P1
+p1 := 1
+
+# METADATA
+# title: P2
+p2 := 2
+`,
+			},
+			entrypoints: []string{"test"},
+			manifest: `
+{
+	"revision":"",
+	"roots":[""],
+	"wasm":[{
+		"entrypoint":"test",
+		"module":"/policy.wasm"
+	}]
+}
+`,
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.note, func(t *testing.T) {
+			test.WithTempFS(tc.files, func(root string) {
+				params := newBuildParams()
+				if err := params.target.Set("wasm"); err != nil {
+					t.Fatal(err)
+				}
+				params.pruneUnused = true
+				params.outputFile = path.Join(root, "bundle.tar.gz")
+				params.entrypoints.v = tc.entrypoints
+
+				// Build should fail if entrypoint is not discovered from annotations.
+				err := dobuild(params, []string{root})
+				if err != nil {
+					t.Fatal(err)
+				}
+
+				_, err = loader.NewFileLoader().AsBundle(params.outputFile)
+				if err != nil {
+					t.Fatal(err)
+				}
+
+				// Check that manifest has expected content
+				f, err := os.Open(params.outputFile)
+				if err != nil {
+					t.Fatal(err)
+				}
+				defer f.Close()
+
+				gr, err := gzip.NewReader(f)
+				if err != nil {
+					t.Fatal(err)
+				}
+
+				tr := tar.NewReader(gr)
+
+				found := false
+				for {
+					f, err := tr.Next()
+					if err == io.EOF {
+						break
+					} else if err != nil {
+						t.Fatal(err)
+					}
+					if f.Name == "/.manifest" {
+						found = true
+						data, err := io.ReadAll(tr)
+						if err != nil {
+							t.Fatal(err)
+						}
+						manifest := util.MustUnmarshalJSON(data)
+						if !reflect.DeepEqual(manifest, util.MustUnmarshalJSON([]byte(tc.manifest))) {
+							t.Fatalf("expected manifest\n\n%v\n\nbut got\n\n%v", tc.manifest, string(util.MustMarshalJSON(manifest)))
+						}
+						break
+					}
+				}
+
+				if !found {
+					t.Fatal("no manifest found in bundle")
+				}
+			})
+		})
+	}
+}
+
 func TestBuildBundleModeIgnoreFlag(t *testing.T) {
 
 	files := map[string]string{

cmd/inspect.go

@@ -216,9 +216,16 @@ func populateAnnotations(out io.Writer, refs []*ast.AnnotationsRef) error {
 			if r := ref.GetRule(); r != nil {
 				fmt.Fprintln(out, "Rule:    ", r.Head.Name)
 			}
-			fmt.Fprintln(out, "Location:", ref.Location.String())
-			if a := ref.Annotations; a != nil && a.Entrypoint {
-				fmt.Fprintln(out, "Entrypoint:", a.Entrypoint)
+			if loc := ref.Location; loc != nil {
+				fmt.Fprintln(out, "Location:", loc.String())
+			}
+			if a := ref.Annotations; a != nil {
+				if len(a.Scope) > 0 {
+					fmt.Fprintln(out, "Scope:", a.Scope)
+				}
+				if a.Entrypoint {
+					fmt.Fprintln(out, "Entrypoint:", a.Entrypoint)
+				}
 			}
 			fmt.Fprintln(out)
 

cmd/inspect_test.go

@@ -288,6 +288,7 @@ pkg-descr
 
 Package:  test
 Location: %[1]s/x.rego:16
+Scope: package
 
 Organizations:
  pkg-org
@@ -313,6 +314,7 @@ doc-descr
 Package:  test
 Rule:     p
 Location: %[1]s/x.rego:50
+Scope: document
 
 Organizations:
  doc-org
@@ -338,6 +340,7 @@ rule-title
 Package:  test
 Rule:     p
 Location: %[1]s/x.rego:50
+Scope: rule
 
 Organizations:
  rule-org
@@ -361,3 +364,203 @@ Custom:
 
 	})
 }
+
+func TestDoInspectTarballPrettyWithAnnotations(t *testing.T) {
+
+	files := [][2]string{
+		{"x.rego", `# METADATA
+# title: pkg-title
+# description: pkg-descr
+# organizations:
+# - pkg-org
+# related_resources:
+# - https://pkg
+# - ref: https://pkg
+#   description: rr-pkg-note
+# authors:
+# - pkg-author
+# schemas:
+# - input: {"type": "boolean"}
+# custom:
+#  pkg: pkg-custom
+package test
+
+# METADATA
+# scope: document
+# title: doc-title
+# description: doc-descr
+# organizations:
+# - doc-org
+# related_resources:
+# - https://doc
+# - ref: https://doc
+#   description: rr-doc-note
+# authors:
+# - doc-author
+# schemas:
+# - input: {"type": "integer"}
+# custom:
+#  doc: doc-custom
+
+# METADATA
+# title: rule-title
+# description: rule-title
+# organizations:
+# - rule-org
+# related_resources:
+# - https://rule
+# - ref: https://rule
+#   description: rr-rule-note
+# authors:
+# - rule-author
+# schemas:
+# - input: {"type": "string"}
+# custom:
+#  rule: rule-custom
+p = 1`},
+		{".manifest", `
+{
+	"revision": "",
+	"roots": [
+		""
+	],
+	"wasm": [
+		{
+			"entrypoint": "test/a",
+			"module": "/policy.wasm"
+		},
+		{
+			"entrypoint": "test/b",
+			"module": "/policy.wasm",
+			"annotations": [
+				{
+					"scope": "rule",
+					"title": "WASM RULE B",
+					"entrypoint": true
+				}
+			]
+		}
+	]
+}`},
+		{"policy.wasm", ""},
+	}
+
+	buf := archive.MustWriteTarGz(files)
+
+	test.WithTempFS(nil, func(rootDir string) {
+		bundleFile := filepath.Join(rootDir, "bundle.tar.gz")
+
+		bf, err := os.Create(bundleFile)
+		if err != nil {
+			t.Fatalf("Unexpected error: %v", err)
+		}
+
+		_, err = bf.Write(buf.Bytes())
+		if err != nil {
+			t.Fatalf("Unexpected error: %v", err)
+		}
+
+		ps := newInspectCommandParams()
+		ps.listAnnotations = true
+		var out bytes.Buffer
+
+		err = doInspect(ps, bundleFile, &out)
+		if err != nil {
+			t.Fatalf("Unexpected error %v", err)
+		}
+
+		bs := out.Bytes()
+		idx := bytes.Index(bs, []byte(`ANNOTATIONS`)) // skip NAMESPACE box
+		output := strings.TrimSpace(string(bs[idx:]))
+		expected := strings.TrimSpace(`
+ANNOTATIONS:
+pkg-title
+=========
+
+pkg-descr
+
+Package:  test
+Location: /x.rego:16
+Scope: package
+
+Organizations:
+ pkg-org
+
+Authors:
+ pkg-author
+
+Schemas:
+ input: {"type":"boolean"}
+
+Related Resources:
+ https://pkg 
+ https://pkg rr-pkg-note
+
+Custom:
+ pkg: "pkg-custom"
+
+WASM RULE B
+===========
+
+Location: /policy.wasm:0
+Scope: rule
+Entrypoint: true
+
+doc-title
+=========
+
+doc-descr
+
+Package:  test
+Rule:     p
+Location: /x.rego:50
+Scope: document
+
+Organizations:
+ doc-org
+
+Authors:
+ doc-author
+
+Schemas:
+ input: {"type":"integer"}
+
+Related Resources:
+ https://doc 
+ https://doc rr-doc-note
+
+Custom:
+ doc: "doc-custom"
+
+rule-title
+==========
+
+rule-title
+
+Package:  test
+Rule:     p
+Location: /x.rego:50
+Scope: rule
+
+Organizations:
+ rule-org
+
+Authors:
+ rule-author
+
+Schemas:
+ input: {"type":"string"}
+
+Related Resources:
+ https://rule 
+ https://rule rr-rule-note
+
+Custom:
+ rule: "rule-custom"`)
+
+		if output != expected {
+			t.Fatalf("Unexpected output. Expected:\n\n%q\n\nGot:\n\n%q", expected, output)
+		}
+
+	})
+}

compile/compile.go

@@ -642,13 +642,21 @@ func (c *Compiler) compileWasm(ctx context.Context) error {
 		Raw:  buf.Bytes(),
 	}}
 
+	flattenedAnnotations := c.compiler.GetAnnotationSet().Flatten()
+
 	// Each entrypoint needs an entry in the manifest
-	for i := range c.entrypointrefs {
+	for i, e := range c.entrypointrefs {
 		entrypointPath := c.entrypoints[i]
 
+		var annotations []*ast.Annotations
+		if !c.isPackage(e) {
+			annotations = findAnnotationsForTerm(e, flattenedAnnotations)
+		}
+
 		c.bundle.Manifest.WasmResolvers = append(c.bundle.Manifest.WasmResolvers, bundle.WasmResolver{
-			Module:     "/" + strings.TrimLeft(modulePath, "/"),
-			Entrypoint: entrypointPath,
+			Module:      "/" + strings.TrimLeft(modulePath, "/"),
+			Entrypoint:  entrypointPath,
+			Annotations: annotations,
 		})
 	}
 
@@ -656,6 +664,33 @@ func (c *Compiler) compileWasm(ctx context.Context) error {
 	return pruneBundleEntrypoints(c.bundle, c.entrypointrefs)
 }
 
+func (c *Compiler) isPackage(term *ast.Term) bool {
+	for _, m := range c.compiler.Modules {
+		if m.Package.Path.Equal(term.Value) {
+			return true
+		}
+	}
+	return false
+}
+
+// findAnnotationsForTerm returns a slice of all annotations directly associated with the given term.
+func findAnnotationsForTerm(term *ast.Term, annotationRefs []*ast.AnnotationsRef) []*ast.Annotations {
+	r, ok := term.Value.(ast.Ref)
+	if !ok {
+		return nil
+	}
+
+	var result []*ast.Annotations
+
+	for _, ar := range annotationRefs {
+		if r.Equal(ar.Path) {
+			result = append(result, ar.Annotations)
+		}
+	}
+
+	return result
+}
+
 // pruneBundleEntrypoints will modify modules in the provided bundle to remove
 // rules matching the entrypoints along with injecting import statements to
 // preserve their ability to compile.
@@ -691,11 +726,43 @@ func pruneBundleEntrypoints(b *bundle.Bundle, entrypointrefs []*ast.Term) error
 				}
 			}
 
-			// If any rules were dropped update the module accordingly
-			if len(rules) != len(mf.Parsed.Rules) {
+			// Drop any Annotations for rules matching the entrypoint path
+			var annotations []*ast.Annotations
+			var prunedAnnotations []*ast.Annotations
+			for _, annotation := range mf.Parsed.Annotations {
+				p := annotation.GetTargetPath()
+				// We prune annotations of dropped rules, but not packages, as the Rego file is always retained
+				if p.Equal(entrypoint.Value) && !mf.Parsed.Package.Path.Equal(entrypoint.Value) {
+					prunedAnnotations = append(prunedAnnotations, annotation)
+				} else {
+					annotations = append(annotations, annotation)
+				}
+			}
+
+			// Drop comments associated with pruned annotations
+			var comments []*ast.Comment
+			for _, comment := range mf.Parsed.Comments {
+				pruned := false
+				for _, annotation := range prunedAnnotations {
+					if comment.Location.Row >= annotation.Location.Row &&
+						comment.Location.Row <= annotation.EndLoc().Row {
+						pruned = true
+						break
+					}
+				}
+
+				if !pruned {
+					comments = append(comments, comment)
+				}
+			}
+
+			// If any rules or annotations were dropped update the module accordingly
+			if len(rules) != len(mf.Parsed.Rules) || len(comments) != len(mf.Parsed.Comments) {
 				mf.Parsed.Rules = rules
+				mf.Parsed.Annotations = annotations
+				mf.Parsed.Comments = comments
 				// Remove the original raw source, we're editing the AST
-				// directly so it wont be in sync anymore.
+				// directly, so it won't be in sync anymore.
 				mf.Raw = nil
 			}
 		}

compile/compile_test.go

@@ -627,6 +627,89 @@ func TestCompilerWasmTargetMultipleEntrypoints(t *testing.T) {
 	})
 }
 
+func TestCompilerWasmTargetAnnotations(t *testing.T) {
+	files := map[string]string{
+		"test.rego": `
+# METADATA
+# title: My test package
+package test
+
+# METADATA
+# title: My P rule
+# entrypoint: true
+p = true`,
+		"policy.rego": `
+package policy
+
+# METADATA
+# title: All my Q rules
+# scope: document
+
+# METADATA
+# title: My Q rule
+q = true`,
+	}
+
+	test.WithTempFS(files, func(root string) {
+
+		compiler := New().WithPaths(root).WithTarget("wasm").
+			WithEntrypoints("test", "policy/q").
+			WithRegoAnnotationEntrypoints(true)
+
+		err := compiler.Build(context.Background())
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		if len(compiler.bundle.WasmModules) != 1 {
+			t.Fatalf("expected 1 Wasm modules, got: %d", len(compiler.bundle.WasmModules))
+		}
+
+		expWasmResolvers := []bundle.WasmResolver{
+			{
+				Entrypoint: "test",
+				Module:     "/policy.wasm",
+			},
+			{
+				Entrypoint: "policy/q",
+				Module:     "/policy.wasm",
+				Annotations: []*ast.Annotations{
+					{
+						Title: "All my Q rules",
+						Scope: "document",
+					},
+					{
+						Title: "My Q rule",
+						Scope: "rule",
+					},
+				},
+			},
+			{
+				Entrypoint: "test/p",
+				Module:     "/policy.wasm",
+				Annotations: []*ast.Annotations{
+					{
+						Title:      "My P rule",
+						Scope:      "rule",
+						Entrypoint: true,
+					},
+				},
+			},
+		}
+
+		if len(expWasmResolvers) != len(compiler.bundle.Manifest.WasmResolvers) {
+			t.Fatalf("\nExpected WasmResolvers:\n  %+v\nGot:\n  %+v\n", expWasmResolvers, compiler.bundle.Manifest.WasmResolvers)
+		}
+
+		for i, expWasmResolver := range expWasmResolvers {
+			if !expWasmResolver.Equal(&compiler.bundle.Manifest.WasmResolvers[i]) {
+				t.Fatalf("WasmResolver at index %v mismatch\n\nExpected WasmResolvers:\n  %+v\nGot:\n  %+v\n",
+					i, expWasmResolvers, compiler.bundle.Manifest.WasmResolvers)
+			}
+		}
+	})
+}
+
 func TestCompilerWasmTargetEntrypointDependents(t *testing.T) {
 	files := map[string]string{
 		"test.rego": `package test

internal/bundle/inspect/inspect.go

@@ -52,7 +52,24 @@ func File(path string, includeAnnotations bool) (*Info, error) {
 		if len(errs) > 0 {
 			return nil, errs
 		}
-		bi.Annotations = as.Flatten()
+		flattened := as.Flatten()
+
+		for _, wr := range bi.Manifest.WasmResolvers {
+			if as := wr.Annotations; len(as) > 0 {
+				path, err := ast.PtrRef(ast.DefaultRootDocument, wr.Entrypoint)
+				if err != nil {
+					return nil, fmt.Errorf("failed to parse Wasm entrypoint in manifest: %s", err)
+				}
+				for _, a := range as {
+					ar := ast.NewAnnotationsRef(a)
+					ar.Path = path
+					ar.Location = ast.NewLocation(nil, wr.Module, 0, 0)
+					flattened = flattened.Insert(ar)
+				}
+			}
+		}
+
+		bi.Annotations = flattened
 	}
 
 	err = bi.getBundleDataWasmAndSignatures(path)