Copy metadata associated with a node in the optimized result

ast/annotations.go

@@ -111,6 +111,11 @@ func (a *Annotations) EndLoc() *Location {
 	return a.comments[count-1].Location
 }
 
+// Comments returns the comments associated with this annotation.
+func (a *Annotations) Comments() []*Comment {
+	return a.comments
+}
+
 // Compare returns an integer indicating if a is less than, equal to, or greater
 // than other.
 func (a *Annotations) Compare(other *Annotations) int {
@@ -181,6 +186,26 @@ func (a *Annotations) GetTargetPath() Ref {
 	}
 }
 
+func (a *Annotations) GetPackage() *Package {
+	switch n := a.node.(type) {
+	case *Package:
+		return n
+	case *Rule:
+		return n.Module.Package
+	default:
+		return nil
+	}
+}
+
+func (a *Annotations) GetRule() *Rule {
+	switch n := a.node.(type) {
+	case *Rule:
+		return n
+	default:
+		return nil
+	}
+}
+
 func (a *Annotations) setJSONOptions(opts astJSON.Options) {
 	a.jsonOptions = opts
 	if a.Location != nil {
@@ -509,6 +534,37 @@ func (a *Annotations) toObject() (*Object, *Error) {
 	return &obj, nil
 }
 
+func attachRuleAnnotations(mod *Module) {
+	// make a copy of the annotations
+	cpy := make([]*Annotations, len(mod.Annotations))
+	for i, a := range mod.Annotations {
+		cpy[i] = a.Copy(a.node)
+	}
+
+	for _, rule := range mod.Rules {
+		var j int
+		var found bool
+		for i, a := range cpy {
+			if rule.Loc().Row > a.Location.Row {
+				if rule.Ref().Equal(a.GetTargetPath()) {
+					rule.Annotations = append(rule.Annotations, a)
+
+					if a.Scope == annotationScopeRule {
+						j = i
+						found = true
+					}
+				}
+			} else {
+				break
+			}
+		}
+
+		if found && j < len(cpy) {
+			cpy = append(cpy[:j], cpy[j+1:]...)
+		}
+	}
+}
+
 func attachAnnotationsNodes(mod *Module) Errors {
 	var errs Errors
 

ast/check_test.go

@@ -2328,6 +2328,7 @@ p { input = "foo" }`}},
 			as, errors := BuildAnnotationSet(modules)
 			typeenv, checkErrors := newTypeChecker().WithSchemaSet(schemaSet).CheckTypes(oldTypeEnv, elems, as)
 			errors = append(errors, checkErrors...)
+
 			if len(errors) > 0 {
 				for _, e := range errors {
 					if tc.err == "" || !strings.Contains(e.Error(), tc.err) {

ast/compile.go

@@ -2191,11 +2191,13 @@ func (c *Compiler) parseMetadataBlocks() {
 
 			if len(mod.Annotations) == 0 {
 				var errs Errors
-				mod.Annotations, errs = parseAnnotations(mod.Comments)
+				mod.Annotations, errs = ParseAnnotations(mod.Comments)
 				errs = append(errs, attachAnnotationsNodes(mod)...)
 				for _, err := range errs {
 					c.err(err)
 				}
+
+				attachRuleAnnotations(mod)
 			}
 		}
 	}

ast/marshal_test.go

@@ -444,6 +444,22 @@ func TestRule_MarshalJSON(t *testing.T) {
 			}(),
 			ExpectedJSON: `{"body":[{"index":0,"terms":{"type":"boolean","value":true}}],"head":{"name":"allow","value":{"type":"boolean","value":true},"ref":[{"type":"var","value":"allow"}]},"location":{"file":"example.rego","row":6,"col":2}}`,
 		},
+		"annotations included": {
+			Rule: func() *Rule {
+				r := rule.Copy()
+				r.Annotations = []*Annotations{{
+					Scope:         "rule",
+					Title:         "My rule",
+					Entrypoint:    true,
+					Organizations: []string{"org1"},
+					Description:   "My desc",
+					Custom: map[string]interface{}{
+						"foo": "bar",
+					}}}
+				return r
+			}(),
+			ExpectedJSON: `{"annotations":[{"custom":{"foo":"bar"},"description":"My desc","entrypoint":true,"organizations":["org1"],"scope":"rule","title":"My rule"}],"body":[{"index":0,"terms":{"type":"boolean","value":true}}],"head":{"name":"allow","value":{"type":"boolean","value":true},"ref":[{"type":"var","value":"allow"}]}}`,
+		},
 	}
 
 	for name, data := range testCases {

ast/parser.go

@@ -433,7 +433,7 @@ func (p *Parser) Parse() ([]Statement, []*Comment, Errors) {
 
 func (p *Parser) parseAnnotations(stmts []Statement) []Statement {
 
-	annotStmts, errs := parseAnnotations(p.s.comments)
+	annotStmts, errs := ParseAnnotations(p.s.comments)
 	for _, err := range errs {
 		p.error(err.Location, err.Message)
 	}
@@ -445,7 +445,7 @@ func (p *Parser) parseAnnotations(stmts []Statement) []Statement {
 	return stmts
 }
 
-func parseAnnotations(comments []*Comment) ([]*Annotations, Errors) {
+func ParseAnnotations(comments []*Comment) ([]*Annotations, Errors) {
 
 	var hint = []byte("METADATA")
 	var curr *metadataParser

ast/parser_ext.go

@@ -713,6 +713,8 @@ func parseModule(filename string, stmts []Statement, comments []*Comment, regoCo
 		return nil, errs
 	}
 
+	attachRuleAnnotations(mod)
+
 	return mod, nil
 }