Merge "rust: Don't produce apex variants if apex_exclude" into main
diff --git a/aconfig/build_flags/build_flags_singleton.go b/aconfig/build_flags/build_flags_singleton.go
index ba27a85..3b40755 100644
--- a/aconfig/build_flags/build_flags_singleton.go
+++ b/aconfig/build_flags/build_flags_singleton.go
@@ -114,7 +114,7 @@
 
 func (this *allBuildFlagDeclarationsSingleton) MakeVars(ctx android.MakeVarsContext) {
 	ctx.DistForGoal("droid", this.flagsBinaryProtoPath)
-	for _, goal := range []string{"docs", "droid", "sdk"} {
+	for _, goal := range []string{"docs", "droid", "sdk", "release_config_metadata"} {
 		ctx.DistForGoalWithFilename(goal, this.flagsBinaryProtoPath, "build_flags/all_flags.pb")
 		ctx.DistForGoalWithFilename(goal, this.flagsTextProtoPath, "build_flags/all_flags.textproto")
 		ctx.DistForGoalWithFilename(goal, this.configsBinaryProtoPath, "build_flags/all_release_config_contributions.pb")
diff --git a/android/config.go b/android/config.go
index 10e43ce..e519760 100644
--- a/android/config.go
+++ b/android/config.go
@@ -2105,3 +2105,10 @@
 func (c *config) BoardAvbSystemAddHashtreeFooterArgs() []string {
 	return c.productVariables.BoardAvbSystemAddHashtreeFooterArgs
 }
+
+// Returns true if RELEASE_INSTALL_APEX_SYSTEMSERVER_DEXPREOPT_SAME_PARTITION is set to true.
+// If true, dexpreopt files of apex system server jars will be installed in the same partition as the parent apex.
+// If false, all these files will be installed in /system partition.
+func (c Config) InstallApexSystemServerDexpreoptSamePartition() bool {
+	return c.config.productVariables.GetBuildFlagBool("RELEASE_INSTALL_APEX_SYSTEMSERVER_DEXPREOPT_SAME_PARTITION")
+}
diff --git a/android/module.go b/android/module.go
index 20caae2..a1a9a4a 100644
--- a/android/module.go
+++ b/android/module.go
@@ -81,6 +81,7 @@
 	InstallInOdm() bool
 	InstallInProduct() bool
 	InstallInVendor() bool
+	InstallInSystemExt() bool
 	InstallForceOS() (*OsType, *ArchType)
 	PartitionTag(DeviceConfig) string
 	HideFromMake()
@@ -1514,6 +1515,10 @@
 	return Bool(m.commonProperties.Vendor) || Bool(m.commonProperties.Soc_specific) || Bool(m.commonProperties.Proprietary)
 }
 
+func (m *ModuleBase) InstallInSystemExt() bool {
+	return Bool(m.commonProperties.System_ext_specific)
+}
+
 func (m *ModuleBase) InstallInRoot() bool {
 	return false
 }
diff --git a/android/mutator.go b/android/mutator.go
index a8b5c7d..8265458 100644
--- a/android/mutator.go
+++ b/android/mutator.go
@@ -32,11 +32,11 @@
 // collateGloballyRegisteredMutators constructs the list of mutators that have been registered
 // with the InitRegistrationContext and will be used at runtime.
 func collateGloballyRegisteredMutators() sortableComponents {
-	return collateRegisteredMutators(preArch, preDeps, postDeps, finalDeps)
+	return collateRegisteredMutators(preArch, preDeps, postDeps, postApex, finalDeps)
 }
 
 // collateRegisteredMutators constructs a single list of mutators from the separate lists.
-func collateRegisteredMutators(preArch, preDeps, postDeps, finalDeps []RegisterMutatorFunc) sortableComponents {
+func collateRegisteredMutators(preArch, preDeps, postDeps, postApex, finalDeps []RegisterMutatorFunc) sortableComponents {
 	mctx := &registerMutatorsContext{}
 
 	register := func(funcs []RegisterMutatorFunc) {
@@ -53,6 +53,8 @@
 
 	register(postDeps)
 
+	register(postApex)
+
 	mctx.finalPhase = true
 	register(finalDeps)
 
@@ -166,6 +168,8 @@
 	RegisterOverridePostDepsMutators,
 }
 
+var postApex = []RegisterMutatorFunc{}
+
 var finalDeps = []RegisterMutatorFunc{}
 
 func PreArchMutators(f RegisterMutatorFunc) {
@@ -180,6 +184,10 @@
 	postDeps = append(postDeps, f)
 }
 
+func PostApexMutators(f RegisterMutatorFunc) {
+	postApex = append(postApex, f)
+}
+
 func FinalDepsMutators(f RegisterMutatorFunc) {
 	finalDeps = append(finalDeps, f)
 }
diff --git a/android/register.go b/android/register.go
index eb6a35e..2ce6025 100644
--- a/android/register.go
+++ b/android/register.go
@@ -235,6 +235,7 @@
 
 	PreDepsMutators(f RegisterMutatorFunc)
 	PostDepsMutators(f RegisterMutatorFunc)
+	PostApexMutators(f RegisterMutatorFunc)
 	FinalDepsMutators(f RegisterMutatorFunc)
 }
 
@@ -326,6 +327,10 @@
 	PostDepsMutators(f)
 }
 
+func (ctx *initRegistrationContext) PostApexMutators(f RegisterMutatorFunc) {
+	PostApexMutators(f)
+}
+
 func (ctx *initRegistrationContext) FinalDepsMutators(f RegisterMutatorFunc) {
 	FinalDepsMutators(f)
 }
diff --git a/android/testing.go b/android/testing.go
index 196b22e..7440869 100644
--- a/android/testing.go
+++ b/android/testing.go
@@ -197,8 +197,8 @@
 
 type TestContext struct {
 	*Context
-	preArch, preDeps, postDeps, finalDeps []RegisterMutatorFunc
-	NameResolver                          *NameResolver
+	preArch, preDeps, postDeps, postApex, finalDeps []RegisterMutatorFunc
+	NameResolver                                    *NameResolver
 
 	// The list of singletons registered for the test.
 	singletons sortableComponents
@@ -229,6 +229,10 @@
 	ctx.postDeps = append(ctx.postDeps, f)
 }
 
+func (ctx *TestContext) PostApexMutators(f RegisterMutatorFunc) {
+	ctx.postApex = append(ctx.postApex, f)
+}
+
 func (ctx *TestContext) FinalDepsMutators(f RegisterMutatorFunc) {
 	ctx.finalDeps = append(ctx.finalDeps, f)
 }
@@ -449,7 +453,7 @@
 func (ctx *TestContext) Register() {
 	globalOrder := globallyRegisteredComponentsOrder()
 
-	mutators := collateRegisteredMutators(ctx.preArch, ctx.preDeps, ctx.postDeps, ctx.finalDeps)
+	mutators := collateRegisteredMutators(ctx.preArch, ctx.preDeps, ctx.postDeps, ctx.postApex, ctx.finalDeps)
 	// Ensure that the mutators used in the test are in the same order as they are used at runtime.
 	globalOrder.mutatorOrder.enforceOrdering(mutators)
 	mutators.registerAll(ctx.Context)
diff --git a/cc/Android.bp b/cc/Android.bp
index 3688c8a..88a793c 100644
--- a/cc/Android.bp
+++ b/cc/Android.bp
@@ -102,6 +102,7 @@
         "orderfile_test.go",
         "prebuilt_test.go",
         "proto_test.go",
+        "sabi_test.go",
         "sanitize_test.go",
         "sdk_test.go",
         "test_data_test.go",
diff --git a/cc/cc.go b/cc/cc.go
index a8ff474..1b7624d 100644
--- a/cc/cc.go
+++ b/cc/cc.go
@@ -76,9 +76,9 @@
 		ctx.BottomUp("double_loadable", checkDoubleLoadableLibraries).Parallel()
 	})
 
-	ctx.FinalDepsMutators(func(ctx android.RegisterMutatorsContext) {
+	ctx.PostApexMutators(func(ctx android.RegisterMutatorsContext) {
 		// sabi mutator needs to be run after apex mutator finishes.
-		ctx.TopDown("sabi_deps", sabiDepsMutator)
+		ctx.Transition("sabi", &sabiTransitionMutator{})
 	})
 
 	ctx.RegisterParallelSingletonType("kythe_extract_all", kytheExtractAllFactory)
diff --git a/cc/compiler.go b/cc/compiler.go
index a6f623f..022b712 100644
--- a/cc/compiler.go
+++ b/cc/compiler.go
@@ -228,9 +228,6 @@
 		Static *bool `android:"arch_variant"`
 	} `android:"arch_variant"`
 
-	// Stores the original list of source files before being cleared by library reuse
-	OriginalSrcs proptools.Configurable[[]string] `blueprint:"mutated"`
-
 	// Build and link with OpenMP
 	Openmp *bool `android:"arch_variant"`
 }
@@ -363,10 +360,20 @@
 	tc := ctx.toolchain()
 	modulePath := ctx.ModuleDir()
 
-	srcs := compiler.Properties.Srcs.GetOrDefault(ctx, nil)
-	exclude_srcs := compiler.Properties.Exclude_srcs.GetOrDefault(ctx, nil)
-	compiler.srcsBeforeGen = android.PathsForModuleSrcExcludes(ctx, srcs, exclude_srcs)
-	compiler.srcsBeforeGen = append(compiler.srcsBeforeGen, deps.GeneratedSources...)
+	reuseObjs := false
+	if len(ctx.GetDirectDepsWithTag(reuseObjTag)) > 0 {
+		reuseObjs = true
+	}
+
+	// If a reuseObjTag dependency exists then this module is reusing the objects (generally the shared variant
+	// reusing objects from the static variant), and doesn't need to compile any sources of its own.
+	var srcs []string
+	if !reuseObjs {
+		srcs = compiler.Properties.Srcs.GetOrDefault(ctx, nil)
+		exclude_srcs := compiler.Properties.Exclude_srcs.GetOrDefault(ctx, nil)
+		compiler.srcsBeforeGen = android.PathsForModuleSrcExcludes(ctx, srcs, exclude_srcs)
+		compiler.srcsBeforeGen = append(compiler.srcsBeforeGen, deps.GeneratedSources...)
+	}
 
 	cflags := compiler.Properties.Cflags.GetOrDefault(ctx, nil)
 	cppflags := compiler.Properties.Cppflags.GetOrDefault(ctx, nil)
@@ -721,11 +728,6 @@
 			return true
 		}
 	}
