ast,format: Allowing keywords in Rego references (#7709)

Updating the parser and formatter to allow keywords in refs.

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

capabilities.json

@@ -4844,6 +4844,7 @@
     }
   ],
   "features": [
+    "keywords_in_refs",
     "rego_v1"
   ]
 }

cmd/bench.go

@@ -57,7 +57,7 @@ func newBenchmarkEvalParams() benchmarkCommandParams {
 			outputFormat: formats.Flag(formats.Pretty, formats.JSON, formats.GoBench),
 			target:       util.NewEnumFlag(compile.TargetRego, []string{compile.TargetRego, compile.TargetWasm}),
 			schema:       &schemaFlags{},
-			capabilities: newcapabilitiesFlag(),
+			capabilities: newCapabilitiesFlag(),
 		},
 		gracefulShutdownPeriod: 10,
 	}

cmd/build.go

@@ -53,7 +53,7 @@ type buildParams struct {
 
 func newBuildParams() buildParams {
 	return buildParams{
-		capabilities: newcapabilitiesFlag(),
+		capabilities: newCapabilitiesFlag(),
 		target:       util.NewEnumFlag(compile.TargetRego, compile.Targets),
 	}
 }

cmd/build_test.go

@@ -203,7 +203,7 @@ import rego.v1`,
 			}
 
 			test.WithTempFS(files, func(root string) {
-				caps := newcapabilitiesFlag()
+				caps := newCapabilitiesFlag()
 				if err := caps.Set(path.Join(root, "capabilities.json")); err != nil {
 					t.Fatal(err)
 				}
@@ -1324,7 +1324,7 @@ p[4] {
 				params.v1Compatible = tc.v1Compatible
 
 				if tc.capabilities != nil {
-					params.capabilities = newcapabilitiesFlag()
+					params.capabilities = newCapabilitiesFlag()
 					params.capabilities.C = tc.capabilities
 				}
 
@@ -2842,7 +2842,7 @@ foo contains __local1__1 if {
 				params.optimizationLevel = 1
 
 				if !tc.regoV1ImportCapable {
-					caps := newcapabilitiesFlag()
+					caps := newCapabilitiesFlag()
 					caps.C = ast.CapabilitiesForThisVersion()
 					caps.C.Features = []string{
 						ast.FeatureRefHeadStringPrefixes,

cmd/capabilities_test.go

@@ -90,6 +90,7 @@ func TestCapabilitiesCurrent(t *testing.T) {
 			note: "current",
 			expFeatures: []string{
 				ast.FeatureRegoV1,
+				ast.FeatureKeywordsInRefs,
 			},
 		},
 		{
@@ -100,6 +101,7 @@ func TestCapabilitiesCurrent(t *testing.T) {
 				ast.FeatureRefHeads,
 				ast.FeatureRegoV1Import,
 				ast.FeatureRegoV1,
+				ast.FeatureKeywordsInRefs,
 			},
 			expFutureKeywords: []string{
 				"in",

cmd/check.go

@@ -38,7 +38,7 @@ type checkParams struct {
 func newCheckParams() checkParams {
 	return checkParams{
 		format:       formats.Flag(formats.Pretty, formats.JSON),
-		capabilities: newcapabilitiesFlag(),
+		capabilities: newCapabilitiesFlag(),
 		schema:       &schemaFlags{},
 	}
 }

cmd/check_test.go

@@ -162,7 +162,7 @@ import rego.v1`,
 			}
 
 			test.WithTempFS(files, func(root string) {
-				caps := newcapabilitiesFlag()
+				caps := newCapabilitiesFlag()
 				if err := caps.Set(path.Join(root, "capabilities.json")); err != nil {
 					t.Fatal(err)
 				}

cmd/eval.go

@@ -92,7 +92,7 @@ func (p *evalCommandParams) regoVersion() ast.RegoVersion {
 
 func newEvalCommandParams() evalCommandParams {
 	return evalCommandParams{
-		capabilities: newcapabilitiesFlag(),
+		capabilities: newCapabilitiesFlag(),
 		outputFormat: formats.Flag(
 			formats.JSON,
 			formats.Values,

cmd/eval_test.go

@@ -2067,7 +2067,7 @@ p contains __local0__1 if __local0__1 = input.v
 						_ = params.outputFormat.Set(format)
 
 						if !tc.regoV1ImportCapable {
-							caps := newcapabilitiesFlag()
+							caps := newCapabilitiesFlag()
 							caps.C = ast.CapabilitiesForThisVersion()
 							caps.C.Features = []string{
 								ast.FeatureRefHeadStringPrefixes,

cmd/flags.go

@@ -201,7 +201,7 @@ type capabilitiesFlag struct {
 	pathOrVersion string
 }
 
-func newcapabilitiesFlag() *capabilitiesFlag {
+func newCapabilitiesFlag() *capabilitiesFlag {
 	return &capabilitiesFlag{
 		// cannot call ast.CapabilitiesForThisVersion here because
 		// custom builtins cannot be registered by this point in execution

cmd/fmt.go

@@ -21,18 +21,30 @@ import (
 )
 
 type fmtCommandParams struct {
-	overwrite     bool
-	list          bool
-	diff          bool
-	fail          bool
-	regoV1        bool
-	v0Compatible  bool
-	v1Compatible  bool
-	checkResult   bool
-	dropV0Imports bool
+	overwrite        bool
+	list             bool
+	diff             bool
+	fail             bool
+	regoV1           bool
+	v0Compatible     bool
+	v1Compatible     bool
+	checkResult      bool
+	dropV0Imports    bool
+	capabilitiesFlag *capabilitiesFlag
 }
 
-var fmtParams = fmtCommandParams{}
+func newFmtCommandParams() *fmtCommandParams {
+	return &fmtCommandParams{
+		capabilitiesFlag: newCapabilitiesFlag(),
+	}
+}
+
+func (p *fmtCommandParams) capabilities() *ast.Capabilities {
+	if p.capabilitiesFlag != nil && p.capabilitiesFlag.C != nil {
+		return p.capabilitiesFlag.C
+	}
+	return ast.CapabilitiesForThisVersion(ast.CapabilitiesRegoVersion(p.regoVersion()))
+}
 
 func (p *fmtCommandParams) regoVersion() ast.RegoVersion {
 	// The '--rego-v1' flag takes precedence over the '--v1-compatible' flag.
@@ -49,58 +61,9 @@ func (p *fmtCommandParams) regoVersion() ast.RegoVersion {
 	return ast.DefaultRegoVersion
 }
 
-var formatCommand = &cobra.Command{
-	Use:   "fmt [path [...]]",
-	Short: "Format Rego source files",
-	Long: `Format Rego source files.
-
-The 'fmt' command takes a Rego source file and outputs a reformatted version. If no file path
-is provided - this tool will use stdin.
-The format of the output is not defined specifically; whatever this tool outputs
-is considered correct format (with the exception of bugs).
-
-If the '-w' option is supplied, the 'fmt' command will overwrite the source file
-instead of printing to stdout.
-
-If the '-d' option is supplied, the 'fmt' command will output a diff between the
-original and formatted source.
-
-If the '-l' option is supplied, the 'fmt' command will output the names of files
-that would change if formatted. The '-l' option will suppress any other output
-to stdout from the 'fmt' command.
-
-If the '--fail' option is supplied, the 'fmt' command will return a non zero exit
-code if a file would be reformatted.
-
-The 'fmt' command can be run in several compatibility modes for consuming and outputting
-different Rego versions:
-
-* ` + "`" + `opa fmt` + "`" + `:
-  * v1 Rego is formatted to v1
-  * ` + "`" + `rego.v1` + "`" + `/` + "`" + `future.keywords` + "`" + ` imports are NOT removed
-  * ` + "`" + `rego.v1` + "`" + `/` + "`" + `future.keywords` + "`" + ` imports are NOT added if missing
-  * v0 rego is rejected
-* ` + "`" + `opa fmt --v0-compatible` + "`" + `:
-  * v0 Rego is formatted to v0
-  * v1 Rego is rejected
-* ` + "`" + `opa fmt --v0-v1` + "`" + `:
-  * v0 Rego is formatted to be compatible with v0 AND v1
-  * v1 Rego is rejected
-* ` + "`" + `opa fmt --v0-v1 --v1-compatible` + "`" + `:
-  * v1 Rego is formatted to be compatible with v0 AND v1
-  * v0 Rego is rejected
-`,
-	PreRunE: func(cmd *cobra.Command, _ []string) error {
-		return env.CmdFlags.CheckEnvironmentVariables(cmd)
-	},
-	Run: func(_ *cobra.Command, args []string) {
-		os.Exit(opaFmt(args))
-	},
-}
-
-func opaFmt(args []string) int {
+func opaFmt(args []string, fmtParams *fmtCommandParams) int {
 	if len(args) == 0 {
-		if err := formatStdin(&fmtParams, os.Stdin, os.Stdout); err != nil {
+		if err := formatStdin(fmtParams, os.Stdin, os.Stdout); err != nil {
 			fmt.Fprintln(os.Stderr, err)
 			return 1
 		}
@@ -116,7 +79,7 @@ func opaFmt(args []string) int {
 			return 1
 		}
 		err = filepath.Walk(filename, func(path string, info os.FileInfo, err error) error {
-			return formatFile(&fmtParams, os.Stdout, path, info, err)
+			return formatFile(fmtParams, os.Stdout, path, info, err)
 		})
 		if err != nil {
 			switch err := err.(type) {
@@ -154,6 +117,7 @@ func formatFile(params *fmtCommandParams, out io.Writer, filename string, info o
 	opts := format.Opts{
 		RegoVersion:   params.regoVersion(),
 		DropV0Imports: params.dropV0Imports,
+		Capabilities:  params.capabilities(),
 	}
 
 	if params.regoV1 {
@@ -235,8 +199,10 @@ func formatStdin(params *fmtCommandParams, r io.Reader, w io.Writer) error {
 		return err
 	}
 
-	opts := format.Opts{}
-	opts.RegoVersion = params.regoVersion()
+	opts := format.Opts{
+		RegoVersion:  params.regoVersion(),
+		Capabilities: params.capabilities(),
+	}
 
 	if params.regoV1 {
 		opts.ParserOptions = &ast.ParserOptions{RegoVersion: ast.RegoV0}
@@ -281,6 +247,57 @@ func newError(msg string, a ...any) fmtError {
 }
 
 func init() {
+	fmtParams := newFmtCommandParams()
+
+	var formatCommand = &cobra.Command{
+		Use:   "fmt [path [...]]",
+		Short: "Format Rego source files",
+		Long: `Format Rego source files.
+
+The 'fmt' command takes a Rego source file and outputs a reformatted version. If no file path
+is provided - this tool will use stdin.
+The format of the output is not defined specifically; whatever this tool outputs
+is considered correct format (with the exception of bugs).
+
+If the '-w' option is supplied, the 'fmt' command will overwrite the source file
+instead of printing to stdout.
+
+If the '-d' option is supplied, the 'fmt' command will output a diff between the
+original and formatted source.
+
+If the '-l' option is supplied, the 'fmt' command will output the names of files
+that would change if formatted. The '-l' option will suppress any other output
+to stdout from the 'fmt' command.
+
+If the '--fail' option is supplied, the 'fmt' command will return a non zero exit
+code if a file would be reformatted.
+
+The 'fmt' command can be run in several compatibility modes for consuming and outputting
+different Rego versions:
+
+* ` + "`" + `opa fmt` + "`" + `:
+  * v1 Rego is formatted to v1
+  * ` + "`" + `rego.v1` + "`" + `/` + "`" + `future.keywords` + "`" + ` imports are NOT removed
+  * ` + "`" + `rego.v1` + "`" + `/` + "`" + `future.keywords` + "`" + ` imports are NOT added if missing
+  * v0 rego is rejected
+* ` + "`" + `opa fmt --v0-compatible` + "`" + `:
+  * v0 Rego is formatted to v0
+  * v1 Rego is rejected
+* ` + "`" + `opa fmt --v0-v1` + "`" + `:
+  * v0 Rego is formatted to be compatible with v0 AND v1
+  * v1 Rego is rejected
+* ` + "`" + `opa fmt --v0-v1 --v1-compatible` + "`" + `:
+  * v1 Rego is formatted to be compatible with v0 AND v1
+  * v0 Rego is rejected
+`,
+		PreRunE: func(cmd *cobra.Command, _ []string) error {
+			return env.CmdFlags.CheckEnvironmentVariables(cmd)
+		},
+		Run: func(_ *cobra.Command, args []string) {
+			os.Exit(opaFmt(args, fmtParams))
+		},
+	}
+
 	formatCommand.Flags().BoolVarP(&fmtParams.overwrite, "write", "w", false, "overwrite the original source file")
 	formatCommand.Flags().BoolVarP(&fmtParams.list, "list", "l", false, "list all files who would change when formatted")
 	formatCommand.Flags().BoolVarP(&fmtParams.diff, "diff", "d", false, "only display a diff of the changes")
@@ -290,6 +307,7 @@ func init() {
 	addV1CompatibleFlag(formatCommand.Flags(), &fmtParams.v1Compatible, false)
 	formatCommand.Flags().BoolVar(&fmtParams.checkResult, "check-result", true, "assert that the formatted code is valid and can be successfully parsed")
 	formatCommand.Flags().BoolVar(&fmtParams.dropV0Imports, "drop-v0-imports", false, "drop v0 imports from the formatted code, such as 'rego.v1' and 'future.keywords'")
+	addCapabilitiesFlag(formatCommand.Flags(), fmtParams.capabilitiesFlag)
 
 	RootCommand.AddCommand(formatCommand)
 }

cmd/fmt_test.go

@@ -1139,3 +1139,274 @@ q := all([true, false])
 		})
 	}
 }
+
+func TestFmtFormatStdin_KeywordsInRefs(t *testing.T) {
+	cases := []struct {
+		note        string
+		params      fmtCommandParams
+		unformatted string
+		formatted   string
+	}{
+		{
+			note:   "v0",
+			params: fmtCommandParams{v0Compatible: true},
+			unformatted: `package test.package.import
+
+p {
+	foo.if.else
+}
+
+foo.if.else {
+    true
+}`,
+			formatted: `package test.package.import
+
+p {
+	foo.if.else
+}
+
+foo.if.else = true
+`,
+		},
+		{
+			note: "v0, no capability",
+			params: func() fmtCommandParams {
+				params := newFmtCommandParams()
+				params.v0Compatible = true
+				params.capabilitiesFlag.C = dropCapabilityFeature(ast.CapabilitiesForThisVersion(ast.CapabilitiesRegoVersion(ast.RegoV0)),
+					ast.FeatureKeywordsInRefs)
+				return *params
+			}(),
+			unformatted: `package test.package.import
+
+p {
+	foo.if.else
+}
+
+foo.if.else {
+	true
+}`,
+			formatted: `package test["package"]["import"]
+
+p {
+	foo.if["else"]
+}
+
+foo.if["else"] = true
+`,
+		},
+
+		{
+			note:   "v1",
+			params: fmtCommandParams{},
+			unformatted: `package test.package.import
+
+p if {
+	foo.if.else
+}
+
+foo.if.else if {
+	true
+}`,
+			formatted: `package test.package.import
+
+p if {
+	foo.if.else
+}
+
+foo.if.else := true
+`,
+		},
+		{
+			note: "v1, no capability",
+			params: func() fmtCommandParams {
+				params := newFmtCommandParams()
+				params.capabilitiesFlag.C = dropCapabilityFeature(ast.CapabilitiesForThisVersion(ast.CapabilitiesRegoVersion(ast.RegoV0)),
+					ast.FeatureKeywordsInRefs)
+				return *params
+			}(),
+			unformatted: `package test.package.import
+
+p if {
+	foo.if.else
+}
+
+foo.if.else if {
+	true
+}`,
+			formatted: `package test["package"]["import"]
+
+p if {
+	foo["if"]["else"]
+}
+
+foo["if"]["else"] := true
+`,
+		},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.note, func(t *testing.T) {
+			var stdout bytes.Buffer
+
+			stdin := bytes.NewBufferString(tc.unformatted)
+
+			files := map[string]string{
+				"policy.rego": tc.unformatted,
+			}
+
+			test.WithTempFS(files, func(path string) {
+				err := formatStdin(&tc.params, stdin, &stdout)
+				if err != nil {
+					t.Fatalf("Unexpected error: %s", err)
+				}
+
+				actual := stdout.String()
+				if actual != tc.formatted {
+					t.Fatalf("Expected:\n%s\n\nGot:\n%s\n\n", tc.formatted, actual)
+				}
+			})
+		})
+	}
+}
+
+func TestFmtFormatFile_KeywordsInRefs(t *testing.T) {
+	cases := []struct {
+		note        string
+		params      fmtCommandParams
+		unformatted string
+		formatted   string
+	}{
+		{
+			note:   "v0",
+			params: fmtCommandParams{v0Compatible: true},
+			unformatted: `package test.package.import
+
+p {
+	foo.if.else
+}
+
+foo.if.else {
+    true
+}`,
+			formatted: `package test.package.import
+
+p {
+	foo.if.else
+}
+
+foo.if.else = true
+`,
+		},
+		{
+			note: "v0, no capability",
+			params: func() fmtCommandParams {
+				params := newFmtCommandParams()
+				params.v0Compatible = true
+				params.capabilitiesFlag.C = dropCapabilityFeature(ast.CapabilitiesForThisVersion(ast.CapabilitiesRegoVersion(ast.RegoV0)),
+					ast.FeatureKeywordsInRefs)
+				return *params
+			}(),
+			unformatted: `package test.package.import
+
+p {
+	foo.if.else
+}
+
+foo.if.else {
+	true
+}`,
+			formatted: `package test["package"]["import"]
+
+p {
+	foo.if["else"]
+}
+
+foo.if["else"] = true
+`,
+		},
+
+		{
+			note:   "v1",
+			params: fmtCommandParams{},
+			unformatted: `package test.package.import
+
+p if {
+	foo.if.else
+}
+
+foo.if.else if {
+	true
+}`,
+			formatted: `package test.package.import
+
+p if {
+	foo.if.else
+}
+
+foo.if.else := true
+`,
+		},
+		{
+			note: "v1, no capability",
+			params: func() fmtCommandParams {
+				params := newFmtCommandParams()
+				params.capabilitiesFlag.C = dropCapabilityFeature(ast.CapabilitiesForThisVersion(ast.CapabilitiesRegoVersion(ast.RegoV0)),
+					ast.FeatureKeywordsInRefs)
+				return *params
+			}(),
+			unformatted: `package test.package.import
+
+p if {
+	foo.if.else
+}
+
+foo.if.else if {
+	true
+}`,
+			formatted: `package test["package"]["import"]
+
+p if {
+	foo["if"]["else"]
+}
+
+foo["if"]["else"] := true
+`,
+		},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.note, func(t *testing.T) {
+			var stdout bytes.Buffer
+
+			files := map[string]string{
+				"policy.rego": tc.unformatted,
+			}
+
+			test.WithTempFS(files, func(path string) {
+				policyFile := filepath.Join(path, "policy.rego")
+				info, err := os.Stat(policyFile)
+				err = formatFile(&tc.params, &stdout, policyFile, info, err)
+				if err != nil {
+					t.Fatalf("Unexpected error: %s", err)
+				}
+
+				actual := stdout.String()
+				if actual != tc.formatted {
+					t.Fatalf("Expected:\n%s\n\nGot:\n%s\n\n", tc.formatted, actual)
+				}
+			})
+		})
+	}
+}
+
+func dropCapabilityFeature(caps *ast.Capabilities, feature string) *ast.Capabilities {
+	feats := make([]string, 0, len(caps.Features))
+	for _, f := range caps.Features {
+		if f != feature {
+			feats = append(feats, f)
+		}
+	}
+	caps.Features = feats
+	return caps
+}

cmd/test.go

@@ -71,7 +71,7 @@ func newTestCommandParams() testCommandParams {
 		outputFormat: formats.Flag(formats.Pretty, formats.JSON, formats.GoBench),
 		explain:      newExplainFlag([]string{explainModeFails, explainModeFull, explainModeNotes, explainModeDebug}),
 		target:       util.NewEnumFlag(compile.TargetRego, []string{compile.TargetRego, compile.TargetWasm}),
-		capabilities: newcapabilitiesFlag(),
+		capabilities: newCapabilitiesFlag(),
 		schema:       &schemaFlags{},
 		output:       os.Stdout,
 		errOutput:    os.Stderr,

v1/ast/capabilities.go

@@ -55,6 +55,7 @@ const FeatureRefHeadStringPrefixes = "rule_head_ref_string_prefixes"
 const FeatureRefHeads = "rule_head_refs"
 const FeatureRegoV1 = "rego_v1"
 const FeatureRegoV1Import = "rego_v1_import"
+const FeatureKeywordsInRefs = "keywords_in_refs"
 
 // Capabilities defines a structure containing data that describes the capabilities
 // or features supported by a particular version of OPA.
@@ -133,6 +134,7 @@ func CapabilitiesForThisVersion(opts ...CapabilitiesOption) *Capabilities {
 			FeatureRefHeads,
 			FeatureRegoV1Import,
 			FeatureRegoV1, // Included in v0 capabilities to allow v1 bundles in --v0-compatible mode
+			FeatureKeywordsInRefs,
 		}
 	default:
 		for kw := range futureKeywords {
@@ -141,6 +143,7 @@ func CapabilitiesForThisVersion(opts ...CapabilitiesOption) *Capabilities {
 
 		f.Features = []string{
 			FeatureRegoV1,
+			FeatureKeywordsInRefs,
 		}
 	}
 

v1/ast/internal/scanner/scanner.go

@@ -116,6 +116,11 @@ func (s *Scanner) HasKeyword(keywords map[string]tokens.Token) bool {
 	return false
 }
 
+func (s *Scanner) IsKeyword(str string) bool {
+	_, ok := s.keywords[str]
+	return ok
+}
+
 func (s *Scanner) SetRegoV1Compatible() {
 	s.regoV1Compatible = true
 }

v1/ast/parser.go

@@ -572,8 +572,19 @@ func (p *Parser) parsePackage() *Package {
 		return nil
 	}
 
-	p.scan()
-	if p.s.tok != tokens.Ident {
+	p.scanWS()
+
+	// Make sure we allow the first term of refs to be the 'package' keyword.
+	if p.s.tok == tokens.Dot || p.s.tok == tokens.LBrack {
+		// This is a ref, not a package declaration.
+		return nil
+	}
+
+	if p.s.tok == tokens.Whitespace {
+		p.scan()
+	}
+
+	if !isIdentOrAllowedRefKeyword(p) {
 		p.illegalToken()
 		return nil
 	}
@@ -630,11 +641,23 @@ func (p *Parser) parseImport() *Import {
 		return nil
 	}
 
-	p.scan()
-	if p.s.tok != tokens.Ident {
-		p.error(p.s.Loc(), "expected ident")
+	p.scanWS()
+
+	// Make sure we allow the first term of refs to be the 'import' keyword.
+	if p.s.tok == tokens.Dot || p.s.tok == tokens.LBrack {
+		// This is a ref, not an import declaration.
 		return nil
 	}
+
+	if p.s.tok == tokens.Whitespace {
+		p.scan()
+	}
+
+	if !isIdentOrAllowedRefKeyword(p) {
+		p.illegalToken()
+		return nil
+	}
+
 	q, prev := p.presentParser()
 	term := q.parseTerm()
 	if term != nil {
@@ -693,17 +716,81 @@ func (p *Parser) parseImport() *Import {
 		return nil
 	}
 
+	if imp.Alias != "" {
+		// Unreachable: parsing the alias var should already have generated an error.
+		name := imp.Alias.String()
+		if IsKeywordInRegoVersion(name, p.po.EffectiveRegoVersion()) {
+			p.errorf(imp.Location, "unexpected import alias, must not be a keyword, got: %s", name)
+		}
+		return &imp
+	}
+
+	r := imp.Path.Value.(Ref)
+
+	// Don't allow keywords in the tail path term unless it's a future import
+	if len(r) == 1 {
+		t := r[0]
+		name := string(t.Value.(Var))
+		if IsKeywordInRegoVersion(name, p.po.EffectiveRegoVersion()) {
+			p.errorf(t.Location, "unexpected import path, must not end with a keyword, got: %s", name)
+			p.hint("import a different path or use an alias")
+		}
+	} else if !FutureRootDocument.Equal(r[0]) {
+		t := r[len(r)-1]
+		name := string(t.Value.(String))
+		if IsKeywordInRegoVersion(name, p.po.EffectiveRegoVersion()) {
+			p.errorf(t.Location, "unexpected import path, must not end with a keyword, got: %s", name)
+			p.hint("import a different path or use an alias")
+		}
+	}
+
 	return &imp
 }
 
+// isIdentOrAllowedRefKeyword checks if the current token is an Ident or a keyword in the active rego-version.
+// If a keyword, sets p.s.token to token.Ident
+func isIdentOrAllowedRefKeyword(p *Parser) bool {
+	if p.s.tok == tokens.Ident {
+		return true
+	}
+
+	if p.isAllowedRefKeyword(p.s.tok) {
+		p.s.tok = tokens.Ident
+		return true
+	}
+
+	return false
+}
+
+func scanAheadRef(p *Parser) bool {
+	if p.isAllowedRefKeyword(p.s.tok) {
+		// scan ahead to check if we're parsing a ref
+		s := p.save()
+		p.scanWS()
+		tok := p.s.tok
+		p.restore(s)
+
+		if tok == tokens.Dot || tok == tokens.LBrack {
+			p.s.tok = tokens.Ident
+			return true
+		}
+	}
+
+	return false
+}
+
 func (p *Parser) parseRules() []*Rule {
 
 	var rule Rule
 	rule.SetLoc(p.s.Loc())
 
+	// This allows keywords in the first var term of the ref
+	_ = scanAheadRef(p)
+
 	if p.s.tok == tokens.Default {
 		p.scan()
 		rule.Default = true
+		_ = scanAheadRef(p)
 	}
 
 	if p.s.tok != tokens.Ident {
@@ -817,17 +904,20 @@ func (p *Parser) parseRules() []*Rule {
 	}
 
 	if p.s.tok == tokens.Else {
-		if r := rule.Head.Ref(); len(r) > 1 && !r.IsGround() {
-			p.error(p.s.Loc(), "else keyword cannot be used on rules with variables in head")
-			return nil
-		}
-		if rule.Head.Key != nil {
-			p.error(p.s.Loc(), "else keyword cannot be used on multi-value rules")
-			return nil
-		}
+		// This might just be a refhead rule with a leading 'else' term.
+		if !scanAheadRef(p) {
+			if r := rule.Head.Ref(); len(r) > 1 && !r.IsGround() {
+				p.error(p.s.Loc(), "else keyword cannot be used on rules with variables in head")
+				return nil
+			}
+			if rule.Head.Key != nil {
+				p.error(p.s.Loc(), "else keyword cannot be used on multi-value rules")
+				return nil
+			}
 
-		if rule.Else = p.parseElse(rule.Head); rule.Else == nil {
-			return nil
+			if rule.Else = p.parseElse(rule.Head); rule.Else == nil {
+				return nil
+			}
 		}
 	}
 
@@ -1102,10 +1192,31 @@ func (p *Parser) parseLiteral() (expr *Expr) {
 		}
 	}()
 
+	// Check that we're not parsing a ref
+	if p.isAllowedRefKeyword(p.s.tok) {
+		// Scan ahead
+		s := p.save()
+		p.scanWS()
+		tok := p.s.tok
+		p.restore(s)
+
+		if tok == tokens.Dot || tok == tokens.LBrack {
+			p.s.tok = tokens.Ident
+			return p.parseLiteralExpr(false)
+		}
+	}
+
 	var negated bool
 	if p.s.tok == tokens.Not {
-		p.scan()
-		negated = true
+		s := p.save()
+		p.scanWS()
+		tok := p.s.tok
+		p.restore(s)
+
+		if tok != tokens.Dot && tok != tokens.LBrack {
+			p.scan()
+			negated = true
+		}
 	}
 
 	switch p.s.tok {
@@ -1122,35 +1233,51 @@ func (p *Parser) parseLiteral() (expr *Expr) {
 		}
 		return p.parseEvery()
 	default:
-		s := p.save()
-		expr := p.parseExpr()
-		if expr != nil {
-			expr.Negated = negated
-			if p.s.tok == tokens.With {
-				if expr.With = p.parseWith(); expr.With == nil {
-					return nil
-				}
-			}
-			// If we find a plain `every` identifier, attempt to parse an every expression,
-			// add hint if it succeeds.
-			if term, ok := expr.Terms.(*Term); ok && Var("every").Equal(term.Value) {
-				var hint bool
-				t := p.save()
-				p.restore(s)
-				if expr := p.futureParser().parseEvery(); expr != nil {
-					_, hint = expr.Terms.(*Every)
-				}
-				p.restore(t)
-				if hint {
-					p.hint("`import future.keywords.every` for `every x in xs { ... }` expressions")
-				}
-			}
-			return expr
-		}
-		return nil
+		return p.parseLiteralExpr(negated)
 	}
 }
 
+func (p *Parser) isAllowedRefKeyword(t tokens.Token) bool {
+	return p.isAllowedRefKeywordStr(t.String())
+}
+
+func (p *Parser) isAllowedRefKeywordStr(s string) bool {
+	if p.po.Capabilities.ContainsFeature(FeatureKeywordsInRefs) {
+		return IsKeywordInRegoVersion(s, p.po.EffectiveRegoVersion()) || p.s.s.IsKeyword(s)
+	}
+
+	return false
+}
+
+func (p *Parser) parseLiteralExpr(negated bool) *Expr {
+	s := p.save()
+	expr := p.parseExpr()
+	if expr != nil {
+		expr.Negated = negated
+		if p.s.tok == tokens.With {
+			if expr.With = p.parseWith(); expr.With == nil {
+				return nil
+			}
+		}
+		// If we find a plain `every` identifier, attempt to parse an every expression,
+		// add hint if it succeeds.
+		if term, ok := expr.Terms.(*Term); ok && Var("every").Equal(term.Value) {
+			var hint bool
+			t := p.save()
+			p.restore(s)
+			if expr := p.futureParser().parseEvery(); expr != nil {
+				_, hint = expr.Terms.(*Every)
+			}
+			p.restore(t)
+			if hint {
+				p.hint("`import future.keywords.every` for `every x in xs { ... }` expressions")
+			}
+		}
+		return expr
+	}
+	return nil
+}
+
 func (p *Parser) parseWith() []*With {
 
 	withs := []*With{}
@@ -1431,6 +1558,9 @@ func (p *Parser) parseTermIn(lhs *Term, keyVal bool, offset int) *Term {
 			}
 			p.restore(s)
 		}
+
+		_ = scanAheadRef(p)
+
 		if op := p.parseTermOpName(memberRef, tokens.In); op != nil {
 			if rhs := p.parseTermRelation(nil, p.s.loc.Offset); rhs != nil {
 				call := p.setLoc(CallTerm(op, lhs, rhs), lhs.Location, offset, p.s.lastEnd)
@@ -1789,7 +1919,7 @@ func (p *Parser) parseRef(head *Term, offset int) (term *Term) {
 		switch p.s.tok {
 		case tokens.Dot:
 			p.scanWS()
-			if p.s.tok != tokens.Ident {
+			if p.s.tok != tokens.Ident && !p.isAllowedRefKeyword(p.s.tok) {
 				p.illegal("expected %v", tokens.Ident)
 				return nil
 			}

v1/ast/rego_v1.go

@@ -191,7 +191,7 @@ func checkRegoV1Rule(rule *Rule, opts RegoCheckOptions) Errors {
 
 	var errs Errors
 
-	if opts.NoKeywordsAsRuleNames && IsKeywordInRegoVersion(rule.Head.Name.String(), RegoV1) {
+	if opts.NoKeywordsAsRuleNames && len(rule.Head.Reference) < 2 && IsKeywordInRegoVersion(rule.Head.Name.String(), RegoV1) {
 		errs = append(errs, NewError(ParseErr, rule.Location, "%s keyword cannot be used for rule name", rule.Head.Name.String()))
 	}
 	if opts.RequireRuleBodyOrValue && rule.generatedBody && rule.Head.generatedValue {

v1/bundle/bundle.go

@@ -1092,31 +1092,55 @@ func (b *Bundle) FormatModules(useModulePath bool) error {
 
 // FormatModulesForRegoVersion formats Rego modules to comply with a given Rego version
 func (b *Bundle) FormatModulesForRegoVersion(version ast.RegoVersion, preserveModuleRegoVersion bool, useModulePath bool) error {
+	return b.FormatModulesWithOptions(BundleFormatOptions{
+		RegoVersion:               version,
+		PreserveModuleRegoVersion: preserveModuleRegoVersion,
+		UseModulePath:             useModulePath,
+	})
+}
+
+type BundleFormatOptions struct {
+	RegoVersion               ast.RegoVersion
+	Capabilities              *ast.Capabilities
+	PreserveModuleRegoVersion bool
+	UseModulePath             bool
+}
+
+// FormatModulesWithOptions formats Rego modules with the given options.
+func (b *Bundle) FormatModulesWithOptions(opts BundleFormatOptions) error {
 	var err error
 
 	for i, module := range b.Modules {
-		opts := format.Opts{}
-		if preserveModuleRegoVersion {
-			opts.RegoVersion = module.Parsed.RegoVersion()
-			opts.ParserOptions = &ast.ParserOptions{
-				RegoVersion: opts.RegoVersion,
+		fmtOpts := format.Opts{
+			RegoVersion:  opts.RegoVersion,
+			Capabilities: opts.Capabilities,
+		}
+
+		if module.Parsed != nil {
+			fmtOpts.ParserOptions = &ast.ParserOptions{
+				RegoVersion: module.Parsed.RegoVersion(),
 			}
-		} else {
-			opts.RegoVersion = version
+			if opts.PreserveModuleRegoVersion {
+				fmtOpts.RegoVersion = module.Parsed.RegoVersion()
+			}
+		}
+
+		if fmtOpts.Capabilities == nil {
+			fmtOpts.Capabilities = ast.CapabilitiesForThisVersion(ast.CapabilitiesRegoVersion(fmtOpts.RegoVersion))
 		}
 
 		if module.Raw == nil {
-			module.Raw, err = format.AstWithOpts(module.Parsed, opts)
+			module.Raw, err = format.AstWithOpts(module.Parsed, fmtOpts)
 			if err != nil {
 				return err
 			}
 		} else {
-			path := module.URL
-			if useModulePath {
-				path = module.Path
+			p := module.URL
+			if opts.UseModulePath {
+				p = module.Path
 			}
 
-			module.Raw, err = format.SourceWithOpts(path, module.Raw, opts)
+			module.Raw, err = format.SourceWithOpts(p, module.Raw, fmtOpts)
 			if err != nil {
 				return err
 			}

v1/compile/compile.go

@@ -381,14 +381,12 @@ func (c *Compiler) Build(ctx context.Context) error {
 		c.bundle.Manifest.Metadata = *c.metadata
 	}
 
-	if c.regoVersion == ast.RegoV1 {
-		if err := c.bundle.FormatModulesForRegoVersion(c.regoVersion, true, false); err != nil {
-			return err
-		}
-	} else {
-		if err := c.bundle.FormatModules(false); err != nil {
-			return err
-		}
+	if err := c.bundle.FormatModulesWithOptions(bundle.BundleFormatOptions{
+		RegoVersion:               c.regoVersion,
+		Capabilities:              c.capabilities,
+		PreserveModuleRegoVersion: true,
+	}); err != nil {
+		return err
 	}
 
 	if c.bsc != nil {

v1/compile/compile_test.go

@@ -1702,12 +1702,13 @@ update {
 
 func TestCompilerOptimizationSupportRegoVersion(t *testing.T) {
 	tests := []struct {
-		note                string
-		modulesRegoVersion  ast.RegoVersion
-		regoV1ImportCapable bool
-		entrypoint          string
-		files               map[string]string
-		expected            []string
+		note                       string
+		modulesRegoVersion         ast.RegoVersion
+		noKeywordsInRefsCapability bool
+		regoV1ImportCapable        bool
+		entrypoint                 string
+		files                      map[string]string
+		expected                   []string
 	}{
 		{
 			note:                "v0 module, rego.v1 capable",
@@ -1775,6 +1776,42 @@ contains := 2 {
 
 import rego.v1
 
+p if {
+	data.foo.contains = input.x
+}
+`,
+				`package foo
+
+contains = 2 {
+	input.a = input.b
+}
+`,
+			},
+		},
+		{
+			note:                       "v0 module, rego.v1 capable, import conflict with keyword, no capability",
+			modulesRegoVersion:         ast.RegoV0,
+			noKeywordsInRefsCapability: true,
+			regoV1ImportCapable:        true,
+			entrypoint:                 "test/p",
+			files: map[string]string{
+				"test.rego": `package test
+import data.foo.contains
+
+p {
+    input.x == contains
+}`,
+				"foo.rego": `package foo
+contains := 2 {
+	input.a == input.b
+}`,
+			},
+			// rego.v1 import used for data.test, since complete ref to data.foo.contains is used locally without original import
+			expected: []string{
+				`package test
+
+import rego.v1
+
 p if {
 	data.foo["contains"] = input.x
 }
@@ -1891,6 +1928,9 @@ p if {
 				if tc.regoV1ImportCapable {
 					capabilities.Features = append(capabilities.Features, ast.FeatureRegoV1Import)
 				}
+				if !tc.noKeywordsInRefsCapability {
+					capabilities.Features = append(capabilities.Features, ast.FeatureKeywordsInRefs)
+				}
 
 				compiler := New().
 					WithCapabilities(capabilities).

v1/format/format.go

@@ -37,6 +37,8 @@ type Opts struct {
 	// DropV0Imports instructs the formatter to drop all v0 imports from the module; i.e. 'rego.v1' and 'future.keywords' imports.
 	// Imports are only removed if [Opts.RegoVersion] makes them redundant.
 	DropV0Imports bool
+
+	Capabilities *ast.Capabilities
 }
 
 func (o Opts) effectiveRegoVersion() ast.RegoVersion {
@@ -146,6 +148,10 @@ type fmtOpts struct {
 	regoV1         bool
 	regoV1Imported bool
 	futureKeywords []string
+
+	// If true, the formatter will retain keywords in refs, e.g. `p.not ` instead of `p["not"]`.
+	// The format of the original ref is preserved, so `p["not"]` will still be formatted as `p["not"]`.
+	allowKeywordsInRefs bool
 }
 
 func (o fmtOpts) keywords() []string {
@@ -179,6 +185,12 @@ func AstWithOpts(x any, opts Opts) ([]byte, error) {
 		o.contains = true
 	}
 
+	capabilities := opts.Capabilities
+	if capabilities == nil {
+		capabilities = ast.CapabilitiesForThisVersion(ast.CapabilitiesRegoVersion(opts.effectiveRegoVersion()))
+	}
+	o.allowKeywordsInRefs = capabilities.ContainsFeature(ast.FeatureKeywordsInRefs)
+
 	memberRef := ast.Member.Ref()
 	memberWithKeyRef := ast.MemberWithKey.Ref()
 
@@ -1070,8 +1082,16 @@ func (w *writer) writeFunctionCall(expr *ast.Expr, comments []*ast.Comment) ([]*
 }
 
 func (w *writer) writeFunctionCallPlain(terms []*ast.Term, comments []*ast.Comment) ([]*ast.Comment, error) {
-	w.write(terms[0].String() + "(")
+	if r, ok := terms[0].Value.(ast.Ref); ok {
+		if c, err := w.writeRef(r, comments); err != nil {
+			return c, err
+		}
+	} else {
+		w.write(terms[0].String())
+	}
+	w.write("(")
 	defer w.write(")")
+
 	args := make([]any, len(terms)-1)
 	for i, t := range terms[1:] {
 		args[i] = t
@@ -1264,7 +1284,7 @@ func (w *writer) writeRef(x ast.Ref, comments []*ast.Comment) ([]*ast.Comment, e
 		for _, t := range path {
 			switch p := t.Value.(type) {
 			case ast.String:
-				w.writeRefStringPath(p)
+				w.writeRefStringPath(p, t.Location)
 			case ast.Var:
 				w.writeBracketed(w.formatVar(p))
 			default:
@@ -1292,15 +1312,33 @@ func (w *writer) writeBracketed(str string) {
 
 var varRegexp = regexp.MustCompile("^[[:alpha:]_][[:alpha:][:digit:]_]*$")
 
-func (w *writer) writeRefStringPath(s ast.String) {
+func (w *writer) writeRefStringPath(s ast.String, l *ast.Location) {
 	str := string(s)
-	if varRegexp.MatchString(str) && !ast.IsInKeywords(str, w.fmtOpts.keywords()) {
-		w.write("." + str)
-	} else {
+	if w.shouldBracketRefTerm(str, l) {
 		w.writeBracketed(s.String())
+	} else {
+		w.write("." + str)
 	}
 }
 
+func (w *writer) shouldBracketRefTerm(s string, l *ast.Location) bool {
+	if !varRegexp.MatchString(s) {
+		return true
+	}
+
+	if ast.IsInKeywords(s, w.fmtOpts.keywords()) {
+		if !w.fmtOpts.allowKeywordsInRefs {
+			return true
+		}
+
+		if l != nil && l.Text[0] == 34 { // If the original term text starts with '"', we preserve the brackets and quotes
+			return true
+		}
+	}
+
+	return false
+}
+
 func (*writer) formatVar(v ast.Var) string {
 	if v.IsWildcard() {
 		return ast.Wildcard.String()

v1/format/format_test.go

@@ -932,6 +932,78 @@ f := 6
 	}
 }
 
+func TestFormatKeywordsInRefs(t *testing.T) {
+	regoVersions := map[string]ast.RegoVersion{
+		"v0": ast.RegoV0,
+		"v1": ast.RegoV1,
+	}
+
+	for versionName, regoVersion := range regoVersions {
+		t.Run(versionName, func(t *testing.T) {
+			regoFiles, err := filepath.Glob(fmt.Sprintf("testfiles/%s/*.rego.formatted_no_keywords_in_refs", versionName))
+			if err != nil {
+				panic(err)
+			}
+
+			for _, expected_rego := range regoFiles {
+				original_rego := strings.TrimSuffix(expected_rego, ".formatted_no_keywords_in_refs")
+				t.Run(original_rego, func(t *testing.T) {
+					expected, err := os.ReadFile(expected_rego)
+					if err != nil {
+						t.Fatalf("Failed to read expected rego source: %v", err)
+					}
+
+					original, err := os.ReadFile(original_rego)
+					if err != nil {
+						t.Fatalf("Failed to read rego source: %v", err)
+					}
+
+					caps := ast.CapabilitiesForThisVersion(ast.CapabilitiesRegoVersion(regoVersion))
+					feats := []string{}
+					for _, f := range caps.Features {
+						if f != ast.FeatureKeywordsInRefs {
+							feats = append(feats, f)
+						}
+					}
+					caps.Features = feats
+
+					popts := ast.ParserOptions{
+						RegoVersion: regoVersion,
+					}
+					opts := Opts{
+						RegoVersion:   regoVersion,
+						ParserOptions: &popts,
+						Capabilities:  caps,
+					}
+
+					formatted, err := SourceWithOpts(original_rego, original, opts)
+					if err != nil {
+						t.Fatalf("Failed to format file: %v", err)
+					}
+
+					if ln, at := differsAt(formatted, expected); ln != 0 {
+						t.Fatalf("Expected formatted bytes to equal expected bytes but differed near line %d / byte %d (got: %q, expected: %q):\n%s", ln, at, formatted[at], expected[at], prefixWithLineNumbers(formatted))
+					}
+
+					if _, err := ast.ParseModuleWithOpts(original_rego+".tmp", string(formatted), popts); err != nil {
+						t.Fatalf("Failed to parse formatted bytes: %v", err)
+					}
+
+					formatted, err = SourceWithOpts(original_rego, formatted, opts)
+					if err != nil {
+						t.Fatalf("Failed to double format file")
+					}
+
+					if ln, at := differsAt(formatted, expected); ln != 0 {
+						t.Fatalf("Expected roundtripped bytes to equal expected bytes but differed near line %d / byte %d:\n%s", ln, at, prefixWithLineNumbers(formatted))
+					}
+
+				})
+			}
+		})
+	}
+}
+
 // 382	   3064960 ns/op	 4573131 B/op	   26266 allocs/op // no optimizations
 // 685	   1737719 ns/op	 1972193 B/op	   14160 allocs/op // pre-allocate partitionComments
 // 708	   1674343 ns/op	 1916700 B/op	   11556 allocs/op // static memberRef & memberWithKeyRef

v1/format/testfiles/v0/test_contains.rego.formatted

@@ -1,4 +1,4 @@
-package test["contains"]
+package test.contains
 
 import future.keywords.contains
 

v1/format/testfiles/v0/test_contains.rego.formatted_no_keywords_in_refs

@@ -0,0 +1,19 @@
+package test["contains"]
+
+import future.keywords.contains
+
+p contains "foo"
+
+deny contains msg {
+	msg := "foo"
+}
+
+deny contains msg {
+	msg := "bar"
+}
+
+# partial objects unchanged
+o[k] = v {
+	k := "ok"
+	v := "nok"
+}

v1/format/testfiles/v0/test_if.rego.formatted

@@ -1,4 +1,4 @@
-package test["if"]
+package test.if
 
 import future.keywords.if
 

v1/format/testfiles/v0/test_if.rego.formatted_no_keywords_in_refs

@@ -0,0 +1,34 @@
+package test["if"]
+
+import future.keywords.if
+
+p if 1 > 0 # shorthand
+
+p if 1 > 0 # longhand one line
+
+p if {
+	1 > 0 # longhand two lines
+}
+
+# same without the comment
+p if {
+	1 > 0
+}
+
+p if { # comment one
+	1 > 0 # comment two
+}
+
+q[x] = y if {
+	y := 10
+	x := "ten"
+}
+
+q[x] = y if { # not using if
+	y := 11
+	x := "eleven"
+}
+
+r[x] { # no if here before
+	x := "set"
+}

v1/format/testfiles/v0/test_if_else.rego.formatted

@@ -1,4 +1,4 @@
-package test["if"]
+package test.if
 
 import future.keywords.if
 

v1/format/testfiles/v0/test_if_else.rego.formatted_no_keywords_in_refs

@@ -0,0 +1,24 @@
+package test["if"]
+
+import future.keywords.if
+
+p := 1 if 1 > 0
+
+else := 2
+
+q := 1 if 1 > 0
+
+else := 2 if 2 > 1
+
+q := 1 if {
+	1 > 0
+	2 > 1
+} else := 2 if 2 > 1
+
+r := 1 if {
+	1 > 0
+}
+
+else := 2 if {
+	2 > 1
+}

v1/format/testfiles/v0/test_in.rego.formatted

@@ -1,4 +1,4 @@
-package test["in"]
+package test.in
 
 import future.keywords.in
 

v1/format/testfiles/v0/test_in.rego.formatted_no_keywords_in_refs

@@ -0,0 +1,13 @@
+package test["in"]
+
+import future.keywords.in
+
+a["in"] := "foo"
+
+b.c["in"] := "bar"
+
+c["in"].d := "baz"
+
+p {
+	input["in"]
+}

v1/format/testfiles/v0/test_keywords_in_refs.rego

@@ -0,0 +1,147 @@
+package test.not.package.import.as.default.else.with.null.true.false.some
+
+p {
+	a.not == 1
+	a.package == 1
+	a.import == 1
+	a.as == 1
+	a.default == 1
+	a.else == 1
+	a.with == 1
+	a.null == 1
+	a.true == 1
+	a.false == 1
+	a.some == 1
+
+	b.c.not == 1
+	b.c.package == 1
+	b.c.import == 1
+	b.c.as == 1
+	b.c.default == 1
+	b.c.else == 1
+	b.c.with == 1
+	b.c.null == 1
+	b.c.true == 1
+	b.c.false == 1
+	b.c.some == 1
+
+	d.not.e == 1
+	d.package.e == 1
+	d.import.e == 1
+	d.as.e == 1
+	d.default.e == 1
+	d.else.e == 1
+	d.with.e == 1
+	d.null.e == 1
+	d.true.e == 1
+	d.false.e == 1
+	d.some.e == 1
+
+	f.not(2)
+	f.package(2)
+	f.import(2)
+	f.as(2)
+	f.default(2)
+	f.else(2)
+	f.with(2)
+	f.null(2)
+	f.true(2)
+	f.false(2)
+	f.some(2)
+
+	g.h.not(2)
+	g.h.package(2)
+	g.h.import(2)
+	g.h.as(2)
+	g.h.default(2)
+	g.h.else(2)
+	g.h.with(2)
+	g.h.null(2)
+	g.h.true(2)
+	g.h.false(2)
+	g.h.some(2)
+
+	i.not.j(2)
+	i.package.j(2)
+	i.import.j(2)
+	i.as.j(2)
+	i.default.j(2)
+	i.else.j(2)
+	i.with.j(2)
+	i.null.j(2)
+	i.true.j(2)
+	i.false.j(2)
+	i.some.j(2)
+}
+
+a.not := 1
+a.package := 1
+a.import := 1
+a.as := 1
+a.default := 1
+a.else := 1
+a.with := 1
+a.null := 1
+a.true := 1
+a.false := 1
+a.some := 1
+
+b.c.not := 1
+b.c.package := 1
+b.c.import := 1
+b.c.as := 1
+b.c.default := 1
+b.c.else := 1
+b.c.with := 1
+b.c.null := 1
+b.c.true := 1
+b.c.false := 1
+b.c.some := 1
+
+d.not.e := 1
+d.package.e := 1
+d.import.e := 1
+d.as.e := 1
+d.default.e := 1
+d.else.e := 1
+d.with.e := 1
+d.null.e := 1
+d.true.e := 1
+d.false.e := 1
+d.some.e := 1
+
+f.not(x) := x
+f.package(x) := x
+f.import(x) := x
+f.as(x) := x
+f.default(x) := x
+f.else(x) := x
+f.with(x) := x
+f.null(x) := x
+f.true(x) := x
+f.false(x) := x
+f.some(x) := x
+
+g.h.not(x) := x
+g.h.package(x) := x
+g.h.import(x) := x
+g.h.as(x) := x
+g.h.default(x) := x
+g.h.else(x) := x
+g.h.with(x) := x
+g.h.null(x) := x
+g.h.true(x) := x
+g.h.false(x) := x
+g.h.some(x) := x
+
+i.not.j(x) := x
+i.package.j(x) := x
+i.import.j(x) := x
+i.as.j(x) := x
+i.default.j(x) := x
+i.else.j(x) := x
+i.with.j(x) := x
+i.null.j(x) := x
+i.true.j(x) := x
+i.false.j(x) := x
+i.some.j(x) := x

v1/format/testfiles/v0/test_keywords_in_refs.rego.formatted

@@ -0,0 +1,207 @@
+package test.not.package.import.as.default.else.with.null.true.false.some
+
+p {
+	a.not == 1
+	a.package == 1
+	a.import == 1
+	a.as == 1
+	a.default == 1
+	a.else == 1
+	a.with == 1
+	a.null == 1
+	a.true == 1
+	a.false == 1
+	a.some == 1
+
+	b.c.not == 1
+	b.c.package == 1
+	b.c.import == 1
+	b.c.as == 1
+	b.c.default == 1
+	b.c.else == 1
+	b.c.with == 1
+	b.c.null == 1
+	b.c.true == 1
+	b.c.false == 1
+	b.c.some == 1
+
+	d.not.e == 1
+	d.package.e == 1
+	d.import.e == 1
+	d.as.e == 1
+	d.default.e == 1
+	d.else.e == 1
+	d.with.e == 1
+	d.null.e == 1
+	d.true.e == 1
+	d.false.e == 1
+	d.some.e == 1
+
+	f.not(2)
+	f.package(2)
+	f.import(2)
+	f.as(2)
+	f.default(2)
+	f.else(2)
+	f.with(2)
+	f.null(2)
+	f.true(2)
+	f.false(2)
+	f.some(2)
+
+	g.h.not(2)
+	g.h.package(2)
+	g.h.import(2)
+	g.h.as(2)
+	g.h.default(2)
+	g.h.else(2)
+	g.h.with(2)
+	g.h.null(2)
+	g.h.true(2)
+	g.h.false(2)
+	g.h.some(2)
+
+	i.not.j(2)
+	i.package.j(2)
+	i.import.j(2)
+	i.as.j(2)
+	i.default.j(2)
+	i.else.j(2)
+	i.with.j(2)
+	i.null.j(2)
+	i.true.j(2)
+	i.false.j(2)
+	i.some.j(2)
+}
+
+a.not := 1
+
+a.package := 1
+
+a.import := 1
+
+a.as := 1
+
+a.default := 1
+
+a.else := 1
+
+a.with := 1
+
+a.null := 1
+
+a.true := 1
+
+a.false := 1
+
+a.some := 1
+
+b.c.not := 1
+
+b.c.package := 1
+
+b.c.import := 1
+
+b.c.as := 1
+
+b.c.default := 1
+
+b.c.else := 1
+
+b.c.with := 1
+
+b.c.null := 1
+
+b.c.true := 1
+
+b.c.false := 1
+
+b.c.some := 1
+
+d.not.e := 1
+
+d.package.e := 1
+
+d.import.e := 1
+
+d.as.e := 1
+
+d.default.e := 1
+
+d.else.e := 1
+
+d.with.e := 1
+
+d.null.e := 1
+
+d.true.e := 1
+
+d.false.e := 1
+
+d.some.e := 1
+
+f.not(x) := x
+
+f.package(x) := x
+
+f.import(x) := x
+
+f.as(x) := x
+
+f.default(x) := x
+
+f.else(x) := x
+
+f.with(x) := x
+
+f.null(x) := x
+
+f.true(x) := x
+
+f.false(x) := x
+
+f.some(x) := x
+
+g.h.not(x) := x
+
+g.h.package(x) := x
+
+g.h.import(x) := x
+
+g.h.as(x) := x
+
+g.h.default(x) := x
+
+g.h.else(x) := x
+
+g.h.with(x) := x
+
+g.h.null(x) := x
+
+g.h.true(x) := x
+
+g.h.false(x) := x
+
+g.h.some(x) := x
+
+i.not.j(x) := x
+
+i.package.j(x) := x
+
+i.import.j(x) := x
+
+i.as.j(x) := x
+
+i.default.j(x) := x
+
+i.else.j(x) := x
+
+i.with.j(x) := x
+
+i.null.j(x) := x
+
+i.true.j(x) := x
+
+i.false.j(x) := x
+
+i.some.j(x) := x

v1/format/testfiles/v0/test_keywords_in_refs.rego.formatted_no_keywords_in_refs

@@ -0,0 +1,207 @@
+package test["not"]["package"]["import"]["as"]["default"]["else"]["with"]["null"]["true"]["false"]["some"]
+
+p {
+	a["not"] == 1
+	a["package"] == 1
+	a["import"] == 1
+	a["as"] == 1
+	a["default"] == 1
+	a["else"] == 1
+	a["with"] == 1
+	a["null"] == 1
+	a["true"] == 1
+	a["false"] == 1
+	a["some"] == 1
+
+	b.c["not"] == 1
+	b.c["package"] == 1
+	b.c["import"] == 1
+	b.c["as"] == 1
+	b.c["default"] == 1
+	b.c["else"] == 1
+	b.c["with"] == 1
+	b.c["null"] == 1
+	b.c["true"] == 1
+	b.c["false"] == 1
+	b.c["some"] == 1
+
+	d["not"].e == 1
+	d["package"].e == 1
+	d["import"].e == 1
+	d["as"].e == 1
+	d["default"].e == 1
+	d["else"].e == 1
+	d["with"].e == 1
+	d["null"].e == 1
+	d["true"].e == 1
+	d["false"].e == 1
+	d["some"].e == 1
+
+	f["not"](2)
+	f["package"](2)
+	f["import"](2)
+	f["as"](2)
+	f["default"](2)
+	f["else"](2)
+	f["with"](2)
+	f["null"](2)
+	f["true"](2)
+	f["false"](2)
+	f["some"](2)
+
+	g.h["not"](2)
+	g.h["package"](2)
+	g.h["import"](2)
+	g.h["as"](2)
+	g.h["default"](2)
+	g.h["else"](2)
+	g.h["with"](2)
+	g.h["null"](2)
+	g.h["true"](2)
+	g.h["false"](2)
+	g.h["some"](2)
+
+	i["not"].j(2)
+	i["package"].j(2)
+	i["import"].j(2)
+	i["as"].j(2)
+	i["default"].j(2)
+	i["else"].j(2)
+	i["with"].j(2)
+	i["null"].j(2)
+	i["true"].j(2)
+	i["false"].j(2)
+	i["some"].j(2)
+}
+
+a["not"] := 1
+
+a["package"] := 1
+
+a["import"] := 1
+
+a["as"] := 1
+
+a["default"] := 1
+
+a["else"] := 1
+
+a["with"] := 1
+
+a["null"] := 1
+
+a["true"] := 1
+
+a["false"] := 1
+
+a["some"] := 1
+
+b.c["not"] := 1
+
+b.c["package"] := 1
+
+b.c["import"] := 1
+
+b.c["as"] := 1
+
+b.c["default"] := 1
+
+b.c["else"] := 1
+
+b.c["with"] := 1
+
+b.c["null"] := 1
+
+b.c["true"] := 1
+
+b.c["false"] := 1
+
+b.c["some"] := 1
+
+d["not"].e := 1
+
+d["package"].e := 1
+
+d["import"].e := 1
+
+d["as"].e := 1
+
+d["default"].e := 1
+
+d["else"].e := 1
+
+d["with"].e := 1
+
+d["null"].e := 1
+
+d["true"].e := 1
+
+d["false"].e := 1
+
+d["some"].e := 1
+
+f["not"](x) := x
+
+f["package"](x) := x
+
+f["import"](x) := x
+
+f["as"](x) := x
+
+f["default"](x) := x
+
+f["else"](x) := x
+
+f["with"](x) := x
+
+f["null"](x) := x
+
+f["true"](x) := x
+
+f["false"](x) := x
+
+f["some"](x) := x
+
+g.h["not"](x) := x
+
+g.h["package"](x) := x
+
+g.h["import"](x) := x
+
+g.h["as"](x) := x
+
+g.h["default"](x) := x
+
+g.h["else"](x) := x
+
+g.h["with"](x) := x
+
+g.h["null"](x) := x
+
+g.h["true"](x) := x
+
+g.h["false"](x) := x
+
+g.h["some"](x) := x
+
+i["not"].j(x) := x
+
+i["package"].j(x) := x
+
+i["import"].j(x) := x
+
+i["as"].j(x) := x
+
+i["default"].j(x) := x
+
+i["else"].j(x) := x
+
+i["with"].j(x) := x
+
+i["null"].j(x) := x
+
+i["true"].j(x) := x
+
+i["false"].j(x) := x
+
+i["some"].j(x) := x

v1/format/testfiles/v0/test_keywords_in_refs_keep_brackets.rego

@@ -0,0 +1,147 @@
+package test["not"]["package"]["import"]["as"]["default"]["else"]["with"]["null"]["true"]["false"]["some"]
+
+p {
+	a["not"] == 1
+	a["package"] == 1
+	a["import"] == 1
+	a["as"] == 1
+	a["default"] == 1
+	a["else"] == 1
+	a["with"] == 1
+	a["null"] == 1
+	a["true"] == 1
+	a["false"] == 1
+	a["some"] == 1
+
+	b.c["not"] == 1
+	b.c["package"] == 1
+	b.c["import"] == 1
+	b.c["as"] == 1
+	b.c["default"] == 1
+	b.c["else"] == 1
+	b.c["with"] == 1
+	b.c["null"] == 1
+	b.c["true"] == 1
+	b.c["false"] == 1
+	b.c["some"] == 1
+
+	d["not"].e == 1
+	d["package"].e == 1
+	d["import"].e == 1
+	d["as"].e == 1
+	d["default"].e == 1
+	d["else"].e == 1
+	d["with"].e == 1
+	d["null"].e == 1
+	d["true"].e == 1
+	d["false"].e == 1
+	d["some"].e == 1
+
+	f["not"](2)
+	f["package"](2)
+	f["import"](2)
+	f["as"](2)
+	f["default"](2)
+	f["else"](2)
+	f["with"](2)
+	f["null"](2)
+	f["true"](2)
+	f["false"](2)
+	f["some"](2)
+
+	g.h["not"](2)
+	g.h["package"](2)
+	g.h["import"](2)
+	g.h["as"](2)
+	g.h["default"](2)
+	g.h["else"](2)
+	g.h["with"](2)
+	g.h["null"](2)
+	g.h["true"](2)
+	g.h["false"](2)
+	g.h["some"](2)
+
+	i["not"].j(2)
+	i["package"].j(2)
+	i["import"].j(2)
+	i["as"].j(2)
+	i["default"].j(2)
+	i["else"].j(2)
+	i["with"].j(2)
+	i["null"].j(2)
+	i["true"].j(2)
+	i["false"].j(2)
+	i["some"].j(2)
+}
+
+a["not"] := 1
+a["package"] := 1
+a["import"] := 1
+a["as"] := 1
+a["default"] := 1
+a["else"] := 1
+a["with"] := 1
+a["null"] := 1
+a["true"] := 1
+a["false"] := 1
+a["some"] := 1
+
+b.c["not"] := 1
+b.c["package"] := 1
+b.c["import"] := 1
+b.c["as"] := 1
+b.c["default"] := 1
+b.c["else"] := 1
+b.c["with"] := 1
+b.c["null"] := 1
+b.c["true"] := 1
+b.c["false"] := 1
+b.c["some"] := 1
+
+d["not"].e := 1
+d["package"].e := 1
+d["import"].e := 1
+d["as"].e := 1
+d["default"].e := 1
+d["else"].e := 1
+d["with"].e := 1
+d["null"].e := 1
+d["true"].e := 1
+d["false"].e := 1
+d["some"].e := 1
+
+f["not"](x) := x
+f["package"](x) := x
+f["import"](x) := x
+f["as"](x) := x
+f["default"](x) := x
+f["else"](x) := x
+f["with"](x) := x
+f["null"](x) := x
+f["true"](x) := x
+f["false"](x) := x
+f["some"](x) := x
+
+g.h["not"](x) := x
+g.h["package"](x) := x
+g.h["import"](x) := x
+g.h["as"](x) := x
+g.h["default"](x) := x
+g.h["else"](x) := x
+g.h["with"](x) := x
+g.h["null"](x) := x
+g.h["true"](x) := x
+g.h["false"](x) := x
+g.h["some"](x) := x
+
+i["not"].j(x) := x
+i["package"].j(x) := x
+i["import"].j(x) := x
+i["as"].j(x) := x
+i["default"].j(x) := x
+i["else"].j(x) := x
+i["with"].j(x) := x
+i["null"].j(x) := x
+i["true"].j(x) := x
+i["false"].j(x) := x
+i["some"].j(x) := x

v1/format/testfiles/v0/test_keywords_in_refs_keep_brackets.rego.formatted

@@ -0,0 +1,207 @@
+package test["not"]["package"]["import"]["as"]["default"]["else"]["with"]["null"]["true"]["false"]["some"]
+
+p {
+	a["not"] == 1
+	a["package"] == 1
+	a["import"] == 1
+	a["as"] == 1
+	a["default"] == 1
+	a["else"] == 1
+	a["with"] == 1
+	a["null"] == 1
+	a["true"] == 1
+	a["false"] == 1
+	a["some"] == 1
+
+	b.c["not"] == 1
+	b.c["package"] == 1
+	b.c["import"] == 1
+	b.c["as"] == 1
+	b.c["default"] == 1
+	b.c["else"] == 1
+	b.c["with"] == 1
+	b.c["null"] == 1
+	b.c["true"] == 1
+	b.c["false"] == 1
+	b.c["some"] == 1
+
+	d["not"].e == 1
+	d["package"].e == 1
+	d["import"].e == 1
+	d["as"].e == 1
+	d["default"].e == 1
+	d["else"].e == 1
+	d["with"].e == 1
+	d["null"].e == 1
+	d["true"].e == 1
+	d["false"].e == 1
+	d["some"].e == 1
+
+	f["not"](2)
+	f["package"](2)
+	f["import"](2)
+	f["as"](2)
+	f["default"](2)
+	f["else"](2)
+	f["with"](2)
+	f["null"](2)
+	f["true"](2)
+	f["false"](2)
+	f["some"](2)
+
+	g.h["not"](2)
+	g.h["package"](2)
+	g.h["import"](2)
+	g.h["as"](2)
+	g.h["default"](2)
+	g.h["else"](2)
+	g.h["with"](2)
+	g.h["null"](2)
+	g.h["true"](2)
+	g.h["false"](2)
+	g.h["some"](2)
+
+	i["not"].j(2)
+	i["package"].j(2)
+	i["import"].j(2)
+	i["as"].j(2)
+	i["default"].j(2)
+	i["else"].j(2)
+	i["with"].j(2)
+	i["null"].j(2)
+	i["true"].j(2)
+	i["false"].j(2)
+	i["some"].j(2)
+}
+
+a["not"] := 1
+
+a["package"] := 1
+
+a["import"] := 1
+
+a["as"] := 1
+
+a["default"] := 1
+
+a["else"] := 1
+
+a["with"] := 1
+
+a["null"] := 1
+
+a["true"] := 1
+
+a["false"] := 1
+
+a["some"] := 1
+
+b.c["not"] := 1
+
+b.c["package"] := 1
+
+b.c["import"] := 1
+
+b.c["as"] := 1
+
+b.c["default"] := 1
+
+b.c["else"] := 1
+
+b.c["with"] := 1
+
+b.c["null"] := 1
+
+b.c["true"] := 1
+
+b.c["false"] := 1
+
+b.c["some"] := 1
+
+d["not"].e := 1
+
+d["package"].e := 1
+
+d["import"].e := 1
+
+d["as"].e := 1
+
+d["default"].e := 1
+
+d["else"].e := 1
+
+d["with"].e := 1
+
+d["null"].e := 1
+
+d["true"].e := 1
+
+d["false"].e := 1
+
+d["some"].e := 1
+
+f["not"](x) := x
+
+f["package"](x) := x
+
+f["import"](x) := x
+
+f["as"](x) := x
+
+f["default"](x) := x
+
+f["else"](x) := x
+
+f["with"](x) := x
+
+f["null"](x) := x
+
+f["true"](x) := x
+
+f["false"](x) := x
+
+f["some"](x) := x
+
+g.h["not"](x) := x
+
+g.h["package"](x) := x
+
+g.h["import"](x) := x
+
+g.h["as"](x) := x
+
+g.h["default"](x) := x
+
+g.h["else"](x) := x
+
+g.h["with"](x) := x
+
+g.h["null"](x) := x
+
+g.h["true"](x) := x
+
+g.h["false"](x) := x
+
+g.h["some"](x) := x
+
+i["not"].j(x) := x
+
+i["package"].j(x) := x
+
+i["import"].j(x) := x
+
+i["as"].j(x) := x
+
+i["default"].j(x) := x
+
+i["else"].j(x) := x
+
+i["with"].j(x) := x
+
+i["null"].j(x) := x
+
+i["true"].j(x) := x
+
+i["false"].j(x) := x
+
+i["some"].j(x) := x

v1/format/testfiles/v0_to_v1/keywords.rego

@@ -1,5 +1,6 @@
 package test.if.contains.in.every
 
 p {
-	data.if.contains.in.every == 42
+	data.if.contains.in.every == 1
+	data["if"]["contains"]["in"]["every"] == 2
 }

v1/format/testfiles/v0_to_v1/keywords.rego.formatted

@@ -1,7 +1,8 @@
-package test["if"]["contains"]["in"]["every"]
+package test.if.contains.in.every
 
 import rego.v1
 
 p if {
-	data["if"]["contains"]["in"]["every"] == 42
+	data.if.contains.in.every == 1
+	data["if"]["contains"]["in"]["every"] == 2
 }

v1/format/testfiles/v1/test_keywords_in_refs.rego

@@ -0,0 +1,195 @@
+package test.not.package.import.as.default.else.with.null.true.false.some.if.contains.in.every
+
+p if {
+	a.not == 1
+	a.package == 1
+	a.import == 1
+	a.as == 1
+	a.default == 1
+	a.else == 1
+	a.with == 1
+	a.null == 1
+	a.true == 1
+	a.false == 1
+	a.some == 1
+	a.if == 1
+	a.contains == 1
+	a.in == 1
+	a.every == 1
+
+	b.c.not == 1
+	b.c.package == 1
+	b.c.import == 1
+	b.c.as == 1
+	b.c.default == 1
+	b.c.else == 1
+	b.c.with == 1
+	b.c.null == 1
+	b.c.true == 1
+	b.c.false == 1
+	b.c.some == 1
+	b.c.if == 1
+	b.c.contains == 1
+	b.c.in == 1
+	b.c.every == 1
+
+	d.not.e == 1
+	d.package.e == 1
+	d.import.e == 1
+	d.as.e == 1
+	d.default.e == 1
+	d.else.e == 1
+	d.with.e == 1
+	d.null.e == 1
+	d.true.e == 1
+	d.false.e == 1
+	d.some.e == 1
+	d.if.e == 1
+	d.contains.e == 1
+	d.in.e == 1
+	d.every.e == 1
+
+	f.not(2)
+	f.package(2)
+	f.import(2)
+	f.as(2)
+	f.default(2)
+	f.else(2)
+	f.with(2)
+	f.null(2)
+	f.true(2)
+	f.false(2)
+	f.some(2)
+	f.if(2)
+	f.contains(2)
+	f.in(2)
+	f.every(2)
+
+	g.h.not(2)
+	g.h.package(2)
+	g.h.import(2)
+	g.h.as(2)
+	g.h.default(2)
+	g.h.else(2)
+	g.h.with(2)
+	g.h.null(2)
+	g.h.true(2)
+	g.h.false(2)
+	g.h.some(2)
+	g.h.if(2)
+	g.h.contains(2)
+	g.h.in(2)
+	g.h.every(2)
+
+	i.not.j(2)
+	i.package.j(2)
+	i.import.j(2)
+	i.as.j(2)
+	i.default.j(2)
+	i.else.j(2)
+	i.with.j(2)
+	i.null.j(2)
+	i.true.j(2)
+	i.false.j(2)
+	i.some.j(2)
+	i.if.j(2)
+	i.contains.j(2)
+	i.in.j(2)
+	i.every.j(2)
+}
+
+a.not := 1
+a.package := 1
+a.import := 1
+a.as := 1
+a.default := 1
+a.else := 1
+a.with := 1
+a.null := 1
+a.true := 1
+a.false := 1
+a.some := 1
+a.if := 1
+a.contains := 1
+a.in := 1
+a.every := 1
+
+b.c.not := 1
+b.c.package := 1
+b.c.import := 1
+b.c.as := 1
+b.c.default := 1
+b.c.else := 1
+b.c.with := 1
+b.c.null := 1
+b.c.true := 1
+b.c.false := 1
+b.c.some := 1
+b.c.if := 1
+b.c.contains := 1
+b.c.in := 1
+b.c.every := 1
+
+d.not.e := 1
+d.package.e := 1
+d.import.e := 1
+d.as.e := 1
+d.default.e := 1
+d.else.e := 1
+d.with.e := 1
+d.null.e := 1
+d.true.e := 1
+d.false.e := 1
+d.some.e := 1
+d.if.e := 1
+d.contains.e := 1
+d.in.e := 1
+d.every.e := 1
+
+f.not(x) := x
+f.package(x) := x
+f.import(x) := x
+f.as(x) := x
+f.default(x) := x
+f.else(x) := x
+f.with(x) := x
+f.null(x) := x
+f.true(x) := x
+f.false(x) := x
+f.some(x) := x
+f.if(x) := x
+f.contains(x) := x
+f.in(x) := x
+f.every(x) := x
+
+g.h.not(x) := x
+g.h.package(x) := x
+g.h.import(x) := x
+g.h.as(x) := x
+g.h.default(x) := x
+g.h.else(x) := x
+g.h.with(x) := x
+g.h.null(x) := x
+g.h.true(x) := x
+g.h.false(x) := x
+g.h.some(x) := x
+g.h.if(x) := x
+g.h.contains(x) := x
+g.h.in(x) := x
+g.h.every(x) := x
+
+i.not.j(x) := x
+i.package.j(x) := x
+i.import.j(x) := x
+i.as.j(x) := x
+i.default.j(x) := x
+i.else.j(x) := x
+i.with.j(x) := x
+i.null.j(x) := x
+i.true.j(x) := x
+i.false.j(x) := x
+i.some.j(x) := x
+i.if.j(x) := x
+i.contains.j(x) := x
+i.in.j(x) := x
+i.every.j(x) := x

v1/format/testfiles/v1/test_keywords_in_refs.rego.formatted

@@ -0,0 +1,195 @@
+package test.not.package.import.as.default.else.with.null.true.false.some.if.contains.in.every
+
+p if {
+	a.not == 1
+	a.package == 1
+	a.import == 1
+	a.as == 1
+	a.default == 1
+	a.else == 1
+	a.with == 1
+	a.null == 1
+	a.true == 1
+	a.false == 1
+	a.some == 1
+	a.if == 1
+	a.contains == 1
+	a.in == 1
+	a.every == 1
+
+	b.c.not == 1
+	b.c.package == 1
+	b.c.import == 1
+	b.c.as == 1
+	b.c.default == 1
+	b.c.else == 1
+	b.c.with == 1
+	b.c.null == 1
+	b.c.true == 1
+	b.c.false == 1
+	b.c.some == 1
+	b.c.if == 1
+	b.c.contains == 1
+	b.c.in == 1
+	b.c.every == 1
+
+	d.not.e == 1
+	d.package.e == 1
+	d.import.e == 1
+	d.as.e == 1
+	d.default.e == 1
+	d.else.e == 1
+	d.with.e == 1
+	d.null.e == 1
+	d.true.e == 1
+	d.false.e == 1
+	d.some.e == 1
+	d.if.e == 1
+	d.contains.e == 1
+	d.in.e == 1
+	d.every.e == 1
+
+	f.not(2)
+	f.package(2)
+	f.import(2)
+	f.as(2)
+	f.default(2)
+	f.else(2)
+	f.with(2)
+	f.null(2)
+	f.true(2)
+	f.false(2)
+	f.some(2)
+	f.if(2)
+	f.contains(2)
+	f.in(2)
+	f.every(2)
+
+	g.h.not(2)
+	g.h.package(2)
+	g.h.import(2)
+	g.h.as(2)
+	g.h.default(2)
+	g.h.else(2)
+	g.h.with(2)
+	g.h.null(2)
+	g.h.true(2)
+	g.h.false(2)
+	g.h.some(2)
+	g.h.if(2)
+	g.h.contains(2)
+	g.h.in(2)
+	g.h.every(2)
+
+	i.not.j(2)
+	i.package.j(2)
+	i.import.j(2)
+	i.as.j(2)
+	i.default.j(2)
+	i.else.j(2)
+	i.with.j(2)
+	i.null.j(2)
+	i.true.j(2)
+	i.false.j(2)
+	i.some.j(2)
+	i.if.j(2)
+	i.contains.j(2)
+	i.in.j(2)
+	i.every.j(2)
+}
+
+a.not := 1
+a.package := 1
+a.import := 1
+a.as := 1
+a.default := 1
+a.else := 1
+a.with := 1
+a.null := 1
+a.true := 1
+a.false := 1
+a.some := 1
+a.if := 1
+a.contains := 1
+a.in := 1
+a.every := 1
+
+b.c.not := 1
+b.c.package := 1
+b.c.import := 1
+b.c.as := 1
+b.c.default := 1
+b.c.else := 1
+b.c.with := 1
+b.c.null := 1
+b.c.true := 1
+b.c.false := 1
+b.c.some := 1
+b.c.if := 1
+b.c.contains := 1
+b.c.in := 1
+b.c.every := 1
+
+d.not.e := 1
+d.package.e := 1
+d.import.e := 1
+d.as.e := 1
+d.default.e := 1
+d.else.e := 1
+d.with.e := 1
+d.null.e := 1
+d.true.e := 1
+d.false.e := 1
+d.some.e := 1
+d.if.e := 1
+d.contains.e := 1
+d.in.e := 1
+d.every.e := 1
+
+f.not(x) := x
+f.package(x) := x
+f.import(x) := x
+f.as(x) := x
+f.default(x) := x
+f.else(x) := x
+f.with(x) := x
+f.null(x) := x
+f.true(x) := x
+f.false(x) := x
+f.some(x) := x
+f.if(x) := x
+f.contains(x) := x
+f.in(x) := x
+f.every(x) := x
+
+g.h.not(x) := x
+g.h.package(x) := x
+g.h.import(x) := x
+g.h.as(x) := x
+g.h.default(x) := x
+g.h.else(x) := x
+g.h.with(x) := x
+g.h.null(x) := x
+g.h.true(x) := x
+g.h.false(x) := x
+g.h.some(x) := x
+g.h.if(x) := x
+g.h.contains(x) := x
+g.h.in(x) := x
+g.h.every(x) := x
+
+i.not.j(x) := x
+i.package.j(x) := x
+i.import.j(x) := x
+i.as.j(x) := x
+i.default.j(x) := x
+i.else.j(x) := x
+i.with.j(x) := x
+i.null.j(x) := x
+i.true.j(x) := x
+i.false.j(x) := x
+i.some.j(x) := x
+i.if.j(x) := x
+i.contains.j(x) := x
+i.in.j(x) := x
+i.every.j(x) := x

v1/format/testfiles/v1/test_keywords_in_refs.rego.formatted_no_keywords_in_refs

@@ -0,0 +1,195 @@
+package test["not"]["package"]["import"]["as"]["default"]["else"]["with"]["null"]["true"]["false"]["some"]["if"]["contains"]["in"]["every"]
+
+p if {
+	a["not"] == 1
+	a["package"] == 1
+	a["import"] == 1
+	a["as"] == 1
+	a["default"] == 1
+	a["else"] == 1
+	a["with"] == 1
+	a["null"] == 1
+	a["true"] == 1
+	a["false"] == 1
+	a["some"] == 1
+	a["if"] == 1
+	a["contains"] == 1
+	a["in"] == 1
+	a["every"] == 1
+
+	b.c["not"] == 1
+	b.c["package"] == 1
+	b.c["import"] == 1
+	b.c["as"] == 1
+	b.c["default"] == 1
+	b.c["else"] == 1
+	b.c["with"] == 1
+	b.c["null"] == 1
+	b.c["true"] == 1
+	b.c["false"] == 1
+	b.c["some"] == 1
+	b.c["if"] == 1
+	b.c["contains"] == 1
+	b.c["in"] == 1
+	b.c["every"] == 1
+
+	d["not"].e == 1
+	d["package"].e == 1
+	d["import"].e == 1
+	d["as"].e == 1
+	d["default"].e == 1
+	d["else"].e == 1
+	d["with"].e == 1
+	d["null"].e == 1
+	d["true"].e == 1
+	d["false"].e == 1
+	d["some"].e == 1
+	d["if"].e == 1
+	d["contains"].e == 1
+	d["in"].e == 1
+	d["every"].e == 1
+
+	f["not"](2)
+	f["package"](2)
+	f["import"](2)
+	f["as"](2)
+	f["default"](2)
+	f["else"](2)
+	f["with"](2)
+	f["null"](2)
+	f["true"](2)
+	f["false"](2)
+	f["some"](2)
+	f["if"](2)
+	f["contains"](2)
+	f["in"](2)
+	f["every"](2)
+
+	g.h["not"](2)
+	g.h["package"](2)
+	g.h["import"](2)
+	g.h["as"](2)
+	g.h["default"](2)
+	g.h["else"](2)
+	g.h["with"](2)
+	g.h["null"](2)
+	g.h["true"](2)
+	g.h["false"](2)
+	g.h["some"](2)
+	g.h["if"](2)
+	g.h["contains"](2)
+	g.h["in"](2)
+	g.h["every"](2)
+
+	i["not"].j(2)
+	i["package"].j(2)
+	i["import"].j(2)
+	i["as"].j(2)
+	i["default"].j(2)
+	i["else"].j(2)
+	i["with"].j(2)
+	i["null"].j(2)
+	i["true"].j(2)
+	i["false"].j(2)
+	i["some"].j(2)
+	i["if"].j(2)
+	i["contains"].j(2)
+	i["in"].j(2)
+	i["every"].j(2)
+}
+
+a["not"] := 1
+a["package"] := 1
+a["import"] := 1
+a["as"] := 1
+a["default"] := 1
+a["else"] := 1
+a["with"] := 1
+a["null"] := 1
+a["true"] := 1
+a["false"] := 1
+a["some"] := 1
+a["if"] := 1
+a["contains"] := 1
+a["in"] := 1
+a["every"] := 1
+
+b.c["not"] := 1
+b.c["package"] := 1
+b.c["import"] := 1
+b.c["as"] := 1
+b.c["default"] := 1
+b.c["else"] := 1
+b.c["with"] := 1
+b.c["null"] := 1
+b.c["true"] := 1
+b.c["false"] := 1
+b.c["some"] := 1
+b.c["if"] := 1
+b.c["contains"] := 1
+b.c["in"] := 1
+b.c["every"] := 1
+
+d["not"].e := 1
+d["package"].e := 1
+d["import"].e := 1
+d["as"].e := 1
+d["default"].e := 1
+d["else"].e := 1
+d["with"].e := 1
+d["null"].e := 1
+d["true"].e := 1
+d["false"].e := 1
+d["some"].e := 1
+d["if"].e := 1
+d["contains"].e := 1
+d["in"].e := 1
+d["every"].e := 1
+
+f["not"](x) := x
+f["package"](x) := x
+f["import"](x) := x
+f["as"](x) := x
+f["default"](x) := x
+f["else"](x) := x
+f["with"](x) := x
+f["null"](x) := x
+f["true"](x) := x
+f["false"](x) := x
+f["some"](x) := x
+f["if"](x) := x
+f["contains"](x) := x
+f["in"](x) := x
+f["every"](x) := x
+
+g.h["not"](x) := x
+g.h["package"](x) := x
+g.h["import"](x) := x
+g.h["as"](x) := x
+g.h["default"](x) := x
+g.h["else"](x) := x
+g.h["with"](x) := x
+g.h["null"](x) := x
+g.h["true"](x) := x
+g.h["false"](x) := x
+g.h["some"](x) := x
+g.h["if"](x) := x
+g.h["contains"](x) := x
+g.h["in"](x) := x
+g.h["every"](x) := x
+
+i["not"].j(x) := x
+i["package"].j(x) := x
+i["import"].j(x) := x
+i["as"].j(x) := x
+i["default"].j(x) := x
+i["else"].j(x) := x
+i["with"].j(x) := x
+i["null"].j(x) := x
+i["true"].j(x) := x
+i["false"].j(x) := x
+i["some"].j(x) := x
+i["if"].j(x) := x
+i["contains"].j(x) := x
+i["in"].j(x) := x
+i["every"].j(x) := x

v1/format/testfiles/v1/test_keywords_in_refs_keep_brackets.rego

@@ -0,0 +1,195 @@
+package test["not"]["package"]["import"]["as"]["default"]["else"]["with"]["null"]["true"]["false"]["some"]["if"]["contains"]["in"]["every"]
+
+p if {
+	a["not"] == 1
+	a["package"] == 1
+	a["import"] == 1
+	a["as"] == 1
+	a["default"] == 1
+	a["else"] == 1
+	a["with"] == 1
+	a["null"] == 1
+	a["true"] == 1
+	a["false"] == 1
+	a["some"] == 1
+	a["if"] == 1
+	a["contains"] == 1
+	a["in"] == 1
+	a["every"] == 1
+
+	b.c["not"] == 1
+	b.c["package"] == 1
+	b.c["import"] == 1
+	b.c["as"] == 1
+	b.c["default"] == 1
+	b.c["else"] == 1
+	b.c["with"] == 1
+	b.c["null"] == 1
+	b.c["true"] == 1
+	b.c["false"] == 1
+	b.c["some"] == 1
+	b.c["if"] == 1
+	b.c["contains"] == 1
+	b.c["in"] == 1
+	b.c["every"] == 1
+
+	d["not"].e == 1
+	d["package"].e == 1
+	d["import"].e == 1
+	d["as"].e == 1
+	d["default"].e == 1
+	d["else"].e == 1
+	d["with"].e == 1
+	d["null"].e == 1
+	d["true"].e == 1
+	d["false"].e == 1
+	d["some"].e == 1
+	d["if"].e == 1
+	d["contains"].e == 1
+	d["in"].e == 1
+	d["every"].e == 1
+
+	f["not"](2)
+	f["package"](2)
+	f["import"](2)
+	f["as"](2)
+	f["default"](2)
+	f["else"](2)
+	f["with"](2)
+	f["null"](2)
+	f["true"](2)
+	f["false"](2)
+	f["some"](2)
+	f["if"](2)
+	f["contains"](2)
+	f["in"](2)
+	f["every"](2)
+
+	g.h["not"](2)
+	g.h["package"](2)
+	g.h["import"](2)
+	g.h["as"](2)
+	g.h["default"](2)
+	g.h["else"](2)
+	g.h["with"](2)
+	g.h["null"](2)
+	g.h["true"](2)
+	g.h["false"](2)
+	g.h["some"](2)
+	g.h["if"](2)
+	g.h["contains"](2)
+	g.h["in"](2)
+	g.h["every"](2)
+
+	i["not"].j(2)
+	i["package"].j(2)
+	i["import"].j(2)
+	i["as"].j(2)
+	i["default"].j(2)
+	i["else"].j(2)
+	i["with"].j(2)
+	i["null"].j(2)
+	i["true"].j(2)
+	i["false"].j(2)
+	i["some"].j(2)
+	i["if"].j(2)
+	i["contains"].j(2)
+	i["in"].j(2)
+	i["every"].j(2)
+}
+
+a["not"] := 1
+a["package"] := 1
+a["import"] := 1
+a["as"] := 1
+a["default"] := 1
+a["else"] := 1
+a["with"] := 1
+a["null"] := 1
+a["true"] := 1
+a["false"] := 1
+a["some"] := 1
+a["if"] := 1
+a["contains"] := 1
+a["in"] := 1
+a["every"] := 1
+
+b.c["not"] := 1
+b.c["package"] := 1
+b.c["import"] := 1
+b.c["as"] := 1
+b.c["default"] := 1
+b.c["else"] := 1
+b.c["with"] := 1
+b.c["null"] := 1
+b.c["true"] := 1
+b.c["false"] := 1
+b.c["some"] := 1
+b.c["if"] := 1
+b.c["contains"] := 1
+b.c["in"] := 1
+b.c["every"] := 1
+
+d["not"].e := 1
+d["package"].e := 1
+d["import"].e := 1
+d["as"].e := 1
+d["default"].e := 1
+d["else"].e := 1
+d["with"].e := 1
+d["null"].e := 1
+d["true"].e := 1
+d["false"].e := 1
+d["some"].e := 1
+d["if"].e := 1
+d["contains"].e := 1
+d["in"].e := 1
+d["every"].e := 1
+
+f["not"](x) := x
+f["package"](x) := x
+f["import"](x) := x
+f["as"](x) := x
+f["default"](x) := x
+f["else"](x) := x
+f["with"](x) := x
+f["null"](x) := x
+f["true"](x) := x
+f["false"](x) := x
+f["some"](x) := x
+f["if"](x) := x
+f["contains"](x) := x
+f["in"](x) := x
+f["every"](x) := x
+
+g.h["not"](x) := x
+g.h["package"](x) := x
+g.h["import"](x) := x
+g.h["as"](x) := x
+g.h["default"](x) := x
+g.h["else"](x) := x
+g.h["with"](x) := x
+g.h["null"](x) := x
+g.h["true"](x) := x
+g.h["false"](x) := x
+g.h["some"](x) := x
+g.h["if"](x) := x
+g.h["contains"](x) := x
+g.h["in"](x) := x
+g.h["every"](x) := x
+
+i["not"].j(x) := x
+i["package"].j(x) := x
+i["import"].j(x) := x
+i["as"].j(x) := x
+i["default"].j(x) := x
+i["else"].j(x) := x
+i["with"].j(x) := x
+i["null"].j(x) := x
+i["true"].j(x) := x
+i["false"].j(x) := x
+i["some"].j(x) := x
+i["if"].j(x) := x
+i["contains"].j(x) := x
+i["in"].j(x) := x
+i["every"].j(x) := x

v1/format/testfiles/v1/test_keywords_in_refs_keep_brackets.rego.formatted

@@ -0,0 +1,195 @@
+package test["not"]["package"]["import"]["as"]["default"]["else"]["with"]["null"]["true"]["false"]["some"]["if"]["contains"]["in"]["every"]
+
+p if {
+	a["not"] == 1
+	a["package"] == 1
+	a["import"] == 1
+	a["as"] == 1
+	a["default"] == 1
+	a["else"] == 1
+	a["with"] == 1
+	a["null"] == 1
+	a["true"] == 1
+	a["false"] == 1
+	a["some"] == 1
+	a["if"] == 1
+	a["contains"] == 1
+	a["in"] == 1
+	a["every"] == 1
+
+	b.c["not"] == 1
+	b.c["package"] == 1
+	b.c["import"] == 1
+	b.c["as"] == 1
+	b.c["default"] == 1
+	b.c["else"] == 1
+	b.c["with"] == 1
+	b.c["null"] == 1
+	b.c["true"] == 1
+	b.c["false"] == 1
+	b.c["some"] == 1
+	b.c["if"] == 1
+	b.c["contains"] == 1
+	b.c["in"] == 1
+	b.c["every"] == 1
+
+	d["not"].e == 1
+	d["package"].e == 1
+	d["import"].e == 1
+	d["as"].e == 1
+	d["default"].e == 1
+	d["else"].e == 1
+	d["with"].e == 1
+	d["null"].e == 1
+	d["true"].e == 1
+	d["false"].e == 1
+	d["some"].e == 1
+	d["if"].e == 1
+	d["contains"].e == 1
+	d["in"].e == 1
+	d["every"].e == 1
+
+	f["not"](2)
+	f["package"](2)
+	f["import"](2)
+	f["as"](2)
+	f["default"](2)
+	f["else"](2)
+	f["with"](2)
+	f["null"](2)
+	f["true"](2)
+	f["false"](2)
+	f["some"](2)
+	f["if"](2)
+	f["contains"](2)
+	f["in"](2)
+	f["every"](2)
+
+	g.h["not"](2)
+	g.h["package"](2)
+	g.h["import"](2)
+	g.h["as"](2)
+	g.h["default"](2)
+	g.h["else"](2)
+	g.h["with"](2)
+	g.h["null"](2)
+	g.h["true"](2)
+	g.h["false"](2)
+	g.h["some"](2)
+	g.h["if"](2)
+	g.h["contains"](2)
+	g.h["in"](2)
+	g.h["every"](2)
+
+	i["not"].j(2)
+	i["package"].j(2)
+	i["import"].j(2)
+	i["as"].j(2)
+	i["default"].j(2)
+	i["else"].j(2)
+	i["with"].j(2)
+	i["null"].j(2)
+	i["true"].j(2)
+	i["false"].j(2)
+	i["some"].j(2)
+	i["if"].j(2)
+	i["contains"].j(2)
+	i["in"].j(2)
+	i["every"].j(2)
+}
+
+a["not"] := 1
+a["package"] := 1
+a["import"] := 1
+a["as"] := 1
+a["default"] := 1
+a["else"] := 1
+a["with"] := 1
+a["null"] := 1
+a["true"] := 1
+a["false"] := 1
+a["some"] := 1
+a["if"] := 1
+a["contains"] := 1
+a["in"] := 1
+a["every"] := 1
+
+b.c["not"] := 1
+b.c["package"] := 1
+b.c["import"] := 1
+b.c["as"] := 1
+b.c["default"] := 1
+b.c["else"] := 1
+b.c["with"] := 1
+b.c["null"] := 1
+b.c["true"] := 1
+b.c["false"] := 1
+b.c["some"] := 1
+b.c["if"] := 1
+b.c["contains"] := 1
+b.c["in"] := 1
+b.c["every"] := 1
+
+d["not"].e := 1
+d["package"].e := 1
+d["import"].e := 1
+d["as"].e := 1
+d["default"].e := 1
+d["else"].e := 1
+d["with"].e := 1
+d["null"].e := 1
+d["true"].e := 1
+d["false"].e := 1
+d["some"].e := 1
+d["if"].e := 1
+d["contains"].e := 1
+d["in"].e := 1
+d["every"].e := 1
+
+f["not"](x) := x
+f["package"](x) := x
+f["import"](x) := x
+f["as"](x) := x
+f["default"](x) := x
+f["else"](x) := x
+f["with"](x) := x
+f["null"](x) := x
+f["true"](x) := x
+f["false"](x) := x
+f["some"](x) := x
+f["if"](x) := x
+f["contains"](x) := x
+f["in"](x) := x
+f["every"](x) := x
+
+g.h["not"](x) := x
+g.h["package"](x) := x
+g.h["import"](x) := x
+g.h["as"](x) := x
+g.h["default"](x) := x
+g.h["else"](x) := x
+g.h["with"](x) := x
+g.h["null"](x) := x
+g.h["true"](x) := x
+g.h["false"](x) := x
+g.h["some"](x) := x
+g.h["if"](x) := x
+g.h["contains"](x) := x
+g.h["in"](x) := x
+g.h["every"](x) := x
+
+i["not"].j(x) := x
+i["package"].j(x) := x
+i["import"].j(x) := x
+i["as"].j(x) := x
+i["default"].j(x) := x
+i["else"].j(x) := x
+i["with"].j(x) := x
+i["null"].j(x) := x
+i["true"].j(x) := x
+i["false"].j(x) := x
+i["some"].j(x) := x
+i["if"].j(x) := x
+i["contains"].j(x) := x
+i["in"].j(x) := x
+i["every"].j(x) := x

v1/test/cases/internal/keywordrefs/main.go

@@ -0,0 +1,65 @@
+// Copyright 2025 The OPA Authors.  All rights reserved.
+// Use of this source code is governed by an Apache2
+// license that can be found in the LICENSE file.
+
+package main
+
+import (
+	"fmt"
+	"os"
+	"strings"
+
+	"github.com/open-policy-agent/opa/v1/ast"
+)
+
+//go:generate go run main.go
+
+func main() {
+	templateFile := "test-keywords-in-ref_v0.yaml.template"
+	templateContent, err := os.ReadFile(templateFile)
+	if err != nil {
+		fmt.Printf("Error reading v0 template file: %v\n", err)
+		os.Exit(1)
+	}
+
+	err = generate(ast.KeywordsV0[:], "../../testdata/v0/keywordrefs", string(templateContent))
+	if err != nil {
+		fmt.Printf("Error:%v\n", err)
+		os.Exit(1)
+	}
+
+	templateFile = "test-keywords-in-ref_v1.yaml.template"
+	templateContent, err = os.ReadFile(templateFile)
+	if err != nil {
+		fmt.Printf("Error reading v1 template file: %v\n", err)
+		os.Exit(1)
+	}
+
+	err = generate(ast.KeywordsV1[:], "../../testdata/v1/keywordrefs", string(templateContent))
+	if err != nil {
+		fmt.Printf("Error:%v\n", err)
+		os.Exit(1)
+	}
+}
+
+func generate(keywords []string, root string, template string) error {
+	err := os.MkdirAll(root, os.ModePerm)
+	if err != nil {
+		return fmt.Errorf("error creating directory %s: %v\n", root, err)
+	}
+
+	// Generate a YAML file for each keyword
+	for _, keyword := range keywords {
+		file := fmt.Sprintf("%s/test-keyword-%s.yaml", root, keyword)
+
+		outputContent := strings.ReplaceAll(template, "%{KW}", keyword)
+
+		err := os.WriteFile(file, []byte(outputContent), 0644)
+		if err != nil {
+			return fmt.Errorf("error writing file %s: %v\n", file, err)
+		}
+		fmt.Printf("Generated file: %s\n", file)
+	}
+
+	return nil
+}

v1/test/cases/internal/keywordrefs/test-keywords-in-ref_v0.yaml.template

@@ -0,0 +1,200 @@
+---
+cases:
+  - note: keywordrefs/%{KW} keyword in package
+    query: data.foo.p = x
+    modules:
+      - |
+        package foo.%{KW}.bar
+
+        baz := 42
+      - |
+        package foo
+        import data.foo.%{KW}.bar
+
+        p {
+          bar.baz == 42
+          data.foo.%{KW}.bar.baz == 42
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/%{KW} keyword in package, import alias
+    query: data.foo.p = x
+    modules:
+      - |
+        package foo.%{KW}
+
+        bar := 42
+      - |
+        package foo
+        import data.foo.%{KW} as my_if
+
+        p {
+          my_if.bar == 42
+          data.foo.%{KW}.bar == 42
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/%{KW} keyword rule refhead
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p {
+          %{KW}.foo == 1
+          foo.%{KW} == 2
+        }
+
+        %{KW}.foo := 1
+
+        foo.%{KW} := 2
+    want_result:
+      - x: true
+  - note: keywordrefs/%{KW} keyword rule refhead, else bodies
+    query: data.test.p = x
+    input:
+      x: 3
+    modules:
+      - |
+        package test
+
+        p {
+          %{KW}.foo.bar == 3
+          foo.bar.%{KW} == 6
+        }
+
+        %{KW}.foo.bar := 1 {
+          input.x == 1
+        } else := 2 {
+          input.x == 2
+        } else := 3
+
+        foo.bar.%{KW} := 4 {
+          input.x == 1
+        } else := 5 {
+          input.x == 2
+        } else := 6
+    want_result:
+      - x: true
+  - note: keywordrefs/%{KW} keyword rule refhead, partial set
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+        import future.keywords.contains
+
+        p {
+          %{KW}.foo.bar == {"a", "c"}
+          foo.bar.%{KW} == {"a", "c"}
+        }
+
+        %{KW}.foo.bar contains "a"
+
+        %{KW}.foo.bar contains "b" {
+          false
+        }
+
+        %{KW}.foo.bar contains "c" {
+          true
+        }
+
+        foo.bar.%{KW} contains "a"
+
+        foo.bar.%{KW} contains "b" {
+          false
+        }
+
+        foo.bar.%{KW} contains "c" {
+          true
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/%{KW} keyword rule refhead, partial object
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p {
+          %{KW}.foo == "a"
+          %{KW}.bar.one == "a"
+          %{KW}.bar.three == "c"
+          foo.%{KW} == "a"
+          bar.baz.%{KW} == "a"
+        }
+
+        %{KW}.foo := "a"
+
+        %{KW}.foo := "b" {
+          false
+        }
+
+        %{KW}.foo := "c" {
+          false
+        }
+
+        %{KW}.bar.one := "a"
+
+        %{KW}.bar.two := "b" {
+          false
+        }
+
+        %{KW}.bar.three := "c" {
+          true
+        }
+
+        foo.%{KW} := "a"
+
+        foo.%{KW} := "b" {
+          false
+        }
+
+        foo.%{KW} := "c" {
+          false
+        }
+
+        bar.baz.%{KW} := "a"
+
+        bar.baz.%{KW} := "b" {
+          false
+        }
+
+        bar.baz.%{KW} := "c" {
+          false
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/%{KW} keyword function refhead
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p {
+          %{KW}.foo(1) == 1
+          %{KW}.foo(11) == 42
+          foo.%{KW}(1) == 1
+          foo.%{KW}(11) == 42
+          bar.%{KW}.baz(1) == 1
+          bar.%{KW}.baz(11) == 42
+        }
+
+        default %{KW}.foo(_) := 42
+
+        %{KW}.foo(x) := x {
+          x < 10
+        }
+
+        default foo.%{KW}(_) := 42
+
+        foo.%{KW}(x) := x {
+          x < 10
+        }
+
+        default bar.%{KW}.baz(_) := 42
+
+        bar.%{KW}.baz(x) := x {
+          x < 10
+        }
+    want_result:
+      - x: true

v1/test/cases/internal/keywordrefs/test-keywords-in-ref_v1.yaml.template

@@ -0,0 +1,199 @@
+---
+cases:
+  - note: keywordrefs/%{KW} keyword in package
+    query: data.foo.p = x
+    modules:
+      - |
+        package foo.%{KW}.bar
+
+        baz := 42
+      - |
+        package foo
+        import data.foo.%{KW}.bar
+
+        p if {
+          bar.baz == 42
+          data.foo.%{KW}.bar.baz == 42
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/%{KW} keyword in package, import alias
+    query: data.foo.p = x
+    modules:
+      - |
+        package foo.%{KW}
+
+        bar := 42
+      - |
+        package foo
+        import data.foo.%{KW} as my_if
+
+        p if {
+          my_if.bar == 42
+          data.foo.%{KW}.bar == 42
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/%{KW} keyword rule refhead
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p if {
+          %{KW}.foo == 1
+          foo.%{KW} == 2
+        }
+
+        %{KW}.foo := 1
+
+        foo.%{KW} := 2
+    want_result:
+      - x: true
+  - note: keywordrefs/%{KW} keyword rule refhead, else bodies
+    query: data.test.p = x
+    input:
+      x: 3
+    modules:
+      - |
+        package test
+
+        p if {
+          %{KW}.foo == 3
+          foo.%{KW} == 6
+        }
+
+        %{KW}.foo := 1 if {
+          input.x == 1
+        } else := 2 if {
+          input.x == 2
+        } else := 3
+
+        foo.%{KW} := 4 if {
+          input.x == 1
+        } else := 5 if {
+          input.x == 2
+        } else := 6
+    want_result:
+      - x: true
+  - note: keywordrefs/%{KW} keyword rule refhead, partial set
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p if {
+          %{KW}.foo == {"a", "c"}
+          foo.%{KW} == {"a", "c"}
+        }
+
+        %{KW}.foo contains "a"
+
+        %{KW}.foo contains "b" if {
+          false
+        }
+
+        %{KW}.foo contains "c" if {
+          true
+        }
+
+        foo.%{KW} contains "a"
+
+        foo.%{KW} contains "b" if {
+          false
+        }
+
+        foo.%{KW} contains "c" if {
+          true
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/%{KW} keyword rule refhead, partial object
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p if {
+          %{KW}.foo == "a"
+          %{KW}.bar.one == "a"
+          %{KW}.bar.three == "c"
+          foo.%{KW} == "a"
+          bar.baz.%{KW} == "a"
+        }
+
+        %{KW}.foo := "a"
+
+        %{KW}.foo := "b" if {
+          false
+        }
+
+        %{KW}.foo := "c" if {
+          false
+        }
+
+        %{KW}.bar.one := "a"
+
+        %{KW}.bar.two := "b" if {
+          false
+        }
+
+        %{KW}.bar.three := "c" if {
+          true
+        }
+
+        foo.%{KW} := "a"
+
+        foo.%{KW} := "b" if {
+          false
+        }
+
+        foo.%{KW} := "c" if {
+          false
+        }
+
+        bar.baz.%{KW} := "a"
+
+        bar.baz.%{KW} := "b" if {
+          false
+        }
+
+        bar.baz.%{KW} := "c" if {
+          false
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/%{KW} keyword function refhead
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p if {
+          %{KW}.foo(1) == 1
+          %{KW}.foo(11) == 42
+          foo.%{KW}(1) == 1
+          foo.%{KW}(11) == 42
+          bar.%{KW}.baz(1) == 1
+          bar.%{KW}.baz(11) == 42
+        }
+
+        default %{KW}.foo(_) := 42
+
+        %{KW}.foo(x) := x if {
+          x < 10
+        }
+
+        default foo.%{KW}(_) := 42
+
+        foo.%{KW}(x) := x if {
+          x < 10
+        }
+
+        default bar.%{KW}.baz(_) := 42
+
+        bar.%{KW}.baz(x) := x if {
+          x < 10
+        }
+    want_result:
+      - x: true

v1/test/cases/testdata/v0/keywordrefs/test-keyword-as.yaml

@@ -0,0 +1,200 @@
+---
+cases:
+  - note: keywordrefs/as keyword in package
+    query: data.foo.p = x
+    modules:
+      - |
+        package foo.as.bar
+
+        baz := 42
+      - |
+        package foo
+        import data.foo.as.bar
+
+        p {
+          bar.baz == 42
+          data.foo.as.bar.baz == 42
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/as keyword in package, import alias
+    query: data.foo.p = x
+    modules:
+      - |
+        package foo.as
+
+        bar := 42
+      - |
+        package foo
+        import data.foo.as as my_if
+
+        p {
+          my_if.bar == 42
+          data.foo.as.bar == 42
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/as keyword rule refhead
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p {
+          as.foo == 1
+          foo.as == 2
+        }
+
+        as.foo := 1
+
+        foo.as := 2
+    want_result:
+      - x: true
+  - note: keywordrefs/as keyword rule refhead, else bodies
+    query: data.test.p = x
+    input:
+      x: 3
+    modules:
+      - |
+        package test
+
+        p {
+          as.foo.bar == 3
+          foo.bar.as == 6
+        }
+
+        as.foo.bar := 1 {
+          input.x == 1
+        } else := 2 {
+          input.x == 2
+        } else := 3
+
+        foo.bar.as := 4 {
+          input.x == 1
+        } else := 5 {
+          input.x == 2
+        } else := 6
+    want_result:
+      - x: true
+  - note: keywordrefs/as keyword rule refhead, partial set
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+        import future.keywords.contains
+
+        p {
+          as.foo.bar == {"a", "c"}
+          foo.bar.as == {"a", "c"}
+        }
+
+        as.foo.bar contains "a"
+
+        as.foo.bar contains "b" {
+          false
+        }
+
+        as.foo.bar contains "c" {
+          true
+        }
+
+        foo.bar.as contains "a"
+
+        foo.bar.as contains "b" {
+          false
+        }
+
+        foo.bar.as contains "c" {
+          true
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/as keyword rule refhead, partial object
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p {
+          as.foo == "a"
+          as.bar.one == "a"
+          as.bar.three == "c"
+          foo.as == "a"
+          bar.baz.as == "a"
+        }
+
+        as.foo := "a"
+
+        as.foo := "b" {
+          false
+        }
+
+        as.foo := "c" {
+          false
+        }
+
+        as.bar.one := "a"
+
+        as.bar.two := "b" {
+          false
+        }
+
+        as.bar.three := "c" {
+          true
+        }
+
+        foo.as := "a"
+
+        foo.as := "b" {
+          false
+        }
+
+        foo.as := "c" {
+          false
+        }
+
+        bar.baz.as := "a"
+
+        bar.baz.as := "b" {
+          false
+        }
+
+        bar.baz.as := "c" {
+          false
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/as keyword function refhead
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p {
+          as.foo(1) == 1
+          as.foo(11) == 42
+          foo.as(1) == 1
+          foo.as(11) == 42
+          bar.as.baz(1) == 1
+          bar.as.baz(11) == 42
+        }
+
+        default as.foo(_) := 42
+
+        as.foo(x) := x {
+          x < 10
+        }
+
+        default foo.as(_) := 42
+
+        foo.as(x) := x {
+          x < 10
+        }
+
+        default bar.as.baz(_) := 42
+
+        bar.as.baz(x) := x {
+          x < 10
+        }
+    want_result:
+      - x: true

v1/test/cases/testdata/v0/keywordrefs/test-keyword-default.yaml

@@ -0,0 +1,200 @@
+---
+cases:
+  - note: keywordrefs/default keyword in package
+    query: data.foo.p = x
+    modules:
+      - |
+        package foo.default.bar
+
+        baz := 42
+      - |
+        package foo
+        import data.foo.default.bar
+
+        p {
+          bar.baz == 42
+          data.foo.default.bar.baz == 42
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/default keyword in package, import alias
+    query: data.foo.p = x
+    modules:
+      - |
+        package foo.default
+
+        bar := 42
+      - |
+        package foo
+        import data.foo.default as my_if
+
+        p {
+          my_if.bar == 42
+          data.foo.default.bar == 42
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/default keyword rule refhead
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p {
+          default.foo == 1
+          foo.default == 2
+        }
+
+        default.foo := 1
+
+        foo.default := 2
+    want_result:
+      - x: true
+  - note: keywordrefs/default keyword rule refhead, else bodies
+    query: data.test.p = x
+    input:
+      x: 3
+    modules:
+      - |
+        package test
+
+        p {
+          default.foo.bar == 3
+          foo.bar.default == 6
+        }
+
+        default.foo.bar := 1 {
+          input.x == 1
+        } else := 2 {
+          input.x == 2
+        } else := 3
+
+        foo.bar.default := 4 {
+          input.x == 1
+        } else := 5 {
+          input.x == 2
+        } else := 6
+    want_result:
+      - x: true
+  - note: keywordrefs/default keyword rule refhead, partial set
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+        import future.keywords.contains
+
+        p {
+          default.foo.bar == {"a", "c"}
+          foo.bar.default == {"a", "c"}
+        }
+
+        default.foo.bar contains "a"
+
+        default.foo.bar contains "b" {
+          false
+        }
+
+        default.foo.bar contains "c" {
+          true
+        }
+
+        foo.bar.default contains "a"
+
+        foo.bar.default contains "b" {
+          false
+        }
+
+        foo.bar.default contains "c" {
+          true
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/default keyword rule refhead, partial object
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p {
+          default.foo == "a"
+          default.bar.one == "a"
+          default.bar.three == "c"
+          foo.default == "a"
+          bar.baz.default == "a"
+        }
+
+        default.foo := "a"
+
+        default.foo := "b" {
+          false
+        }
+
+        default.foo := "c" {
+          false
+        }
+
+        default.bar.one := "a"
+
+        default.bar.two := "b" {
+          false
+        }
+
+        default.bar.three := "c" {
+          true
+        }
+
+        foo.default := "a"
+
+        foo.default := "b" {
+          false
+        }
+
+        foo.default := "c" {
+          false
+        }
+
+        bar.baz.default := "a"
+
+        bar.baz.default := "b" {
+          false
+        }
+
+        bar.baz.default := "c" {
+          false
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/default keyword function refhead
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p {
+          default.foo(1) == 1
+          default.foo(11) == 42
+          foo.default(1) == 1
+          foo.default(11) == 42
+          bar.default.baz(1) == 1
+          bar.default.baz(11) == 42
+        }
+
+        default default.foo(_) := 42
+
+        default.foo(x) := x {
+          x < 10
+        }
+
+        default foo.default(_) := 42
+
+        foo.default(x) := x {
+          x < 10
+        }
+
+        default bar.default.baz(_) := 42
+
+        bar.default.baz(x) := x {
+          x < 10
+        }
+    want_result:
+      - x: true

v1/test/cases/testdata/v0/keywordrefs/test-keyword-else.yaml

@@ -0,0 +1,200 @@
+---
+cases:
+  - note: keywordrefs/else keyword in package
+    query: data.foo.p = x
+    modules:
+      - |
+        package foo.else.bar
+
+        baz := 42
+      - |
+        package foo
+        import data.foo.else.bar
+
+        p {
+          bar.baz == 42
+          data.foo.else.bar.baz == 42
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/else keyword in package, import alias
+    query: data.foo.p = x
+    modules:
+      - |
+        package foo.else
+
+        bar := 42
+      - |
+        package foo
+        import data.foo.else as my_if
+
+        p {
+          my_if.bar == 42
+          data.foo.else.bar == 42
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/else keyword rule refhead
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p {
+          else.foo == 1
+          foo.else == 2
+        }
+
+        else.foo := 1
+
+        foo.else := 2
+    want_result:
+      - x: true
+  - note: keywordrefs/else keyword rule refhead, else bodies
+    query: data.test.p = x
+    input:
+      x: 3
+    modules:
+      - |
+        package test
+
+        p {
+          else.foo.bar == 3
+          foo.bar.else == 6
+        }
+
+        else.foo.bar := 1 {
+          input.x == 1
+        } else := 2 {
+          input.x == 2
+        } else := 3
+
+        foo.bar.else := 4 {
+          input.x == 1
+        } else := 5 {
+          input.x == 2
+        } else := 6
+    want_result:
+      - x: true
+  - note: keywordrefs/else keyword rule refhead, partial set
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+        import future.keywords.contains
+
+        p {
+          else.foo.bar == {"a", "c"}
+          foo.bar.else == {"a", "c"}
+        }
+
+        else.foo.bar contains "a"
+
+        else.foo.bar contains "b" {
+          false
+        }
+
+        else.foo.bar contains "c" {
+          true
+        }
+
+        foo.bar.else contains "a"
+
+        foo.bar.else contains "b" {
+          false
+        }
+
+        foo.bar.else contains "c" {
+          true
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/else keyword rule refhead, partial object
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p {
+          else.foo == "a"
+          else.bar.one == "a"
+          else.bar.three == "c"
+          foo.else == "a"
+          bar.baz.else == "a"
+        }
+
+        else.foo := "a"
+
+        else.foo := "b" {
+          false
+        }
+
+        else.foo := "c" {
+          false
+        }
+
+        else.bar.one := "a"
+
+        else.bar.two := "b" {
+          false
+        }
+
+        else.bar.three := "c" {
+          true
+        }
+
+        foo.else := "a"
+
+        foo.else := "b" {
+          false
+        }
+
+        foo.else := "c" {
+          false
+        }
+
+        bar.baz.else := "a"
+
+        bar.baz.else := "b" {
+          false
+        }
+
+        bar.baz.else := "c" {
+          false
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/else keyword function refhead
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p {
+          else.foo(1) == 1
+          else.foo(11) == 42
+          foo.else(1) == 1
+          foo.else(11) == 42
+          bar.else.baz(1) == 1
+          bar.else.baz(11) == 42
+        }
+
+        default else.foo(_) := 42
+
+        else.foo(x) := x {
+          x < 10
+        }
+
+        default foo.else(_) := 42
+
+        foo.else(x) := x {
+          x < 10
+        }
+
+        default bar.else.baz(_) := 42
+
+        bar.else.baz(x) := x {
+          x < 10
+        }
+    want_result:
+      - x: true

v1/test/cases/testdata/v0/keywordrefs/test-keyword-false.yaml

@@ -0,0 +1,200 @@
+---
+cases:
+  - note: keywordrefs/false keyword in package
+    query: data.foo.p = x
+    modules:
+      - |
+        package foo.false.bar
+
+        baz := 42
+      - |
+        package foo
+        import data.foo.false.bar
+
+        p {
+          bar.baz == 42
+          data.foo.false.bar.baz == 42
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/false keyword in package, import alias
+    query: data.foo.p = x
+    modules:
+      - |
+        package foo.false
+
+        bar := 42
+      - |
+        package foo
+        import data.foo.false as my_if
+
+        p {
+          my_if.bar == 42
+          data.foo.false.bar == 42
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/false keyword rule refhead
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p {
+          false.foo == 1
+          foo.false == 2
+        }
+
+        false.foo := 1
+
+        foo.false := 2
+    want_result:
+      - x: true
+  - note: keywordrefs/false keyword rule refhead, else bodies
+    query: data.test.p = x
+    input:
+      x: 3
+    modules:
+      - |
+        package test
+
+        p {
+          false.foo.bar == 3
+          foo.bar.false == 6
+        }
+
+        false.foo.bar := 1 {
+          input.x == 1
+        } else := 2 {
+          input.x == 2
+        } else := 3
+
+        foo.bar.false := 4 {
+          input.x == 1
+        } else := 5 {
+          input.x == 2
+        } else := 6
+    want_result:
+      - x: true
+  - note: keywordrefs/false keyword rule refhead, partial set
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+        import future.keywords.contains
+
+        p {
+          false.foo.bar == {"a", "c"}
+          foo.bar.false == {"a", "c"}
+        }
+
+        false.foo.bar contains "a"
+
+        false.foo.bar contains "b" {
+          false
+        }
+
+        false.foo.bar contains "c" {
+          true
+        }
+
+        foo.bar.false contains "a"
+
+        foo.bar.false contains "b" {
+          false
+        }
+
+        foo.bar.false contains "c" {
+          true
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/false keyword rule refhead, partial object
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p {
+          false.foo == "a"
+          false.bar.one == "a"
+          false.bar.three == "c"
+          foo.false == "a"
+          bar.baz.false == "a"
+        }
+
+        false.foo := "a"
+
+        false.foo := "b" {
+          false
+        }
+
+        false.foo := "c" {
+          false
+        }
+
+        false.bar.one := "a"
+
+        false.bar.two := "b" {
+          false
+        }
+
+        false.bar.three := "c" {
+          true
+        }
+
+        foo.false := "a"
+
+        foo.false := "b" {
+          false
+        }
+
+        foo.false := "c" {
+          false
+        }
+
+        bar.baz.false := "a"
+
+        bar.baz.false := "b" {
+          false
+        }
+
+        bar.baz.false := "c" {
+          false
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/false keyword function refhead
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p {
+          false.foo(1) == 1
+          false.foo(11) == 42
+          foo.false(1) == 1
+          foo.false(11) == 42
+          bar.false.baz(1) == 1
+          bar.false.baz(11) == 42
+        }
+
+        default false.foo(_) := 42
+
+        false.foo(x) := x {
+          x < 10
+        }
+
+        default foo.false(_) := 42
+
+        foo.false(x) := x {
+          x < 10
+        }
+
+        default bar.false.baz(_) := 42
+
+        bar.false.baz(x) := x {
+          x < 10
+        }
+    want_result:
+      - x: true

v1/test/cases/testdata/v0/keywordrefs/test-keyword-import.yaml

@@ -0,0 +1,200 @@
+---
+cases:
+  - note: keywordrefs/import keyword in package
+    query: data.foo.p = x
+    modules:
+      - |
+        package foo.import.bar
+
+        baz := 42
+      - |
+        package foo
+        import data.foo.import.bar
+
+        p {
+          bar.baz == 42
+          data.foo.import.bar.baz == 42
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/import keyword in package, import alias
+    query: data.foo.p = x
+    modules:
+      - |
+        package foo.import
+
+        bar := 42
+      - |
+        package foo
+        import data.foo.import as my_if
+
+        p {
+          my_if.bar == 42
+          data.foo.import.bar == 42
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/import keyword rule refhead
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p {
+          import.foo == 1
+          foo.import == 2
+        }
+
+        import.foo := 1
+
+        foo.import := 2
+    want_result:
+      - x: true
+  - note: keywordrefs/import keyword rule refhead, else bodies
+    query: data.test.p = x
+    input:
+      x: 3
+    modules:
+      - |
+        package test
+
+        p {
+          import.foo.bar == 3
+          foo.bar.import == 6
+        }
+
+        import.foo.bar := 1 {
+          input.x == 1
+        } else := 2 {
+          input.x == 2
+        } else := 3
+
+        foo.bar.import := 4 {
+          input.x == 1
+        } else := 5 {
+          input.x == 2
+        } else := 6
+    want_result:
+      - x: true
+  - note: keywordrefs/import keyword rule refhead, partial set
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+        import future.keywords.contains
+
+        p {
+          import.foo.bar == {"a", "c"}
+          foo.bar.import == {"a", "c"}
+        }
+
+        import.foo.bar contains "a"
+
+        import.foo.bar contains "b" {
+          false
+        }
+
+        import.foo.bar contains "c" {
+          true
+        }
+
+        foo.bar.import contains "a"
+
+        foo.bar.import contains "b" {
+          false
+        }
+
+        foo.bar.import contains "c" {
+          true
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/import keyword rule refhead, partial object
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p {
+          import.foo == "a"
+          import.bar.one == "a"
+          import.bar.three == "c"
+          foo.import == "a"
+          bar.baz.import == "a"
+        }
+
+        import.foo := "a"
+
+        import.foo := "b" {
+          false
+        }
+
+        import.foo := "c" {
+          false
+        }
+
+        import.bar.one := "a"
+
+        import.bar.two := "b" {
+          false
+        }
+
+        import.bar.three := "c" {
+          true
+        }
+
+        foo.import := "a"
+
+        foo.import := "b" {
+          false
+        }
+
+        foo.import := "c" {
+          false
+        }
+
+        bar.baz.import := "a"
+
+        bar.baz.import := "b" {
+          false
+        }
+
+        bar.baz.import := "c" {
+          false
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/import keyword function refhead
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p {
+          import.foo(1) == 1
+          import.foo(11) == 42
+          foo.import(1) == 1
+          foo.import(11) == 42
+          bar.import.baz(1) == 1
+          bar.import.baz(11) == 42
+        }
+
+        default import.foo(_) := 42
+
+        import.foo(x) := x {
+          x < 10
+        }
+
+        default foo.import(_) := 42
+
+        foo.import(x) := x {
+          x < 10
+        }
+
+        default bar.import.baz(_) := 42
+
+        bar.import.baz(x) := x {
+          x < 10
+        }
+    want_result:
+      - x: true

v1/test/cases/testdata/v0/keywordrefs/test-keyword-not.yaml

@@ -0,0 +1,200 @@
+---
+cases:
+  - note: keywordrefs/not keyword in package
+    query: data.foo.p = x
+    modules:
+      - |
+        package foo.not.bar
+
+        baz := 42
+      - |
+        package foo
+        import data.foo.not.bar
+
+        p {
+          bar.baz == 42
+          data.foo.not.bar.baz == 42
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/not keyword in package, import alias
+    query: data.foo.p = x
+    modules:
+      - |
+        package foo.not
+
+        bar := 42
+      - |
+        package foo
+        import data.foo.not as my_if
+
+        p {
+          my_if.bar == 42
+          data.foo.not.bar == 42
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/not keyword rule refhead
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p {
+          not.foo == 1
+          foo.not == 2
+        }
+
+        not.foo := 1
+
+        foo.not := 2
+    want_result:
+      - x: true
+  - note: keywordrefs/not keyword rule refhead, else bodies
+    query: data.test.p = x
+    input:
+      x: 3
+    modules:
+      - |
+        package test
+
+        p {
+          not.foo.bar == 3
+          foo.bar.not == 6
+        }
+
+        not.foo.bar := 1 {
+          input.x == 1
+        } else := 2 {
+          input.x == 2
+        } else := 3
+
+        foo.bar.not := 4 {
+          input.x == 1
+        } else := 5 {
+          input.x == 2
+        } else := 6
+    want_result:
+      - x: true
+  - note: keywordrefs/not keyword rule refhead, partial set
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+        import future.keywords.contains
+
+        p {
+          not.foo.bar == {"a", "c"}
+          foo.bar.not == {"a", "c"}
+        }
+
+        not.foo.bar contains "a"
+
+        not.foo.bar contains "b" {
+          false
+        }
+
+        not.foo.bar contains "c" {
+          true
+        }
+
+        foo.bar.not contains "a"
+
+        foo.bar.not contains "b" {
+          false
+        }
+
+        foo.bar.not contains "c" {
+          true
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/not keyword rule refhead, partial object
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p {
+          not.foo == "a"
+          not.bar.one == "a"
+          not.bar.three == "c"
+          foo.not == "a"
+          bar.baz.not == "a"
+        }
+
+        not.foo := "a"
+
+        not.foo := "b" {
+          false
+        }
+
+        not.foo := "c" {
+          false
+        }
+
+        not.bar.one := "a"
+
+        not.bar.two := "b" {
+          false
+        }
+
+        not.bar.three := "c" {
+          true
+        }
+
+        foo.not := "a"
+
+        foo.not := "b" {
+          false
+        }
+
+        foo.not := "c" {
+          false
+        }
+
+        bar.baz.not := "a"
+
+        bar.baz.not := "b" {
+          false
+        }
+
+        bar.baz.not := "c" {
+          false
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/not keyword function refhead
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p {
+          not.foo(1) == 1
+          not.foo(11) == 42
+          foo.not(1) == 1
+          foo.not(11) == 42
+          bar.not.baz(1) == 1
+          bar.not.baz(11) == 42
+        }
+
+        default not.foo(_) := 42
+
+        not.foo(x) := x {
+          x < 10
+        }
+
+        default foo.not(_) := 42
+
+        foo.not(x) := x {
+          x < 10
+        }
+
+        default bar.not.baz(_) := 42
+
+        bar.not.baz(x) := x {
+          x < 10
+        }
+    want_result:
+      - x: true

v1/test/cases/testdata/v0/keywordrefs/test-keyword-null.yaml

@@ -0,0 +1,200 @@
+---
+cases:
+  - note: keywordrefs/null keyword in package
+    query: data.foo.p = x
+    modules:
+      - |
+        package foo.null.bar
+
+        baz := 42
+      - |
+        package foo
+        import data.foo.null.bar
+
+        p {
+          bar.baz == 42
+          data.foo.null.bar.baz == 42
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/null keyword in package, import alias
+    query: data.foo.p = x
+    modules:
+      - |
+        package foo.null
+
+        bar := 42
+      - |
+        package foo
+        import data.foo.null as my_if
+
+        p {
+          my_if.bar == 42
+          data.foo.null.bar == 42
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/null keyword rule refhead
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p {
+          null.foo == 1
+          foo.null == 2
+        }
+
+        null.foo := 1
+
+        foo.null := 2
+    want_result:
+      - x: true
+  - note: keywordrefs/null keyword rule refhead, else bodies
+    query: data.test.p = x
+    input:
+      x: 3
+    modules:
+      - |
+        package test
+
+        p {
+          null.foo.bar == 3
+          foo.bar.null == 6
+        }
+
+        null.foo.bar := 1 {
+          input.x == 1
+        } else := 2 {
+          input.x == 2
+        } else := 3
+
+        foo.bar.null := 4 {
+          input.x == 1
+        } else := 5 {
+          input.x == 2
+        } else := 6
+    want_result:
+      - x: true
+  - note: keywordrefs/null keyword rule refhead, partial set
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+        import future.keywords.contains
+
+        p {
+          null.foo.bar == {"a", "c"}
+          foo.bar.null == {"a", "c"}
+        }
+
+        null.foo.bar contains "a"
+
+        null.foo.bar contains "b" {
+          false
+        }
+
+        null.foo.bar contains "c" {
+          true
+        }
+
+        foo.bar.null contains "a"
+
+        foo.bar.null contains "b" {
+          false
+        }
+
+        foo.bar.null contains "c" {
+          true
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/null keyword rule refhead, partial object
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p {
+          null.foo == "a"
+          null.bar.one == "a"
+          null.bar.three == "c"
+          foo.null == "a"
+          bar.baz.null == "a"
+        }
+
+        null.foo := "a"
+
+        null.foo := "b" {
+          false
+        }
+
+        null.foo := "c" {
+          false
+        }
+
+        null.bar.one := "a"
+
+        null.bar.two := "b" {
+          false
+        }
+
+        null.bar.three := "c" {
+          true
+        }
+
+        foo.null := "a"
+
+        foo.null := "b" {
+          false
+        }
+
+        foo.null := "c" {
+          false
+        }
+
+        bar.baz.null := "a"
+
+        bar.baz.null := "b" {
+          false
+        }
+
+        bar.baz.null := "c" {
+          false
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/null keyword function refhead
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p {
+          null.foo(1) == 1
+          null.foo(11) == 42
+          foo.null(1) == 1
+          foo.null(11) == 42
+          bar.null.baz(1) == 1
+          bar.null.baz(11) == 42
+        }
+
+        default null.foo(_) := 42
+
+        null.foo(x) := x {
+          x < 10
+        }
+
+        default foo.null(_) := 42
+
+        foo.null(x) := x {
+          x < 10
+        }
+
+        default bar.null.baz(_) := 42
+
+        bar.null.baz(x) := x {
+          x < 10
+        }
+    want_result:
+      - x: true

v1/test/cases/testdata/v0/keywordrefs/test-keyword-package.yaml

@@ -0,0 +1,200 @@
+---
+cases:
+  - note: keywordrefs/package keyword in package
+    query: data.foo.p = x
+    modules:
+      - |
+        package foo.package.bar
+
+        baz := 42
+      - |
+        package foo
+        import data.foo.package.bar
+
+        p {
+          bar.baz == 42
+          data.foo.package.bar.baz == 42
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/package keyword in package, import alias
+    query: data.foo.p = x
+    modules:
+      - |
+        package foo.package
+
+        bar := 42
+      - |
+        package foo
+        import data.foo.package as my_if
+
+        p {
+          my_if.bar == 42
+          data.foo.package.bar == 42
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/package keyword rule refhead
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p {
+          package.foo == 1
+          foo.package == 2
+        }
+
+        package.foo := 1
+
+        foo.package := 2
+    want_result:
+      - x: true
+  - note: keywordrefs/package keyword rule refhead, else bodies
+    query: data.test.p = x
+    input:
+      x: 3
+    modules:
+      - |
+        package test
+
+        p {
+          package.foo.bar == 3
+          foo.bar.package == 6
+        }
+
+        package.foo.bar := 1 {
+          input.x == 1
+        } else := 2 {
+          input.x == 2
+        } else := 3
+
+        foo.bar.package := 4 {
+          input.x == 1
+        } else := 5 {
+          input.x == 2
+        } else := 6
+    want_result:
+      - x: true
+  - note: keywordrefs/package keyword rule refhead, partial set
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+        import future.keywords.contains
+
+        p {
+          package.foo.bar == {"a", "c"}
+          foo.bar.package == {"a", "c"}
+        }
+
+        package.foo.bar contains "a"
+
+        package.foo.bar contains "b" {
+          false
+        }
+
+        package.foo.bar contains "c" {
+          true
+        }
+
+        foo.bar.package contains "a"
+
+        foo.bar.package contains "b" {
+          false
+        }
+
+        foo.bar.package contains "c" {
+          true
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/package keyword rule refhead, partial object
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p {
+          package.foo == "a"
+          package.bar.one == "a"
+          package.bar.three == "c"
+          foo.package == "a"
+          bar.baz.package == "a"
+        }
+
+        package.foo := "a"
+
+        package.foo := "b" {
+          false
+        }
+
+        package.foo := "c" {
+          false
+        }
+
+        package.bar.one := "a"
+
+        package.bar.two := "b" {
+          false
+        }
+
+        package.bar.three := "c" {
+          true
+        }
+
+        foo.package := "a"
+
+        foo.package := "b" {
+          false
+        }
+
+        foo.package := "c" {
+          false
+        }
+
+        bar.baz.package := "a"
+
+        bar.baz.package := "b" {
+          false
+        }
+
+        bar.baz.package := "c" {
+          false
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/package keyword function refhead
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p {
+          package.foo(1) == 1
+          package.foo(11) == 42
+          foo.package(1) == 1
+          foo.package(11) == 42
+          bar.package.baz(1) == 1
+          bar.package.baz(11) == 42
+        }
+
+        default package.foo(_) := 42
+
+        package.foo(x) := x {
+          x < 10
+        }
+
+        default foo.package(_) := 42
+
+        foo.package(x) := x {
+          x < 10
+        }
+
+        default bar.package.baz(_) := 42
+
+        bar.package.baz(x) := x {
+          x < 10
+        }
+    want_result:
+      - x: true

v1/test/cases/testdata/v0/keywordrefs/test-keyword-some.yaml

@@ -0,0 +1,200 @@
+---
+cases:
+  - note: keywordrefs/some keyword in package
+    query: data.foo.p = x
+    modules:
+      - |
+        package foo.some.bar
+
+        baz := 42
+      - |
+        package foo
+        import data.foo.some.bar
+
+        p {
+          bar.baz == 42
+          data.foo.some.bar.baz == 42
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/some keyword in package, import alias
+    query: data.foo.p = x
+    modules:
+      - |
+        package foo.some
+
+        bar := 42
+      - |
+        package foo
+        import data.foo.some as my_if
+
+        p {
+          my_if.bar == 42
+          data.foo.some.bar == 42
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/some keyword rule refhead
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p {
+          some.foo == 1
+          foo.some == 2
+        }
+
+        some.foo := 1
+
+        foo.some := 2
+    want_result:
+      - x: true
+  - note: keywordrefs/some keyword rule refhead, else bodies
+    query: data.test.p = x
+    input:
+      x: 3
+    modules:
+      - |
+        package test
+
+        p {
+          some.foo.bar == 3
+          foo.bar.some == 6
+        }
+
+        some.foo.bar := 1 {
+          input.x == 1
+        } else := 2 {
+          input.x == 2
+        } else := 3
+
+        foo.bar.some := 4 {
+          input.x == 1
+        } else := 5 {
+          input.x == 2
+        } else := 6
+    want_result:
+      - x: true
+  - note: keywordrefs/some keyword rule refhead, partial set
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+        import future.keywords.contains
+
+        p {
+          some.foo.bar == {"a", "c"}
+          foo.bar.some == {"a", "c"}
+        }
+
+        some.foo.bar contains "a"
+
+        some.foo.bar contains "b" {
+          false
+        }
+
+        some.foo.bar contains "c" {
+          true
+        }
+
+        foo.bar.some contains "a"
+
+        foo.bar.some contains "b" {
+          false
+        }
+
+        foo.bar.some contains "c" {
+          true
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/some keyword rule refhead, partial object
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p {
+          some.foo == "a"
+          some.bar.one == "a"
+          some.bar.three == "c"
+          foo.some == "a"
+          bar.baz.some == "a"
+        }
+
+        some.foo := "a"
+
+        some.foo := "b" {
+          false
+        }
+
+        some.foo := "c" {
+          false
+        }
+
+        some.bar.one := "a"
+
+        some.bar.two := "b" {
+          false
+        }
+
+        some.bar.three := "c" {
+          true
+        }
+
+        foo.some := "a"
+
+        foo.some := "b" {
+          false
+        }
+
+        foo.some := "c" {
+          false
+        }
+
+        bar.baz.some := "a"
+
+        bar.baz.some := "b" {
+          false
+        }
+
+        bar.baz.some := "c" {
+          false
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/some keyword function refhead
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p {
+          some.foo(1) == 1
+          some.foo(11) == 42
+          foo.some(1) == 1
+          foo.some(11) == 42
+          bar.some.baz(1) == 1
+          bar.some.baz(11) == 42
+        }
+
+        default some.foo(_) := 42
+
+        some.foo(x) := x {
+          x < 10
+        }
+
+        default foo.some(_) := 42
+
+        foo.some(x) := x {
+          x < 10
+        }
+
+        default bar.some.baz(_) := 42
+
+        bar.some.baz(x) := x {
+          x < 10
+        }
+    want_result:
+      - x: true

v1/test/cases/testdata/v0/keywordrefs/test-keyword-true.yaml

@@ -0,0 +1,200 @@
+---
+cases:
+  - note: keywordrefs/true keyword in package
+    query: data.foo.p = x
+    modules:
+      - |
+        package foo.true.bar
+
+        baz := 42
+      - |
+        package foo
+        import data.foo.true.bar
+
+        p {
+          bar.baz == 42
+          data.foo.true.bar.baz == 42
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/true keyword in package, import alias
+    query: data.foo.p = x
+    modules:
+      - |
+        package foo.true
+
+        bar := 42
+      - |
+        package foo
+        import data.foo.true as my_if
+
+        p {
+          my_if.bar == 42
+          data.foo.true.bar == 42
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/true keyword rule refhead
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p {
+          true.foo == 1
+          foo.true == 2
+        }
+
+        true.foo := 1
+
+        foo.true := 2
+    want_result:
+      - x: true
+  - note: keywordrefs/true keyword rule refhead, else bodies
+    query: data.test.p = x
+    input:
+      x: 3
+    modules:
+      - |
+        package test
+
+        p {
+          true.foo.bar == 3
+          foo.bar.true == 6
+        }
+
+        true.foo.bar := 1 {
+          input.x == 1
+        } else := 2 {
+          input.x == 2
+        } else := 3
+
+        foo.bar.true := 4 {
+          input.x == 1
+        } else := 5 {
+          input.x == 2
+        } else := 6
+    want_result:
+      - x: true
+  - note: keywordrefs/true keyword rule refhead, partial set
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+        import future.keywords.contains
+
+        p {
+          true.foo.bar == {"a", "c"}
+          foo.bar.true == {"a", "c"}
+        }
+
+        true.foo.bar contains "a"
+
+        true.foo.bar contains "b" {
+          false
+        }
+
+        true.foo.bar contains "c" {
+          true
+        }
+
+        foo.bar.true contains "a"
+
+        foo.bar.true contains "b" {
+          false
+        }
+
+        foo.bar.true contains "c" {
+          true
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/true keyword rule refhead, partial object
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p {
+          true.foo == "a"
+          true.bar.one == "a"
+          true.bar.three == "c"
+          foo.true == "a"
+          bar.baz.true == "a"
+        }
+
+        true.foo := "a"
+
+        true.foo := "b" {
+          false
+        }
+
+        true.foo := "c" {
+          false
+        }
+
+        true.bar.one := "a"
+
+        true.bar.two := "b" {
+          false
+        }
+
+        true.bar.three := "c" {
+          true
+        }
+
+        foo.true := "a"
+
+        foo.true := "b" {
+          false
+        }
+
+        foo.true := "c" {
+          false
+        }
+
+        bar.baz.true := "a"
+
+        bar.baz.true := "b" {
+          false
+        }
+
+        bar.baz.true := "c" {
+          false
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/true keyword function refhead
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p {
+          true.foo(1) == 1
+          true.foo(11) == 42
+          foo.true(1) == 1
+          foo.true(11) == 42
+          bar.true.baz(1) == 1
+          bar.true.baz(11) == 42
+        }
+
+        default true.foo(_) := 42
+
+        true.foo(x) := x {
+          x < 10
+        }
+
+        default foo.true(_) := 42
+
+        foo.true(x) := x {
+          x < 10
+        }
+
+        default bar.true.baz(_) := 42
+
+        bar.true.baz(x) := x {
+          x < 10
+        }
+    want_result:
+      - x: true

v1/test/cases/testdata/v0/keywordrefs/test-keyword-with.yaml

@@ -0,0 +1,200 @@
+---
+cases:
+  - note: keywordrefs/with keyword in package
+    query: data.foo.p = x
+    modules:
+      - |
+        package foo.with.bar
+
+        baz := 42
+      - |
+        package foo
+        import data.foo.with.bar
+
+        p {
+          bar.baz == 42
+          data.foo.with.bar.baz == 42
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/with keyword in package, import alias
+    query: data.foo.p = x
+    modules:
+      - |
+        package foo.with
+
+        bar := 42
+      - |
+        package foo
+        import data.foo.with as my_if
+
+        p {
+          my_if.bar == 42
+          data.foo.with.bar == 42
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/with keyword rule refhead
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p {
+          with.foo == 1
+          foo.with == 2
+        }
+
+        with.foo := 1
+
+        foo.with := 2
+    want_result:
+      - x: true
+  - note: keywordrefs/with keyword rule refhead, else bodies
+    query: data.test.p = x
+    input:
+      x: 3
+    modules:
+      - |
+        package test
+
+        p {
+          with.foo.bar == 3
+          foo.bar.with == 6
+        }
+
+        with.foo.bar := 1 {
+          input.x == 1
+        } else := 2 {
+          input.x == 2
+        } else := 3
+
+        foo.bar.with := 4 {
+          input.x == 1
+        } else := 5 {
+          input.x == 2
+        } else := 6
+    want_result:
+      - x: true
+  - note: keywordrefs/with keyword rule refhead, partial set
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+        import future.keywords.contains
+
+        p {
+          with.foo.bar == {"a", "c"}
+          foo.bar.with == {"a", "c"}
+        }
+
+        with.foo.bar contains "a"
+
+        with.foo.bar contains "b" {
+          false
+        }
+
+        with.foo.bar contains "c" {
+          true
+        }
+
+        foo.bar.with contains "a"
+
+        foo.bar.with contains "b" {
+          false
+        }
+
+        foo.bar.with contains "c" {
+          true
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/with keyword rule refhead, partial object
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p {
+          with.foo == "a"
+          with.bar.one == "a"
+          with.bar.three == "c"
+          foo.with == "a"
+          bar.baz.with == "a"
+        }
+
+        with.foo := "a"
+
+        with.foo := "b" {
+          false
+        }
+
+        with.foo := "c" {
+          false
+        }
+
+        with.bar.one := "a"
+
+        with.bar.two := "b" {
+          false
+        }
+
+        with.bar.three := "c" {
+          true
+        }
+
+        foo.with := "a"
+
+        foo.with := "b" {
+          false
+        }
+
+        foo.with := "c" {
+          false
+        }
+
+        bar.baz.with := "a"
+
+        bar.baz.with := "b" {
+          false
+        }
+
+        bar.baz.with := "c" {
+          false
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/with keyword function refhead
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p {
+          with.foo(1) == 1
+          with.foo(11) == 42
+          foo.with(1) == 1
+          foo.with(11) == 42
+          bar.with.baz(1) == 1
+          bar.with.baz(11) == 42
+        }
+
+        default with.foo(_) := 42
+
+        with.foo(x) := x {
+          x < 10
+        }
+
+        default foo.with(_) := 42
+
+        foo.with(x) := x {
+          x < 10
+        }
+
+        default bar.with.baz(_) := 42
+
+        bar.with.baz(x) := x {
+          x < 10
+        }
+    want_result:
+      - x: true

v1/test/cases/testdata/v1/keywordrefs/test-keyword-as.yaml

@@ -0,0 +1,199 @@
+---
+cases:
+  - note: keywordrefs/as keyword in package
+    query: data.foo.p = x
+    modules:
+      - |
+        package foo.as.bar
+
+        baz := 42
+      - |
+        package foo
+        import data.foo.as.bar
+
+        p if {
+          bar.baz == 42
+          data.foo.as.bar.baz == 42
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/as keyword in package, import alias
+    query: data.foo.p = x
+    modules:
+      - |
+        package foo.as
+
+        bar := 42
+      - |
+        package foo
+        import data.foo.as as my_if
+
+        p if {
+          my_if.bar == 42
+          data.foo.as.bar == 42
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/as keyword rule refhead
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p if {
+          as.foo == 1
+          foo.as == 2
+        }
+
+        as.foo := 1
+
+        foo.as := 2
+    want_result:
+      - x: true
+  - note: keywordrefs/as keyword rule refhead, else bodies
+    query: data.test.p = x
+    input:
+      x: 3
+    modules:
+      - |
+        package test
+
+        p if {
+          as.foo == 3
+          foo.as == 6
+        }
+
+        as.foo := 1 if {
+          input.x == 1
+        } else := 2 if {
+          input.x == 2
+        } else := 3
+
+        foo.as := 4 if {
+          input.x == 1
+        } else := 5 if {
+          input.x == 2
+        } else := 6
+    want_result:
+      - x: true
+  - note: keywordrefs/as keyword rule refhead, partial set
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p if {
+          as.foo == {"a", "c"}
+          foo.as == {"a", "c"}
+        }
+
+        as.foo contains "a"
+
+        as.foo contains "b" if {
+          false
+        }
+
+        as.foo contains "c" if {
+          true
+        }
+
+        foo.as contains "a"
+
+        foo.as contains "b" if {
+          false
+        }
+
+        foo.as contains "c" if {
+          true
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/as keyword rule refhead, partial object
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p if {
+          as.foo == "a"
+          as.bar.one == "a"
+          as.bar.three == "c"
+          foo.as == "a"
+          bar.baz.as == "a"
+        }
+
+        as.foo := "a"
+
+        as.foo := "b" if {
+          false
+        }
+
+        as.foo := "c" if {
+          false
+        }
+
+        as.bar.one := "a"
+
+        as.bar.two := "b" if {
+          false
+        }
+
+        as.bar.three := "c" if {
+          true
+        }
+
+        foo.as := "a"
+
+        foo.as := "b" if {
+          false
+        }
+
+        foo.as := "c" if {
+          false
+        }
+
+        bar.baz.as := "a"
+
+        bar.baz.as := "b" if {
+          false
+        }
+
+        bar.baz.as := "c" if {
+          false
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/as keyword function refhead
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p if {
+          as.foo(1) == 1
+          as.foo(11) == 42
+          foo.as(1) == 1
+          foo.as(11) == 42
+          bar.as.baz(1) == 1
+          bar.as.baz(11) == 42
+        }
+
+        default as.foo(_) := 42
+
+        as.foo(x) := x if {
+          x < 10
+        }
+
+        default foo.as(_) := 42
+
+        foo.as(x) := x if {
+          x < 10
+        }
+
+        default bar.as.baz(_) := 42
+
+        bar.as.baz(x) := x if {
+          x < 10
+        }
+    want_result:
+      - x: true

v1/test/cases/testdata/v1/keywordrefs/test-keyword-contains.yaml

@@ -0,0 +1,199 @@
+---
+cases:
+  - note: keywordrefs/contains keyword in package
+    query: data.foo.p = x
+    modules:
+      - |
+        package foo.contains.bar
+
+        baz := 42
+      - |
+        package foo
+        import data.foo.contains.bar
+
+        p if {
+          bar.baz == 42
+          data.foo.contains.bar.baz == 42
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/contains keyword in package, import alias
+    query: data.foo.p = x
+    modules:
+      - |
+        package foo.contains
+
+        bar := 42
+      - |
+        package foo
+        import data.foo.contains as my_if
+
+        p if {
+          my_if.bar == 42
+          data.foo.contains.bar == 42
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/contains keyword rule refhead
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p if {
+          contains.foo == 1
+          foo.contains == 2
+        }
+
+        contains.foo := 1
+
+        foo.contains := 2
+    want_result:
+      - x: true
+  - note: keywordrefs/contains keyword rule refhead, else bodies
+    query: data.test.p = x
+    input:
+      x: 3
+    modules:
+      - |
+        package test
+
+        p if {
+          contains.foo == 3
+          foo.contains == 6
+        }
+
+        contains.foo := 1 if {
+          input.x == 1
+        } else := 2 if {
+          input.x == 2
+        } else := 3
+
+        foo.contains := 4 if {
+          input.x == 1
+        } else := 5 if {
+          input.x == 2
+        } else := 6
+    want_result:
+      - x: true
+  - note: keywordrefs/contains keyword rule refhead, partial set
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p if {
+          contains.foo == {"a", "c"}
+          foo.contains == {"a", "c"}
+        }
+
+        contains.foo contains "a"
+
+        contains.foo contains "b" if {
+          false
+        }
+
+        contains.foo contains "c" if {
+          true
+        }
+
+        foo.contains contains "a"
+
+        foo.contains contains "b" if {
+          false
+        }
+
+        foo.contains contains "c" if {
+          true
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/contains keyword rule refhead, partial object
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p if {
+          contains.foo == "a"
+          contains.bar.one == "a"
+          contains.bar.three == "c"
+          foo.contains == "a"
+          bar.baz.contains == "a"
+        }
+
+        contains.foo := "a"
+
+        contains.foo := "b" if {
+          false
+        }
+
+        contains.foo := "c" if {
+          false
+        }
+
+        contains.bar.one := "a"
+
+        contains.bar.two := "b" if {
+          false
+        }
+
+        contains.bar.three := "c" if {
+          true
+        }
+
+        foo.contains := "a"
+
+        foo.contains := "b" if {
+          false
+        }
+
+        foo.contains := "c" if {
+          false
+        }
+
+        bar.baz.contains := "a"
+
+        bar.baz.contains := "b" if {
+          false
+        }
+
+        bar.baz.contains := "c" if {
+          false
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/contains keyword function refhead
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p if {
+          contains.foo(1) == 1
+          contains.foo(11) == 42
+          foo.contains(1) == 1
+          foo.contains(11) == 42
+          bar.contains.baz(1) == 1
+          bar.contains.baz(11) == 42
+        }
+
+        default contains.foo(_) := 42
+
+        contains.foo(x) := x if {
+          x < 10
+        }
+
+        default foo.contains(_) := 42
+
+        foo.contains(x) := x if {
+          x < 10
+        }
+
+        default bar.contains.baz(_) := 42
+
+        bar.contains.baz(x) := x if {
+          x < 10
+        }
+    want_result:
+      - x: true

v1/test/cases/testdata/v1/keywordrefs/test-keyword-default.yaml

@@ -0,0 +1,199 @@
+---
+cases:
+  - note: keywordrefs/default keyword in package
+    query: data.foo.p = x
+    modules:
+      - |
+        package foo.default.bar
+
+        baz := 42
+      - |
+        package foo
+        import data.foo.default.bar
+
+        p if {
+          bar.baz == 42
+          data.foo.default.bar.baz == 42
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/default keyword in package, import alias
+    query: data.foo.p = x
+    modules:
+      - |
+        package foo.default
+
+        bar := 42
+      - |
+        package foo
+        import data.foo.default as my_if
+
+        p if {
+          my_if.bar == 42
+          data.foo.default.bar == 42
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/default keyword rule refhead
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p if {
+          default.foo == 1
+          foo.default == 2
+        }
+
+        default.foo := 1
+
+        foo.default := 2
+    want_result:
+      - x: true
+  - note: keywordrefs/default keyword rule refhead, else bodies
+    query: data.test.p = x
+    input:
+      x: 3
+    modules:
+      - |
+        package test
+
+        p if {
+          default.foo == 3
+          foo.default == 6
+        }
+
+        default.foo := 1 if {
+          input.x == 1
+        } else := 2 if {
+          input.x == 2
+        } else := 3
+
+        foo.default := 4 if {
+          input.x == 1
+        } else := 5 if {
+          input.x == 2
+        } else := 6
+    want_result:
+      - x: true
+  - note: keywordrefs/default keyword rule refhead, partial set
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p if {
+          default.foo == {"a", "c"}
+          foo.default == {"a", "c"}
+        }
+
+        default.foo contains "a"
+
+        default.foo contains "b" if {
+          false
+        }
+
+        default.foo contains "c" if {
+          true
+        }
+
+        foo.default contains "a"
+
+        foo.default contains "b" if {
+          false
+        }
+
+        foo.default contains "c" if {
+          true
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/default keyword rule refhead, partial object
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p if {
+          default.foo == "a"
+          default.bar.one == "a"
+          default.bar.three == "c"
+          foo.default == "a"
+          bar.baz.default == "a"
+        }
+
+        default.foo := "a"
+
+        default.foo := "b" if {
+          false
+        }
+
+        default.foo := "c" if {
+          false
+        }
+
+        default.bar.one := "a"
+
+        default.bar.two := "b" if {
+          false
+        }
+
+        default.bar.three := "c" if {
+          true
+        }
+
+        foo.default := "a"
+
+        foo.default := "b" if {
+          false
+        }
+
+        foo.default := "c" if {
+          false
+        }
+
+        bar.baz.default := "a"
+
+        bar.baz.default := "b" if {
+          false
+        }
+
+        bar.baz.default := "c" if {
+          false
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/default keyword function refhead
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p if {
+          default.foo(1) == 1
+          default.foo(11) == 42
+          foo.default(1) == 1
+          foo.default(11) == 42
+          bar.default.baz(1) == 1
+          bar.default.baz(11) == 42
+        }
+
+        default default.foo(_) := 42
+
+        default.foo(x) := x if {
+          x < 10
+        }
+
+        default foo.default(_) := 42
+
+        foo.default(x) := x if {
+          x < 10
+        }
+
+        default bar.default.baz(_) := 42
+
+        bar.default.baz(x) := x if {
+          x < 10
+        }
+    want_result:
+      - x: true

v1/test/cases/testdata/v1/keywordrefs/test-keyword-else.yaml

@@ -0,0 +1,199 @@
+---
+cases:
+  - note: keywordrefs/else keyword in package
+    query: data.foo.p = x
+    modules:
+      - |
+        package foo.else.bar
+
+        baz := 42
+      - |
+        package foo
+        import data.foo.else.bar
+
+        p if {
+          bar.baz == 42
+          data.foo.else.bar.baz == 42
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/else keyword in package, import alias
+    query: data.foo.p = x
+    modules:
+      - |
+        package foo.else
+
+        bar := 42
+      - |
+        package foo
+        import data.foo.else as my_if
+
+        p if {
+          my_if.bar == 42
+          data.foo.else.bar == 42
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/else keyword rule refhead
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p if {
+          else.foo == 1
+          foo.else == 2
+        }
+
+        else.foo := 1
+
+        foo.else := 2
+    want_result:
+      - x: true
+  - note: keywordrefs/else keyword rule refhead, else bodies
+    query: data.test.p = x
+    input:
+      x: 3
+    modules:
+      - |
+        package test
+
+        p if {
+          else.foo == 3
+          foo.else == 6
+        }
+
+        else.foo := 1 if {
+          input.x == 1
+        } else := 2 if {
+          input.x == 2
+        } else := 3
+
+        foo.else := 4 if {
+          input.x == 1
+        } else := 5 if {
+          input.x == 2
+        } else := 6
+    want_result:
+      - x: true
+  - note: keywordrefs/else keyword rule refhead, partial set
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p if {
+          else.foo == {"a", "c"}
+          foo.else == {"a", "c"}
+        }
+
+        else.foo contains "a"
+
+        else.foo contains "b" if {
+          false
+        }
+
+        else.foo contains "c" if {
+          true
+        }
+
+        foo.else contains "a"
+
+        foo.else contains "b" if {
+          false
+        }
+
+        foo.else contains "c" if {
+          true
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/else keyword rule refhead, partial object
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p if {
+          else.foo == "a"
+          else.bar.one == "a"
+          else.bar.three == "c"
+          foo.else == "a"
+          bar.baz.else == "a"
+        }
+
+        else.foo := "a"
+
+        else.foo := "b" if {
+          false
+        }
+
+        else.foo := "c" if {
+          false
+        }
+
+        else.bar.one := "a"
+
+        else.bar.two := "b" if {
+          false
+        }
+
+        else.bar.three := "c" if {
+          true
+        }
+
+        foo.else := "a"
+
+        foo.else := "b" if {
+          false
+        }
+
+        foo.else := "c" if {
+          false
+        }
+
+        bar.baz.else := "a"
+
+        bar.baz.else := "b" if {
+          false
+        }
+
+        bar.baz.else := "c" if {
+          false
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/else keyword function refhead
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p if {
+          else.foo(1) == 1
+          else.foo(11) == 42
+          foo.else(1) == 1
+          foo.else(11) == 42
+          bar.else.baz(1) == 1
+          bar.else.baz(11) == 42
+        }
+
+        default else.foo(_) := 42
+
+        else.foo(x) := x if {
+          x < 10
+        }
+
+        default foo.else(_) := 42
+
+        foo.else(x) := x if {
+          x < 10
+        }
+
+        default bar.else.baz(_) := 42
+
+        bar.else.baz(x) := x if {
+          x < 10
+        }
+    want_result:
+      - x: true

v1/test/cases/testdata/v1/keywordrefs/test-keyword-every.yaml

@@ -0,0 +1,199 @@
+---
+cases:
+  - note: keywordrefs/every keyword in package
+    query: data.foo.p = x
+    modules:
+      - |
+        package foo.every.bar
+
+        baz := 42
+      - |
+        package foo
+        import data.foo.every.bar
+
+        p if {
+          bar.baz == 42
+          data.foo.every.bar.baz == 42
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/every keyword in package, import alias
+    query: data.foo.p = x
+    modules:
+      - |
+        package foo.every
+
+        bar := 42
+      - |
+        package foo
+        import data.foo.every as my_if
+
+        p if {
+          my_if.bar == 42
+          data.foo.every.bar == 42
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/every keyword rule refhead
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p if {
+          every.foo == 1
+          foo.every == 2
+        }
+
+        every.foo := 1
+
+        foo.every := 2
+    want_result:
+      - x: true
+  - note: keywordrefs/every keyword rule refhead, else bodies
+    query: data.test.p = x
+    input:
+      x: 3
+    modules:
+      - |
+        package test
+
+        p if {
+          every.foo == 3
+          foo.every == 6
+        }
+
+        every.foo := 1 if {
+          input.x == 1
+        } else := 2 if {
+          input.x == 2
+        } else := 3
+
+        foo.every := 4 if {
+          input.x == 1
+        } else := 5 if {
+          input.x == 2
+        } else := 6
+    want_result:
+      - x: true
+  - note: keywordrefs/every keyword rule refhead, partial set
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p if {
+          every.foo == {"a", "c"}
+          foo.every == {"a", "c"}
+        }
+
+        every.foo contains "a"
+
+        every.foo contains "b" if {
+          false
+        }
+
+        every.foo contains "c" if {
+          true
+        }
+
+        foo.every contains "a"
+
+        foo.every contains "b" if {
+          false
+        }
+
+        foo.every contains "c" if {
+          true
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/every keyword rule refhead, partial object
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p if {
+          every.foo == "a"
+          every.bar.one == "a"
+          every.bar.three == "c"
+          foo.every == "a"
+          bar.baz.every == "a"
+        }
+
+        every.foo := "a"
+
+        every.foo := "b" if {
+          false
+        }
+
+        every.foo := "c" if {
+          false
+        }
+
+        every.bar.one := "a"
+
+        every.bar.two := "b" if {
+          false
+        }
+
+        every.bar.three := "c" if {
+          true
+        }
+
+        foo.every := "a"
+
+        foo.every := "b" if {
+          false
+        }
+
+        foo.every := "c" if {
+          false
+        }
+
+        bar.baz.every := "a"
+
+        bar.baz.every := "b" if {
+          false
+        }
+
+        bar.baz.every := "c" if {
+          false
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/every keyword function refhead
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p if {
+          every.foo(1) == 1
+          every.foo(11) == 42
+          foo.every(1) == 1
+          foo.every(11) == 42
+          bar.every.baz(1) == 1
+          bar.every.baz(11) == 42
+        }
+
+        default every.foo(_) := 42
+
+        every.foo(x) := x if {
+          x < 10
+        }
+
+        default foo.every(_) := 42
+
+        foo.every(x) := x if {
+          x < 10
+        }
+
+        default bar.every.baz(_) := 42
+
+        bar.every.baz(x) := x if {
+          x < 10
+        }
+    want_result:
+      - x: true

v1/test/cases/testdata/v1/keywordrefs/test-keyword-false.yaml

@@ -0,0 +1,199 @@
+---
+cases:
+  - note: keywordrefs/false keyword in package
+    query: data.foo.p = x
+    modules:
+      - |
+        package foo.false.bar
+
+        baz := 42
+      - |
+        package foo
+        import data.foo.false.bar
+
+        p if {
+          bar.baz == 42
+          data.foo.false.bar.baz == 42
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/false keyword in package, import alias
+    query: data.foo.p = x
+    modules:
+      - |
+        package foo.false
+
+        bar := 42
+      - |
+        package foo
+        import data.foo.false as my_if
+
+        p if {
+          my_if.bar == 42
+          data.foo.false.bar == 42
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/false keyword rule refhead
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p if {
+          false.foo == 1
+          foo.false == 2
+        }
+
+        false.foo := 1
+
+        foo.false := 2
+    want_result:
+      - x: true
+  - note: keywordrefs/false keyword rule refhead, else bodies
+    query: data.test.p = x
+    input:
+      x: 3
+    modules:
+      - |
+        package test
+
+        p if {
+          false.foo == 3
+          foo.false == 6
+        }
+
+        false.foo := 1 if {
+          input.x == 1
+        } else := 2 if {
+          input.x == 2
+        } else := 3
+
+        foo.false := 4 if {
+          input.x == 1
+        } else := 5 if {
+          input.x == 2
+        } else := 6
+    want_result:
+      - x: true
+  - note: keywordrefs/false keyword rule refhead, partial set
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p if {
+          false.foo == {"a", "c"}
+          foo.false == {"a", "c"}
+        }
+
+        false.foo contains "a"
+
+        false.foo contains "b" if {
+          false
+        }
+
+        false.foo contains "c" if {
+          true
+        }
+
+        foo.false contains "a"
+
+        foo.false contains "b" if {
+          false
+        }
+
+        foo.false contains "c" if {
+          true
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/false keyword rule refhead, partial object
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p if {
+          false.foo == "a"
+          false.bar.one == "a"
+          false.bar.three == "c"
+          foo.false == "a"
+          bar.baz.false == "a"
+        }
+
+        false.foo := "a"
+
+        false.foo := "b" if {
+          false
+        }
+
+        false.foo := "c" if {
+          false
+        }
+
+        false.bar.one := "a"
+
+        false.bar.two := "b" if {
+          false
+        }
+
+        false.bar.three := "c" if {
+          true
+        }
+
+        foo.false := "a"
+
+        foo.false := "b" if {
+          false
+        }
+
+        foo.false := "c" if {
+          false
+        }
+
+        bar.baz.false := "a"
+
+        bar.baz.false := "b" if {
+          false
+        }
+
+        bar.baz.false := "c" if {
+          false
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/false keyword function refhead
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p if {
+          false.foo(1) == 1
+          false.foo(11) == 42
+          foo.false(1) == 1
+          foo.false(11) == 42
+          bar.false.baz(1) == 1
+          bar.false.baz(11) == 42
+        }
+
+        default false.foo(_) := 42
+
+        false.foo(x) := x if {
+          x < 10
+        }
+
+        default foo.false(_) := 42
+
+        foo.false(x) := x if {
+          x < 10
+        }
+
+        default bar.false.baz(_) := 42
+
+        bar.false.baz(x) := x if {
+          x < 10
+        }
+    want_result:
+      - x: true

v1/test/cases/testdata/v1/keywordrefs/test-keyword-if.yaml

@@ -0,0 +1,199 @@
+---
+cases:
+  - note: keywordrefs/if keyword in package
+    query: data.foo.p = x
+    modules:
+      - |
+        package foo.if.bar
+
+        baz := 42
+      - |
+        package foo
+        import data.foo.if.bar
+
+        p if {
+          bar.baz == 42
+          data.foo.if.bar.baz == 42
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/if keyword in package, import alias
+    query: data.foo.p = x
+    modules:
+      - |
+        package foo.if
+
+        bar := 42
+      - |
+        package foo
+        import data.foo.if as my_if
+
+        p if {
+          my_if.bar == 42
+          data.foo.if.bar == 42
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/if keyword rule refhead
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p if {
+          if.foo == 1
+          foo.if == 2
+        }
+
+        if.foo := 1
+
+        foo.if := 2
+    want_result:
+      - x: true
+  - note: keywordrefs/if keyword rule refhead, else bodies
+    query: data.test.p = x
+    input:
+      x: 3
+    modules:
+      - |
+        package test
+
+        p if {
+          if.foo == 3
+          foo.if == 6
+        }
+
+        if.foo := 1 if {
+          input.x == 1
+        } else := 2 if {
+          input.x == 2
+        } else := 3
+
+        foo.if := 4 if {
+          input.x == 1
+        } else := 5 if {
+          input.x == 2
+        } else := 6
+    want_result:
+      - x: true
+  - note: keywordrefs/if keyword rule refhead, partial set
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p if {
+          if.foo == {"a", "c"}
+          foo.if == {"a", "c"}
+        }
+
+        if.foo contains "a"
+
+        if.foo contains "b" if {
+          false
+        }
+
+        if.foo contains "c" if {
+          true
+        }
+
+        foo.if contains "a"
+
+        foo.if contains "b" if {
+          false
+        }
+
+        foo.if contains "c" if {
+          true
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/if keyword rule refhead, partial object
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p if {
+          if.foo == "a"
+          if.bar.one == "a"
+          if.bar.three == "c"
+          foo.if == "a"
+          bar.baz.if == "a"
+        }
+
+        if.foo := "a"
+
+        if.foo := "b" if {
+          false
+        }
+
+        if.foo := "c" if {
+          false
+        }
+
+        if.bar.one := "a"
+
+        if.bar.two := "b" if {
+          false
+        }
+
+        if.bar.three := "c" if {
+          true
+        }
+
+        foo.if := "a"
+
+        foo.if := "b" if {
+          false
+        }
+
+        foo.if := "c" if {
+          false
+        }
+
+        bar.baz.if := "a"
+
+        bar.baz.if := "b" if {
+          false
+        }
+
+        bar.baz.if := "c" if {
+          false
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/if keyword function refhead
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p if {
+          if.foo(1) == 1
+          if.foo(11) == 42
+          foo.if(1) == 1
+          foo.if(11) == 42
+          bar.if.baz(1) == 1
+          bar.if.baz(11) == 42
+        }
+
+        default if.foo(_) := 42
+
+        if.foo(x) := x if {
+          x < 10
+        }
+
+        default foo.if(_) := 42
+
+        foo.if(x) := x if {
+          x < 10
+        }
+
+        default bar.if.baz(_) := 42
+
+        bar.if.baz(x) := x if {
+          x < 10
+        }
+    want_result:
+      - x: true

v1/test/cases/testdata/v1/keywordrefs/test-keyword-import.yaml

@@ -0,0 +1,199 @@
+---
+cases:
+  - note: keywordrefs/import keyword in package
+    query: data.foo.p = x
+    modules:
+      - |
+        package foo.import.bar
+
+        baz := 42
+      - |
+        package foo
+        import data.foo.import.bar
+
+        p if {
+          bar.baz == 42
+          data.foo.import.bar.baz == 42
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/import keyword in package, import alias
+    query: data.foo.p = x
+    modules:
+      - |
+        package foo.import
+
+        bar := 42
+      - |
+        package foo
+        import data.foo.import as my_if
+
+        p if {
+          my_if.bar == 42
+          data.foo.import.bar == 42
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/import keyword rule refhead
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p if {
+          import.foo == 1
+          foo.import == 2
+        }
+
+        import.foo := 1
+
+        foo.import := 2
+    want_result:
+      - x: true
+  - note: keywordrefs/import keyword rule refhead, else bodies
+    query: data.test.p = x
+    input:
+      x: 3
+    modules:
+      - |
+        package test
+
+        p if {
+          import.foo == 3
+          foo.import == 6
+        }
+
+        import.foo := 1 if {
+          input.x == 1
+        } else := 2 if {
+          input.x == 2
+        } else := 3
+
+        foo.import := 4 if {
+          input.x == 1
+        } else := 5 if {
+          input.x == 2
+        } else := 6
+    want_result:
+      - x: true
+  - note: keywordrefs/import keyword rule refhead, partial set
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p if {
+          import.foo == {"a", "c"}
+          foo.import == {"a", "c"}
+        }
+
+        import.foo contains "a"
+
+        import.foo contains "b" if {
+          false
+        }
+
+        import.foo contains "c" if {
+          true
+        }
+
+        foo.import contains "a"
+
+        foo.import contains "b" if {
+          false
+        }
+
+        foo.import contains "c" if {
+          true
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/import keyword rule refhead, partial object
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p if {
+          import.foo == "a"
+          import.bar.one == "a"
+          import.bar.three == "c"
+          foo.import == "a"
+          bar.baz.import == "a"
+        }
+
+        import.foo := "a"
+
+        import.foo := "b" if {
+          false
+        }
+
+        import.foo := "c" if {
+          false
+        }
+
+        import.bar.one := "a"
+
+        import.bar.two := "b" if {
+          false
+        }
+
+        import.bar.three := "c" if {
+          true
+        }
+
+        foo.import := "a"
+
+        foo.import := "b" if {
+          false
+        }
+
+        foo.import := "c" if {
+          false
+        }
+
+        bar.baz.import := "a"
+
+        bar.baz.import := "b" if {
+          false
+        }
+
+        bar.baz.import := "c" if {
+          false
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/import keyword function refhead
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p if {
+          import.foo(1) == 1
+          import.foo(11) == 42
+          foo.import(1) == 1
+          foo.import(11) == 42
+          bar.import.baz(1) == 1
+          bar.import.baz(11) == 42
+        }
+
+        default import.foo(_) := 42
+
+        import.foo(x) := x if {
+          x < 10
+        }
+
+        default foo.import(_) := 42
+
+        foo.import(x) := x if {
+          x < 10
+        }
+
+        default bar.import.baz(_) := 42
+
+        bar.import.baz(x) := x if {
+          x < 10
+        }
+    want_result:
+      - x: true

v1/test/cases/testdata/v1/keywordrefs/test-keyword-in.yaml

@@ -0,0 +1,199 @@
+---
+cases:
+  - note: keywordrefs/in keyword in package
+    query: data.foo.p = x
+    modules:
+      - |
+        package foo.in.bar
+
+        baz := 42
+      - |
+        package foo
+        import data.foo.in.bar
+
+        p if {
+          bar.baz == 42
+          data.foo.in.bar.baz == 42
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/in keyword in package, import alias
+    query: data.foo.p = x
+    modules:
+      - |
+        package foo.in
+
+        bar := 42
+      - |
+        package foo
+        import data.foo.in as my_if
+
+        p if {
+          my_if.bar == 42
+          data.foo.in.bar == 42
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/in keyword rule refhead
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p if {
+          in.foo == 1
+          foo.in == 2
+        }
+
+        in.foo := 1
+
+        foo.in := 2
+    want_result:
+      - x: true
+  - note: keywordrefs/in keyword rule refhead, else bodies
+    query: data.test.p = x
+    input:
+      x: 3
+    modules:
+      - |
+        package test
+
+        p if {
+          in.foo == 3
+          foo.in == 6
+        }
+
+        in.foo := 1 if {
+          input.x == 1
+        } else := 2 if {
+          input.x == 2
+        } else := 3
+
+        foo.in := 4 if {
+          input.x == 1
+        } else := 5 if {
+          input.x == 2
+        } else := 6
+    want_result:
+      - x: true
+  - note: keywordrefs/in keyword rule refhead, partial set
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p if {
+          in.foo == {"a", "c"}
+          foo.in == {"a", "c"}
+        }
+
+        in.foo contains "a"
+
+        in.foo contains "b" if {
+          false
+        }
+
+        in.foo contains "c" if {
+          true
+        }
+
+        foo.in contains "a"
+
+        foo.in contains "b" if {
+          false
+        }
+
+        foo.in contains "c" if {
+          true
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/in keyword rule refhead, partial object
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p if {
+          in.foo == "a"
+          in.bar.one == "a"
+          in.bar.three == "c"
+          foo.in == "a"
+          bar.baz.in == "a"
+        }
+
+        in.foo := "a"
+
+        in.foo := "b" if {
+          false
+        }
+
+        in.foo := "c" if {
+          false
+        }
+
+        in.bar.one := "a"
+
+        in.bar.two := "b" if {
+          false
+        }
+
+        in.bar.three := "c" if {
+          true
+        }
+
+        foo.in := "a"
+
+        foo.in := "b" if {
+          false
+        }
+
+        foo.in := "c" if {
+          false
+        }
+
+        bar.baz.in := "a"
+
+        bar.baz.in := "b" if {
+          false
+        }
+
+        bar.baz.in := "c" if {
+          false
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/in keyword function refhead
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p if {
+          in.foo(1) == 1
+          in.foo(11) == 42
+          foo.in(1) == 1
+          foo.in(11) == 42
+          bar.in.baz(1) == 1
+          bar.in.baz(11) == 42
+        }
+
+        default in.foo(_) := 42
+
+        in.foo(x) := x if {
+          x < 10
+        }
+
+        default foo.in(_) := 42
+
+        foo.in(x) := x if {
+          x < 10
+        }
+
+        default bar.in.baz(_) := 42
+
+        bar.in.baz(x) := x if {
+          x < 10
+        }
+    want_result:
+      - x: true

v1/test/cases/testdata/v1/keywordrefs/test-keyword-not.yaml

@@ -0,0 +1,199 @@
+---
+cases:
+  - note: keywordrefs/not keyword in package
+    query: data.foo.p = x
+    modules:
+      - |
+        package foo.not.bar
+
+        baz := 42
+      - |
+        package foo
+        import data.foo.not.bar
+
+        p if {
+          bar.baz == 42
+          data.foo.not.bar.baz == 42
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/not keyword in package, import alias
+    query: data.foo.p = x
+    modules:
+      - |
+        package foo.not
+
+        bar := 42
+      - |
+        package foo
+        import data.foo.not as my_if
+
+        p if {
+          my_if.bar == 42
+          data.foo.not.bar == 42
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/not keyword rule refhead
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p if {
+          not.foo == 1
+          foo.not == 2
+        }
+
+        not.foo := 1
+
+        foo.not := 2
+    want_result:
+      - x: true
+  - note: keywordrefs/not keyword rule refhead, else bodies
+    query: data.test.p = x
+    input:
+      x: 3
+    modules:
+      - |
+        package test
+
+        p if {
+          not.foo == 3
+          foo.not == 6
+        }
+
+        not.foo := 1 if {
+          input.x == 1
+        } else := 2 if {
+          input.x == 2
+        } else := 3
+
+        foo.not := 4 if {
+          input.x == 1
+        } else := 5 if {
+          input.x == 2
+        } else := 6
+    want_result:
+      - x: true
+  - note: keywordrefs/not keyword rule refhead, partial set
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p if {
+          not.foo == {"a", "c"}
+          foo.not == {"a", "c"}
+        }
+
+        not.foo contains "a"
+
+        not.foo contains "b" if {
+          false
+        }
+
+        not.foo contains "c" if {
+          true
+        }
+
+        foo.not contains "a"
+
+        foo.not contains "b" if {
+          false
+        }
+
+        foo.not contains "c" if {
+          true
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/not keyword rule refhead, partial object
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p if {
+          not.foo == "a"
+          not.bar.one == "a"
+          not.bar.three == "c"
+          foo.not == "a"
+          bar.baz.not == "a"
+        }
+
+        not.foo := "a"
+
+        not.foo := "b" if {
+          false
+        }
+
+        not.foo := "c" if {
+          false
+        }
+
+        not.bar.one := "a"
+
+        not.bar.two := "b" if {
+          false
+        }
+
+        not.bar.three := "c" if {
+          true
+        }
+
+        foo.not := "a"
+
+        foo.not := "b" if {
+          false
+        }
+
+        foo.not := "c" if {
+          false
+        }
+
+        bar.baz.not := "a"
+
+        bar.baz.not := "b" if {
+          false
+        }
+
+        bar.baz.not := "c" if {
+          false
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/not keyword function refhead
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p if {
+          not.foo(1) == 1
+          not.foo(11) == 42
+          foo.not(1) == 1
+          foo.not(11) == 42
+          bar.not.baz(1) == 1
+          bar.not.baz(11) == 42
+        }
+
+        default not.foo(_) := 42
+
+        not.foo(x) := x if {
+          x < 10
+        }
+
+        default foo.not(_) := 42
+
+        foo.not(x) := x if {
+          x < 10
+        }
+
+        default bar.not.baz(_) := 42
+
+        bar.not.baz(x) := x if {
+          x < 10
+        }
+    want_result:
+      - x: true

v1/test/cases/testdata/v1/keywordrefs/test-keyword-null.yaml

@@ -0,0 +1,199 @@
+---
+cases:
+  - note: keywordrefs/null keyword in package
+    query: data.foo.p = x
+    modules:
+      - |
+        package foo.null.bar
+
+        baz := 42
+      - |
+        package foo
+        import data.foo.null.bar
+
+        p if {
+          bar.baz == 42
+          data.foo.null.bar.baz == 42
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/null keyword in package, import alias
+    query: data.foo.p = x
+    modules:
+      - |
+        package foo.null
+
+        bar := 42
+      - |
+        package foo
+        import data.foo.null as my_if
+
+        p if {
+          my_if.bar == 42
+          data.foo.null.bar == 42
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/null keyword rule refhead
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p if {
+          null.foo == 1
+          foo.null == 2
+        }
+
+        null.foo := 1
+
+        foo.null := 2
+    want_result:
+      - x: true
+  - note: keywordrefs/null keyword rule refhead, else bodies
+    query: data.test.p = x
+    input:
+      x: 3
+    modules:
+      - |
+        package test
+
+        p if {
+          null.foo == 3
+          foo.null == 6
+        }
+
+        null.foo := 1 if {
+          input.x == 1
+        } else := 2 if {
+          input.x == 2
+        } else := 3
+
+        foo.null := 4 if {
+          input.x == 1
+        } else := 5 if {
+          input.x == 2
+        } else := 6
+    want_result:
+      - x: true
+  - note: keywordrefs/null keyword rule refhead, partial set
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p if {
+          null.foo == {"a", "c"}
+          foo.null == {"a", "c"}
+        }
+
+        null.foo contains "a"
+
+        null.foo contains "b" if {
+          false
+        }
+
+        null.foo contains "c" if {
+          true
+        }
+
+        foo.null contains "a"
+
+        foo.null contains "b" if {
+          false
+        }
+
+        foo.null contains "c" if {
+          true
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/null keyword rule refhead, partial object
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p if {
+          null.foo == "a"
+          null.bar.one == "a"
+          null.bar.three == "c"
+          foo.null == "a"
+          bar.baz.null == "a"
+        }
+
+        null.foo := "a"
+
+        null.foo := "b" if {
+          false
+        }
+
+        null.foo := "c" if {
+          false
+        }
+
+        null.bar.one := "a"
+
+        null.bar.two := "b" if {
+          false
+        }
+
+        null.bar.three := "c" if {
+          true
+        }
+
+        foo.null := "a"
+
+        foo.null := "b" if {
+          false
+        }
+
+        foo.null := "c" if {
+          false
+        }
+
+        bar.baz.null := "a"
+
+        bar.baz.null := "b" if {
+          false
+        }
+
+        bar.baz.null := "c" if {
+          false
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/null keyword function refhead
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p if {
+          null.foo(1) == 1
+          null.foo(11) == 42
+          foo.null(1) == 1
+          foo.null(11) == 42
+          bar.null.baz(1) == 1
+          bar.null.baz(11) == 42
+        }
+
+        default null.foo(_) := 42
+
+        null.foo(x) := x if {
+          x < 10
+        }
+
+        default foo.null(_) := 42
+
+        foo.null(x) := x if {
+          x < 10
+        }
+
+        default bar.null.baz(_) := 42
+
+        bar.null.baz(x) := x if {
+          x < 10
+        }
+    want_result:
+      - x: true

v1/test/cases/testdata/v1/keywordrefs/test-keyword-package.yaml

@@ -0,0 +1,199 @@
+---
+cases:
+  - note: keywordrefs/package keyword in package
+    query: data.foo.p = x
+    modules:
+      - |
+        package foo.package.bar
+
+        baz := 42
+      - |
+        package foo
+        import data.foo.package.bar
+
+        p if {
+          bar.baz == 42
+          data.foo.package.bar.baz == 42
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/package keyword in package, import alias
+    query: data.foo.p = x
+    modules:
+      - |
+        package foo.package
+
+        bar := 42
+      - |
+        package foo
+        import data.foo.package as my_if
+
+        p if {
+          my_if.bar == 42
+          data.foo.package.bar == 42
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/package keyword rule refhead
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p if {
+          package.foo == 1
+          foo.package == 2
+        }
+
+        package.foo := 1
+
+        foo.package := 2
+    want_result:
+      - x: true
+  - note: keywordrefs/package keyword rule refhead, else bodies
+    query: data.test.p = x
+    input:
+      x: 3
+    modules:
+      - |
+        package test
+
+        p if {
+          package.foo == 3
+          foo.package == 6
+        }
+
+        package.foo := 1 if {
+          input.x == 1
+        } else := 2 if {
+          input.x == 2
+        } else := 3
+
+        foo.package := 4 if {
+          input.x == 1
+        } else := 5 if {
+          input.x == 2
+        } else := 6
+    want_result:
+      - x: true
+  - note: keywordrefs/package keyword rule refhead, partial set
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p if {
+          package.foo == {"a", "c"}
+          foo.package == {"a", "c"}
+        }
+
+        package.foo contains "a"
+
+        package.foo contains "b" if {
+          false
+        }
+
+        package.foo contains "c" if {
+          true
+        }
+
+        foo.package contains "a"
+
+        foo.package contains "b" if {
+          false
+        }
+
+        foo.package contains "c" if {
+          true
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/package keyword rule refhead, partial object
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p if {
+          package.foo == "a"
+          package.bar.one == "a"
+          package.bar.three == "c"
+          foo.package == "a"
+          bar.baz.package == "a"
+        }
+
+        package.foo := "a"
+
+        package.foo := "b" if {
+          false
+        }
+
+        package.foo := "c" if {
+          false
+        }
+
+        package.bar.one := "a"
+
+        package.bar.two := "b" if {
+          false
+        }
+
+        package.bar.three := "c" if {
+          true
+        }
+
+        foo.package := "a"
+
+        foo.package := "b" if {
+          false
+        }
+
+        foo.package := "c" if {
+          false
+        }
+
+        bar.baz.package := "a"
+
+        bar.baz.package := "b" if {
+          false
+        }
+
+        bar.baz.package := "c" if {
+          false
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/package keyword function refhead
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p if {
+          package.foo(1) == 1
+          package.foo(11) == 42
+          foo.package(1) == 1
+          foo.package(11) == 42
+          bar.package.baz(1) == 1
+          bar.package.baz(11) == 42
+        }
+
+        default package.foo(_) := 42
+
+        package.foo(x) := x if {
+          x < 10
+        }
+
+        default foo.package(_) := 42
+
+        foo.package(x) := x if {
+          x < 10
+        }
+
+        default bar.package.baz(_) := 42
+
+        bar.package.baz(x) := x if {
+          x < 10
+        }
+    want_result:
+      - x: true

v1/test/cases/testdata/v1/keywordrefs/test-keyword-some.yaml

@@ -0,0 +1,199 @@
+---
+cases:
+  - note: keywordrefs/some keyword in package
+    query: data.foo.p = x
+    modules:
+      - |
+        package foo.some.bar
+
+        baz := 42
+      - |
+        package foo
+        import data.foo.some.bar
+
+        p if {
+          bar.baz == 42
+          data.foo.some.bar.baz == 42
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/some keyword in package, import alias
+    query: data.foo.p = x
+    modules:
+      - |
+        package foo.some
+
+        bar := 42
+      - |
+        package foo
+        import data.foo.some as my_if
+
+        p if {
+          my_if.bar == 42
+          data.foo.some.bar == 42
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/some keyword rule refhead
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p if {
+          some.foo == 1
+          foo.some == 2
+        }
+
+        some.foo := 1
+
+        foo.some := 2
+    want_result:
+      - x: true
+  - note: keywordrefs/some keyword rule refhead, else bodies
+    query: data.test.p = x
+    input:
+      x: 3
+    modules:
+      - |
+        package test
+
+        p if {
+          some.foo == 3
+          foo.some == 6
+        }
+
+        some.foo := 1 if {
+          input.x == 1
+        } else := 2 if {
+          input.x == 2
+        } else := 3
+
+        foo.some := 4 if {
+          input.x == 1
+        } else := 5 if {
+          input.x == 2
+        } else := 6
+    want_result:
+      - x: true
+  - note: keywordrefs/some keyword rule refhead, partial set
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p if {
+          some.foo == {"a", "c"}
+          foo.some == {"a", "c"}
+        }
+
+        some.foo contains "a"
+
+        some.foo contains "b" if {
+          false
+        }
+
+        some.foo contains "c" if {
+          true
+        }
+
+        foo.some contains "a"
+
+        foo.some contains "b" if {
+          false
+        }
+
+        foo.some contains "c" if {
+          true
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/some keyword rule refhead, partial object
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p if {
+          some.foo == "a"
+          some.bar.one == "a"
+          some.bar.three == "c"
+          foo.some == "a"
+          bar.baz.some == "a"
+        }
+
+        some.foo := "a"
+
+        some.foo := "b" if {
+          false
+        }
+
+        some.foo := "c" if {
+          false
+        }
+
+        some.bar.one := "a"
+
+        some.bar.two := "b" if {
+          false
+        }
+
+        some.bar.three := "c" if {
+          true
+        }
+
+        foo.some := "a"
+
+        foo.some := "b" if {
+          false
+        }
+
+        foo.some := "c" if {
+          false
+        }
+
+        bar.baz.some := "a"
+
+        bar.baz.some := "b" if {
+          false
+        }
+
+        bar.baz.some := "c" if {
+          false
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/some keyword function refhead
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p if {
+          some.foo(1) == 1
+          some.foo(11) == 42
+          foo.some(1) == 1
+          foo.some(11) == 42
+          bar.some.baz(1) == 1
+          bar.some.baz(11) == 42
+        }
+
+        default some.foo(_) := 42
+
+        some.foo(x) := x if {
+          x < 10
+        }
+
+        default foo.some(_) := 42
+
+        foo.some(x) := x if {
+          x < 10
+        }
+
+        default bar.some.baz(_) := 42
+
+        bar.some.baz(x) := x if {
+          x < 10
+        }
+    want_result:
+      - x: true

v1/test/cases/testdata/v1/keywordrefs/test-keyword-true.yaml

@@ -0,0 +1,199 @@
+---
+cases:
+  - note: keywordrefs/true keyword in package
+    query: data.foo.p = x
+    modules:
+      - |
+        package foo.true.bar
+
+        baz := 42
+      - |
+        package foo
+        import data.foo.true.bar
+
+        p if {
+          bar.baz == 42
+          data.foo.true.bar.baz == 42
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/true keyword in package, import alias
+    query: data.foo.p = x
+    modules:
+      - |
+        package foo.true
+
+        bar := 42
+      - |
+        package foo
+        import data.foo.true as my_if
+
+        p if {
+          my_if.bar == 42
+          data.foo.true.bar == 42
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/true keyword rule refhead
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p if {
+          true.foo == 1
+          foo.true == 2
+        }
+
+        true.foo := 1
+
+        foo.true := 2
+    want_result:
+      - x: true
+  - note: keywordrefs/true keyword rule refhead, else bodies
+    query: data.test.p = x
+    input:
+      x: 3
+    modules:
+      - |
+        package test
+
+        p if {
+          true.foo == 3
+          foo.true == 6
+        }
+
+        true.foo := 1 if {
+          input.x == 1
+        } else := 2 if {
+          input.x == 2
+        } else := 3
+
+        foo.true := 4 if {
+          input.x == 1
+        } else := 5 if {
+          input.x == 2
+        } else := 6
+    want_result:
+      - x: true
+  - note: keywordrefs/true keyword rule refhead, partial set
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p if {
+          true.foo == {"a", "c"}
+          foo.true == {"a", "c"}
+        }
+
+        true.foo contains "a"
+
+        true.foo contains "b" if {
+          false
+        }
+
+        true.foo contains "c" if {
+          true
+        }
+
+        foo.true contains "a"
+
+        foo.true contains "b" if {
+          false
+        }
+
+        foo.true contains "c" if {
+          true
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/true keyword rule refhead, partial object
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p if {
+          true.foo == "a"
+          true.bar.one == "a"
+          true.bar.three == "c"
+          foo.true == "a"
+          bar.baz.true == "a"
+        }
+
+        true.foo := "a"
+
+        true.foo := "b" if {
+          false
+        }
+
+        true.foo := "c" if {
+          false
+        }
+
+        true.bar.one := "a"
+
+        true.bar.two := "b" if {
+          false
+        }
+
+        true.bar.three := "c" if {
+          true
+        }
+
+        foo.true := "a"
+
+        foo.true := "b" if {
+          false
+        }
+
+        foo.true := "c" if {
+          false
+        }
+
+        bar.baz.true := "a"
+
+        bar.baz.true := "b" if {
+          false
+        }
+
+        bar.baz.true := "c" if {
+          false
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/true keyword function refhead
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p if {
+          true.foo(1) == 1
+          true.foo(11) == 42
+          foo.true(1) == 1
+          foo.true(11) == 42
+          bar.true.baz(1) == 1
+          bar.true.baz(11) == 42
+        }
+
+        default true.foo(_) := 42
+
+        true.foo(x) := x if {
+          x < 10
+        }
+
+        default foo.true(_) := 42
+
+        foo.true(x) := x if {
+          x < 10
+        }
+
+        default bar.true.baz(_) := 42
+
+        bar.true.baz(x) := x if {
+          x < 10
+        }
+    want_result:
+      - x: true

v1/test/cases/testdata/v1/keywordrefs/test-keyword-with.yaml

@@ -0,0 +1,199 @@
+---
+cases:
+  - note: keywordrefs/with keyword in package
+    query: data.foo.p = x
+    modules:
+      - |
+        package foo.with.bar
+
+        baz := 42
+      - |
+        package foo
+        import data.foo.with.bar
+
+        p if {
+          bar.baz == 42
+          data.foo.with.bar.baz == 42
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/with keyword in package, import alias
+    query: data.foo.p = x
+    modules:
+      - |
+        package foo.with
+
+        bar := 42
+      - |
+        package foo
+        import data.foo.with as my_if
+
+        p if {
+          my_if.bar == 42
+          data.foo.with.bar == 42
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/with keyword rule refhead
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p if {
+          with.foo == 1
+          foo.with == 2
+        }
+
+        with.foo := 1
+
+        foo.with := 2
+    want_result:
+      - x: true
+  - note: keywordrefs/with keyword rule refhead, else bodies
+    query: data.test.p = x
+    input:
+      x: 3
+    modules:
+      - |
+        package test
+
+        p if {
+          with.foo == 3
+          foo.with == 6
+        }
+
+        with.foo := 1 if {
+          input.x == 1
+        } else := 2 if {
+          input.x == 2
+        } else := 3
+
+        foo.with := 4 if {
+          input.x == 1
+        } else := 5 if {
+          input.x == 2
+        } else := 6
+    want_result:
+      - x: true
+  - note: keywordrefs/with keyword rule refhead, partial set
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p if {
+          with.foo == {"a", "c"}
+          foo.with == {"a", "c"}
+        }
+
+        with.foo contains "a"
+
+        with.foo contains "b" if {
+          false
+        }
+
+        with.foo contains "c" if {
+          true
+        }
+
+        foo.with contains "a"
+
+        foo.with contains "b" if {
+          false
+        }
+
+        foo.with contains "c" if {
+          true
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/with keyword rule refhead, partial object
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p if {
+          with.foo == "a"
+          with.bar.one == "a"
+          with.bar.three == "c"
+          foo.with == "a"
+          bar.baz.with == "a"
+        }
+
+        with.foo := "a"
+
+        with.foo := "b" if {
+          false
+        }
+
+        with.foo := "c" if {
+          false
+        }
+
+        with.bar.one := "a"
+
+        with.bar.two := "b" if {
+          false
+        }
+
+        with.bar.three := "c" if {
+          true
+        }
+
+        foo.with := "a"
+
+        foo.with := "b" if {
+          false
+        }
+
+        foo.with := "c" if {
+          false
+        }
+
+        bar.baz.with := "a"
+
+        bar.baz.with := "b" if {
+          false
+        }
+
+        bar.baz.with := "c" if {
+          false
+        }
+    want_result:
+      - x: true
+  - note: keywordrefs/with keyword function refhead
+    query: data.test.p = x
+    modules:
+      - |
+        package test
+
+        p if {
+          with.foo(1) == 1
+          with.foo(11) == 42
+          foo.with(1) == 1
+          foo.with(11) == 42
+          bar.with.baz(1) == 1
+          bar.with.baz(11) == 42
+        }
+
+        default with.foo(_) := 42
+
+        with.foo(x) := x if {
+          x < 10
+        }
+
+        default foo.with(_) := 42
+
+        foo.with(x) := x if {
+          x < 10
+        }
+
+        default bar.with.baz(_) := 42
+
+        bar.with.baz(x) := x if {
+          x < 10
+        }
+    want_result:
+      - x: true