Add dex_import module type

Add a module type for importing a prebuilt jar that contains
classes.dex files.

Test: m with a prebuilt jar that contains classes.dex files in
      PRODUCT_BOOT_JARS
Test: java_test.go, dexpreopt_test.go, dexpreopt_bootjars_test.go
Bug: 124804356
Bug: 125517186

Change-Id: I496848f9dca11f758d49b1cb68168cec7f8e1718
diff --git a/java/androidmk.go b/java/androidmk.go
index 533b82e..0c3b1c7 100644
--- a/java/androidmk.go
+++ b/java/androidmk.go
@@ -131,6 +131,28 @@
 	}
 }
 
+func (prebuilt *DexImport) AndroidMk() android.AndroidMkData {
+	return android.AndroidMkData{
+		Class:      "JAVA_LIBRARIES",
+		OutputFile: android.OptionalPathForPath(prebuilt.maybeStrippedDexJarFile),
+		Include:    "$(BUILD_SYSTEM)/soong_java_prebuilt.mk",
+		Extra: []android.AndroidMkExtraFunc{
+			func(w io.Writer, outputFile android.Path) {
+				if prebuilt.dexJarFile != nil {
+					fmt.Fprintln(w, "LOCAL_SOONG_DEX_JAR :=", prebuilt.dexJarFile.String())
+					// TODO(b/125517186): export the dex jar as a classes jar to match some mis-uses in Make until
+					// boot_jars_package_check.mk can check dex jars.
+					fmt.Fprintln(w, "LOCAL_SOONG_HEADER_JAR :=", prebuilt.dexJarFile.String())
+					fmt.Fprintln(w, "LOCAL_SOONG_CLASSES_JAR :=", prebuilt.dexJarFile.String())
+				}
+				if len(prebuilt.dexpreopter.builtInstalled) > 0 {
+					fmt.Fprintln(w, "LOCAL_SOONG_BUILT_INSTALLED :=", prebuilt.dexpreopter.builtInstalled)
+				}
+			},
+		},
+	}
+}
+
 func (prebuilt *AARImport) AndroidMk() android.AndroidMkData {
 	return android.AndroidMkData{
 		Class:      "JAVA_LIBRARIES",
diff --git a/java/dexpreopt_bootjars.go b/java/dexpreopt_bootjars.go
index 64de21a..a23b477 100644
--- a/java/dexpreopt_bootjars.go
+++ b/java/dexpreopt_bootjars.go
@@ -152,7 +152,7 @@
 
 	ctx.VisitAllModules(func(module android.Module) {
 		// Collect dex jar paths for the modules listed above.
-		if j, ok := module.(Dependency); ok {
+		if j, ok := module.(interface{ DexJar() android.Path }); ok {
 			name := ctx.ModuleName(module)
 			if i := android.IndexList(name, image.modules); i != -1 {
 				bootDexJars[i] = j.DexJar()
diff --git a/java/dexpreopt_bootjars_test.go b/java/dexpreopt_bootjars_test.go
index c99540d..141f7ba 100644
--- a/java/dexpreopt_bootjars_test.go
+++ b/java/dexpreopt_bootjars_test.go
@@ -37,13 +37,18 @@
 			srcs: ["b.java"],
 			installable: true,
 		}
+
+		dex_import {
+			name: "baz",
+			jars: ["a.jar"],
+		}
 	`
 
 	config := testConfig(nil)
 
 	pathCtx := android.PathContextForTesting(config, nil)
 	dexpreoptConfig := dexpreopt.GlobalConfigForTests(pathCtx)
-	dexpreoptConfig.RuntimeApexJars = []string{"foo", "bar"}
+	dexpreoptConfig.RuntimeApexJars = []string{"foo", "bar", "baz"}
 	setDexpreoptTestGlobalConfig(config, dexpreoptConfig)
 
 	ctx := testContext(config, bp, nil)
@@ -59,6 +64,7 @@
 	expectedInputs := []string{
 		"dex_bootjars_input/foo.jar",
 		"dex_bootjars_input/bar.jar",
+		"dex_bootjars_input/baz.jar",
 	}
 
 	for i := range expectedInputs {
@@ -78,15 +84,19 @@
 
 		"dex_bootjars/system/framework/arm64/boot.art",
 		"dex_bootjars/system/framework/arm64/boot-bar.art",
+		"dex_bootjars/system/framework/arm64/boot-baz.art",
 
 		"dex_bootjars/system/framework/arm64/boot.oat",
 		"dex_bootjars/system/framework/arm64/boot-bar.oat",
+		"dex_bootjars/system/framework/arm64/boot-baz.oat",
 
 		"dex_bootjars/system/framework/arm64/boot.vdex",
 		"dex_bootjars/system/framework/arm64/boot-bar.vdex",
+		"dex_bootjars/system/framework/arm64/boot-baz.vdex",
 
 		"dex_bootjars_unstripped/system/framework/arm64/boot.oat",
 		"dex_bootjars_unstripped/system/framework/arm64/boot-bar.oat",
+		"dex_bootjars_unstripped/system/framework/arm64/boot-baz.oat",
 	}
 
 	for i := range expectedOutputs {
diff --git a/java/dexpreopt_test.go b/java/dexpreopt_test.go
index 6838bd2..4af2f5c 100644
--- a/java/dexpreopt_test.go
+++ b/java/dexpreopt_test.go
@@ -119,6 +119,15 @@
 				}`,
 			enabled: false,
 		},
+		{
+			name: "dex_import",
+			bp: `
+				dex_import {
+					name: "foo",
+					jars: ["a.jar"],
+				}`,
+			enabled: true,
+		},
 	}
 
 	for _, test := range tests {
diff --git a/java/java.go b/java/java.go
index f088d86..986e425 100644
--- a/java/java.go
+++ b/java/java.go
@@ -46,6 +46,7 @@
 	android.RegisterModuleType("java_import_host", ImportFactoryHost)
 	android.RegisterModuleType("java_device_for_host", DeviceForHostFactory)
 	android.RegisterModuleType("java_host_for_device", HostForDeviceFactory)
+	android.RegisterModuleType("dex_import", DexImportFactory)
 
 	android.RegisterSingletonType("logtags", LogtagsSingleton)
 }
@@ -1431,14 +1432,14 @@
 	Module
 }
 
-func (j *Library) shouldUncompressDex(ctx android.ModuleContext) bool {
+func shouldUncompressDex(ctx android.ModuleContext, dexpreopter *dexpreopter) bool {
 	// Store uncompressed (and do not strip) dex files from boot class path jars.
 	if inList(ctx.ModuleName(), ctx.Config().BootJars()) {
 		return true
 	}
 
 	// Store uncompressed dex files that are preopted on /system.
-	if !j.dexpreopter.dexpreoptDisabled(ctx) && (ctx.Host() || !odexOnSystemOther(ctx, j.dexpreopter.installPath)) {
+	if !dexpreopter.dexpreoptDisabled(ctx) && (ctx.Host() || !odexOnSystemOther(ctx, dexpreopter.installPath)) {
 		return true
 	}
 	if ctx.Config().UncompressPrivAppDex() &&
@@ -1453,7 +1454,7 @@
 	j.dexpreopter.installPath = android.PathForModuleInstall(ctx, "framework", ctx.ModuleName()+".jar")
 	j.dexpreopter.isSDKLibrary = j.deviceProperties.IsSDKLibrary
 	j.dexpreopter.isInstallable = Bool(j.properties.Installable)
-	j.dexpreopter.uncompressedDex = j.shouldUncompressDex(ctx)
+	j.dexpreopter.uncompressedDex = shouldUncompressDex(ctx, &j.dexpreopter)
 	j.deviceProperties.UncompressDex = j.dexpreopter.uncompressedDex
 	j.compile(ctx)
 
@@ -1901,6 +1902,113 @@
 	return module
 }
 
+// dex_import module
+
+type DexImportProperties struct {
+	Jars []string
+}
+
+type DexImport struct {
+	android.ModuleBase
+	android.DefaultableModuleBase
+	prebuilt android.Prebuilt
+
+	properties DexImportProperties
+
+	dexJarFile              android.Path
+	maybeStrippedDexJarFile android.Path
+
+	dexpreopter
+}
+
+func (j *DexImport) Prebuilt() *android.Prebuilt {
+	return &j.prebuilt
+}
+
+func (j *DexImport) PrebuiltSrcs() []string {
+	return j.properties.Jars
+}
+
+func (j *DexImport) Name() string {
+	return j.prebuilt.Name(j.ModuleBase.Name())
+}
+
+func (j *DexImport) DepsMutator(ctx android.BottomUpMutatorContext) {
+	android.ExtractSourcesDeps(ctx, j.properties.Jars)
+}
+
+func (j *DexImport) GenerateAndroidBuildActions(ctx android.ModuleContext) {
+	if len(j.properties.Jars) != 1 {
+		ctx.PropertyErrorf("jars", "exactly one jar must be provided")
+	}
+
+	j.dexpreopter.installPath = android.PathForModuleInstall(ctx, "framework", ctx.ModuleName()+".jar")
+	j.dexpreopter.isInstallable = true
+	j.dexpreopter.uncompressedDex = shouldUncompressDex(ctx, &j.dexpreopter)
+
+	inputJar := ctx.ExpandSource(j.properties.Jars[0], "jars")
+	dexOutputFile := android.PathForModuleOut(ctx, ctx.ModuleName()+".jar")
+
+	if j.dexpreopter.uncompressedDex {
+		rule := android.NewRuleBuilder()
+
+		temporary := android.PathForModuleOut(ctx, ctx.ModuleName()+".jar.unaligned")
+		rule.Temporary(temporary)
+
+		// use zip2zip to uncompress classes*.dex files
+		rule.Command().
+			Tool(ctx.Config().HostToolPath(ctx, "zip2zip")).
+			FlagWithInput("-i ", inputJar).
+			FlagWithOutput("-o ", temporary).
+			FlagWithArg("-0 ", "'classes*.dex'")
+
+		// use zipalign to align uncompressed classes*.dex files
+		rule.Command().
+			Tool(ctx.Config().HostToolPath(ctx, "zipalign")).
+			Flag("-f").
+			Text("4").
+			Input(temporary).
+			Output(dexOutputFile)
+
+		rule.DeleteTemporaryFiles()
+
+		rule.Build(pctx, ctx, "uncompress_dex", "uncompress dex")
+	} else {
+		ctx.Build(pctx, android.BuildParams{
+			Rule:   android.Cp,
+			Input:  inputJar,
+			Output: dexOutputFile,
+		})
+	}
+
+	j.dexJarFile = dexOutputFile
+
+	dexOutputFile = j.dexpreopt(ctx, dexOutputFile)
+
+	j.maybeStrippedDexJarFile = dexOutputFile
+
+	ctx.InstallFile(android.PathForModuleInstall(ctx, "framework"),
+		ctx.ModuleName()+".jar", dexOutputFile)
+}
+
+func (j *DexImport) DexJar() android.Path {
+	return j.dexJarFile
+}
+
+// dex_import imports a `.jar` file containing classes.dex files.
+//
+// A dex_import module cannot be used as a dependency of a java_* or android_* module, it can only be installed
+// to the device.
+func DexImportFactory() android.Module {
+	module := &DexImport{}
+
+	module.AddProperties(&module.properties)
+
+	android.InitPrebuiltModule(module, &module.properties.Jars)
+	InitJavaModule(module, android.DeviceSupported)
+	return module
+}
+
 //
 // Defaults
 //
@@ -1963,6 +2071,7 @@
 		&ImportProperties{},
 		&AARImportProperties{},
 		&sdkLibraryProperties{},
+		&DexImportProperties{},
 	)
 
 	android.InitDefaultsModule(module)
diff --git a/java/java_test.go b/java/java_test.go
index 952da11..35dd696 100644
--- a/java/java_test.go
+++ b/java/java_test.go
@@ -79,6 +79,7 @@
 	ctx.RegisterModuleType("java_system_modules", android.ModuleFactoryAdaptor(SystemModulesFactory))
 	ctx.RegisterModuleType("java_genrule", android.ModuleFactoryAdaptor(genRuleFactory))
 	ctx.RegisterModuleType("java_plugin", android.ModuleFactoryAdaptor(PluginFactory))
+	ctx.RegisterModuleType("dex_import", android.ModuleFactoryAdaptor(DexImportFactory))
 	ctx.RegisterModuleType("filegroup", android.ModuleFactoryAdaptor(android.FileGroupFactory))
 	ctx.RegisterModuleType("genrule", android.ModuleFactoryAdaptor(genrule.GenRuleFactory))
 	ctx.RegisterModuleType("droiddoc", android.ModuleFactoryAdaptor(DroiddocFactory))
@@ -330,6 +331,11 @@
 			name: "baz",
 			jars: ["b.jar"],
 		}
+
+		dex_import {
+			name: "qux",
+			jars: ["b.jar"],
+		}
 		`)
 
 	javac := ctx.ModuleForTests("foo", "android_common").Rule("javac")
@@ -344,6 +350,8 @@
 	if len(combineJar.Inputs) != 2 || combineJar.Inputs[1].String() != bazJar.String() {
 		t.Errorf("foo combineJar inputs %v does not contain %q", combineJar.Inputs, bazJar.String())
 	}
+
+	ctx.ModuleForTests("qux", "android_common").Rule("Cp")
 }
 
 func TestDefaults(t *testing.T) {