-	for _, src := range compiler.Properties.OriginalSrcs.GetOrDefault(ctx, nil) {
-		if filepath.Ext(src) == ext {
-			return true
-		}
-	}
 
 	return false
 }
diff --git a/cc/library.go b/cc/library.go
index 3833b98..988a7fa 100644
--- a/cc/library.go
+++ b/cc/library.go
@@ -548,8 +548,7 @@
 	return flags
 }
 
-func (library *libraryDecorator) getHeaderAbiCheckerProperties(ctx android.BaseModuleContext) headerAbiCheckerProperties {
-	m := ctx.Module().(*Module)
+func (library *libraryDecorator) getHeaderAbiCheckerProperties(m *Module) headerAbiCheckerProperties {
 	variantProps := &library.Properties.Target.Platform.Header_abi_checker
 	if m.InVendor() {
 		variantProps = &library.Properties.Target.Vendor.Header_abi_checker
@@ -559,7 +558,7 @@
 	props := library.Properties.Header_abi_checker
 	err := proptools.AppendProperties(&props, variantProps, nil)
 	if err != nil {
-		ctx.ModuleErrorf("Cannot merge headerAbiCheckerProperties: %s", err.Error())
+		panic(fmt.Errorf("Cannot merge headerAbiCheckerProperties: %s", err.Error()))
 	}
 	return props
 }
@@ -718,7 +717,7 @@
 	setShared()
 
 	// Gets the ABI properties for vendor, product, or platform variant
-	getHeaderAbiCheckerProperties(ctx android.BaseModuleContext) headerAbiCheckerProperties
+	getHeaderAbiCheckerProperties(m *Module) headerAbiCheckerProperties
 
 	// Write LOCAL_ADDITIONAL_DEPENDENCIES for ABI diff
 	androidMkWriteAdditionalDependenciesForSourceAbiDiff(w io.Writer)
@@ -1365,7 +1364,7 @@
 	sourceVersion, errorMessage string) {
 
 	extraFlags := []string{"-target-version", sourceVersion}
-	headerAbiChecker := library.getHeaderAbiCheckerProperties(ctx)
+	headerAbiChecker := library.getHeaderAbiCheckerProperties(ctx.Module().(*Module))
 	if Bool(headerAbiChecker.Check_all_apis) {
 		extraFlags = append(extraFlags, "-check-all-apis")
 	} else {
@@ -1437,7 +1436,7 @@
 func (library *libraryDecorator) linkSAbiDumpFiles(ctx ModuleContext, deps PathDeps, objs Objects, fileName string, soFile android.Path) {
 	if library.sabi.shouldCreateSourceAbiDump() {
 		exportedIncludeDirs := library.exportedIncludeDirsForAbiCheck(ctx)
-		headerAbiChecker := library.getHeaderAbiCheckerProperties(ctx)
+		headerAbiChecker := library.getHeaderAbiCheckerProperties(ctx.Module().(*Module))
 		currSdkVersion := currRefAbiDumpSdkVersion(ctx)
 		currVendorVersion := ctx.Config().VendorApiLevel()
 
@@ -1451,7 +1450,7 @@
 			[]string{} /* includeSymbolTags */, currSdkVersion, false /* isLlndk */)
 
 		var llndkDump, apexVariantDump android.Path
-		tags := classifySourceAbiDump(ctx)
+		tags := classifySourceAbiDump(ctx.Module().(*Module))
 		optInTags := []lsdumpTag{}
 		for _, tag := range tags {
 			if tag == llndkLsdumpTag && currVendorVersion != "" {
@@ -1868,7 +1867,7 @@
 }
 
 func (library *libraryDecorator) symbolFileForAbiCheck(ctx ModuleContext) *string {
-	if props := library.getHeaderAbiCheckerProperties(ctx); props.Symbol_file != nil {
+	if props := library.getHeaderAbiCheckerProperties(ctx.Module().(*Module)); props.Symbol_file != nil {
 		return props.Symbol_file
 	}
 	if library.hasStubsVariants() && library.Properties.Stubs.Symbol_file != nil {
@@ -2071,12 +2070,7 @@
 			sharedCompiler.StaticProperties.Static.System_shared_libs == nil &&
 			sharedCompiler.SharedProperties.Shared.System_shared_libs == nil {
 
-			// TODO: namespaces?
 			ctx.AddVariationDependencies([]blueprint.Variation{{"link", "static"}}, reuseObjTag, ctx.ModuleName())
-			sharedCompiler.baseCompiler.Properties.OriginalSrcs =
-				sharedCompiler.baseCompiler.Properties.Srcs
-			sharedCompiler.baseCompiler.Properties.Srcs = proptools.NewConfigurable[[]string](nil, nil)
-			sharedCompiler.baseCompiler.Properties.Generated_sources = nil
 		}
 
 		// This dep is just to reference static variant from shared variant
diff --git a/cc/sabi.go b/cc/sabi.go
index 64eab41..2caf0d4 100644
--- a/cc/sabi.go
+++ b/cc/sabi.go
@@ -84,8 +84,8 @@
 
 type SAbiProperties struct {
 	// Whether ABI dump should be created for this module.
-	// Set by `sabiDepsMutator` if this module is a shared library that needs ABI check, or a static
-	// library that is depended on by an ABI checked library.
+	// Set by `sabiTransitionMutator` if this module is a shared library that needs ABI check,
+	// or a static library that is depended on by an ABI checked library.
 	ShouldCreateSourceAbiDump bool `blueprint:"mutated"`
 
 	// Include directories that may contain ABI information exported by a library.
@@ -121,10 +121,9 @@
 }
 
 // Returns a slice of strings that represent the ABI dumps generated for this module.
-func classifySourceAbiDump(ctx android.BaseModuleContext) []lsdumpTag {
+func classifySourceAbiDump(m *Module) []lsdumpTag {
 	result := []lsdumpTag{}
-	m := ctx.Module().(*Module)
-	headerAbiChecker := m.library.getHeaderAbiCheckerProperties(ctx)
+	headerAbiChecker := m.library.getHeaderAbiCheckerProperties(m)
 	if headerAbiChecker.explicitlyDisabled() {
 		return result
 	}
@@ -149,24 +148,37 @@
 	return result
 }
 
-// Called from sabiDepsMutator to check whether ABI dumps should be created for this module.
+type shouldCreateAbiDumpContext interface {
+	android.ModuleProviderContext
+	Module() android.Module
+	Config() android.Config
+}
+
+var _ shouldCreateAbiDumpContext = android.ModuleContext(nil)
+var _ shouldCreateAbiDumpContext = android.OutgoingTransitionContext(nil)
+
+// Called from sabiTransitionMutator to check whether ABI dumps should be created for this module.
 // ctx should be wrapping a native library type module.
-func shouldCreateSourceAbiDumpForLibrary(ctx android.BaseModuleContext) bool {
-	// Only generate ABI dump for device modules.
-	if !ctx.Device() {
+func shouldCreateSourceAbiDumpForLibrary(ctx shouldCreateAbiDumpContext) bool {
+	m, ok := ctx.Module().(*Module)
+	if !ok {
 		return false
 	}
 
-	m := ctx.Module().(*Module)
+	// Only generate ABI dump for device modules.
+	if !m.Device() {
+		return false
+	}
 
 	// Only create ABI dump for native library module types.
 	if m.library == nil {
 		return false
 	}
 
-	// Create ABI dump for static libraries only if they are dependencies of ABI checked libraries.
+	// Don't create ABI dump for static libraries
+	// The sabi variant will be propagated to dependencies of ABI checked libraries.
 	if m.library.static() {
-		return m.sabi.shouldCreateSourceAbiDump()
+		return false
 	}
 
 	// Module is shared library type.
@@ -215,31 +227,64 @@
 			return false
 		}
 	}
-	return len(classifySourceAbiDump(ctx)) > 0
+	return len(classifySourceAbiDump(m)) > 0
 }
 
 // Mark the direct and transitive dependencies of libraries that need ABI check, so that ABI dumps
 // of their dependencies would be generated.
-func sabiDepsMutator(mctx android.TopDownMutatorContext) {
+type sabiTransitionMutator struct{}
+
+func (s *sabiTransitionMutator) Split(ctx android.BaseModuleContext) []string {
+	return []string{""}
+}
+
+func (s *sabiTransitionMutator) OutgoingTransition(ctx android.OutgoingTransitionContext, sourceVariation string) string {
 	// Escape hatch to not check any ABI dump.
-	if mctx.Config().IsEnvTrue("SKIP_ABI_CHECKS") {
-		return
+	if ctx.Config().IsEnvTrue("SKIP_ABI_CHECKS") {
+		return ""
 	}
+
 	// Only create ABI dump for native shared libraries and their static library dependencies.
-	if m, ok := mctx.Module().(*Module); ok && m.sabi != nil {
-		if shouldCreateSourceAbiDumpForLibrary(mctx) {
-			// Mark this module so that .sdump / .lsdump for this library can be generated.
+	if m, ok := ctx.Module().(*Module); ok && m.sabi != nil {
+		if shouldCreateSourceAbiDumpForLibrary(ctx) {
+			if IsStaticDepTag(ctx.DepTag()) || ctx.DepTag() == reuseObjTag {
+				return "sabi"
+			}
+		} else if sourceVariation == "sabi" {
+			if IsWholeStaticLib(ctx.DepTag()) || ctx.DepTag() == reuseObjTag {
+				return "sabi"
+			}
+		}
+	}
+
+	return ""
+}
+
+func (s *sabiTransitionMutator) IncomingTransition(ctx android.IncomingTransitionContext, incomingVariation string) string {
+	if incomingVariation == "" {
+		return ""
+	}
+
+	if incomingVariation == "sabi" {
+		if m, ok := ctx.Module().(*Module); ok && m.sabi != nil {
+			return "sabi"
+		}
+	}
+
+	return ""
+}
+
+func (s *sabiTransitionMutator) Mutate(ctx android.BottomUpMutatorContext, variation string) {
+	if m, ok := ctx.Module().(*Module); ok && m.sabi != nil {
+		if variation == "sabi" {
 			m.sabi.Properties.ShouldCreateSourceAbiDump = true
-			// Mark all of its static library dependencies.
-			mctx.VisitDirectDeps(func(child android.Module) {
-				depTag := mctx.OtherModuleDependencyTag(child)
-				if IsStaticDepTag(depTag) || depTag == reuseObjTag {
-					if c, ok := child.(*Module); ok && c.sabi != nil {
-						// Mark this module so that .sdump for this static library can be generated.
-						c.sabi.Properties.ShouldCreateSourceAbiDump = true
-					}
-				}
-			})
+			m.HideFromMake()
+			m.Properties.PreventInstall = true
+		} else if shouldCreateSourceAbiDumpForLibrary(ctx) {
+			// Escape hatch to not check any ABI dump.
+			if !ctx.Config().IsEnvTrue("SKIP_ABI_CHECKS") {
+				m.sabi.Properties.ShouldCreateSourceAbiDump = true
+			}
 		}
 	}
 }
diff --git a/cc/sabi_test.go b/cc/sabi_test.go
new file mode 100644
index 0000000..6b8cc17
--- /dev/null
+++ b/cc/sabi_test.go
@@ -0,0 +1,66 @@
+// Copyright 2024 Google Inc. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package cc
+
+import (
+	"android/soong/android"
+	"testing"
+)
+
+func TestSabi(t *testing.T) {
+	bp := `
+		cc_library {
+			name: "libsabi",
+			srcs: ["sabi.cpp"],
+			static_libs: ["libdirect"],
+			header_abi_checker: {
+				enabled: true,
+				symbol_file: "libsabi.map.txt",
+                ref_dump_dirs: ["abi-dumps"],
+			},
+		}
+
+		cc_library {
+			name: "libdirect",
+			srcs: ["direct.cpp"],
+			whole_static_libs: ["libtransitive"],
+		}
+
+		cc_library {
+			name: "libtransitive",
+			srcs: ["transitive.cpp"],
+		}
+	`
+
+	result := android.GroupFixturePreparers(
+		PrepareForTestWithCcDefaultModules,
+	).RunTestWithBp(t, bp)
+
+	libsabiStatic := result.ModuleForTests("libsabi", "android_arm64_armv8-a_static_sabi")
+	sabiObjSDump := libsabiStatic.Output("obj/sabi.sdump")
+
+	libDirect := result.ModuleForTests("libdirect", "android_arm64_armv8-a_static_sabi")
+	directObjSDump := libDirect.Output("obj/direct.sdump")
+
+	libTransitive := result.ModuleForTests("libtransitive", "android_arm64_armv8-a_static_sabi")
+	transitiveObjSDump := libTransitive.Output("obj/transitive.sdump")
+
+	libsabiShared := result.ModuleForTests("libsabi", "android_arm64_armv8-a_shared")
+	sabiLink := libsabiShared.Rule("sAbiLink")
+
+	android.AssertStringListContains(t, "sabi link inputs", sabiLink.Inputs.Strings(), sabiObjSDump.Output.String())
+	android.AssertStringListContains(t, "sabi link inputs", sabiLink.Inputs.Strings(), directObjSDump.Output.String())
+	android.AssertStringListContains(t, "sabi link inputs", sabiLink.Inputs.Strings(), transitiveObjSDump.Output.String())
+}
diff --git a/dexpreopt/config.go b/dexpreopt/config.go
index fe6317c..84d4f10 100644
--- a/dexpreopt/config.go
+++ b/dexpreopt/config.go
@@ -191,6 +191,10 @@
 	ForceCreateAppImage bool
 
 	PresignedPrebuilt bool
+
+	// ApexPartition is the partition in which the dexpreopt files of apex system server jars (if any) are installed.
+	// This is a noop unless the module is apex system server jar.
+	ApexPartition string
 }
 
 type globalSoongConfigSingleton struct{}
diff --git a/dexpreopt/dexpreopt.go b/dexpreopt/dexpreopt.go
index 5616483..7a39fa1 100644
--- a/dexpreopt/dexpreopt.go
+++ b/dexpreopt/dexpreopt.go
@@ -219,9 +219,9 @@
 }
 
 // Returns the location to the odex file for the dex file at `path`.
-func ToOdexPath(path string, arch android.ArchType) string {
+func ToOdexPath(path string, arch android.ArchType, partition string) string {
 	if strings.HasPrefix(path, "/apex/") {
-		return filepath.Join("/system/framework/oat", arch.String(),
+		return filepath.Join(partition, "framework/oat", arch.String(),
 			strings.ReplaceAll(path[1:], "/", "@")+"@classes.odex")
 	}
 
@@ -245,7 +245,7 @@
 
 	odexPath := module.BuildPath.InSameDir(ctx, "oat", arch.String(), pathtools.ReplaceExtension(base, "odex"))
 	odexSymbolsPath := odexPath.ReplaceExtension(ctx, "symbols.odex")
-	odexInstallPath := ToOdexPath(module.DexLocation, arch)
+	odexInstallPath := ToOdexPath(module.DexLocation, arch, module.ApexPartition)
 	if odexOnSystemOther(module, global) {
 		odexInstallPath = filepath.Join(SystemOtherPartition, odexInstallPath)
 	}
diff --git a/dexpreopt/dexpreopt_test.go b/dexpreopt/dexpreopt_test.go
index 6f7d3bb..7b0f51f 100644
--- a/dexpreopt/dexpreopt_test.go
+++ b/dexpreopt/dexpreopt_test.go
@@ -42,12 +42,14 @@
 }
 
 func testApexModuleConfig(ctx android.PathContext, name, apexName string) *ModuleConfig {
-	return createTestModuleConfig(
+	ret := createTestModuleConfig(
 		name,
 		fmt.Sprintf("/apex/%s/javalib/%s.jar", apexName, name),
 		android.PathForOutput(ctx, fmt.Sprintf("%s/dexpreopt/%s.jar", name, name)),
 		android.PathForOutput(ctx, fmt.Sprintf("%s/aligned/%s.jar", name, name)),
 		android.PathForOutput(ctx, fmt.Sprintf("%s/enforce_uses_libraries.status", name)))
+	ret.ApexPartition = "/system"
+	return ret
 }
 
 func testPlatformSystemServerModuleConfig(ctx android.PathContext, name string) *ModuleConfig {
@@ -221,6 +223,49 @@
 	DexpreoptRunningInSoong = oldDexpreoptRunningInSoong
 }
 
+// Same as `TestDexPreoptApexSystemServerJars`, but the apex jar is in /system_ext
+func TestDexPreoptApexSystemServerJarsSystemExt(t *testing.T) {
+	// modify the global variable for test
+	var oldDexpreoptRunningInSoong = DexpreoptRunningInSoong
+	DexpreoptRunningInSoong = true
+
+	// test begin
+	config := android.TestConfig("out", nil, "", nil)
+	ctx := android.BuilderContextForTesting(config)
+	globalSoong := globalSoongConfigForTests(ctx)
+	global := GlobalConfigForTests(ctx)
+	module := testApexModuleConfig(ctx, "service-A", "com.android.apex1")
+	module.ApexPartition = "/system_ext"
+	productPackages := android.PathForTesting("product_packages.txt")
+
+	global.ApexSystemServerJars = android.CreateTestConfiguredJarList(
+		[]string{"com.android.apex1:service-A"})
+
+	rule, err := GenerateDexpreoptRule(ctx, globalSoong, global, module, productPackages, true)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	wantInstalls := android.RuleBuilderInstalls{
+		{android.PathForOutput(ctx, "service-A/dexpreopt/oat/arm/javalib.odex"), "/system_ext/framework/oat/arm/apex@com.android.apex1@javalib@service-A.jar@classes.odex"},
+		{android.PathForOutput(ctx, "service-A/dexpreopt/oat/arm/javalib.vdex"), "/system_ext/framework/oat/arm/apex@com.android.apex1@javalib@service-A.jar@classes.vdex"},
+	}
+
+	android.AssertStringEquals(t, "installs", wantInstalls.String(), rule.Installs().String())
+
+	android.AssertStringListContains(t, "apex sscp jar copy", rule.Outputs().Strings(), "out/soong/system_server_dexjars/service-A.jar")
+
+	// rule with apex sscp cp as false
+	rule, err = GenerateDexpreoptRule(ctx, globalSoong, global, module, productPackages, false)
+	if err != nil {
+		t.Fatal(err)
+	}
+	android.AssertStringListDoesNotContain(t, "apex sscp jar copy", rule.Outputs().Strings(), "out/soong/system_server_dexjars/service-A.jar")
+
+	// cleanup the global variable for test
+	DexpreoptRunningInSoong = oldDexpreoptRunningInSoong
+}
+
 func TestDexPreoptStandaloneSystemServerJars(t *testing.T) {
 	config := android.TestConfig("out", nil, "", nil)
 	ctx := android.BuilderContextForTesting(config)
diff --git a/filesystem/filesystem.go b/filesystem/filesystem.go
index a26fac7..09d8fba 100644
--- a/filesystem/filesystem.go
+++ b/filesystem/filesystem.go
@@ -35,7 +35,7 @@
 }
 
 func registerBuildComponents(ctx android.RegistrationContext) {
-	ctx.RegisterModuleType("android_filesystem", filesystemFactory)
+	ctx.RegisterModuleType("android_filesystem", FilesystemFactory)
 	ctx.RegisterModuleType("android_filesystem_defaults", filesystemDefaultsFactory)
 	ctx.RegisterModuleType("android_system_image", SystemImageFactory)
 	ctx.RegisterModuleType("avb_add_hash_footer", avbAddHashFooterFactory)
@@ -137,6 +137,12 @@
 	Gen_aconfig_flags_pb *bool
 
 	Fsverity fsverityProperties
+
+	// If this property is set to true, the filesystem will call ctx.UncheckedModule(), causing
+	// it to not be built on checkbuilds. Used for the automatic migration from make to soong
+	// build modules, where we want to emit some not-yet-working filesystems and we don't want them
+	// to be built.
+	Unchecked_module *bool `blueprint:"mutated"`
 }
 
 // android_filesystem packages a set of modules and their transitive dependencies into a filesystem
@@ -144,7 +150,7 @@
 // modules in the filesystem image are built for the target device (i.e. Android, not Linux host).
 // The modules are placed in the filesystem image just like they are installed to the ordinary
 // partitions like system.img. For example, cc_library modules are placed under ./lib[64] directory.
-func filesystemFactory() android.Module {
+func FilesystemFactory() android.Module {
 	module := &filesystem{}
 	module.filterPackagingSpec = module.filterInstallablePackagingSpec
 	initFilesystemModule(module, module)
@@ -177,6 +183,13 @@
 	unknown
 )
 
+type FilesystemInfo struct {
+	// A text file containing the list of paths installed on the partition.
+	FileListFile android.Path
+}
+
+var FilesystemProvider = blueprint.NewProvider[FilesystemInfo]()
+
 func (f *filesystem) fsType(ctx android.ModuleContext) fsType {
 	typeStr := proptools.StringDefault(f.properties.Type, "ext4")
 	switch typeStr {
@@ -227,6 +240,14 @@
 
 	f.fileListFile = android.PathForModuleOut(ctx, "fileList").OutputPath
 	android.WriteFileRule(ctx, f.fileListFile, f.installedFilesList())
+
+	android.SetProvider(ctx, FilesystemProvider, FilesystemInfo{
+		FileListFile: f.fileListFile,
+	})
+
+	if proptools.Bool(f.properties.Unchecked_module) {
+		ctx.UncheckedModule()
+	}
 }
 
 func (f *filesystem) appendToEntry(ctx android.ModuleContext, installedFile android.OutputPath) {
diff --git a/fsgen/Android.bp b/fsgen/Android.bp
index aa8881f..9fa9557 100644
--- a/fsgen/Android.bp
+++ b/fsgen/Android.bp
@@ -19,3 +19,7 @@
     ],
     pluginFor: ["soong_build"],
 }
+
+soong_filesystem_creator {
+    name: "soong_filesystem_creator",
+}
diff --git a/fsgen/filesystem_creator.go b/fsgen/filesystem_creator.go
index ca948f4..eb4f318 100644
--- a/fsgen/filesystem_creator.go
+++ b/fsgen/filesystem_creator.go
@@ -17,12 +17,16 @@
 import (
 	"android/soong/android"
 	"android/soong/filesystem"
+	"crypto/sha256"
 	"fmt"
 	"strconv"
 
+	"github.com/google/blueprint"
 	"github.com/google/blueprint/proptools"
 )
 
+var pctx = android.NewPackageContext("android/soong/fsgen")
+
 func init() {
 	registerBuildComponents(android.InitRegistrationContext)
 }
@@ -31,14 +35,22 @@
 	ctx.RegisterModuleType("soong_filesystem_creator", filesystemCreatorFactory)
 }
 
+type filesystemCreatorProps struct {
+	Generated_partition_types   []string `blueprint:"mutated"`
+	Unsupported_partition_types []string `blueprint:"mutated"`
+}
+
 type filesystemCreator struct {
 	android.ModuleBase
+
+	properties filesystemCreatorProps
 }
 
 func filesystemCreatorFactory() android.Module {
 	module := &filesystemCreator{}
 
 	android.InitAndroidModule(module)
+	module.AddProperties(&module.properties)
 	android.AddLoadHook(module, func(ctx android.LoadHookContext) {
 		module.createInternalModules(ctx)
 	})
@@ -47,36 +59,62 @@
 }
 
 func (f *filesystemCreator) createInternalModules(ctx android.LoadHookContext) {
-	f.createSystemImage(ctx)
+	for _, partitionType := range []string{"system"} {
+		if f.createPartition(ctx, partitionType) {
+			f.properties.Generated_partition_types = append(f.properties.Generated_partition_types, partitionType)
+		} else {
+			f.properties.Unsupported_partition_types = append(f.properties.Unsupported_partition_types, partitionType)
+		}
+	}
 }
 
-func (f *filesystemCreator) createSystemImage(ctx android.LoadHookContext) {
+func (f *filesystemCreator) generatedModuleNameForPartition(cfg android.Config, partitionType string) string {
+	prefix := "soong"
+	if cfg.HasDeviceProduct() {
+		prefix = cfg.DeviceProduct()
+	}
+	return fmt.Sprintf("%s_generated_%s_image", prefix, partitionType)
+}
+
+// Creates a soong module to build the given partition. Returns false if we can't support building
+// it.
+func (f *filesystemCreator) createPartition(ctx android.LoadHookContext, partitionType string) bool {
 	baseProps := &struct {
 		Name *string
 	}{
-		Name: proptools.StringPtr(fmt.Sprintf("%s_generated_system_image", ctx.Config().DeviceProduct())),
+		Name: proptools.StringPtr(f.generatedModuleNameForPartition(ctx.Config(), partitionType)),
 	}
 
-	fsProps := &(filesystem.FilesystemProperties{})
+	fsProps := &filesystem.FilesystemProperties{}
+
+	// Don't build this module on checkbuilds, the soong-built partitions are still in-progress
+	// and sometimes don't build.
+	fsProps.Unchecked_module = proptools.BoolPtr(true)
+
 	partitionVars := ctx.Config().ProductVariables().PartitionVarsForSoongMigrationOnlyDoNotUse
-	systemPartitionVars := partitionVars.PartitionQualifiedVariables["system"]
+	specificPartitionVars := partitionVars.PartitionQualifiedVariables[partitionType]
 
 	// BOARD_AVB_ENABLE
 	fsProps.Use_avb = proptools.BoolPtr(partitionVars.BoardAvbEnable)
 	// BOARD_AVB_KEY_PATH
-	fsProps.Avb_private_key = proptools.StringPtr(systemPartitionVars.BoardAvbKeyPath)
+	fsProps.Avb_private_key = proptools.StringPtr(specificPartitionVars.BoardAvbKeyPath)
 	// BOARD_AVB_ALGORITHM
-	fsProps.Avb_algorithm = proptools.StringPtr(systemPartitionVars.BoardAvbAlgorithm)
+	fsProps.Avb_algorithm = proptools.StringPtr(specificPartitionVars.BoardAvbAlgorithm)
 	// BOARD_AVB_SYSTEM_ROLLBACK_INDEX
-	if rollbackIndex, err := strconv.ParseInt(systemPartitionVars.BoardAvbRollbackIndex, 10, 64); err == nil {
+	if rollbackIndex, err := strconv.ParseInt(specificPartitionVars.BoardAvbRollbackIndex, 10, 64); err == nil {
 		fsProps.Rollback_index = proptools.Int64Ptr(rollbackIndex)
 	}
 
-	fsProps.Partition_name = proptools.StringPtr("system")
+	fsProps.Partition_name = proptools.StringPtr(partitionType)
 	// BOARD_SYSTEMIMAGE_FILE_SYSTEM_TYPE
-	fsProps.Type = proptools.StringPtr(systemPartitionVars.BoardFileSystemType)
+	fsProps.Type = proptools.StringPtr(specificPartitionVars.BoardFileSystemType)
+	if *fsProps.Type != "ext4" {
+		// Currently the android_filesystem module type only supports ext4:
+		// https://cs.android.com/android/platform/superproject/main/+/main:build/soong/filesystem/filesystem.go;l=416;drc=98047cfd07944b297a12d173453bc984806760d2
+		return false
+	}
 
-	fsProps.Base_dir = proptools.StringPtr("system")
+	fsProps.Base_dir = proptools.StringPtr(partitionType)
 
 	fsProps.Gen_aconfig_flags_pb = proptools.BoolPtr(true)
 
@@ -103,9 +141,77 @@
 	// - filesystemProperties.Build_logtags
 	// - filesystemProperties.Fsverity.Libs
 	// - systemImageProperties.Linker_config_src
-	ctx.CreateModule(filesystem.SystemImageFactory, baseProps, fsProps)
+	var module android.Module
+	if partitionType == "system" {
+		module = ctx.CreateModule(filesystem.SystemImageFactory, baseProps, fsProps)
+	} else {
+		module = ctx.CreateModule(filesystem.FilesystemFactory, baseProps, fsProps)
+	}
+	module.HideFromMake()
+	return true
+}
+
+func (f *filesystemCreator) createDiffTest(ctx android.ModuleContext, partitionType string) android.Path {
+	partitionModuleName := f.generatedModuleNameForPartition(ctx.Config(), partitionType)
+	systemImage := ctx.GetDirectDepWithTag(partitionModuleName, generatedFilesystemDepTag)
+	filesystemInfo, ok := android.OtherModuleProvider(ctx, systemImage, filesystem.FilesystemProvider)
+	if !ok {
+		ctx.ModuleErrorf("Expected module %s to provide FileysystemInfo", partitionModuleName)
+	}
+	makeFileList := android.PathForArbitraryOutput(ctx, fmt.Sprintf("target/product/%s/obj/PACKAGING/%s_intermediates/file_list.txt", ctx.Config().DeviceName(), partitionType))
+	// For now, don't allowlist anything. The test will fail, but that's fine in the current
+	// early stages where we're just figuring out what we need
+	emptyAllowlistFile := android.PathForModuleOut(ctx, "allowlist_%s.txt", partitionModuleName)
+	android.WriteFileRule(ctx, emptyAllowlistFile, "")
+	diffTestResultFile := android.PathForModuleOut(ctx, "diff_test_%s.txt", partitionModuleName)
+
+	builder := android.NewRuleBuilder(pctx, ctx)
+	builder.Command().BuiltTool("file_list_diff").
+		Input(makeFileList).
+		Input(filesystemInfo.FileListFile).
+		Input(emptyAllowlistFile).
+		Text(partitionModuleName)
+	builder.Command().Text("touch").Output(diffTestResultFile)
+	builder.Build(partitionModuleName+" diff test", partitionModuleName+" diff test")
+	return diffTestResultFile
+}
+
+func createFailingCommand(ctx android.ModuleContext, message string) android.Path {
+	hasher := sha256.New()
+	hasher.Write([]byte(message))
+	filename := fmt.Sprintf("failing_command_%x.txt", hasher.Sum(nil))
+	file := android.PathForModuleOut(ctx, filename)
+	builder := android.NewRuleBuilder(pctx, ctx)
+	builder.Command().Textf("echo %s", proptools.NinjaAndShellEscape(message))
+	builder.Command().Text("exit 1 #").Output(file)
+	builder.Build("failing command "+filename, "failing command "+filename)
+	return file
+}
+
+type systemImageDepTagType struct {
+	blueprint.BaseDependencyTag
+}
+
+var generatedFilesystemDepTag systemImageDepTagType
+
+func (f *filesystemCreator) DepsMutator(ctx android.BottomUpMutatorContext) {
+	for _, partitionType := range f.properties.Generated_partition_types {
+		ctx.AddDependency(ctx.Module(), generatedFilesystemDepTag, f.generatedModuleNameForPartition(ctx.Config(), partitionType))
+	}
 }
 
 func (f *filesystemCreator) GenerateAndroidBuildActions(ctx android.ModuleContext) {
+	if ctx.ModuleDir() != "build/soong/fsgen" {
+		ctx.ModuleErrorf("There can only be one soong_filesystem_creator in build/soong/fsgen")
+	}
+	f.HideFromMake()
 
+	var diffTestFiles []android.Path
+	for _, partitionType := range f.properties.Generated_partition_types {
+		diffTestFiles = append(diffTestFiles, f.createDiffTest(ctx, partitionType))
+	}
+	for _, partitionType := range f.properties.Unsupported_partition_types {
+		diffTestFiles = append(diffTestFiles, createFailingCommand(ctx, fmt.Sprintf("Couldn't build %s partition", partitionType)))
+	}
+	ctx.Phony("soong_generated_filesystem_tests", diffTestFiles...)
 }
diff --git a/fsgen/filesystem_creator_test.go b/fsgen/filesystem_creator_test.go
index 6bb0e77..554b66b 100644
--- a/fsgen/filesystem_creator_test.go
+++ b/fsgen/filesystem_creator_test.go
@@ -44,12 +44,13 @@
 		}),
 		android.FixtureMergeMockFs(android.MockFS{
 			"external/avb/test/data/testkey_rsa4096.pem": nil,
+			"build/soong/fsgen/Android.bp": []byte(`
+			soong_filesystem_creator {
+				name: "foo",
+			}
+			`),
 		}),
-	).RunTestWithBp(t, `
-	soong_filesystem_creator {
-		name: "foo",
-	}
-	`)
+	).RunTest(t)
 
 	fooSystem := result.ModuleForTests("test_product_generated_system_image", "android_common").Module().(interface {
 		FsProps() filesystem.FilesystemProperties
diff --git a/java/app.go b/java/app.go
index 1358faf..dd99675 100644
--- a/java/app.go
+++ b/java/app.go
@@ -172,7 +172,7 @@
 	RotationMinSdkVersion *string
 
 	// the package name of this app. The package name in the manifest file is used if one was not given.
-	Package_name *string
+	Package_name proptools.Configurable[string]
 
 	// the logging parent of this app.
 	Logging_parent *string
@@ -386,7 +386,8 @@
 	checkMinSdkVersionMts(ctx, a.MinSdkVersion(ctx))
 	applicationId := a.appTestHelperAppProperties.Manifest_values.ApplicationId
 	if applicationId != nil {
-		if a.overridableAppProperties.Package_name != nil {
+		packageName := a.overridableAppProperties.Package_name.Get(ctx)
+		if packageName.IsPresent() {
 			ctx.PropertyErrorf("manifest_values.applicationId", "property is not supported when property package_name is set.")
 		}
 		a.aapt.manifestValues.applicationId = *applicationId
@@ -586,10 +587,11 @@
 	}
 
 	manifestPackageName, overridden := ctx.DeviceConfig().OverrideManifestPackageNameFor(ctx.ModuleName())
-	if overridden || a.overridableAppProperties.Package_name != nil {
+	packageNameProp := a.overridableAppProperties.Package_name.Get(ctx)
+	if overridden || packageNameProp.IsPresent() {
 		// The product override variable has a priority over the package_name property.
 		if !overridden {
-			manifestPackageName = *a.overridableAppProperties.Package_name
+			manifestPackageName = packageNameProp.Get()
 		}
 		aaptLinkFlags = append(aaptLinkFlags, generateAaptRenamePackageFlags(manifestPackageName, a.renameResourcesPackage())...)
 		a.overriddenManifestPackageName = manifestPackageName
@@ -829,11 +831,12 @@
 		return android.PathForModuleSrc(ctx, *a.appProperties.Privapp_allowlist)
 	}
 
-	if a.overridableAppProperties.Package_name == nil {
+	packageNameProp := a.overridableAppProperties.Package_name.Get(ctx)
+	if packageNameProp.IsEmpty() {
 		ctx.PropertyErrorf("privapp_allowlist", "package_name must be set to use privapp_allowlist")
 	}
 
-	packageName := *a.overridableAppProperties.Package_name
+	packageName := packageNameProp.Get()
 	fileName := "privapp_allowlist_" + packageName + ".xml"
 	outPath := android.PathForModuleOut(ctx, fileName).OutputPath
 	ctx.Build(pctx, android.BuildParams{
@@ -1418,7 +1421,8 @@
 	}
 	applicationId := a.appTestProperties.Manifest_values.ApplicationId
 	if applicationId != nil {
-		if a.overridableAppProperties.Package_name != nil {
+		packageNameProp := a.overridableAppProperties.Package_name.Get(ctx)
+		if packageNameProp.IsPresent() {
 			ctx.PropertyErrorf("manifest_values.applicationId", "property is not supported when property package_name is set.")
 		}
 		a.aapt.manifestValues.applicationId = *applicationId
@@ -1469,10 +1473,11 @@
 		command.FlagWithArg("--test-file-name ", a.installApkName+".apk")
 	}
 
-	if a.overridableAppProperties.Package_name != nil {
+	packageNameProp := a.overridableAppProperties.Package_name.Get(ctx)
+	if packageNameProp.IsPresent() {
 		fixNeeded = true
 		command.FlagWithInput("--manifest ", a.manifestPath).
-			FlagWithArg("--package-name ", *a.overridableAppProperties.Package_name)
+			FlagWithArg("--package-name ", packageNameProp.Get())
 	}
 
 	if a.appTestProperties.Mainline_package_name != nil {
diff --git a/java/dex.go b/java/dex.go
index e16b052..a3f699b 100644
--- a/java/dex.go
+++ b/java/dex.go
@@ -133,7 +133,7 @@
 			`$d8Template${config.D8Cmd} ${config.D8Flags} $d8Flags --output $outDir --no-dex-input-jar $in && ` +
 			`$zipTemplate${config.SoongZipCmd} $zipFlags -o $outDir/classes.dex.jar -C $outDir -f "$outDir/classes*.dex" && ` +
 			`${config.MergeZipsCmd} -D -stripFile "**/*.class" $mergeZipsFlags $out $outDir/classes.dex.jar $in && ` +
-			`rm -f "$outDir/classes*.dex" "$outDir/classes.dex.jar"`,
+			`rm -f "$outDir"/classes*.dex "$outDir/classes.dex.jar"`,
 		CommandDeps: []string{
 			"${config.D8Cmd}",
 			"${config.SoongZipCmd}",
@@ -172,7 +172,7 @@
 			`rm -rf ${outUsageDir} && ` +
 			`$zipTemplate${config.SoongZipCmd} $zipFlags -o $outDir/classes.dex.jar -C $outDir -f "$outDir/classes*.dex" && ` +
 			`${config.MergeZipsCmd} -D -stripFile "**/*.class" $mergeZipsFlags $out $outDir/classes.dex.jar $in && ` +
-			`rm -f "$outDir/classes*.dex" "$outDir/classes.dex.jar"`,
+			`rm -f "$outDir"/classes*.dex "$outDir/classes.dex.jar"`,
 		Depfile: "${out}.d",
 		Deps:    blueprint.DepsGCC,
 		CommandDeps: []string{
diff --git a/java/dexpreopt.go b/java/dexpreopt.go
index 63a8634..637da36 100644
--- a/java/dexpreopt.go
+++ b/java/dexpreopt.go
@@ -494,6 +494,12 @@
 		PresignedPrebuilt: d.isPresignedPrebuilt,
 	}
 
+	if ctx.Config().InstallApexSystemServerDexpreoptSamePartition() {
+		dexpreoptConfig.ApexPartition = android.PathForModuleInstall(ctx).Partition()
+	} else {
+		dexpreoptConfig.ApexPartition = "system"
+	}
+
 	d.configPath = android.PathForModuleOut(ctx, "dexpreopt", dexJarStem, "dexpreopt.config")
 	dexpreopt.WriteModuleConfig(ctx, dexpreoptConfig, d.configPath)
 	ctx.CheckbuildFile(d.configPath)
diff --git a/java/dexpreopt_check.go b/java/dexpreopt_check.go
index 33be603..c971565 100644
--- a/java/dexpreopt_check.go
+++ b/java/dexpreopt_check.go
@@ -17,6 +17,8 @@
 import (
 	"strings"
 
+	"github.com/google/blueprint"
+
 	"android/soong/android"
 	"android/soong/dexpreopt"
 
@@ -43,16 +45,12 @@
 type dexpreoptSystemserverCheck struct {
 	android.SingletonModuleBase
 
-	// Mapping from the module name to the install paths to the compilation artifacts.
-	artifactsByModuleName map[string][]string
-
 	// The install paths to the compilation artifacts.
 	artifacts []string
 }
 
 func dexpreoptSystemserverCheckFactory() android.SingletonModule {
 	m := &dexpreoptSystemserverCheck{}
-	m.artifactsByModuleName = make(map[string][]string)
 	android.InitAndroidArchModule(m, android.DeviceSupported, android.MultilibCommon)
 	return m
 }
@@ -62,7 +60,25 @@
 		ctx, "", strings.TrimPrefix(location, "/"))
 }
 
-func (m *dexpreoptSystemserverCheck) GenerateAndroidBuildActions(ctx android.ModuleContext) {
+type systemServerDependencyTag struct {
+	blueprint.BaseDependencyTag
+}
+
+// systemServerJarDepTag willl be used for validation. Skip visiblility.
+func (b systemServerDependencyTag) ExcludeFromVisibilityEnforcement() {
+}
+
+var (
+	// dep tag for platform and apex system server jars
+	systemServerJarDepTag = systemServerDependencyTag{}
+)
+
+var _ android.ExcludeFromVisibilityEnforcementTag = systemServerJarDepTag
+
+// Add a depenendency on the system server jars. The dexpreopt files of those will be emitted to make.
+// The kati packaging system will verify that those files appear in installed files.
+// Adding the dependency allows the singleton module to determine whether an apex system server jar is system_ext specific.
+func (m *dexpreoptSystemserverCheck) DepsMutator(ctx android.BottomUpMutatorContext) {
 	global := dexpreopt.GetGlobalConfig(ctx)
 	targets := ctx.Config().Targets[android.Android]
 
@@ -72,23 +88,27 @@
 		return
 	}
 
-	systemServerJars := global.AllSystemServerJars(ctx)
-	for _, jar := range systemServerJars.CopyOfJars() {
-		dexLocation := dexpreopt.GetSystemServerDexLocation(ctx, global, jar)
-		odexLocation := dexpreopt.ToOdexPath(dexLocation, targets[0].Arch.ArchType)
+	ctx.AddDependency(ctx.Module(), systemServerJarDepTag, global.AllSystemServerJars(ctx).CopyOfJars()...)
+}
+
+func (m *dexpreoptSystemserverCheck) GenerateAndroidBuildActions(ctx android.ModuleContext) {
+	global := dexpreopt.GetGlobalConfig(ctx)
+	targets := ctx.Config().Targets[android.Android]
+
+	ctx.VisitDirectDepsWithTag(systemServerJarDepTag, func(systemServerJar android.Module) {
+		partition := "system"
+		if systemServerJar.InstallInSystemExt() && ctx.Config().InstallApexSystemServerDexpreoptSamePartition() {
+			partition = ctx.DeviceConfig().SystemExtPath() // system_ext
+		}
+		dexLocation := dexpreopt.GetSystemServerDexLocation(ctx, global, systemServerJar.Name())
+		odexLocation := dexpreopt.ToOdexPath(dexLocation, targets[0].Arch.ArchType, partition)
 		odexPath := getInstallPath(ctx, odexLocation)
 		vdexPath := getInstallPath(ctx, pathtools.ReplaceExtension(odexLocation, "vdex"))
-		m.artifactsByModuleName[jar] = []string{odexPath.String(), vdexPath.String()}
-	}
+		m.artifacts = append(m.artifacts, odexPath.String(), vdexPath.String())
+	})
 }
 
 func (m *dexpreoptSystemserverCheck) GenerateSingletonBuildActions(ctx android.SingletonContext) {
-	// Only keep modules defined in Soong.
-	ctx.VisitAllModules(func(module android.Module) {
-		if artifacts, ok := m.artifactsByModuleName[module.Name()]; ok {
-			m.artifacts = append(m.artifacts, artifacts...)
-		}
-	})
 }
 
 func (m *dexpreoptSystemserverCheck) MakeVars(ctx android.MakeVarsContext) {
diff --git a/ui/build/config.go b/ui/build/config.go
index bd20442..75edfcd 100644
--- a/ui/build/config.go
+++ b/ui/build/config.go
@@ -1754,12 +1754,10 @@
 }
 
 func (c *configImpl) PrebuiltBuildTool(name string) string {
-	if v, ok := c.environ.Get("SANITIZE_HOST"); ok {
-		if sanitize := strings.Fields(v); inList("address", sanitize) {
-			asan := filepath.Join("prebuilts/build-tools", c.HostPrebuiltTag(), "asan/bin", name)
-			if _, err := os.Stat(asan); err == nil {
-				return asan
-			}
+	if c.environ.IsEnvTrue("SANITIZE_BUILD_TOOL_PREBUILTS") {
+		asan := filepath.Join("prebuilts/build-tools", c.HostPrebuiltTag(), "asan/bin", name)
+		if _, err := os.Stat(asan); err == nil {
+			return asan
 		}
 	}
 	return filepath.Join("prebuilts/build-tools", c.HostPrebuiltTag(), "bin", name)