diff --git a/android/androidmk.go b/android/androidmk.go
index 5856851..32d7712 100644
--- a/android/androidmk.go
+++ b/android/androidmk.go
@@ -141,7 +141,20 @@
 	entryOrder []string
 }
 
-type AndroidMkExtraEntriesFunc func(entries *AndroidMkEntries)
+type AndroidMkExtraEntriesContext interface {
+	Provider(provider blueprint.ProviderKey) interface{}
+}
+
+type androidMkExtraEntriesContext struct {
+	ctx fillInEntriesContext
+	mod blueprint.Module
+}
+
+func (a *androidMkExtraEntriesContext) Provider(provider blueprint.ProviderKey) interface{} {
+	return a.ctx.ModuleProvider(a.mod, provider)
+}
+
+type AndroidMkExtraEntriesFunc func(ctx AndroidMkExtraEntriesContext, entries *AndroidMkEntries)
 type AndroidMkExtraFootersFunc func(w io.Writer, name, prefix, moduleDir string)
 
 // Utility funcs to manipulate Android.mk variable entries.
@@ -448,7 +461,13 @@
 
 // fillInEntries goes through the common variable processing and calls the extra data funcs to
 // generate and fill in AndroidMkEntries's in-struct data, ready to be flushed to a file.
-func (a *AndroidMkEntries) fillInEntries(config Config, bpPath string, mod blueprint.Module) {
+type fillInEntriesContext interface {
+	ModuleDir(module blueprint.Module) string
+	Config() Config
+	ModuleProvider(module blueprint.Module, provider blueprint.ProviderKey) interface{}
+}
+
+func (a *AndroidMkEntries) fillInEntries(ctx fillInEntriesContext, mod blueprint.Module) {
 	a.EntryMap = make(map[string][]string)
 	amod := mod.(Module).base()
 	name := amod.BaseModuleName()
@@ -470,7 +489,7 @@
 	fmt.Fprintln(&a.header, "\ninclude $(CLEAR_VARS)")
 
 	// Collect make variable assignment entries.
-	a.SetString("LOCAL_PATH", filepath.Dir(bpPath))
+	a.SetString("LOCAL_PATH", ctx.ModuleDir(mod))
 	a.SetString("LOCAL_MODULE", name+a.SubName)
 	a.AddStrings("LOCAL_LICENSE_KINDS", amod.commonProperties.Effective_license_kinds...)
 	a.AddStrings("LOCAL_LICENSE_CONDITIONS", amod.commonProperties.Effective_license_conditions...)
@@ -561,17 +580,23 @@
 
 		}
 
-		if amod.Arch().ArchType != config.Targets[amod.Os()][0].Arch.ArchType {
+		if amod.Arch().ArchType != ctx.Config().Targets[amod.Os()][0].Arch.ArchType {
 			prefix = "2ND_" + prefix
 		}
 	}
+
+	extraCtx := &androidMkExtraEntriesContext{
+		ctx: ctx,
+		mod: mod,
+	}
+
 	for _, extra := range a.ExtraEntries {
-		extra(a)
+		extra(extraCtx, a)
 	}
 
 	// Write to footer.
 	fmt.Fprintln(&a.footer, "include "+a.Include)
-	blueprintDir := filepath.Dir(bpPath)
+	blueprintDir := ctx.ModuleDir(mod)
 	for _, footerFunc := range a.ExtraFooters {
 		footerFunc(&a.footer, name, prefix, blueprintDir)
 	}
@@ -732,7 +757,7 @@
 	return nil
 }
 
-func (data *AndroidMkData) fillInData(config Config, bpPath string, mod blueprint.Module) {
+func (data *AndroidMkData) fillInData(ctx fillInEntriesContext, mod blueprint.Module) {
 	// Get the preamble content through AndroidMkEntries logic.
 	data.Entries = AndroidMkEntries{
 		Class:           data.Class,
@@ -745,7 +770,7 @@
 		Host_required:   data.Host_required,
 		Target_required: data.Target_required,
 	}
-	data.Entries.fillInEntries(config, bpPath, mod)
+	data.Entries.fillInEntries(ctx, mod)
 
 	// copy entries back to data since it is used in Custom
 	data.Required = data.Entries.Required
@@ -768,7 +793,7 @@
 		data.Include = "$(BUILD_PREBUILT)"
 	}
 
-	data.fillInData(ctx.Config(), ctx.BlueprintFile(mod), mod)
+	data.fillInData(ctx, mod)
 
 	prefix := ""
 	if amod.ArchSpecific() {
@@ -795,7 +820,7 @@
 	if data.Custom != nil {
 		// List of module types allowed to use .Custom(...)
 		// Additions to the list require careful review for proper license handling.
-		switch reflect.TypeOf(mod).String() {  // ctx.ModuleType(mod) doesn't work: aidl_interface creates phony without type
+		switch reflect.TypeOf(mod).String() { // ctx.ModuleType(mod) doesn't work: aidl_interface creates phony without type
 		case "*aidl.aidlApi": // writes non-custom before adding .phony
 		case "*aidl.aidlMapping": // writes non-custom before adding .phony
 		case "*android.customModule": // appears in tests only
@@ -850,7 +875,7 @@
 
 	// Any new or special cases here need review to verify correct propagation of license information.
 	for _, entries := range provider.AndroidMkEntries() {
-		entries.fillInEntries(ctx.Config(), ctx.BlueprintFile(mod), mod)
+		entries.fillInEntries(ctx, mod)
 		entries.write(w)
 	}
 
diff --git a/android/androidmk_test.go b/android/androidmk_test.go
index 347b92e..2f568fb 100644
--- a/android/androidmk_test.go
+++ b/android/androidmk_test.go
@@ -137,9 +137,9 @@
 	return module
 }
 
-// buildConfigAndCustomModuleFoo creates a config object, processes the supplied
+// buildContextAndCustomModuleFoo creates a config object, processes the supplied
 // bp module and then returns the config and the custom module called "foo".
-func buildConfigAndCustomModuleFoo(t *testing.T, bp string) (Config, *customModule) {
+func buildContextAndCustomModuleFoo(t *testing.T, bp string) (*TestContext, *customModule) {
 	t.Helper()
 	config := TestConfig(buildDir, nil, bp, nil)
 	config.katiEnabled = true // Enable androidmk Singleton
@@ -155,7 +155,7 @@
 	FailIfErrored(t, errs)
 
 	module := ctx.ModuleForTests("foo", "").Module().(*customModule)
-	return config, module
+	return ctx, module
 }
 
 func TestAndroidMkSingleton_PassesUpdatedAndroidMkDataToCustomCallback(t *testing.T) {
@@ -168,7 +168,7 @@
 	}
 	`
 
-	_, m := buildConfigAndCustomModuleFoo(t, bp)
+	_, m := buildContextAndCustomModuleFoo(t, bp)
 
 	assertEqual := func(expected interface{}, actual interface{}) {
 		if !reflect.DeepEqual(expected, actual) {
@@ -253,8 +253,8 @@
 		"$(call dist-for-goals,my_goal my_other_goal,three/four.out:four.out)\n",
 	}
 
-	config, module := buildConfigAndCustomModuleFoo(t, bp)
-	entries := AndroidMkEntriesForTest(t, config, "", module)
+	ctx, module := buildContextAndCustomModuleFoo(t, bp)
+	entries := AndroidMkEntriesForTest(t, ctx, module)
 	if len(entries) != 1 {
 		t.Errorf("Expected a single AndroidMk entry, got %d", len(entries))
 	}
@@ -343,8 +343,8 @@
 		t.Run(name, func(t *testing.T) {
 			t.Helper()
 
-			config, module := buildConfigAndCustomModuleFoo(t, bp)
-			entries := AndroidMkEntriesForTest(t, config, "", module)
+			ctx, module := buildContextAndCustomModuleFoo(t, bp)
+			entries := AndroidMkEntriesForTest(t, ctx, module)
 			if len(entries) != 1 {
 				t.Errorf("Expected a single AndroidMk entry, got %d", len(entries))
 			}
diff --git a/android/config.go b/android/config.go
index 095b1e4..f0bba81 100644
--- a/android/config.go
+++ b/android/config.go
@@ -913,6 +913,25 @@
 	return "json"
 }
 
+// XrefCuJavaSourceMax returns the maximum number of the Java source files
+// in a single compilation unit
+const xrefJavaSourceFileMaxDefault = "1000"
+
+func (c Config) XrefCuJavaSourceMax() string {
+	v := c.Getenv("KYTHE_JAVA_SOURCE_BATCH_SIZE")
+	if v == "" {
+		return xrefJavaSourceFileMaxDefault
+	}
+	if _, err := strconv.ParseUint(v, 0, 0); err != nil {
+		fmt.Fprintf(os.Stderr,
+			"bad KYTHE_JAVA_SOURCE_BATCH_SIZE value: %s, will use %s",
+			err, xrefJavaSourceFileMaxDefault)
+		return xrefJavaSourceFileMaxDefault
+	}
+	return v
+
+}
+
 func (c *config) EmitXrefRules() bool {
 	return c.XrefCorpusName() != ""
 }
diff --git a/android/csuite_config.go b/android/csuite_config.go
index bf24d98..56d2408 100644
--- a/android/csuite_config.go
+++ b/android/csuite_config.go
@@ -40,7 +40,7 @@
 		OutputFile: OptionalPathForPath(me.OutputFilePath),
 	}
 	androidMkEntries.ExtraEntries = []AndroidMkExtraEntriesFunc{
-		func(entries *AndroidMkEntries) {
+		func(ctx AndroidMkExtraEntriesContext, entries *AndroidMkEntries) {
 			if me.properties.Test_config != nil {
 				entries.SetString("LOCAL_TEST_CONFIG", *me.properties.Test_config)
 			}
diff --git a/android/testing.go b/android/testing.go
index de338bf..7852b60 100644
--- a/android/testing.go
+++ b/android/testing.go
@@ -201,6 +201,10 @@
 		"\nall singletons: %v", name, allSingletonNames))
 }
 
+func (ctx *TestContext) Config() Config {
+	return ctx.config
+}
+
 type testBuildProvider interface {
 	BuildParamsForTests() []BuildParams
 	RuleParamsForTests() map[blueprint.Rule]blueprint.RuleParams
@@ -461,7 +465,7 @@
 	config.katiEnabled = true
 }
 
-func AndroidMkEntriesForTest(t *testing.T, config Config, bpPath string, mod blueprint.Module) []AndroidMkEntries {
+func AndroidMkEntriesForTest(t *testing.T, ctx *TestContext, mod blueprint.Module) []AndroidMkEntries {
 	var p AndroidMkEntriesProvider
 	var ok bool
 	if p, ok = mod.(AndroidMkEntriesProvider); !ok {
@@ -470,19 +474,19 @@
 
 	entriesList := p.AndroidMkEntries()
 	for i, _ := range entriesList {
-		entriesList[i].fillInEntries(config, bpPath, mod)
+		entriesList[i].fillInEntries(ctx, mod)
 	}
 	return entriesList
 }
 
-func AndroidMkDataForTest(t *testing.T, config Config, bpPath string, mod blueprint.Module) AndroidMkData {
+func AndroidMkDataForTest(t *testing.T, ctx *TestContext, mod blueprint.Module) AndroidMkData {
 	var p AndroidMkDataProvider
 	var ok bool
 	if p, ok = mod.(AndroidMkDataProvider); !ok {
 		t.Errorf("module does not implement AndroidMkDataProvider: " + mod.Name())
 	}
 	data := p.AndroidMk()
-	data.fillInData(config, bpPath, mod)
+	data.fillInData(ctx, mod)
 	return data
 }
 
diff --git a/apex/allowed_deps.txt b/apex/allowed_deps.txt
index 90edca5..72394db 100644
--- a/apex/allowed_deps.txt
+++ b/apex/allowed_deps.txt
@@ -478,6 +478,7 @@
 mediaswcodec(minSdkVersion:29)
 metrics-constants-protos(minSdkVersion:29)
 modules-utils-build(minSdkVersion:29)
+modules-utils-build_system(minSdkVersion:29)
 modules-utils-os(minSdkVersion:30)
 ndk_crtbegin_so.19(minSdkVersion:(no version))
 ndk_crtbegin_so.21(minSdkVersion:(no version))
diff --git a/apex/apex_test.go b/apex/apex_test.go
index bbc4b93..3f56047 100644
--- a/apex/apex_test.go
+++ b/apex/apex_test.go
@@ -354,7 +354,7 @@
 
 // Minimal test
 func TestBasicApex(t *testing.T) {
-	ctx, config := testApex(t, `
+	ctx, _ := testApex(t, `
 		apex_defaults {
 			name: "myapex-defaults",
 			manifest: ":myapex.manifest",
@@ -562,7 +562,7 @@
 
 	// Make sure that Android.mk is created
 	ab := ctx.ModuleForTests("myapex", "android_common_myapex_image").Module().(*apexBundle)
-	data := android.AndroidMkDataForTest(t, config, "", ab)
+	data := android.AndroidMkDataForTest(t, ctx, ab)
 	var builder strings.Builder
 	data.Custom(&builder, ab.BaseModuleName(), "TARGET_", "", data)
 
@@ -2508,7 +2508,7 @@
 }
 
 func TestVendorApex(t *testing.T) {
-	ctx, config := testApex(t, `
+	ctx, _ := testApex(t, `
 		apex {
 			name: "myapex",
 			key: "myapex.key",
@@ -2539,7 +2539,7 @@
 	})
 
 	apexBundle := ctx.ModuleForTests("myapex", "android_common_myapex_image").Module().(*apexBundle)
-	data := android.AndroidMkDataForTest(t, config, "", apexBundle)
+	data := android.AndroidMkDataForTest(t, ctx, apexBundle)
 	name := apexBundle.BaseModuleName()
 	prefix := "TARGET_"
 	var builder strings.Builder
@@ -2643,7 +2643,7 @@
 }
 
 func TestAndroidMk_UseVendorRequired(t *testing.T) {
-	ctx, config := testApex(t, `
+	ctx, _ := testApex(t, `
 		apex {
 			name: "myapex",
 			key: "myapex.key",
@@ -2667,7 +2667,7 @@
 	})
 
 	apexBundle := ctx.ModuleForTests("myapex", "android_common_myapex_image").Module().(*apexBundle)
-	data := android.AndroidMkDataForTest(t, config, "", apexBundle)
+	data := android.AndroidMkDataForTest(t, ctx, apexBundle)
 	name := apexBundle.BaseModuleName()
 	prefix := "TARGET_"
 	var builder strings.Builder
@@ -2677,7 +2677,7 @@
 }
 
 func TestAndroidMk_VendorApexRequired(t *testing.T) {
-	ctx, config := testApex(t, `
+	ctx, _ := testApex(t, `
 		apex {
 			name: "myapex",
 			key: "myapex.key",
@@ -2698,7 +2698,7 @@
 	`)
 
 	apexBundle := ctx.ModuleForTests("myapex", "android_common_myapex_image").Module().(*apexBundle)
-	data := android.AndroidMkDataForTest(t, config, "", apexBundle)
+	data := android.AndroidMkDataForTest(t, ctx, apexBundle)
 	name := apexBundle.BaseModuleName()
 	prefix := "TARGET_"
 	var builder strings.Builder
@@ -2708,7 +2708,7 @@
 }
 
 func TestAndroidMkWritesCommonProperties(t *testing.T) {
-	ctx, config := testApex(t, `
+	ctx, _ := testApex(t, `
 		apex {
 			name: "myapex",
 			key: "myapex.key",
@@ -2726,7 +2726,7 @@
 	`)
 
 	apexBundle := ctx.ModuleForTests("myapex", "android_common_myapex_image").Module().(*apexBundle)
-	data := android.AndroidMkDataForTest(t, config, "", apexBundle)
+	data := android.AndroidMkDataForTest(t, ctx, apexBundle)
 	name := apexBundle.BaseModuleName()
 	prefix := "TARGET_"
 	var builder strings.Builder
@@ -3799,7 +3799,7 @@
 }
 
 func TestApexName(t *testing.T) {
-	ctx, config := testApex(t, `
+	ctx, _ := testApex(t, `
 		apex {
 			name: "myapex",
 			key: "myapex.key",
@@ -3832,7 +3832,7 @@
 	ensureContains(t, apexRule.Args["opt_flags"], "--do_not_check_keyname")
 
 	apexBundle := ctx.ModuleForTests("myapex", "android_common_myapex_image").Module().(*apexBundle)
-	data := android.AndroidMkDataForTest(t, config, "", apexBundle)
+	data := android.AndroidMkDataForTest(t, ctx, apexBundle)
 	name := apexBundle.BaseModuleName()
 	prefix := "TARGET_"
 	var builder strings.Builder
@@ -4336,7 +4336,7 @@
 }
 
 func TestPrebuiltOverrides(t *testing.T) {
-	ctx, config := testApex(t, `
+	ctx, _ := testApex(t, `
 		prebuilt_apex {
 			name: "myapex.prebuilt",
 			src: "myapex-arm.apex",
@@ -4349,7 +4349,7 @@
 	p := ctx.ModuleForTests("myapex.prebuilt", "android_common").Module().(*Prebuilt)
 
 	expected := []string{"myapex"}
-	actual := android.AndroidMkEntriesForTest(t, config, "", p)[0].EntryMap["LOCAL_OVERRIDES_MODULES"]
+	actual := android.AndroidMkEntriesForTest(t, ctx, p)[0].EntryMap["LOCAL_OVERRIDES_MODULES"]
 	if !reflect.DeepEqual(actual, expected) {
 		t.Errorf("Incorrect LOCAL_OVERRIDES_MODULES value '%s', expected '%s'", actual, expected)
 	}
@@ -4708,7 +4708,7 @@
 }
 
 func TestApexWithTests(t *testing.T) {
-	ctx, config := testApex(t, `
+	ctx, _ := testApex(t, `
 		apex_test {
 			name: "myapex",
 			key: "myapex.key",
@@ -4796,7 +4796,7 @@
 
 	// Ensure the module is correctly translated.
 	bundle := ctx.ModuleForTests("myapex", "android_common_myapex_image").Module().(*apexBundle)
-	data := android.AndroidMkDataForTest(t, config, "", bundle)
+	data := android.AndroidMkDataForTest(t, ctx, bundle)
 	name := bundle.BaseModuleName()
 	prefix := "TARGET_"
 	var builder strings.Builder
@@ -4811,7 +4811,7 @@
 	ensureContains(t, androidMk, "LOCAL_MODULE := myapex\n")
 
 	flatBundle := ctx.ModuleForTests("myapex", "android_common_myapex_flattened").Module().(*apexBundle)
-	data = android.AndroidMkDataForTest(t, config, "", flatBundle)
+	data = android.AndroidMkDataForTest(t, ctx, flatBundle)
 	data.Custom(&builder, name, prefix, "", data)
 	flatAndroidMk := builder.String()
 	ensureContainsOnce(t, flatAndroidMk, "LOCAL_TEST_DATA := :baz :bar/baz\n")
@@ -4819,7 +4819,7 @@
 }
 
 func TestInstallExtraFlattenedApexes(t *testing.T) {
-	ctx, config := testApex(t, `
+	ctx, _ := testApex(t, `
 		apex {
 			name: "myapex",
 			key: "myapex.key",
@@ -4834,7 +4834,7 @@
 	})
 	ab := ctx.ModuleForTests("myapex", "android_common_myapex_image").Module().(*apexBundle)
 	ensureListContains(t, ab.requiredDeps, "myapex.flattened")
-	mk := android.AndroidMkDataForTest(t, config, "", ab)
+	mk := android.AndroidMkDataForTest(t, ctx, ab)
 	var builder strings.Builder
 	mk.Custom(&builder, ab.Name(), "TARGET_", "", mk)
 	androidMk := builder.String()
@@ -5382,7 +5382,7 @@
 }
 
 func TestOverrideApex(t *testing.T) {
-	ctx, config := testApex(t, `
+	ctx, _ := testApex(t, `
 		apex {
 			name: "myapex",
 			key: "myapex.key",
@@ -5450,7 +5450,7 @@
 	optFlags := apexRule.Args["opt_flags"]
 	ensureContains(t, optFlags, "--override_apk_package_name test.overridden.package")
 
-	data := android.AndroidMkDataForTest(t, config, "", apexBundle)
+	data := android.AndroidMkDataForTest(t, ctx, apexBundle)
 	var builder strings.Builder
 	data.Custom(&builder, name, "TARGET_", "", data)
 	androidMk := builder.String()
@@ -5838,7 +5838,7 @@
 }
 
 func TestCarryRequiredModuleNames(t *testing.T) {
-	ctx, config := testApex(t, `
+	ctx, _ := testApex(t, `
 		apex {
 			name: "myapex",
 			key: "myapex.key",
@@ -5864,7 +5864,7 @@
 	`)
 
 	apexBundle := ctx.ModuleForTests("myapex", "android_common_myapex_image").Module().(*apexBundle)
-	data := android.AndroidMkDataForTest(t, config, "", apexBundle)
+	data := android.AndroidMkDataForTest(t, ctx, apexBundle)
 	name := apexBundle.BaseModuleName()
 	prefix := "TARGET_"
 	var builder strings.Builder
@@ -6005,7 +6005,7 @@
 }
 
 func TestSymlinksFromApexToSystemRequiredModuleNames(t *testing.T) {
-	ctx, config := testApex(t, `
+	ctx, _ := testApex(t, `
 		apex {
 			name: "myapex",
 			key: "myapex.key",
@@ -6043,7 +6043,7 @@
 	`)
 
 	apexBundle := ctx.ModuleForTests("myapex", "android_common_myapex_image").Module().(*apexBundle)
-	data := android.AndroidMkDataForTest(t, config, "", apexBundle)
+	data := android.AndroidMkDataForTest(t, ctx, apexBundle)
 	var builder strings.Builder
 	data.Custom(&builder, apexBundle.BaseModuleName(), "TARGET_", "", data)
 	androidMk := builder.String()
@@ -6728,7 +6728,7 @@
 }
 
 func TestApexSet(t *testing.T) {
-	ctx, config := testApex(t, `
+	ctx, _ := testApex(t, `
 		apex_set {
 			name: "myapex",
 			set: "myapex.apks",
@@ -6760,7 +6760,7 @@
 
 	a := m.Module().(*ApexSet)
 	expectedOverrides := []string{"foo"}
-	actualOverrides := android.AndroidMkEntriesForTest(t, config, "", a)[0].EntryMap["LOCAL_OVERRIDES_MODULES"]
+	actualOverrides := android.AndroidMkEntriesForTest(t, ctx, a)[0].EntryMap["LOCAL_OVERRIDES_MODULES"]
 	if !reflect.DeepEqual(actualOverrides, expectedOverrides) {
 		t.Errorf("Incorrect LOCAL_OVERRIDES_MODULES - expected %q vs actual %q", expectedOverrides, actualOverrides)
 	}
@@ -6933,7 +6933,7 @@
 }
 
 func TestCompressedApex(t *testing.T) {
-	ctx, config := testApex(t, `
+	ctx, _ := testApex(t, `
 		apex {
 			name: "myapex",
 			key: "myapex.key",
@@ -6959,7 +6959,7 @@
 	ensureContains(t, ab.outputFile.String(), "myapex.capex")
 
 	// Verify android.mk rules
-	data := android.AndroidMkDataForTest(t, config, "", ab)
+	data := android.AndroidMkDataForTest(t, ctx, ab)
 	var builder strings.Builder
 	data.Custom(&builder, ab.BaseModuleName(), "TARGET_", "", data)
 	androidMk := builder.String()
@@ -6967,7 +6967,7 @@
 }
 
 func TestPreferredPrebuiltSharedLibDep(t *testing.T) {
-	ctx, config := testApex(t, `
+	ctx, _ := testApex(t, `
 		apex {
 			name: "myapex",
 			key: "myapex.key",
@@ -7007,7 +7007,7 @@
 	`)
 
 	ab := ctx.ModuleForTests("myapex", "android_common_myapex_image").Module().(*apexBundle)
-	data := android.AndroidMkDataForTest(t, config, "", ab)
+	data := android.AndroidMkDataForTest(t, ctx, ab)
 	var builder strings.Builder
 	data.Custom(&builder, ab.BaseModuleName(), "TARGET_", "", data)
 	androidMk := builder.String()
@@ -7159,7 +7159,7 @@
 		t.Run(test.name, func(t *testing.T) {
 			for _, otherApexEnabled := range test.otherApexEnabled {
 				t.Run("otherapex_enabled_"+otherApexEnabled, func(t *testing.T) {
-					ctx, config := testApex(t, fmt.Sprintf(bpBase, otherApexEnabled)+test.stublibBp)
+					ctx, _ := testApex(t, fmt.Sprintf(bpBase, otherApexEnabled)+test.stublibBp)
 
 					type modAndMkEntries struct {
 						mod       *cc.Module
@@ -7177,7 +7177,7 @@
 							if !mod.Enabled() || mod.IsHideFromMake() {
 								continue
 							}
-							for _, ent := range android.AndroidMkEntriesForTest(t, config, "", mod) {
+							for _, ent := range android.AndroidMkEntriesForTest(t, ctx, mod) {
 								if ent.Disabled {
 									continue
 								}
diff --git a/apex/prebuilt.go b/apex/prebuilt.go
index 041afb3..ec7f253 100644
--- a/apex/prebuilt.go
+++ b/apex/prebuilt.go
@@ -414,7 +414,7 @@
 		OutputFile: android.OptionalPathForPath(p.inputApex),
 		Include:    "$(BUILD_PREBUILT)",
 		ExtraEntries: []android.AndroidMkExtraEntriesFunc{
-			func(entries *android.AndroidMkEntries) {
+			func(ctx android.AndroidMkExtraEntriesContext, entries *android.AndroidMkEntries) {
 				entries.SetString("LOCAL_MODULE_PATH", p.installDir.ToMakePath().String())
 				entries.SetString("LOCAL_MODULE_STEM", p.installFilename)
 				entries.SetBoolIfTrue("LOCAL_UNINSTALLABLE_MODULE", !p.installable())
@@ -609,7 +609,7 @@
 		Include:       "$(BUILD_PREBUILT)",
 		Host_required: a.hostRequired,
 		ExtraEntries: []android.AndroidMkExtraEntriesFunc{
-			func(entries *android.AndroidMkEntries) {
+			func(ctx android.AndroidMkExtraEntriesContext, entries *android.AndroidMkEntries) {
 				entries.SetString("LOCAL_MODULE_PATH", a.installDir.ToMakePath().String())
 				entries.SetString("LOCAL_MODULE_STEM", a.installFilename)
 				entries.SetBoolIfTrue("LOCAL_UNINSTALLABLE_MODULE", !a.installable())
diff --git a/bp2build/build_conversion.go b/bp2build/build_conversion.go
index fcad686..41ad409 100644
--- a/bp2build/build_conversion.go
+++ b/bp2build/build_conversion.go
@@ -208,12 +208,6 @@
 	return buildFileToTargets, metrics
 }
 
-// Helper method for tests to easily access the targets in a dir.
-func GenerateBazelTargetsForDir(codegenCtx CodegenContext, dir string) BazelTargets {
-	bazelTargetsMap, _ := GenerateBazelTargets(codegenCtx)
-	return bazelTargetsMap[dir]
-}
-
 // Helper method to trim quotes around strings.
 func trimQuotes(s string) string {
 	if s == "" {
diff --git a/bp2build/build_conversion_test.go b/bp2build/build_conversion_test.go
index fce6433..aa4fc1d 100644
--- a/bp2build/build_conversion_test.go
+++ b/bp2build/build_conversion_test.go
@@ -204,7 +204,7 @@
 		android.FailIfErrored(t, errs)
 
 		codegenCtx := NewCodegenContext(config, *ctx.Context, QueryView)
-		bazelTargets := GenerateBazelTargetsForDir(codegenCtx, dir)
+		bazelTargets := generateBazelTargetsForDir(codegenCtx, dir)
 		if actualCount, expectedCount := len(bazelTargets), 1; actualCount != expectedCount {
 			t.Fatalf("Expected %d bazel target, got %d", expectedCount, actualCount)
 		}
@@ -262,7 +262,7 @@
 		}
 
 		codegenCtx := NewCodegenContext(config, *ctx.Context, Bp2Build)
-		bazelTargets := GenerateBazelTargetsForDir(codegenCtx, dir)
+		bazelTargets := generateBazelTargetsForDir(codegenCtx, dir)
 
 		if actualCount, expectedCount := len(bazelTargets), 1; actualCount != expectedCount {
 			t.Errorf("Expected %d bazel target, got %d", expectedCount, actualCount)
@@ -421,7 +421,7 @@
 		android.FailIfErrored(t, errs)
 
 		codegenCtx := NewCodegenContext(config, *ctx.Context, Bp2Build)
-		bazelTargets := GenerateBazelTargetsForDir(codegenCtx, dir)
+		bazelTargets := generateBazelTargetsForDir(codegenCtx, dir)
 		if actualCount := len(bazelTargets); actualCount != testCase.expectedBazelTargetCount {
 			t.Fatalf("Expected %d bazel target, got %d", testCase.expectedBazelTargetCount, actualCount)
 		}
@@ -912,7 +912,7 @@
 		}
 
 		codegenCtx := NewCodegenContext(config, *ctx.Context, Bp2Build)
-		bazelTargets := GenerateBazelTargetsForDir(codegenCtx, checkDir)
+		bazelTargets := generateBazelTargetsForDir(codegenCtx, checkDir)
 		if actualCount, expectedCount := len(bazelTargets), len(testCase.expectedBazelTargets); actualCount != expectedCount {
 			t.Errorf("%s: Expected %d bazel target, got %d", testCase.description, expectedCount, actualCount)
 		} else {
@@ -1127,7 +1127,7 @@
 		android.FailIfErrored(t, errs)
 
 		codegenCtx := NewCodegenContext(config, *ctx.Context, Bp2Build)
-		bazelTargets := GenerateBazelTargetsForDir(codegenCtx, dir)
+		bazelTargets := generateBazelTargetsForDir(codegenCtx, dir)
 		if actualCount := len(bazelTargets); actualCount != 1 {
 			t.Fatalf("%s: Expected 1 bazel target, got %d", testCase.description, actualCount)
 		}
@@ -1215,7 +1215,7 @@
 		android.FailIfErrored(t, errs)
 
 		codegenCtx := NewCodegenContext(config, *ctx.Context, Bp2Build)
-		bazelTargets := GenerateBazelTargetsForDir(codegenCtx, dir)
+		bazelTargets := generateBazelTargetsForDir(codegenCtx, dir)
 		if actualCount := len(bazelTargets); actualCount != testCase.expectedCount {
 			t.Fatalf("%s: Expected %d bazel target, got %d", testCase.description, testCase.expectedCount, actualCount)
 		}
diff --git a/bp2build/cc_library_headers_conversion_test.go b/bp2build/cc_library_headers_conversion_test.go
index b2d8941..5bf5c80 100644
--- a/bp2build/cc_library_headers_conversion_test.go
+++ b/bp2build/cc_library_headers_conversion_test.go
@@ -203,7 +203,7 @@
 			checkDir = testCase.dir
 		}
 		codegenCtx := NewCodegenContext(config, *ctx.Context, Bp2Build)
-		bazelTargets := GenerateBazelTargetsForDir(codegenCtx, checkDir)
+		bazelTargets := generateBazelTargetsForDir(codegenCtx, checkDir)
 		if actualCount, expectedCount := len(bazelTargets), len(testCase.expectedBazelTargets); actualCount != expectedCount {
 			t.Errorf("%s: Expected %d bazel target, got %d", testCase.description, expectedCount, actualCount)
 		} else {
diff --git a/bp2build/cc_object_conversion_test.go b/bp2build/cc_object_conversion_test.go
index 816edb8..f094102 100644
--- a/bp2build/cc_object_conversion_test.go
+++ b/bp2build/cc_object_conversion_test.go
@@ -167,7 +167,7 @@
 		}
 
 		codegenCtx := NewCodegenContext(config, *ctx.Context, Bp2Build)
-		bazelTargets := GenerateBazelTargetsForDir(codegenCtx, dir)
+		bazelTargets := generateBazelTargetsForDir(codegenCtx, dir)
 		if actualCount, expectedCount := len(bazelTargets), len(testCase.expectedBazelTargets); actualCount != expectedCount {
 			fmt.Println(bazelTargets)
 			t.Errorf("%s: Expected %d bazel target, got %d", testCase.description, expectedCount, actualCount)
diff --git a/bp2build/sh_conversion_test.go b/bp2build/sh_conversion_test.go
index 3e15dfc..2aa373c 100644
--- a/bp2build/sh_conversion_test.go
+++ b/bp2build/sh_conversion_test.go
@@ -116,7 +116,7 @@
 			checkDir = testCase.dir
 		}
 		codegenCtx := NewCodegenContext(config, *ctx.Context, Bp2Build)
-		bazelTargets := GenerateBazelTargetsForDir(codegenCtx, checkDir)
+		bazelTargets := generateBazelTargetsForDir(codegenCtx, checkDir)
 		if actualCount, expectedCount := len(bazelTargets), len(testCase.expectedBazelTargets); actualCount != expectedCount {
 			t.Errorf("%s: Expected %d bazel target, got %d", testCase.description, expectedCount, actualCount)
 		} else {
diff --git a/bp2build/testing.go b/bp2build/testing.go
index 2e59999..9774915 100644
--- a/bp2build/testing.go
+++ b/bp2build/testing.go
@@ -175,3 +175,9 @@
 		ctx.CreateBazelTargetModule(customBazelModuleFactory, myProtoLibraryProps, attrs)
 	}
 }
+
+// Helper method for tests to easily access the targets in a dir.
+func generateBazelTargetsForDir(codegenCtx CodegenContext, dir string) BazelTargets {
+	buildFileToTargets, _ := GenerateBazelTargets(codegenCtx)
+	return buildFileToTargets[dir]
+}
diff --git a/build_kzip.bash b/build_kzip.bash
index 0018ea9..9564723 100755
--- a/build_kzip.bash
+++ b/build_kzip.bash
@@ -7,14 +7,16 @@
 #   BUILD_NUMBER          build number, used to generate unique ID (will use UUID if not set)
 #   DIST_DIR              where the resulting all.kzip will be placed
 #   KYTHE_KZIP_ENCODING   proto or json (proto is default)
+#   KYTHE_JAVA_SOURCE_BATCH_SIZE maximum number of the Java source files in a compilation unit
 #   OUT_DIR               output directory (out if not specified})
 #   TARGET_BUILD_VARIANT  variant, e.g., `userdebug`
 #   TARGET_PRODUCT        target device name, e.g., 'aosp_blueline'
 #   XREF_CORPUS           source code repository URI, e.g., 'android.googlesource.com/platform/superproject'
 
 : ${BUILD_NUMBER:=$(uuidgen)}
+: ${KYTHE_JAVA_SOURCE_BATCH_SIZE:=500}
 : ${KYTHE_KZIP_ENCODING:=proto}
-export KYTHE_KZIP_ENCODING
+export KYTHE_JAVA_SOURCE_BATCH_SIZE KYTHE_KZIP_ENCODING
 
 # The extraction might fail for some source files, so run with -k and then check that
 # sufficiently many files were generated.
diff --git a/cc/androidmk.go b/cc/androidmk.go
index 8652c10..536efa4 100644
--- a/cc/androidmk.go
+++ b/cc/androidmk.go
@@ -83,7 +83,7 @@
 		Include:  "$(BUILD_SYSTEM)/soong_cc_prebuilt.mk",
 
 		ExtraEntries: []android.AndroidMkExtraEntriesFunc{
-			func(entries *android.AndroidMkEntries) {
+			func(ctx android.AndroidMkExtraEntriesContext, entries *android.AndroidMkEntries) {
 				if len(c.Properties.Logtags) > 0 {
 					entries.AddStrings("LOCAL_LOGTAGS_FILES", c.Properties.Logtags...)
 				}
@@ -162,16 +162,17 @@
 
 func androidMkWriteExtraTestConfigs(extraTestConfigs android.Paths, entries *android.AndroidMkEntries) {
 	if len(extraTestConfigs) > 0 {
-		entries.ExtraEntries = append(entries.ExtraEntries, func(entries *android.AndroidMkEntries) {
-			entries.AddStrings("LOCAL_EXTRA_FULL_TEST_CONFIGS", extraTestConfigs.Strings()...)
-		})
+		entries.ExtraEntries = append(entries.ExtraEntries,
+			func(ctx android.AndroidMkExtraEntriesContext, entries *android.AndroidMkEntries) {
+				entries.AddStrings("LOCAL_EXTRA_FULL_TEST_CONFIGS", extraTestConfigs.Strings()...)
+			})
 	}
 }
 
 func AndroidMkWriteTestData(data []android.DataPath, entries *android.AndroidMkEntries) {
 	testFiles := android.AndroidMkDataPaths(data)
 	if len(testFiles) > 0 {
-		entries.ExtraEntries = append(entries.ExtraEntries, func(entries *android.AndroidMkEntries) {
+		entries.ExtraEntries = append(entries.ExtraEntries, func(ctx android.AndroidMkExtraEntriesContext, entries *android.AndroidMkEntries) {
 			entries.AddStrings("LOCAL_TEST_DATA", testFiles...)
 		})
 	}
@@ -224,7 +225,7 @@
 		entries.Class = "STATIC_LIBRARIES"
 	} else if library.shared() {
 		entries.Class = "SHARED_LIBRARIES"
-		entries.ExtraEntries = append(entries.ExtraEntries, func(entries *android.AndroidMkEntries) {
+		entries.ExtraEntries = append(entries.ExtraEntries, func(_ android.AndroidMkExtraEntriesContext, entries *android.AndroidMkEntries) {
 			entries.SetString("LOCAL_SOONG_TOC", library.toc().String())
 			if !library.buildStubs() {
 				entries.SetString("LOCAL_SOONG_UNSTRIPPED_BINARY", library.unstrippedOutputFile.String())
@@ -244,7 +245,7 @@
 		entries.DistFiles = android.MakeDefaultDistFiles(library.distFile)
 	}
 
-	entries.ExtraEntries = append(entries.ExtraEntries, func(entries *android.AndroidMkEntries) {
+	entries.ExtraEntries = append(entries.ExtraEntries, func(ctx android.AndroidMkExtraEntriesContext, entries *android.AndroidMkEntries) {
 		library.androidMkWriteExportedFlags(entries)
 		library.androidMkEntriesWriteAdditionalDependenciesForSourceAbiDiff(entries)
 
@@ -272,7 +273,7 @@
 		if library.buildStubs() && library.stubsVersion() != "" {
 			entries.SubName = "." + library.stubsVersion()
 		}
-		entries.ExtraEntries = append(entries.ExtraEntries, func(entries *android.AndroidMkEntries) {
+		entries.ExtraEntries = append(entries.ExtraEntries, func(ctx android.AndroidMkExtraEntriesContext, entries *android.AndroidMkEntries) {
 			// library.makeUninstallable() depends on this to bypass HideFromMake() for
 			// static libraries.
 			entries.SetBool("LOCAL_UNINSTALLABLE_MODULE", true)
@@ -315,7 +316,7 @@
 
 	entries.Class = "EXECUTABLES"
 	entries.DistFiles = binary.distFiles
-	entries.ExtraEntries = append(entries.ExtraEntries, func(entries *android.AndroidMkEntries) {
+	entries.ExtraEntries = append(entries.ExtraEntries, func(_ android.AndroidMkExtraEntriesContext, entries *android.AndroidMkEntries) {
 		entries.SetString("LOCAL_SOONG_UNSTRIPPED_BINARY", binary.unstrippedOutputFile.String())
 		if len(binary.symlinks) > 0 {
 			entries.AddStrings("LOCAL_MODULE_SYMLINKS", binary.symlinks...)
@@ -337,7 +338,7 @@
 func (benchmark *benchmarkDecorator) AndroidMkEntries(ctx AndroidMkContext, entries *android.AndroidMkEntries) {
 	ctx.subAndroidMk(entries, benchmark.binaryDecorator)
 	entries.Class = "NATIVE_TESTS"
-	entries.ExtraEntries = append(entries.ExtraEntries, func(entries *android.AndroidMkEntries) {
+	entries.ExtraEntries = append(entries.ExtraEntries, func(ctx android.AndroidMkExtraEntriesContext, entries *android.AndroidMkEntries) {
 		if len(benchmark.Properties.Test_suites) > 0 {
 			entries.AddCompatibilityTestSuites(benchmark.Properties.Test_suites...)
 		}
@@ -362,7 +363,7 @@
 	if Bool(test.Properties.Test_per_src) {
 		entries.SubName = "_" + String(test.binaryDecorator.Properties.Stem)
 	}
-	entries.ExtraEntries = append(entries.ExtraEntries, func(entries *android.AndroidMkEntries) {
+	entries.ExtraEntries = append(entries.ExtraEntries, func(ctx android.AndroidMkExtraEntriesContext, entries *android.AndroidMkEntries) {
 		if len(test.Properties.Test_suites) > 0 {
 			entries.AddCompatibilityTestSuites(test.Properties.Test_suites...)
 		}
@@ -406,7 +407,7 @@
 			filepath.Dir(fuzz.config.String())+":config.json")
 	}
 
-	entries.ExtraEntries = append(entries.ExtraEntries, func(entries *android.AndroidMkEntries) {
+	entries.ExtraEntries = append(entries.ExtraEntries, func(ctx android.AndroidMkExtraEntriesContext, entries *android.AndroidMkEntries) {
 		entries.SetBool("LOCAL_IS_FUZZ_TARGET", true)
 		if len(fuzzFiles) > 0 {
 			entries.AddStrings("LOCAL_TEST_DATA", fuzzFiles...)
@@ -423,7 +424,7 @@
 
 func (library *toolchainLibraryDecorator) AndroidMkEntries(ctx AndroidMkContext, entries *android.AndroidMkEntries) {
 	entries.Class = "STATIC_LIBRARIES"
-	entries.ExtraEntries = append(entries.ExtraEntries, func(entries *android.AndroidMkEntries) {
+	entries.ExtraEntries = append(entries.ExtraEntries, func(ctx android.AndroidMkExtraEntriesContext, entries *android.AndroidMkEntries) {
 		_, suffix, _ := android.SplitFileExt(entries.OutputFile.Path().Base())
 		entries.SetString("LOCAL_MODULE_SUFFIX", suffix)
 	})
@@ -439,7 +440,7 @@
 		entries.OutputFile = android.OptionalPathForPath(installer.path)
 	}
 
-	entries.ExtraEntries = append(entries.ExtraEntries, func(entries *android.AndroidMkEntries) {
+	entries.ExtraEntries = append(entries.ExtraEntries, func(ctx android.AndroidMkExtraEntriesContext, entries *android.AndroidMkEntries) {
 		path, file := filepath.Split(installer.path.ToMakePath().String())
 		stem, suffix, _ := android.SplitFileExt(file)
 		entries.SetString("LOCAL_MODULE_SUFFIX", suffix)
@@ -457,7 +458,7 @@
 		return
 	}
 
-	entries.ExtraEntries = append(entries.ExtraEntries, func(entries *android.AndroidMkEntries) {
+	entries.ExtraEntries = append(entries.ExtraEntries, func(ctx android.AndroidMkExtraEntriesContext, entries *android.AndroidMkEntries) {
 		path, file := filepath.Split(c.installPath.String())
 		stem, suffix, _ := android.SplitFileExt(file)
 		entries.SetString("LOCAL_MODULE_SUFFIX", suffix)
@@ -481,7 +482,7 @@
 
 	entries.SubName = c.androidMkSuffix
 
-	entries.ExtraEntries = append(entries.ExtraEntries, func(entries *android.AndroidMkEntries) {
+	entries.ExtraEntries = append(entries.ExtraEntries, func(ctx android.AndroidMkExtraEntriesContext, entries *android.AndroidMkEntries) {
 		c.libraryDecorator.androidMkWriteExportedFlags(entries)
 
 		// Specifying stem is to pass check_elf_files when vendor modules link against vndk prebuilt.
@@ -521,7 +522,7 @@
 
 	entries.SubName += c.baseProperties.Androidmk_suffix
 
-	entries.ExtraEntries = append(entries.ExtraEntries, func(entries *android.AndroidMkEntries) {
+	entries.ExtraEntries = append(entries.ExtraEntries, func(ctx android.AndroidMkExtraEntriesContext, entries *android.AndroidMkEntries) {
 		c.libraryDecorator.androidMkWriteExportedFlags(entries)
 
 		if c.shared() || c.static() {
@@ -548,7 +549,7 @@
 	entries.Class = "EXECUTABLES"
 	entries.SubName = c.baseProperties.Androidmk_suffix
 
-	entries.ExtraEntries = append(entries.ExtraEntries, func(entries *android.AndroidMkEntries) {
+	entries.ExtraEntries = append(entries.ExtraEntries, func(ctx android.AndroidMkExtraEntriesContext, entries *android.AndroidMkEntries) {
 		entries.AddStrings("LOCAL_MODULE_SYMLINKS", c.Properties.Symlinks...)
 	})
 }
@@ -575,7 +576,7 @@
 	entries.Class = "SHARED_LIBRARIES"
 	entries.SubName = vendorPublicLibrarySuffix
 
-	entries.ExtraEntries = append(entries.ExtraEntries, func(entries *android.AndroidMkEntries) {
+	entries.ExtraEntries = append(entries.ExtraEntries, func(ctx android.AndroidMkExtraEntriesContext, entries *android.AndroidMkEntries) {
 		c.libraryDecorator.androidMkWriteExportedFlags(entries)
 		_, _, ext := android.SplitFileExt(entries.OutputFile.Path().Base())
 
@@ -586,7 +587,7 @@
 }
 
 func (p *prebuiltLinker) AndroidMkEntries(ctx AndroidMkContext, entries *android.AndroidMkEntries) {
-	entries.ExtraEntries = append(entries.ExtraEntries, func(entries *android.AndroidMkEntries) {
+	entries.ExtraEntries = append(entries.ExtraEntries, func(ctx android.AndroidMkExtraEntriesContext, entries *android.AndroidMkEntries) {
 		if p.properties.Check_elf_files != nil {
 			entries.SetBool("LOCAL_CHECK_ELF_FILES", *p.properties.Check_elf_files)
 		} else {
@@ -616,7 +617,7 @@
 func androidMkWriteAllowUndefinedSymbols(linker *baseLinker, entries *android.AndroidMkEntries) {
 	allow := linker.Properties.Allow_undefined_symbols
 	if allow != nil {
-		entries.ExtraEntries = append(entries.ExtraEntries, func(entries *android.AndroidMkEntries) {
+		entries.ExtraEntries = append(entries.ExtraEntries, func(ctx android.AndroidMkExtraEntriesContext, entries *android.AndroidMkEntries) {
 			entries.SetBool("LOCAL_ALLOW_UNDEFINED_SYMBOLS", *allow)
 		})
 	}
diff --git a/cc/cc_test.go b/cc/cc_test.go
index 942e397..e640f12 100644
--- a/cc/cc_test.go
+++ b/cc/cc_test.go
@@ -567,7 +567,7 @@
 	ctx := testCcWithConfig(t, config)
 
 	module := ctx.ModuleForTests("llndk.libraries.txt", "")
-	entries := android.AndroidMkEntriesForTest(t, config, "", module.Module())[0]
+	entries := android.AndroidMkEntriesForTest(t, ctx, module.Module())[0]
 	assertArrayString(t, entries.EntryMap["LOCAL_MODULE_STEM"], []string{"llndk.libraries.VER.txt"})
 }
 
@@ -713,7 +713,7 @@
 	if !strings.HasSuffix(outputPath, "/main_test") {
 		t.Errorf("expected test output file to be 'main_test', but was '%s'", outputPath)
 	}
-	entries := android.AndroidMkEntriesForTest(t, config, "", module)[0]
+	entries := android.AndroidMkEntriesForTest(t, ctx, module)[0]
 	if !strings.HasSuffix(entries.EntryMap["LOCAL_TEST_DATA"][0], ":test_lib.so:foo/bar/baz") {
 		t.Errorf("expected LOCAL_TEST_DATA to end with `:test_lib.so:foo/bar/baz`,"+
 			" but was '%s'", entries.EntryMap["LOCAL_TEST_DATA"][0])
@@ -3105,7 +3105,7 @@
 	if !strings.HasSuffix(outputPath, "/main_test") {
 		t.Errorf("expected test output file to be 'main_test', but was '%s'", outputPath)
 	}
-	entries := android.AndroidMkEntriesForTest(t, config, "", module)[0]
+	entries := android.AndroidMkEntriesForTest(t, ctx, module)[0]
 	if !strings.HasSuffix(entries.EntryMap["LOCAL_TEST_DATA"][0], ":test_lib.so:foo/bar/baz") {
 		t.Errorf("expected LOCAL_TEST_DATA to end with `:test_lib.so:foo/bar/baz`,"+
 			" but was '%s'", entries.EntryMap["LOCAL_TEST_DATA"][0])
@@ -4092,7 +4092,6 @@
 		)
 	})
 
-	// TODO: fix this test as it exports all generated headers.
 	t.Run("ensure only aidl headers are exported", func(t *testing.T) {
 		ctx := testCc(t, genRuleModules+`
 		cc_library_shared {
@@ -4117,18 +4116,15 @@
 				.intermediates/libfoo/android_arm64_armv8-a_shared/gen/aidl/b.h
 				.intermediates/libfoo/android_arm64_armv8-a_shared/gen/aidl/Bnb.h
 				.intermediates/libfoo/android_arm64_armv8-a_shared/gen/aidl/Bpb.h
-				.intermediates/libfoo/android_arm64_armv8-a_shared/gen/proto/a.pb.h
 			`),
 			expectedOrderOnlyDeps(`
 				.intermediates/libfoo/android_arm64_armv8-a_shared/gen/aidl/b.h
 				.intermediates/libfoo/android_arm64_armv8-a_shared/gen/aidl/Bnb.h
 				.intermediates/libfoo/android_arm64_armv8-a_shared/gen/aidl/Bpb.h
-				.intermediates/libfoo/android_arm64_armv8-a_shared/gen/proto/a.pb.h
 			`),
 		)
 	})
 
-	// TODO: fix this test as it exports all generated headers.
 	t.Run("ensure only proto headers are exported", func(t *testing.T) {
 		ctx := testCc(t, genRuleModules+`
 		cc_library_shared {
@@ -4150,22 +4146,15 @@
 			`),
 			expectedSystemIncludeDirs(``),
 			expectedGeneratedHeaders(`
-				.intermediates/libfoo/android_arm64_armv8-a_shared/gen/aidl/b.h
-				.intermediates/libfoo/android_arm64_armv8-a_shared/gen/aidl/Bnb.h
-				.intermediates/libfoo/android_arm64_armv8-a_shared/gen/aidl/Bpb.h
 				.intermediates/libfoo/android_arm64_armv8-a_shared/gen/proto/a.pb.h
 			`),
 			expectedOrderOnlyDeps(`
-				.intermediates/libfoo/android_arm64_armv8-a_shared/gen/aidl/b.h
-				.intermediates/libfoo/android_arm64_armv8-a_shared/gen/aidl/Bnb.h
-				.intermediates/libfoo/android_arm64_armv8-a_shared/gen/aidl/Bpb.h
 				.intermediates/libfoo/android_arm64_armv8-a_shared/gen/proto/a.pb.h
 			`),
 		)
 	})
 
-	// TODO: fix this test as it exports all generated headers including public and non-public.
-	t.Run("ensure only non-public sysprop headers are exported", func(t *testing.T) {
+	t.Run("ensure only sysprop headers are exported", func(t *testing.T) {
 		ctx := testCc(t, genRuleModules+`
 		cc_library_shared {
 			name: "libfoo",
@@ -4185,19 +4174,10 @@
 			expectedSystemIncludeDirs(``),
 			expectedGeneratedHeaders(`
 				.intermediates/libfoo/android_arm64_armv8-a_shared/gen/sysprop/include/a.sysprop.h
-				.intermediates/libfoo/android_arm64_armv8-a_shared/gen/sysprop/public/include/a.sysprop.h
-				.intermediates/libfoo/android_arm64_armv8-a_shared/gen/aidl/b.h
-				.intermediates/libfoo/android_arm64_armv8-a_shared/gen/aidl/Bnb.h
-				.intermediates/libfoo/android_arm64_armv8-a_shared/gen/aidl/Bpb.h
-				.intermediates/libfoo/android_arm64_armv8-a_shared/gen/proto/a.pb.h
 			`),
 			expectedOrderOnlyDeps(`
 				.intermediates/libfoo/android_arm64_armv8-a_shared/gen/sysprop/include/a.sysprop.h
 				.intermediates/libfoo/android_arm64_armv8-a_shared/gen/sysprop/public/include/a.sysprop.h
-				.intermediates/libfoo/android_arm64_armv8-a_shared/gen/aidl/b.h
-				.intermediates/libfoo/android_arm64_armv8-a_shared/gen/aidl/Bnb.h
-				.intermediates/libfoo/android_arm64_armv8-a_shared/gen/aidl/Bpb.h
-				.intermediates/libfoo/android_arm64_armv8-a_shared/gen/proto/a.pb.h
 			`),
 		)
 	})
diff --git a/cc/compiler.go b/cc/compiler.go
index 5f30d3d..e96295c 100644
--- a/cc/compiler.go
+++ b/cc/compiler.go
@@ -230,6 +230,8 @@
 	// other modules and filegroups. May include source files that have not yet been translated to
 	// C/C++ (.aidl, .proto, etc.)
 	srcsBeforeGen android.Paths
+
+	generatedSourceInfo
 }
 
 var _ compiler = (*baseCompiler)(nil)
@@ -634,10 +636,11 @@
 
 	srcs := append(android.Paths(nil), compiler.srcsBeforeGen...)
 
-	srcs, genDeps := genSources(ctx, srcs, buildFlags)
+	srcs, genDeps, info := genSources(ctx, srcs, buildFlags)
 	pathDeps = append(pathDeps, genDeps...)
 
 	compiler.pathDeps = pathDeps
+	compiler.generatedSourceInfo = info
 	compiler.cFlagsDeps = flags.CFlagsDeps
 
 	// Save src, buildFlags and context
diff --git a/cc/gen.go b/cc/gen.go
index 75b9e49..83c019c 100644
--- a/cc/gen.go
+++ b/cc/gen.go
@@ -220,8 +220,35 @@
 	return rcFile, headerFile
 }
 
+// Used to communicate information from the genSources method back to the library code that uses
+// it.
+type generatedSourceInfo struct {
+	// The headers created from .proto files
+	protoHeaders android.Paths
+
+	// The files that can be used as order only dependencies in order to ensure that the proto header
+	// files are up to date.
+	protoOrderOnlyDeps android.Paths
+
+	// The headers created from .aidl files
+	aidlHeaders android.Paths
+
+	// The files that can be used as order only dependencies in order to ensure that the aidl header
+	// files are up to date.
+	aidlOrderOnlyDeps android.Paths
+
+	// The headers created from .sysprop files
+	syspropHeaders android.Paths
+
+	// The files that can be used as order only dependencies in order to ensure that the sysprop
+	// header files are up to date.
+	syspropOrderOnlyDeps android.Paths
+}
+
 func genSources(ctx android.ModuleContext, srcFiles android.Paths,
-	buildFlags builderFlags) (android.Paths, android.Paths) {
+	buildFlags builderFlags) (android.Paths, android.Paths, generatedSourceInfo) {
+
+	var info generatedSourceInfo
 
 	var deps android.Paths
 	var rsFiles android.Paths
@@ -258,7 +285,9 @@
 		case ".proto":
 			ccFile, headerFile := genProto(ctx, srcFile, buildFlags)
 			srcFiles[i] = ccFile
-			deps = append(deps, headerFile)
+			info.protoHeaders = append(info.protoHeaders, headerFile)
+			// Use the generated header as an order only dep to ensure that it is up to date when needed.
+			info.protoOrderOnlyDeps = append(info.protoOrderOnlyDeps, headerFile)
 		case ".aidl":
 			if aidlRule == nil {
 				aidlRule = android.NewRuleBuilder(pctx, ctx).Sbox(android.PathForModuleGen(ctx, "aidl"),
@@ -267,7 +296,12 @@
 			cppFile := android.GenPathWithExt(ctx, "aidl", srcFile, "cpp")
 			depFile := android.GenPathWithExt(ctx, "aidl", srcFile, "cpp.d")
 			srcFiles[i] = cppFile
-			deps = append(deps, genAidl(ctx, aidlRule, srcFile, cppFile, depFile, buildFlags.aidlFlags)...)
+			aidlHeaders := genAidl(ctx, aidlRule, srcFile, cppFile, depFile, buildFlags.aidlFlags)
+			info.aidlHeaders = append(info.aidlHeaders, aidlHeaders...)
+			// Use the generated headers as order only deps to ensure that they are up to date when
+			// needed.
+			// TODO: Reduce the size of the ninja file by using one order only dep for the whole rule
+			info.aidlOrderOnlyDeps = append(info.aidlOrderOnlyDeps, aidlHeaders...)
 		case ".rscript", ".fs":
 			cppFile := rsGeneratedCppFile(ctx, srcFile)
 			rsFiles = append(rsFiles, srcFiles[i])
@@ -279,7 +313,10 @@
 		case ".sysprop":
 			cppFile, headerFiles := genSysprop(ctx, srcFile)
 			srcFiles[i] = cppFile
-			deps = append(deps, headerFiles...)
+			info.syspropHeaders = append(info.syspropHeaders, headerFiles...)
+			// Use the generated headers as order only deps to ensure that they are up to date when
+			// needed.
+			info.syspropOrderOnlyDeps = append(info.syspropOrderOnlyDeps, headerFiles...)
 		}
 	}
 
@@ -291,9 +328,13 @@
 		yaccRule_.Build("yacc", "gen yacc")
 	}
 
+	deps = append(deps, info.protoOrderOnlyDeps...)
+	deps = append(deps, info.aidlOrderOnlyDeps...)
+	deps = append(deps, info.syspropOrderOnlyDeps...)
+
 	if len(rsFiles) > 0 {
 		deps = append(deps, rsGenerateCpp(ctx, rsFiles, buildFlags.rsFlags)...)
 	}
 
-	return srcFiles, deps
+	return srcFiles, deps, info
 }
diff --git a/cc/library.go b/cc/library.go
index 65533bc..0e6e107 100644
--- a/cc/library.go
+++ b/cc/library.go
@@ -1313,9 +1313,8 @@
 			dir := android.PathForModuleGen(ctx, "aidl")
 			library.reexportDirs(dir)
 
-			// TODO: restrict to aidl deps
-			library.reexportDeps(library.baseCompiler.pathDeps...)
-			library.addExportedGeneratedHeaders(library.baseCompiler.pathDeps...)
+			library.reexportDeps(library.baseCompiler.aidlOrderOnlyDeps...)
+			library.addExportedGeneratedHeaders(library.baseCompiler.aidlHeaders...)
 		}
 	}
 
@@ -1329,9 +1328,8 @@
 			includes = append(includes, flags.proto.Dir)
 			library.reexportDirs(includes...)
 
-			// TODO: restrict to proto deps
-			library.reexportDeps(library.baseCompiler.pathDeps...)
-			library.addExportedGeneratedHeaders(library.baseCompiler.pathDeps...)
+			library.reexportDeps(library.baseCompiler.protoOrderOnlyDeps...)
+			library.addExportedGeneratedHeaders(library.baseCompiler.protoHeaders...)
 		}
 	}
 
@@ -1350,10 +1348,16 @@
 			}
 		}
 
+		// Make sure to only export headers which are within the include directory.
+		_, headers := android.FilterPathListPredicate(library.baseCompiler.syspropHeaders, func(path android.Path) bool {
+			_, isRel := android.MaybeRel(ctx, dir.String(), path.String())
+			return isRel
+		})
+
 		// Add sysprop-related directories to the exported directories of this library.
 		library.reexportDirs(dir)
-		library.reexportDeps(library.baseCompiler.pathDeps...)
-		library.addExportedGeneratedHeaders(library.baseCompiler.pathDeps...)
+		library.reexportDeps(library.baseCompiler.syspropOrderOnlyDeps...)
+		library.addExportedGeneratedHeaders(headers...)
 	}
 
 	// Add stub-related flags if this library is a stub library.
diff --git a/cc/vndk.go b/cc/vndk.go
index ff38db5..85028d0 100644
--- a/cc/vndk.go
+++ b/cc/vndk.go
@@ -561,7 +561,7 @@
 		Class:      "ETC",
 		OutputFile: android.OptionalPathForPath(txt.outputFile),
 		ExtraEntries: []android.AndroidMkExtraEntriesFunc{
-			func(entries *android.AndroidMkEntries) {
+			func(ctx android.AndroidMkExtraEntriesContext, entries *android.AndroidMkEntries) {
 				entries.SetString("LOCAL_MODULE_STEM", txt.outputFile.Base())
 			},
 		},
diff --git a/etc/prebuilt_etc.go b/etc/prebuilt_etc.go
index e15324b..b07ad91 100644
--- a/etc/prebuilt_etc.go
+++ b/etc/prebuilt_etc.go
@@ -309,7 +309,7 @@
 		SubName:    nameSuffix,
 		OutputFile: android.OptionalPathForPath(p.outputFilePath),
 		ExtraEntries: []android.AndroidMkExtraEntriesFunc{
-			func(entries *android.AndroidMkEntries) {
+			func(ctx android.AndroidMkExtraEntriesContext, entries *android.AndroidMkEntries) {
 				entries.SetString("LOCAL_MODULE_TAGS", "optional")
 				entries.SetString("LOCAL_MODULE_PATH", p.installDirPath.ToMakePath().String())
 				entries.SetString("LOCAL_INSTALLED_MODULE_STEM", p.outputFilePath.Base())
diff --git a/etc/prebuilt_etc_test.go b/etc/prebuilt_etc_test.go
index 6c4c0b6..585760d 100644
--- a/etc/prebuilt_etc_test.go
+++ b/etc/prebuilt_etc_test.go
@@ -175,7 +175,7 @@
 }
 
 func TestPrebuiltEtcAndroidMk(t *testing.T) {
-	ctx, config := testPrebuiltEtc(t, `
+	ctx, _ := testPrebuiltEtc(t, `
 		prebuilt_etc {
 			name: "foo",
 			src: "foo.conf",
@@ -198,7 +198,7 @@
 	}
 
 	mod := ctx.ModuleForTests("foo", "android_arm64_armv8-a").Module().(*PrebuiltEtc)
-	entries := android.AndroidMkEntriesForTest(t, config, "", mod)[0]
+	entries := android.AndroidMkEntriesForTest(t, ctx, mod)[0]
 	for k, expectedValue := range expected {
 		if value, ok := entries.EntryMap[k]; ok {
 			if !reflect.DeepEqual(value, expectedValue) {
diff --git a/filesystem/bootimg.go b/filesystem/bootimg.go
index c90929a..764f045 100644
--- a/filesystem/bootimg.go
+++ b/filesystem/bootimg.go
@@ -107,14 +107,8 @@
 }
 
 func (b *bootimg) GenerateAndroidBuildActions(ctx android.ModuleContext) {
-	var unsignedOutput android.OutputPath
-	if proptools.Bool(b.properties.Vendor_boot) {
-		unsignedOutput = b.buildVendorBootImage(ctx)
-	} else {
-		// TODO(jiyong): fix this
-		ctx.PropertyErrorf("vendor_boot", "only vendor_boot:true is supported")
-		return
-	}
+	vendor := proptools.Bool(b.properties.Vendor_boot)
+	unsignedOutput := b.buildBootImage(ctx, vendor)
 
 	if proptools.Bool(b.properties.Use_avb) {
 		b.output = b.signImage(ctx, unsignedOutput)
@@ -126,17 +120,24 @@
 	ctx.InstallFile(b.installDir, b.installFileName(), b.output)
 }
 
-func (b *bootimg) buildVendorBootImage(ctx android.ModuleContext) android.OutputPath {
+func (b *bootimg) buildBootImage(ctx android.ModuleContext, vendor bool) android.OutputPath {
 	output := android.PathForModuleOut(ctx, "unsigned", b.installFileName()).OutputPath
 
 	builder := android.NewRuleBuilder(pctx, ctx)
 	cmd := builder.Command().BuiltTool("mkbootimg")
 
-	kernel := android.OptionalPathForModuleSrc(ctx, b.properties.Kernel_prebuilt)
-	if kernel.Valid() {
+	kernel := proptools.String(b.properties.Kernel_prebuilt)
+	if vendor && kernel != "" {
 		ctx.PropertyErrorf("kernel_prebuilt", "vendor_boot partition can't have kernel")
 		return output
 	}
+	if !vendor && kernel == "" {
+		ctx.PropertyErrorf("kernel_prebuilt", "boot partition must have kernel")
+		return output
+	}
+	if kernel != "" {
+		cmd.FlagWithInput("--kernel ", android.PathForModuleSrc(ctx, kernel))
+	}
 
 	dtbName := proptools.String(b.properties.Dtb_prebuilt)
 	if dtbName == "" {
@@ -148,7 +149,11 @@
 
 	cmdline := proptools.String(b.properties.Cmdline)
 	if cmdline != "" {
-		cmd.FlagWithArg("--vendor_cmdline ", "\""+cmdline+"\"")
+		flag := "--cmdline "
+		if vendor {
+			flag = "--vendor_cmdline "
+		}
+		cmd.FlagWithArg(flag, "\""+proptools.ShellEscape(cmdline)+"\"")
 	}
 
 	headerVersion := proptools.String(b.properties.Header_version)
@@ -174,15 +179,23 @@
 	}
 	ramdisk := ctx.GetDirectDepWithTag(ramdiskName, bootimgRamdiskDep)
 	if filesystem, ok := ramdisk.(*filesystem); ok {
-		cmd.FlagWithInput("--vendor_ramdisk ", filesystem.OutputPath())
+		flag := "--ramdisk "
+		if vendor {
+			flag = "--vendor_ramdisk "
+		}
+		cmd.FlagWithInput(flag, filesystem.OutputPath())
 	} else {
 		ctx.PropertyErrorf("ramdisk", "%q is not android_filesystem module", ramdisk.Name())
 		return output
 	}
 
-	cmd.FlagWithOutput("--vendor_boot ", output)
+	flag := "--output "
+	if vendor {
+		flag = "--vendor_boot "
+	}
+	cmd.FlagWithOutput(flag, output)
 
-	builder.Build("build_vendor_bootimg", fmt.Sprintf("Creating %s", b.BaseModuleName()))
+	builder.Build("build_bootimg", fmt.Sprintf("Creating %s", b.BaseModuleName()))
 	return output
 }
 
@@ -211,7 +224,7 @@
 		Class:      "ETC",
 		OutputFile: android.OptionalPathForPath(b.output),
 		ExtraEntries: []android.AndroidMkExtraEntriesFunc{
-			func(entries *android.AndroidMkEntries) {
+			func(ctx android.AndroidMkExtraEntriesContext, entries *android.AndroidMkEntries) {
 				entries.SetString("LOCAL_MODULE_PATH", b.installDir.ToMakePath().String())
 				entries.SetString("LOCAL_INSTALLED_MODULE_STEM", b.installFileName())
 			},
diff --git a/filesystem/filesystem.go b/filesystem/filesystem.go
index 3bccde9..3b0a7ae 100644
--- a/filesystem/filesystem.go
+++ b/filesystem/filesystem.go
@@ -358,7 +358,7 @@
 		Class:      "ETC",
 		OutputFile: android.OptionalPathForPath(f.output),
 		ExtraEntries: []android.AndroidMkExtraEntriesFunc{
-			func(entries *android.AndroidMkEntries) {
+			func(ctx android.AndroidMkExtraEntriesContext, entries *android.AndroidMkEntries) {
 				entries.SetString("LOCAL_MODULE_PATH", f.installDir.ToMakePath().String())
 				entries.SetString("LOCAL_INSTALLED_MODULE_STEM", f.installFileName())
 			},
diff --git a/java/androidmk.go b/java/androidmk.go
index 6e7c437..e7261f8 100644
--- a/java/androidmk.go
+++ b/java/androidmk.go
@@ -41,7 +41,7 @@
 			Required:   library.deviceProperties.Target.Hostdex.Required,
 			Include:    "$(BUILD_SYSTEM)/soong_java_prebuilt.mk",
 			ExtraEntries: []android.AndroidMkExtraEntriesFunc{
-				func(entries *android.AndroidMkEntries) {
+				func(ctx android.AndroidMkExtraEntriesContext, entries *android.AndroidMkEntries) {
 					entries.SetBool("LOCAL_IS_HOST_MODULE", true)
 					entries.SetPath("LOCAL_PREBUILT_MODULE_FILE", output)
 					if library.dexJarFile != nil {
@@ -74,7 +74,7 @@
 				OutputFile: android.OptionalPathForPath(checkedModulePaths[0]),
 				Include:    "$(BUILD_PHONY_PACKAGE)",
 				ExtraEntries: []android.AndroidMkExtraEntriesFunc{
-					func(entries *android.AndroidMkEntries) {
+					func(ctx android.AndroidMkExtraEntriesContext, entries *android.AndroidMkEntries) {
 						entries.AddStrings("LOCAL_ADDITIONAL_CHECKED_MODULE", checkedModulePaths.Strings()...)
 					},
 				},
@@ -88,7 +88,7 @@
 			OutputFile: android.OptionalPathForPath(library.outputFile),
 			Include:    "$(BUILD_SYSTEM)/soong_java_prebuilt.mk",
 			ExtraEntries: []android.AndroidMkExtraEntriesFunc{
-				func(entries *android.AndroidMkEntries) {
+				func(ctx android.AndroidMkExtraEntriesContext, entries *android.AndroidMkEntries) {
 					if len(library.logtagsSrcs) > 0 {
 						var logtags []string
 						for _, l := range library.logtagsSrcs {
@@ -152,7 +152,7 @@
 func (j *Test) AndroidMkEntries() []android.AndroidMkEntries {
 	entriesList := j.Library.AndroidMkEntries()
 	entries := &entriesList[0]
-	entries.ExtraEntries = append(entries.ExtraEntries, func(entries *android.AndroidMkEntries) {
+	entries.ExtraEntries = append(entries.ExtraEntries, func(ctx android.AndroidMkExtraEntriesContext, entries *android.AndroidMkEntries) {
 		testSuiteComponent(entries, j.testProperties.Test_suites)
 		if j.testConfig != nil {
 			entries.SetPath("LOCAL_FULL_TEST_CONFIG", j.testConfig)
@@ -180,7 +180,7 @@
 func (j *TestHelperLibrary) AndroidMkEntries() []android.AndroidMkEntries {
 	entriesList := j.Library.AndroidMkEntries()
 	entries := &entriesList[0]
-	entries.ExtraEntries = append(entries.ExtraEntries, func(entries *android.AndroidMkEntries) {
+	entries.ExtraEntries = append(entries.ExtraEntries, func(ctx android.AndroidMkExtraEntriesContext, entries *android.AndroidMkEntries) {
 		testSuiteComponent(entries, j.testHelperLibraryProperties.Test_suites)
 	})
 
@@ -198,7 +198,7 @@
 		OutputFile: android.OptionalPathForPath(prebuilt.combinedClasspathFile),
 		Include:    "$(BUILD_SYSTEM)/soong_java_prebuilt.mk",
 		ExtraEntries: []android.AndroidMkExtraEntriesFunc{
-			func(entries *android.AndroidMkEntries) {
+			func(ctx android.AndroidMkExtraEntriesContext, entries *android.AndroidMkEntries) {
 				entries.SetBool("LOCAL_UNINSTALLABLE_MODULE", !Bool(prebuilt.properties.Installable))
 				if prebuilt.dexJarFile != nil {
 					entries.SetPath("LOCAL_SOONG_DEX_JAR", prebuilt.dexJarFile)
@@ -223,7 +223,7 @@
 		OutputFile: android.OptionalPathForPath(prebuilt.maybeStrippedDexJarFile),
 		Include:    "$(BUILD_SYSTEM)/soong_java_prebuilt.mk",
 		ExtraEntries: []android.AndroidMkExtraEntriesFunc{
-			func(entries *android.AndroidMkEntries) {
+			func(ctx android.AndroidMkExtraEntriesContext, entries *android.AndroidMkEntries) {
 				if prebuilt.dexJarFile != nil {
 					entries.SetPath("LOCAL_SOONG_DEX_JAR", prebuilt.dexJarFile)
 				}
@@ -247,7 +247,7 @@
 		OutputFile: android.OptionalPathForPath(prebuilt.classpathFile),
 		Include:    "$(BUILD_SYSTEM)/soong_java_prebuilt.mk",
 		ExtraEntries: []android.AndroidMkExtraEntriesFunc{
-			func(entries *android.AndroidMkEntries) {
+			func(ctx android.AndroidMkExtraEntriesContext, entries *android.AndroidMkEntries) {
 				entries.SetBool("LOCAL_UNINSTALLABLE_MODULE", true)
 				entries.SetPath("LOCAL_SOONG_HEADER_JAR", prebuilt.classpathFile)
 				entries.SetPath("LOCAL_SOONG_CLASSES_JAR", prebuilt.classpathFile)
@@ -269,7 +269,7 @@
 			OutputFile: android.OptionalPathForPath(binary.outputFile),
 			Include:    "$(BUILD_SYSTEM)/soong_java_prebuilt.mk",
 			ExtraEntries: []android.AndroidMkExtraEntriesFunc{
-				func(entries *android.AndroidMkEntries) {
+				func(ctx android.AndroidMkExtraEntriesContext, entries *android.AndroidMkEntries) {
 					entries.SetPath("LOCAL_SOONG_HEADER_JAR", binary.headerJarFile)
 					entries.SetPath("LOCAL_SOONG_CLASSES_JAR", binary.implementationAndResourcesJar)
 					if binary.dexJarFile != nil {
@@ -291,7 +291,7 @@
 			Class:      "EXECUTABLES",
 			OutputFile: android.OptionalPathForPath(binary.wrapperFile),
 			ExtraEntries: []android.AndroidMkExtraEntriesFunc{
-				func(entries *android.AndroidMkEntries) {
+				func(ctx android.AndroidMkExtraEntriesContext, entries *android.AndroidMkEntries) {
 					entries.SetBool("LOCAL_STRIP_MODULE", false)
 				},
 			},
@@ -317,7 +317,7 @@
 		OutputFile: android.OptionalPathForPath(app.outputFile),
 		Include:    "$(BUILD_SYSTEM)/soong_app_prebuilt.mk",
 		ExtraEntries: []android.AndroidMkExtraEntriesFunc{
-			func(entries *android.AndroidMkEntries) {
+			func(ctx android.AndroidMkExtraEntriesContext, entries *android.AndroidMkEntries) {
 				// App module names can be overridden.
 				entries.SetString("LOCAL_MODULE", app.installApkName)
 				entries.SetBoolIfTrue("LOCAL_UNINSTALLABLE_MODULE", app.appProperties.PreventInstall)
@@ -432,7 +432,7 @@
 func (a *AndroidTest) AndroidMkEntries() []android.AndroidMkEntries {
 	entriesList := a.AndroidApp.AndroidMkEntries()
 	entries := &entriesList[0]
-	entries.ExtraEntries = append(entries.ExtraEntries, func(entries *android.AndroidMkEntries) {
+	entries.ExtraEntries = append(entries.ExtraEntries, func(ctx android.AndroidMkExtraEntriesContext, entries *android.AndroidMkEntries) {
 		testSuiteComponent(entries, a.testProperties.Test_suites)
 		if a.testConfig != nil {
 			entries.SetPath("LOCAL_FULL_TEST_CONFIG", a.testConfig)
@@ -448,7 +448,7 @@
 func (a *AndroidTestHelperApp) AndroidMkEntries() []android.AndroidMkEntries {
 	entriesList := a.AndroidApp.AndroidMkEntries()
 	entries := &entriesList[0]
-	entries.ExtraEntries = append(entries.ExtraEntries, func(entries *android.AndroidMkEntries) {
+	entries.ExtraEntries = append(entries.ExtraEntries, func(ctx android.AndroidMkExtraEntriesContext, entries *android.AndroidMkEntries) {
 		testSuiteComponent(entries, a.appTestHelperAppProperties.Test_suites)
 	})
 
@@ -464,7 +464,7 @@
 	entriesList := a.Library.AndroidMkEntries()
 	entries := &entriesList[0]
 
-	entries.ExtraEntries = append(entries.ExtraEntries, func(entries *android.AndroidMkEntries) {
+	entries.ExtraEntries = append(entries.ExtraEntries, func(ctx android.AndroidMkExtraEntriesContext, entries *android.AndroidMkEntries) {
 		if a.aarFile != nil {
 			entries.SetPath("LOCAL_SOONG_AAR", a.aarFile)
 		}
@@ -492,7 +492,7 @@
 		OutputFile: android.OptionalPathForPath(jd.stubsSrcJar),
 		Include:    "$(BUILD_SYSTEM)/soong_droiddoc_prebuilt.mk",
 		ExtraEntries: []android.AndroidMkExtraEntriesFunc{
-			func(entries *android.AndroidMkEntries) {
+			func(ctx android.AndroidMkExtraEntriesContext, entries *android.AndroidMkEntries) {
 				if BoolDefault(jd.properties.Installable, true) {
 					entries.SetPath("LOCAL_DROIDDOC_DOC_ZIP", jd.docZip)
 				}
@@ -510,7 +510,7 @@
 		OutputFile: android.OptionalPathForPath(ddoc.Javadoc.docZip),
 		Include:    "$(BUILD_SYSTEM)/soong_droiddoc_prebuilt.mk",
 		ExtraEntries: []android.AndroidMkExtraEntriesFunc{
-			func(entries *android.AndroidMkEntries) {
+			func(ctx android.AndroidMkExtraEntriesContext, entries *android.AndroidMkEntries) {
 				if ddoc.Javadoc.docZip != nil {
 					entries.SetPath("LOCAL_DROIDDOC_DOC_ZIP", ddoc.Javadoc.docZip)
 				}
@@ -539,7 +539,7 @@
 		OutputFile: outputFile,
 		Include:    "$(BUILD_SYSTEM)/soong_droiddoc_prebuilt.mk",
 		ExtraEntries: []android.AndroidMkExtraEntriesFunc{
-			func(entries *android.AndroidMkEntries) {
+			func(ctx android.AndroidMkExtraEntriesContext, entries *android.AndroidMkEntries) {
 				if dstubs.Javadoc.stubsSrcJar != nil {
 					entries.SetPath("LOCAL_DROIDDOC_STUBS_SRCJAR", dstubs.Javadoc.stubsSrcJar)
 				}
@@ -638,7 +638,7 @@
 		OutputFile: android.OptionalPathForPath(a.outputFile),
 		Include:    "$(BUILD_SYSTEM)/soong_app_prebuilt.mk",
 		ExtraEntries: []android.AndroidMkExtraEntriesFunc{
-			func(entries *android.AndroidMkEntries) {
+			func(ctx android.AndroidMkExtraEntriesContext, entries *android.AndroidMkEntries) {
 				entries.SetBoolIfTrue("LOCAL_PRIVILEGED_MODULE", a.Privileged())
 				entries.SetString("LOCAL_CERTIFICATE", a.certificate.AndroidMkString())
 				entries.AddStrings("LOCAL_OVERRIDES_PACKAGES", a.properties.Overrides...)
@@ -657,7 +657,7 @@
 func (a *AndroidTestImport) AndroidMkEntries() []android.AndroidMkEntries {
 	entriesList := a.AndroidAppImport.AndroidMkEntries()
 	entries := &entriesList[0]
-	entries.ExtraEntries = append(entries.ExtraEntries, func(entries *android.AndroidMkEntries) {
+	entries.ExtraEntries = append(entries.ExtraEntries, func(ctx android.AndroidMkExtraEntriesContext, entries *android.AndroidMkEntries) {
 		testSuiteComponent(entries, a.testProperties.Test_suites)
 		androidMkWriteTestData(a.data, entries)
 	})
@@ -678,7 +678,7 @@
 		OutputFile: android.OptionalPathForPath(r.outputFile),
 		Include:    "$(BUILD_SYSTEM)/soong_app_prebuilt.mk",
 		ExtraEntries: []android.AndroidMkExtraEntriesFunc{
-			func(entries *android.AndroidMkEntries) {
+			func(ctx android.AndroidMkExtraEntriesContext, entries *android.AndroidMkEntries) {
 				entries.SetString("LOCAL_CERTIFICATE", r.certificate.AndroidMkString())
 				entries.SetPath("LOCAL_MODULE_PATH", r.installDir.ToMakePath())
 				entries.AddStrings("LOCAL_OVERRIDES_PACKAGES", r.properties.Overrides...)
@@ -694,7 +694,7 @@
 			OutputFile: android.OptionalPathForPath(apkSet.packedOutput),
 			Include:    "$(BUILD_SYSTEM)/soong_android_app_set.mk",
 			ExtraEntries: []android.AndroidMkExtraEntriesFunc{
-				func(entries *android.AndroidMkEntries) {
+				func(ctx android.AndroidMkExtraEntriesContext, entries *android.AndroidMkEntries) {
 					entries.SetBoolIfTrue("LOCAL_PRIVILEGED_MODULE", apkSet.Privileged())
 					entries.SetString("LOCAL_APK_SET_INSTALL_FILE", apkSet.InstallFile())
 					entries.SetPath("LOCAL_APKCERTS_FILE", apkSet.apkcertsFile)
diff --git a/java/androidmk_test.go b/java/androidmk_test.go
index e2647cf..e758a92 100644
--- a/java/androidmk_test.go
+++ b/java/androidmk_test.go
@@ -22,7 +22,7 @@
 )
 
 func TestRequired(t *testing.T) {
-	ctx, config := testJava(t, `
+	ctx, _ := testJava(t, `
 		java_library {
 			name: "foo",
 			srcs: ["a.java"],
@@ -31,7 +31,7 @@
 	`)
 
 	mod := ctx.ModuleForTests("foo", "android_common").Module()
-	entries := android.AndroidMkEntriesForTest(t, config, "", mod)[0]
+	entries := android.AndroidMkEntriesForTest(t, ctx, mod)[0]
 
 	expected := []string{"libfoo"}
 	actual := entries.EntryMap["LOCAL_REQUIRED_MODULES"]
@@ -41,7 +41,7 @@
 }
 
 func TestHostdex(t *testing.T) {
-	ctx, config := testJava(t, `
+	ctx, _ := testJava(t, `
 		java_library {
 			name: "foo",
 			srcs: ["a.java"],
@@ -50,7 +50,7 @@
 	`)
 
 	mod := ctx.ModuleForTests("foo", "android_common").Module()
-	entriesList := android.AndroidMkEntriesForTest(t, config, "", mod)
+	entriesList := android.AndroidMkEntriesForTest(t, ctx, mod)
 	if len(entriesList) != 2 {
 		t.Errorf("two entries are expected, but got %d", len(entriesList))
 	}
@@ -71,7 +71,7 @@
 }
 
 func TestHostdexRequired(t *testing.T) {
-	ctx, config := testJava(t, `
+	ctx, _ := testJava(t, `
 		java_library {
 			name: "foo",
 			srcs: ["a.java"],
@@ -81,7 +81,7 @@
 	`)
 
 	mod := ctx.ModuleForTests("foo", "android_common").Module()
-	entriesList := android.AndroidMkEntriesForTest(t, config, "", mod)
+	entriesList := android.AndroidMkEntriesForTest(t, ctx, mod)
 	if len(entriesList) != 2 {
 		t.Errorf("two entries are expected, but got %d", len(entriesList))
 	}
@@ -102,7 +102,7 @@
 }
 
 func TestHostdexSpecificRequired(t *testing.T) {
-	ctx, config := testJava(t, `
+	ctx, _ := testJava(t, `
 		java_library {
 			name: "foo",
 			srcs: ["a.java"],
@@ -116,7 +116,7 @@
 	`)
 
 	mod := ctx.ModuleForTests("foo", "android_common").Module()
-	entriesList := android.AndroidMkEntriesForTest(t, config, "", mod)
+	entriesList := android.AndroidMkEntriesForTest(t, ctx, mod)
 	if len(entriesList) != 2 {
 		t.Errorf("two entries are expected, but got %d", len(entriesList))
 	}
@@ -135,7 +135,7 @@
 }
 
 func TestJavaSdkLibrary_RequireXmlPermissionFile(t *testing.T) {
-	ctx, config := testJava(t, `
+	ctx, _ := testJava(t, `
 		java_sdk_library {
 			name: "foo-shared_library",
 			srcs: ["a.java"],
@@ -159,7 +159,7 @@
 	}
 	for _, tc := range testCases {
 		mod := ctx.ModuleForTests(tc.moduleName, "android_common").Module()
-		entries := android.AndroidMkEntriesForTest(t, config, "", mod)[0]
+		entries := android.AndroidMkEntriesForTest(t, ctx, mod)[0]
 		actual := entries.EntryMap["LOCAL_REQUIRED_MODULES"]
 		if !reflect.DeepEqual(tc.expected, actual) {
 			t.Errorf("Unexpected required modules - expected: %q, actual: %q", tc.expected, actual)
@@ -168,7 +168,7 @@
 }
 
 func TestImportSoongDexJar(t *testing.T) {
-	ctx, config := testJava(t, `
+	ctx, _ := testJava(t, `
 		java_import {
 			name: "my-java-import",
 			jars: ["a.jar"],
@@ -178,7 +178,7 @@
 	`)
 
 	mod := ctx.ModuleForTests("my-java-import", "android_common").Module()
-	entries := android.AndroidMkEntriesForTest(t, config, "", mod)[0]
+	entries := android.AndroidMkEntriesForTest(t, ctx, mod)[0]
 	expectedSoongDexJar := buildDir + "/.intermediates/my-java-import/android_common/dex/my-java-import.jar"
 	actualSoongDexJar := entries.EntryMap["LOCAL_SOONG_DEX_JAR"]
 
diff --git a/java/app_import_test.go b/java/app_import_test.go
index d7f69eb..dc31d07 100644
--- a/java/app_import_test.go
+++ b/java/app_import_test.go
@@ -232,7 +232,7 @@
 }
 
 func TestAndroidAppImport_Filename(t *testing.T) {
-	ctx, config := testJava(t, `
+	ctx, _ := testJava(t, `
 		android_app_import {
 			name: "foo",
 			apk: "prebuilts/apk/app.apk",
@@ -269,8 +269,7 @@
 
 		a := variant.Module().(*AndroidAppImport)
 		expectedValues := []string{test.expected}
-		actualValues := android.AndroidMkEntriesForTest(
-			t, config, "", a)[0].EntryMap["LOCAL_INSTALLED_MODULE_STEM"]
+		actualValues := android.AndroidMkEntriesForTest(t, ctx, a)[0].EntryMap["LOCAL_INSTALLED_MODULE_STEM"]
 		if !reflect.DeepEqual(actualValues, expectedValues) {
 			t.Errorf("Incorrect LOCAL_INSTALLED_MODULE_STEM value '%s', expected '%s'",
 				actualValues, expectedValues)
@@ -394,7 +393,7 @@
 }
 
 func TestAndroidAppImport_frameworkRes(t *testing.T) {
-	ctx, config := testJava(t, `
+	ctx, _ := testJava(t, `
 		android_app_import {
 			name: "framework-res",
 			certificate: "platform",
@@ -424,7 +423,7 @@
 
 	}
 
-	entries := android.AndroidMkEntriesForTest(t, config, "", mod)[0]
+	entries := android.AndroidMkEntriesForTest(t, ctx, mod)[0]
 
 	expectedPath := "."
 	// From apk property above, in the root of the source tree.
@@ -457,7 +456,7 @@
 }
 
 func TestAndroidTestImport(t *testing.T) {
-	ctx, config := testJava(t, `
+	ctx, _ := testJava(t, `
 		android_test_import {
 			name: "foo",
 			apk: "prebuilts/apk/app.apk",
@@ -471,7 +470,7 @@
 	test := ctx.ModuleForTests("foo", "android_common").Module().(*AndroidTestImport)
 
 	// Check android mks.
-	entries := android.AndroidMkEntriesForTest(t, config, "", test)[0]
+	entries := android.AndroidMkEntriesForTest(t, ctx, test)[0]
 	expected := []string{"tests"}
 	actual := entries.EntryMap["LOCAL_MODULE_TAGS"]
 	if !reflect.DeepEqual(expected, actual) {
diff --git a/java/app_set_test.go b/java/app_set_test.go
index d31900d..ab55758 100644
--- a/java/app_set_test.go
+++ b/java/app_set_test.go
@@ -22,7 +22,7 @@
 )
 
 func TestAndroidAppSet(t *testing.T) {
-	ctx, config := testJava(t, `
+	ctx, _ := testJava(t, `
 		android_app_set {
 			name: "foo",
 			set: "prebuilts/apks/app.apks",
@@ -40,7 +40,7 @@
 	if s := params.Args["partition"]; s != "system" {
 		t.Errorf("wrong partition value: '%s', expected 'system'", s)
 	}
-	mkEntries := android.AndroidMkEntriesForTest(t, config, "", module.Module())[0]
+	mkEntries := android.AndroidMkEntriesForTest(t, ctx, module.Module())[0]
 	actualInstallFile := mkEntries.EntryMap["LOCAL_APK_SET_INSTALL_FILE"]
 	expectedInstallFile := []string{"foo.apk"}
 	if !reflect.DeepEqual(actualInstallFile, expectedInstallFile) {
diff --git a/java/builder.go b/java/builder.go
index 22a891a..33206ce 100644
--- a/java/builder.go
+++ b/java/builder.go
@@ -80,6 +80,8 @@
 		func(ctx android.PackageVarContext) string { return ctx.Config().XrefCorpusName() })
 	_ = pctx.VariableFunc("kytheCuEncoding",
 		func(ctx android.PackageVarContext) string { return ctx.Config().XrefCuEncoding() })
+	_ = pctx.VariableFunc("kytheCuJavaSourceMax",
+		func(ctx android.PackageVarContext) string { return ctx.Config().XrefCuJavaSourceMax() })
 	_ = pctx.SourcePathVariable("kytheVnames", "build/soong/vnames.json")
 	// Run it with -add-opens=java.base/java.nio=ALL-UNNAMED to avoid JDK9's warning about
 	// "Illegal reflective access by com.google.protobuf.Utf8$UnsafeProcessor ...
@@ -93,6 +95,7 @@
 				`KYTHE_CORPUS=${kytheCorpus} ` +
 				`KYTHE_VNAMES=${kytheVnames} ` +
 				`KYTHE_KZIP_ENCODING=${kytheCuEncoding} ` +
+				`KYTHE_JAVA_SOURCE_BATCH_SIZE=${kytheCuJavaSourceMax} ` +
 				`${config.SoongJavacWrapper} ${config.JavaCmd} ` +
 				`--add-opens=java.base/java.nio=ALL-UNNAMED ` +
 				`-jar ${config.JavaKytheExtractorJar} ` +
diff --git a/java/hiddenapi.go b/java/hiddenapi.go
index f8e41c4..da2c48f 100644
--- a/java/hiddenapi.go
+++ b/java/hiddenapi.go
@@ -15,8 +15,6 @@
 package java
 
 import (
-	"strings"
-
 	"github.com/google/blueprint"
 
 	"android/soong/android"
@@ -29,8 +27,8 @@
 
 type hiddenAPI struct {
 	// The name of the module as it would be used in the boot jars configuration, e.g. without any
-	// prebuilt_ prefix (if it is a prebuilt), without any "-hiddenapi" suffix if it just provides
-	// annotations and without any ".impl" suffix if it is a java_sdk_library implementation library.
+	// prebuilt_ prefix (if it is a prebuilt) and without any ".impl" suffix if it is a
+	// java_sdk_library implementation library.
 	configurationName string
 
 	// True if the module containing this structure contributes to the hiddenapi information or has
@@ -49,11 +47,6 @@
 	// annotation information.
 	primary bool
 
-	// True if the module only contains additional annotations and so does not require hiddenapi
-	// information to be encoded in its dex file and should not be used to generate the
-	// hiddenAPISingletonPathsStruct.stubFlags file.
-	annotationsOnly bool
-
 	// The path to the dex jar that is in the boot class path. If this is nil then the associated
 	// module is not a boot jar, but could be one of the <x>-hiddenapi modules that provide additional
 	// annotations for the <x> boot dex jar but which do not actually provide a boot dex jar
@@ -119,48 +112,40 @@
 var _ hiddenAPIIntf = (*hiddenAPI)(nil)
 
 // Initialize the hiddenapi structure
-func (h *hiddenAPI) initHiddenAPI(ctx android.BaseModuleContext, name string) {
+func (h *hiddenAPI) initHiddenAPI(ctx android.BaseModuleContext, configurationName string) {
 	// If hiddenapi processing is disabled treat this as inactive.
 	if ctx.Config().IsEnvTrue("UNSAFE_DISABLE_HIDDENAPI_FLAGS") {
 		return
 	}
 
-	// Modules whose names are of the format <x>-hiddenapi provide hiddenapi information for the boot
-	// jar module <x>. Otherwise, the module provides information for itself. Either way extract the
-	// configurationName of the boot jar module.
-	configurationName := strings.TrimSuffix(name, "-hiddenapi")
 	h.configurationName = configurationName
 
 	// It is important that hiddenapi information is only gathered for/from modules that are actually
 	// on the boot jars list because the runtime only enforces access to the hidden API for the
 	// bootclassloader. If information is gathered for modules not on the list then that will cause
 	// failures in the CtsHiddenApiBlocklist... tests.
-	h.active = inList(configurationName, ctx.Config().BootJars())
+	module := ctx.Module()
+	h.active = isModuleInBootClassPath(ctx, module)
 	if !h.active {
 		// The rest of the properties will be ignored if active is false.
 		return
 	}
 
-	// If this module has a suffix of -hiddenapi then it only provides additional annotation
-	// information for a module on the boot jars list.
-	h.annotationsOnly = strings.HasSuffix(name, "-hiddenapi")
-
 	// Determine whether this module is the primary module or not.
 	primary := true
 
 	// A prebuilt module is only primary if it is preferred and conversely a source module is only
 	// primary if it has not been replaced by a prebuilt module.
-	module := ctx.Module()
 	if pi, ok := module.(android.PrebuiltInterface); ok {
 		if p := pi.Prebuilt(); p != nil {
 			primary = p.UsePrebuilt()
 		}
 	} else {
-		// The only module that will pass a different name to its module name to this method is the
-		// implementation library of a java_sdk_library. It has a configuration name of <x> the same
-		// as its parent java_sdk_library but a module name of <x>.impl. It is not the primary module,
-		// the java_sdk_library with the name of <x> is.
-		primary = name == ctx.ModuleName()
+		// The only module that will pass a different configurationName to its module name to this
+		// method is the implementation library of a java_sdk_library. It has a configuration name of
+		// <x> the same as its parent java_sdk_library but a module name of <x>.impl. It is not the
+		// primary module, the java_sdk_library with the name of <x> is.
+		primary = configurationName == ctx.ModuleName()
 
 		// A source module that has been replaced by a prebuilt can never be the primary module.
 		primary = primary && !module.IsReplacedByPrebuilt()
@@ -168,6 +153,15 @@
 	h.primary = primary
 }
 
+func isModuleInBootClassPath(ctx android.BaseModuleContext, module android.Module) bool {
+	// Get the configured non-updatable and updatable boot jars.
+	nonUpdatableBootJars := ctx.Config().NonUpdatableBootJars()
+	updatableBootJars := ctx.Config().UpdatableBootJars()
+	active := isModuleInConfiguredList(ctx, module, nonUpdatableBootJars) ||
+		isModuleInConfiguredList(ctx, module, updatableBootJars)
+	return active
+}
+
 // hiddenAPIExtractAndEncode is called by any module that could contribute to the hiddenapi
 // processing.
 //
@@ -191,15 +185,13 @@
 
 	h.hiddenAPIExtractInformation(ctx, dexJar, implementationJar)
 
-	if !h.annotationsOnly {
-		hiddenAPIJar := android.PathForModuleOut(ctx, "hiddenapi", h.configurationName+".jar").OutputPath
+	hiddenAPIJar := android.PathForModuleOut(ctx, "hiddenapi", h.configurationName+".jar").OutputPath
 
-		// Create a copy of the dex jar which has been encoded with hiddenapi flags.
-		hiddenAPIEncodeDex(ctx, hiddenAPIJar, dexJar, uncompressDex)
+	// Create a copy of the dex jar which has been encoded with hiddenapi flags.
+	hiddenAPIEncodeDex(ctx, hiddenAPIJar, dexJar, uncompressDex)
 
-		// Use the encoded dex jar from here onwards.
-		dexJar = hiddenAPIJar
-	}
+	// Use the encoded dex jar from here onwards.
+	dexJar = hiddenAPIJar
 
 	return dexJar
 }
@@ -262,6 +254,7 @@
 	rule.Command().
 		BuiltTool("merge_csv").
 		Flag("--zip_input").
+		Flag("--key_field signature").
 		FlagWithOutput("--output=", indexCSV).
 		Inputs(classesJars)
 	rule.Build("merged-hiddenapi-index", "Merged Hidden API index")
diff --git a/java/hiddenapi_singleton.go b/java/hiddenapi_singleton.go
index 6341a34..82e8b3f 100644
--- a/java/hiddenapi_singleton.go
+++ b/java/hiddenapi_singleton.go
@@ -217,10 +217,6 @@
 
 	var bootDexJars android.Paths
 
-	// Get the configured non-updatable and updatable boot jars.
-	nonUpdatableBootJars := ctx.Config().NonUpdatableBootJars()
-	updatableBootJars := ctx.Config().UpdatableBootJars()
-
 	ctx.VisitAllModules(func(module android.Module) {
 		// Collect dex jar paths for the modules listed above.
 		if j, ok := module.(UsesLibraryDependency); ok {
@@ -235,11 +231,6 @@
 		// Collect dex jar paths for modules that had hiddenapi encode called on them.
 		if h, ok := module.(hiddenAPIIntf); ok {
 			if jar := h.bootDexJar(); jar != nil {
-				if !isModuleInConfiguredList(ctx, module, nonUpdatableBootJars) &&
-					!isModuleInConfiguredList(ctx, module, updatableBootJars) {
-					return
-				}
-
 				bootDexJars = append(bootDexJars, jar)
 			}
 		}
@@ -291,8 +282,8 @@
 // there too.
 //
 // TODO(b/179354495): Avoid having to perform this type of check or if necessary dedup it.
-func isModuleInConfiguredList(ctx android.SingletonContext, module android.Module, configuredBootJars android.ConfiguredJarList) bool {
-	name := ctx.ModuleName(module)
+func isModuleInConfiguredList(ctx android.BaseModuleContext, module android.Module, configuredBootJars android.ConfiguredJarList) bool {
+	name := ctx.OtherModuleName(module)
 
 	// Strip a prebuilt_ prefix so that this can match a prebuilt module that has not been renamed.
 	name = android.RemoveOptionalPrebuiltPrefix(name)
@@ -305,11 +296,11 @@
 
 	// It is an error if the module is not an ApexModule.
 	if _, ok := module.(android.ApexModule); !ok {
-		ctx.Errorf("module %q configured in boot jars does not support being added to an apex", module)
+		ctx.ModuleErrorf("is configured in boot jars but does not support being added to an apex")
 		return false
 	}
 
-	apexInfo := ctx.ModuleProvider(module, android.ApexInfoProvider).(android.ApexInfo)
+	apexInfo := ctx.OtherModuleProvider(module, android.ApexInfoProvider).(android.ApexInfo)
 
 	// Now match the apex part of the boot image configuration.
 	requiredApex := configuredBootJars.Apex(index)
@@ -433,6 +424,7 @@
 
 	rule.Command().
 		BuiltTool("merge_csv").
+		Flag("--key_field signature").
 		FlagWithOutput("--output=", outputPath).
 		Inputs(metadataCSV)
 
@@ -544,6 +536,7 @@
 	rule := android.NewRuleBuilder(pctx, ctx)
 	rule.Command().
 		BuiltTool("merge_csv").
+		Flag("--key_field signature").
 		FlagWithArg("--header=", "signature,file,startline,startcol,endline,endcol,properties").
 		FlagWithOutput("--output=", hiddenAPISingletonPaths(ctx).index).
 		Inputs(indexes)
diff --git a/java/hiddenapi_singleton_test.go b/java/hiddenapi_singleton_test.go
index 4670d03..c0f0e38 100644
--- a/java/hiddenapi_singleton_test.go
+++ b/java/hiddenapi_singleton_test.go
@@ -89,12 +89,6 @@
 		}
 
 		java_library {
-			name: "foo-hiddenapi",
-			srcs: ["a.java"],
-			compile_dex: true,
-		}
-
-		java_library {
 			name: "foo-hiddenapi-annotations",
 			srcs: ["a.java"],
 			compile_dex: true,
@@ -118,7 +112,6 @@
 	indexRule := hiddenAPIIndex.Rule("singleton-merged-hiddenapi-index")
 	CheckHiddenAPIRuleInputs(t, `
 .intermediates/bar/android_common/hiddenapi/index.csv
-.intermediates/foo-hiddenapi/android_common/hiddenapi/index.csv
 .intermediates/foo/android_common/hiddenapi/index.csv
 `,
 		indexRule)
diff --git a/java/java.go b/java/java.go
index dbfad02..78cd362 100644
--- a/java/java.go
+++ b/java/java.go
@@ -1038,7 +1038,7 @@
 	case javaPlatform:
 		return "private API"
 	default:
-		panic(fmt.Errorf("unrecognized linktype: %v", lt))
+		panic(fmt.Errorf("unrecognized linktype: %d", lt))
 	}
 }
 
diff --git a/java/java_test.go b/java/java_test.go
index 0ef4db6..11f6a7c 100644
--- a/java/java_test.go
+++ b/java/java_test.go
@@ -2509,7 +2509,7 @@
 }
 
 func TestDataNativeBinaries(t *testing.T) {
-	ctx, config := testJava(t, `
+	ctx, _ := testJava(t, `
 		java_test_host {
 			name: "foo",
 			srcs: ["a.java"],
@@ -2525,7 +2525,7 @@
 	buildOS := android.BuildOs.String()
 
 	test := ctx.ModuleForTests("foo", buildOS+"_common").Module().(*TestHost)
-	entries := android.AndroidMkEntriesForTest(t, config, "", test)[0]
+	entries := android.AndroidMkEntriesForTest(t, ctx, test)[0]
 	expected := []string{buildDir + "/.intermediates/bin/" + buildOS + "_x86_64_PY3/bin:bin"}
 	actual := entries.EntryMap["LOCAL_COMPATIBILITY_SUPPORT_FILES"]
 	if !reflect.DeepEqual(expected, actual) {
diff --git a/java/platform_compat_config.go b/java/platform_compat_config.go
index 9bc821d..2c47b0a 100644
--- a/java/platform_compat_config.go
+++ b/java/platform_compat_config.go
@@ -131,7 +131,7 @@
 		OutputFile: android.OptionalPathForPath(p.configFile),
 		Include:    "$(BUILD_PREBUILT)",
 		ExtraEntries: []android.AndroidMkExtraEntriesFunc{
-			func(entries *android.AndroidMkEntries) {
+			func(ctx android.AndroidMkExtraEntriesContext, entries *android.AndroidMkEntries) {
 				entries.SetString("LOCAL_MODULE_PATH", p.installDirPath.ToMakePath().String())
 				entries.SetString("LOCAL_INSTALLED_MODULE_STEM", p.configFile.Base())
 			},
diff --git a/java/rro_test.go b/java/rro_test.go
index 0a2f848..edbf170 100644
--- a/java/rro_test.go
+++ b/java/rro_test.go
@@ -96,7 +96,7 @@
 	if expected != signingFlag {
 		t.Errorf("Incorrect signing flags, expected: %q, got: %q", expected, signingFlag)
 	}
-	androidMkEntries := android.AndroidMkEntriesForTest(t, config, "", m.Module())[0]
+	androidMkEntries := android.AndroidMkEntriesForTest(t, ctx, m.Module())[0]
 	path := androidMkEntries.EntryMap["LOCAL_CERTIFICATE"]
 	expectedPath := []string{"build/make/target/product/security/platform.x509.pem"}
 	if !reflect.DeepEqual(path, expectedPath) {
@@ -112,7 +112,7 @@
 
 	// A themed module has a different device location
 	m = ctx.ModuleForTests("foo_themed", "android_common")
-	androidMkEntries = android.AndroidMkEntriesForTest(t, config, "", m.Module())[0]
+	androidMkEntries = android.AndroidMkEntriesForTest(t, ctx, m.Module())[0]
 	path = androidMkEntries.EntryMap["LOCAL_MODULE_PATH"]
 	expectedPath = []string{"/tmp/target/product/test_device/product/overlay/faza"}
 	if !reflect.DeepEqual(path, expectedPath) {
@@ -127,7 +127,7 @@
 }
 
 func TestRuntimeResourceOverlay_JavaDefaults(t *testing.T) {
-	ctx, config := testJava(t, `
+	ctx, _ := testJava(t, `
 		java_defaults {
 			name: "rro_defaults",
 			theme: "default_theme",
@@ -159,7 +159,7 @@
 	}
 
 	// Check device location.
-	path := android.AndroidMkEntriesForTest(t, config, "", m.Module())[0].EntryMap["LOCAL_MODULE_PATH"]
+	path := android.AndroidMkEntriesForTest(t, ctx, m.Module())[0].EntryMap["LOCAL_MODULE_PATH"]
 	expectedPath := []string{"/tmp/target/product/test_device/product/overlay/default_theme"}
 	if !reflect.DeepEqual(path, expectedPath) {
 		t.Errorf("Unexpected LOCAL_MODULE_PATH value: %q, expected: %q", path, expectedPath)
@@ -178,7 +178,7 @@
 	}
 
 	// Check device location.
-	path = android.AndroidMkEntriesForTest(t, config, "", m.Module())[0].EntryMap["LOCAL_MODULE_PATH"]
+	path = android.AndroidMkEntriesForTest(t, ctx, m.Module())[0].EntryMap["LOCAL_MODULE_PATH"]
 	expectedPath = []string{"/tmp/target/product/test_device/system/overlay"}
 	if !reflect.DeepEqual(path, expectedPath) {
 		t.Errorf("Unexpected LOCAL_MODULE_PATH value: %v, expected: %v", path, expectedPath)
@@ -345,7 +345,7 @@
 			modules := []string{"foo", "bar"}
 			for _, moduleName := range modules {
 				module := ctx.ModuleForTests(moduleName, "android_common")
-				mkEntries := android.AndroidMkEntriesForTest(t, config, "", module.Module())[0]
+				mkEntries := android.AndroidMkEntriesForTest(t, ctx, module.Module())[0]
 				actualRRODirs := mkEntries.EntryMap["LOCAL_SOONG_PRODUCT_RRO_DIRS"]
 				if !reflect.DeepEqual(actualRRODirs, testCase.rroDirs[moduleName]) {
 					t.Errorf("exected %s LOCAL_SOONG_PRODUCT_RRO_DIRS entry: %v\ngot:%q",
diff --git a/java/sdk_library.go b/java/sdk_library.go
index aa96e0d..90823a0 100644
--- a/java/sdk_library.go
+++ b/java/sdk_library.go
@@ -2238,7 +2238,7 @@
 		Class:      "ETC",
 		OutputFile: android.OptionalPathForPath(module.outputFilePath),
 		ExtraEntries: []android.AndroidMkExtraEntriesFunc{
-			func(entries *android.AndroidMkEntries) {
+			func(ctx android.AndroidMkExtraEntriesContext, entries *android.AndroidMkEntries) {
 				entries.SetString("LOCAL_MODULE_TAGS", "optional")
 				entries.SetString("LOCAL_MODULE_PATH", module.installDirPath.ToMakePath().String())
 				entries.SetString("LOCAL_INSTALLED_MODULE_STEM", module.outputFilePath.Base())
diff --git a/linkerconfig/linkerconfig.go b/linkerconfig/linkerconfig.go
index d538ce4..ff548e5 100644
--- a/linkerconfig/linkerconfig.go
+++ b/linkerconfig/linkerconfig.go
@@ -98,7 +98,7 @@
 		Class:      "ETC",
 		OutputFile: android.OptionalPathForPath(l.outputFilePath),
 		ExtraEntries: []android.AndroidMkExtraEntriesFunc{
-			func(entries *android.AndroidMkEntries) {
+			func(ctx android.AndroidMkExtraEntriesContext, entries *android.AndroidMkEntries) {
 				entries.SetString("LOCAL_MODULE_PATH", l.installDirPath.ToMakePath().String())
 				entries.SetString("LOCAL_INSTALLED_MODULE_STEM", l.outputFilePath.Base())
 				entries.SetBoolIfTrue("LOCAL_UNINSTALLABLE_MODULE", !installable)
diff --git a/linkerconfig/linkerconfig_test.go b/linkerconfig/linkerconfig_test.go
index 01f4657..8eed4b5 100644
--- a/linkerconfig/linkerconfig_test.go
+++ b/linkerconfig/linkerconfig_test.go
@@ -47,7 +47,7 @@
 	os.Exit(run())
 }
 
-func testContext(t *testing.T, bp string) (*android.TestContext, android.Config) {
+func testContext(t *testing.T, bp string) *android.TestContext {
 	t.Helper()
 
 	fs := map[string][]byte{
@@ -65,11 +65,11 @@
 	_, errs = ctx.PrepareBuildActions(config)
 	android.FailIfErrored(t, errs)
 
-	return ctx, config
+	return ctx
 }
 
 func TestBaseLinkerConfig(t *testing.T) {
-	ctx, config := testContext(t, `
+	ctx := testContext(t, `
 	linker_config {
 		name: "linker-config-base",
 		src: "linker.config.json",
@@ -88,7 +88,7 @@
 		t.Errorf("expected linker.config.pb, got %q", p.outputFilePath.Base())
 	}
 
-	entries := android.AndroidMkEntriesForTest(t, config, "", p)[0]
+	entries := android.AndroidMkEntriesForTest(t, ctx, p)[0]
 	for k, expectedValue := range expected {
 		if value, ok := entries.EntryMap[k]; ok {
 			if !reflect.DeepEqual(value, expectedValue) {
@@ -105,7 +105,7 @@
 }
 
 func TestUninstallableLinkerConfig(t *testing.T) {
-	ctx, config := testContext(t, `
+	ctx := testContext(t, `
 	linker_config {
 		name: "linker-config-base",
 		src: "linker.config.json",
@@ -116,7 +116,7 @@
 	expected := []string{"true"}
 
 	p := ctx.ModuleForTests("linker-config-base", "android_arm64_armv8-a").Module().(*linkerConfig)
-	entries := android.AndroidMkEntriesForTest(t, config, "", p)[0]
+	entries := android.AndroidMkEntriesForTest(t, ctx, p)[0]
 	if value, ok := entries.EntryMap["LOCAL_UNINSTALLABLE_MODULE"]; ok {
 		if !reflect.DeepEqual(value, expected) {
 			t.Errorf("LOCAL_UNINSTALLABLE_MODULE is expected to be true but %s", value)
diff --git a/python/androidmk.go b/python/androidmk.go
index 60637d3..13b4172 100644
--- a/python/androidmk.go
+++ b/python/androidmk.go
@@ -48,27 +48,29 @@
 func (p *binaryDecorator) AndroidMk(base *Module, entries *android.AndroidMkEntries) {
 	entries.Class = "EXECUTABLES"
 
-	entries.ExtraEntries = append(entries.ExtraEntries, func(entries *android.AndroidMkEntries) {
-		entries.AddCompatibilityTestSuites(p.binaryProperties.Test_suites...)
-	})
+	entries.ExtraEntries = append(entries.ExtraEntries,
+		func(ctx android.AndroidMkExtraEntriesContext, entries *android.AndroidMkEntries) {
+			entries.AddCompatibilityTestSuites(p.binaryProperties.Test_suites...)
+		})
 	base.subAndroidMk(entries, p.pythonInstaller)
 }
 
 func (p *testDecorator) AndroidMk(base *Module, entries *android.AndroidMkEntries) {
 	entries.Class = "NATIVE_TESTS"
 
-	entries.ExtraEntries = append(entries.ExtraEntries, func(entries *android.AndroidMkEntries) {
-		entries.AddCompatibilityTestSuites(p.binaryDecorator.binaryProperties.Test_suites...)
-		if p.testConfig != nil {
-			entries.SetString("LOCAL_FULL_TEST_CONFIG", p.testConfig.String())
-		}
+	entries.ExtraEntries = append(entries.ExtraEntries,
+		func(ctx android.AndroidMkExtraEntriesContext, entries *android.AndroidMkEntries) {
+			entries.AddCompatibilityTestSuites(p.binaryDecorator.binaryProperties.Test_suites...)
+			if p.testConfig != nil {
+				entries.SetString("LOCAL_FULL_TEST_CONFIG", p.testConfig.String())
+			}
 
-		entries.SetBoolIfTrue("LOCAL_DISABLE_AUTO_GENERATE_TEST_CONFIG", !BoolDefault(p.binaryProperties.Auto_gen_config, true))
+			entries.SetBoolIfTrue("LOCAL_DISABLE_AUTO_GENERATE_TEST_CONFIG", !BoolDefault(p.binaryProperties.Auto_gen_config, true))
 
-		entries.AddStrings("LOCAL_TEST_DATA", android.AndroidMkDataPaths(p.data)...)
+			entries.AddStrings("LOCAL_TEST_DATA", android.AndroidMkDataPaths(p.data)...)
 
-		entries.SetBoolIfTrue("LOCAL_IS_UNIT_TEST", Bool(p.testProperties.Test_options.Unit_test))
-	})
+			entries.SetBoolIfTrue("LOCAL_IS_UNIT_TEST", Bool(p.testProperties.Test_options.Unit_test))
+		})
 	base.subAndroidMk(entries, p.binaryDecorator.pythonInstaller)
 }
 
@@ -80,14 +82,15 @@
 	}
 
 	entries.Required = append(entries.Required, "libc++")
-	entries.ExtraEntries = append(entries.ExtraEntries, func(entries *android.AndroidMkEntries) {
-		path, file := filepath.Split(installer.path.ToMakePath().String())
-		stem := strings.TrimSuffix(file, filepath.Ext(file))
+	entries.ExtraEntries = append(entries.ExtraEntries,
+		func(ctx android.AndroidMkExtraEntriesContext, entries *android.AndroidMkEntries) {
+			path, file := filepath.Split(installer.path.ToMakePath().String())
+			stem := strings.TrimSuffix(file, filepath.Ext(file))
 
-		entries.SetString("LOCAL_MODULE_SUFFIX", filepath.Ext(file))
-		entries.SetString("LOCAL_MODULE_PATH", path)
-		entries.SetString("LOCAL_MODULE_STEM", stem)
-		entries.AddStrings("LOCAL_SHARED_LIBRARIES", installer.androidMkSharedLibs...)
-		entries.SetBool("LOCAL_CHECK_ELF_FILES", false)
-	})
+			entries.SetString("LOCAL_MODULE_SUFFIX", filepath.Ext(file))
+			entries.SetString("LOCAL_MODULE_PATH", path)
+			entries.SetString("LOCAL_MODULE_STEM", stem)
+			entries.AddStrings("LOCAL_SHARED_LIBRARIES", installer.androidMkSharedLibs...)
+			entries.SetBool("LOCAL_CHECK_ELF_FILES", false)
+		})
 }
diff --git a/rust/androidmk.go b/rust/androidmk.go
index 1a286f7..0f9a17d 100644
--- a/rust/androidmk.go
+++ b/rust/androidmk.go
@@ -53,7 +53,7 @@
 		OutputFile: mod.outputFile,
 		Include:    "$(BUILD_SYSTEM)/soong_rust_prebuilt.mk",
 		ExtraEntries: []android.AndroidMkExtraEntriesFunc{
-			func(entries *android.AndroidMkEntries) {
+			func(ctx android.AndroidMkExtraEntriesContext, entries *android.AndroidMkEntries) {
 				entries.AddStrings("LOCAL_RLIB_LIBRARIES", mod.Properties.AndroidMkRlibs...)
 				entries.AddStrings("LOCAL_DYLIB_LIBRARIES", mod.Properties.AndroidMkDylibs...)
 				entries.AddStrings("LOCAL_PROC_MACRO_LIBRARIES", mod.Properties.AndroidMkProcMacroLibs...)
@@ -89,14 +89,15 @@
 	ctx.SubAndroidMk(ret, test.binaryDecorator)
 
 	ret.Class = "NATIVE_TESTS"
-	ret.ExtraEntries = append(ret.ExtraEntries, func(entries *android.AndroidMkEntries) {
-		entries.AddCompatibilityTestSuites(test.Properties.Test_suites...)
-		if test.testConfig != nil {
-			entries.SetString("LOCAL_FULL_TEST_CONFIG", test.testConfig.String())
-		}
-		entries.SetBoolIfTrue("LOCAL_DISABLE_AUTO_GENERATE_TEST_CONFIG", !BoolDefault(test.Properties.Auto_gen_config, true))
-		entries.SetBoolIfTrue("LOCAL_IS_UNIT_TEST", Bool(test.Properties.Test_options.Unit_test))
-	})
+	ret.ExtraEntries = append(ret.ExtraEntries,
+		func(ctx android.AndroidMkExtraEntriesContext, entries *android.AndroidMkEntries) {
+			entries.AddCompatibilityTestSuites(test.Properties.Test_suites...)
+			if test.testConfig != nil {
+				entries.SetString("LOCAL_FULL_TEST_CONFIG", test.testConfig.String())
+			}
+			entries.SetBoolIfTrue("LOCAL_DISABLE_AUTO_GENERATE_TEST_CONFIG", !BoolDefault(test.Properties.Auto_gen_config, true))
+			entries.SetBoolIfTrue("LOCAL_IS_UNIT_TEST", Bool(test.Properties.Test_options.Unit_test))
+		})
 
 	cc.AndroidMkWriteTestData(test.data, ret)
 }
@@ -134,13 +135,14 @@
 	ret.Class = "ETC"
 	ret.OutputFile = android.OptionalPathForPath(outFile)
 	ret.SubName += sourceProvider.subName
-	ret.ExtraEntries = append(ret.ExtraEntries, func(entries *android.AndroidMkEntries) {
-		_, file := filepath.Split(outFile.String())
-		stem, suffix, _ := android.SplitFileExt(file)
-		entries.SetString("LOCAL_MODULE_SUFFIX", suffix)
-		entries.SetString("LOCAL_MODULE_STEM", stem)
-		entries.SetBool("LOCAL_UNINSTALLABLE_MODULE", true)
-	})
+	ret.ExtraEntries = append(ret.ExtraEntries,
+		func(ctx android.AndroidMkExtraEntriesContext, entries *android.AndroidMkEntries) {
+			_, file := filepath.Split(outFile.String())
+			stem, suffix, _ := android.SplitFileExt(file)
+			entries.SetString("LOCAL_MODULE_SUFFIX", suffix)
+			entries.SetString("LOCAL_MODULE_STEM", stem)
+			entries.SetBool("LOCAL_UNINSTALLABLE_MODULE", true)
+		})
 }
 
 func (bindgen *bindgenDecorator) AndroidMk(ctx AndroidMkContext, ret *android.AndroidMkEntries) {
@@ -165,12 +167,13 @@
 		unstrippedOutputFile = ret.OutputFile
 		ret.OutputFile = compiler.strippedOutputFile
 	}
-	ret.ExtraEntries = append(ret.ExtraEntries, func(entries *android.AndroidMkEntries) {
-		entries.SetOptionalPath("LOCAL_SOONG_UNSTRIPPED_BINARY", unstrippedOutputFile)
-		path, file := filepath.Split(compiler.path.ToMakePath().String())
-		stem, suffix, _ := android.SplitFileExt(file)
-		entries.SetString("LOCAL_MODULE_SUFFIX", suffix)
-		entries.SetString("LOCAL_MODULE_PATH", path)
-		entries.SetString("LOCAL_MODULE_STEM", stem)
-	})
+	ret.ExtraEntries = append(ret.ExtraEntries,
+		func(ctx android.AndroidMkExtraEntriesContext, entries *android.AndroidMkEntries) {
+			entries.SetOptionalPath("LOCAL_SOONG_UNSTRIPPED_BINARY", unstrippedOutputFile)
+			path, file := filepath.Split(compiler.path.ToMakePath().String())
+			stem, suffix, _ := android.SplitFileExt(file)
+			entries.SetString("LOCAL_MODULE_SUFFIX", suffix)
+			entries.SetString("LOCAL_MODULE_PATH", path)
+			entries.SetString("LOCAL_MODULE_STEM", stem)
+		})
 }
diff --git a/rust/builder.go b/rust/builder.go
index 547d705..6326124 100644
--- a/rust/builder.go
+++ b/rust/builder.go
@@ -230,6 +230,8 @@
 		envVars = append(envVars, "OUT_DIR="+filepath.Join(outDirPrefix, moduleGenDir.String()))
 	}
 
+	envVars = append(envVars, "ANDROID_RUST_VERSION="+config.RustDefaultVersion)
+
 	if flags.Clippy {
 		clippyFile := android.PathForModuleOut(ctx, outputFile.Base()+".clippy")
 		ctx.Build(pctx, android.BuildParams{
diff --git a/scripts/hiddenapi/merge_csv.py b/scripts/hiddenapi/merge_csv.py
index 5ad61b2..b047aab 100755
--- a/scripts/hiddenapi/merge_csv.py
+++ b/scripts/hiddenapi/merge_csv.py
@@ -20,6 +20,9 @@
 import argparse
 import csv
 import io
+import heapq
+import itertools
+import operator
 
 from zipfile import ZipFile
 
@@ -28,6 +31,10 @@
                                           'if missing determines the header from input files.')
 args_parser.add_argument('--zip_input', help='Treat files as ZIP archives containing CSV files to merge.',
                          action="store_true")
+args_parser.add_argument('--key_field', help='The name of the field by which the rows should be sorted. '
+                                             'Must be in the field names. '
+                                             'Will be the first field in the output. '
+                                             'All input files must be sorted by that field.')
 args_parser.add_argument('--output', help='Output file for merged CSV.',
                          default='-', type=argparse.FileType('w'))
 args_parser.add_argument('files', nargs=argparse.REMAINDER)
@@ -57,10 +64,29 @@
         headers = headers.union(reader.fieldnames)
     fieldnames = sorted(headers)
 
-# Concatenate all files to output:
+# By default chain the csv readers together so that the resulting output is
+# the concatenation of the rows from each of them:
+all_rows = itertools.chain.from_iterable(csv_readers)
+
+if len(csv_readers) > 0:
+    keyField = args.key_field
+    if keyField:
+        assert keyField in fieldnames, (
+            "--key_field {} not found, must be one of {}\n").format(
+            keyField, ",".join(fieldnames))
+        # Make the key field the first field in the output
+        keyFieldIndex = fieldnames.index(args.key_field)
+        fieldnames.insert(0, fieldnames.pop(keyFieldIndex))
+        # Create an iterable that performs a lazy merge sort on the csv readers
+        # sorting the rows by the key field.
+        all_rows = heapq.merge(*csv_readers, key=operator.itemgetter(keyField))
+
+# Write all rows from the input files to the output:
 writer = csv.DictWriter(args.output, delimiter=',', quotechar='|', quoting=csv.QUOTE_MINIMAL,
                         dialect='unix', fieldnames=fieldnames)
 writer.writeheader()
-for reader in csv_readers:
-    for row in reader:
-        writer.writerow(row)
+
+# Read all the rows from the input and write them to the output in the correct
+# order:
+for row in all_rows:
+  writer.writerow(row)
diff --git a/sh/sh_binary.go b/sh/sh_binary.go
index 54dfc24..58f8cf6 100644
--- a/sh/sh_binary.go
+++ b/sh/sh_binary.go
@@ -246,7 +246,7 @@
 		OutputFile: android.OptionalPathForPath(s.outputFilePath),
 		Include:    "$(BUILD_SYSTEM)/soong_cc_prebuilt.mk",
 		ExtraEntries: []android.AndroidMkExtraEntriesFunc{
-			func(entries *android.AndroidMkEntries) {
+			func(ctx android.AndroidMkExtraEntriesContext, entries *android.AndroidMkEntries) {
 				s.customAndroidMkEntries(entries)
 				entries.SetString("LOCAL_MODULE_RELATIVE_PATH", proptools.String(s.properties.Sub_dir))
 			},
@@ -395,7 +395,7 @@
 		OutputFile: android.OptionalPathForPath(s.outputFilePath),
 		Include:    "$(BUILD_SYSTEM)/soong_cc_prebuilt.mk",
 		ExtraEntries: []android.AndroidMkExtraEntriesFunc{
-			func(entries *android.AndroidMkEntries) {
+			func(ctx android.AndroidMkExtraEntriesContext, entries *android.AndroidMkEntries) {
 				s.customAndroidMkEntries(entries)
 				entries.SetPath("LOCAL_MODULE_PATH", s.installDir.ToMakePath())
 				entries.AddCompatibilityTestSuites(s.testProperties.Test_suites...)
diff --git a/sh/sh_binary_test.go b/sh/sh_binary_test.go
index c664461..fb7ab13 100644
--- a/sh/sh_binary_test.go
+++ b/sh/sh_binary_test.go
@@ -61,7 +61,7 @@
 }
 
 func TestShTestSubDir(t *testing.T) {
-	ctx, config := testShBinary(t, `
+	ctx, _ := testShBinary(t, `
 		sh_test {
 			name: "foo",
 			src: "test.sh",
@@ -71,7 +71,7 @@
 
 	mod := ctx.ModuleForTests("foo", "android_arm64_armv8-a").Module().(*ShTest)
 
-	entries := android.AndroidMkEntriesForTest(t, config, "", mod)[0]
+	entries := android.AndroidMkEntriesForTest(t, ctx, mod)[0]
 
 	expectedPath := "/tmp/target/product/test_device/data/nativetest64/foo_test"
 	actualPath := entries.EntryMap["LOCAL_MODULE_PATH"][0]
@@ -81,7 +81,7 @@
 }
 
 func TestShTest(t *testing.T) {
-	ctx, config := testShBinary(t, `
+	ctx, _ := testShBinary(t, `
 		sh_test {
 			name: "foo",
 			src: "test.sh",
@@ -95,7 +95,7 @@
 
 	mod := ctx.ModuleForTests("foo", "android_arm64_armv8-a").Module().(*ShTest)
 
-	entries := android.AndroidMkEntriesForTest(t, config, "", mod)[0]
+	entries := android.AndroidMkEntriesForTest(t, ctx, mod)[0]
 
 	expectedPath := "/tmp/target/product/test_device/data/nativetest64/foo"
 	actualPath := entries.EntryMap["LOCAL_MODULE_PATH"][0]
@@ -111,7 +111,7 @@
 }
 
 func TestShTest_dataModules(t *testing.T) {
-	ctx, config := testShBinary(t, `
+	ctx, _ := testShBinary(t, `
 		sh_test {
 			name: "foo",
 			src: "test.sh",
@@ -157,7 +157,7 @@
 		}
 
 		mod := variant.Module().(*ShTest)
-		entries := android.AndroidMkEntriesForTest(t, config, "", mod)[0]
+		entries := android.AndroidMkEntriesForTest(t, ctx, mod)[0]
 		expectedData := []string{
 			filepath.Join(buildDir, ".intermediates/bar", arch, ":bar"),
 			filepath.Join(buildDir, ".intermediates/foo", arch, "relocated/:lib64/libbar"+libExt),
@@ -190,7 +190,7 @@
 }
 
 func TestShTestHost_dataDeviceModules(t *testing.T) {
-	ctx, config := testShBinary(t, `
+	ctx, _ := testShBinary(t, `
 		sh_test_host {
 			name: "foo",
 			src: "test.sh",
@@ -227,7 +227,7 @@
 	}
 
 	mod := variant.Module().(*ShTest)
-	entries := android.AndroidMkEntriesForTest(t, config, "", mod)[0]
+	entries := android.AndroidMkEntriesForTest(t, ctx, mod)[0]
 	expectedData := []string{
 		filepath.Join(buildDir, ".intermediates/bar/android_arm64_armv8-a/:bar"),
 		// libbar has been relocated, and so has a variant that matches the host arch.