Attach rust_benchmark to atest and tradefed.

Automatically generate required tradefed configs for rust benchmarks so
that they're available in atest.

Test: atest <module with rust_benchmark defined>
Bug: 155309706
Change-Id: I6002100367a66b6b0555614acc6cebb00dbf435d
diff --git a/rust/androidmk.go b/rust/androidmk.go
index 0f9a17d..eda2bb4 100644
--- a/rust/androidmk.go
+++ b/rust/androidmk.go
@@ -102,6 +102,20 @@
 	cc.AndroidMkWriteTestData(test.data, ret)
 }
 
+func (benchmark *benchmarkDecorator) AndroidMk(ctx AndroidMkContext, ret *android.AndroidMkEntries) {
+	benchmark.binaryDecorator.AndroidMk(ctx, ret)
+	ret.Class = "NATIVE_TESTS"
+	ret.ExtraEntries = append(ret.ExtraEntries,
+		func(ctx android.AndroidMkExtraEntriesContext, entries *android.AndroidMkEntries) {
+			entries.AddCompatibilityTestSuites(benchmark.Properties.Test_suites...)
+			if benchmark.testConfig != nil {
+				entries.SetString("LOCAL_FULL_TEST_CONFIG", benchmark.testConfig.String())
+			}
+			entries.SetBool("LOCAL_NATIVE_BENCHMARK", true)
+			entries.SetBoolIfTrue("LOCAL_DISABLE_AUTO_GENERATE_TEST_CONFIG", !BoolDefault(benchmark.Properties.Auto_gen_config, true))
+		})
+}
+
 func (library *libraryDecorator) AndroidMk(ctx AndroidMkContext, ret *android.AndroidMkEntries) {
 	ctx.SubAndroidMk(ret, library.baseCompiler)
 
diff --git a/rust/benchmark.go b/rust/benchmark.go
index 6e4fb85..b89f5cd 100644
--- a/rust/benchmark.go
+++ b/rust/benchmark.go
@@ -16,9 +16,31 @@
 
 import (
 	"android/soong/android"
+	"android/soong/tradefed"
 )
 
 type BenchmarkProperties struct {
+	// Disables the creation of a test-specific directory when used with
+	// relative_install_path. Useful if several tests need to be in the same
+	// directory, but test_per_src doesn't work.
+	No_named_install_directory *bool
+
+	// the name of the test configuration (for example "AndroidBenchmark.xml") that should be
+	// installed with the module.
+	Test_config *string `android:"path,arch_variant"`
+
+	// the name of the test configuration template (for example "AndroidBenchmarkTemplate.xml") that
+	// should be installed with the module.
+	Test_config_template *string `android:"path,arch_variant"`
+
+	// list of compatibility suites (for example "cts", "vts") that the module should be
+	// installed into.
+	Test_suites []string `android:"arch_variant"`
+
+	// Flag to indicate whether or not to create test config automatically. If AndroidTest.xml
+	// doesn't exist next to the Android.bp, this attribute doesn't need to be set to true
+	// explicitly.
+	Auto_gen_config *bool
 }
 
 type benchmarkDecorator struct {
@@ -87,3 +109,21 @@
 func (benchmark *benchmarkDecorator) compilerProps() []interface{} {
 	return append(benchmark.binaryDecorator.compilerProps(), &benchmark.Properties)
 }
+
+func (benchmark *benchmarkDecorator) install(ctx ModuleContext) {
+	benchmark.testConfig = tradefed.AutoGenRustBenchmarkConfig(ctx,
+		benchmark.Properties.Test_config,
+		benchmark.Properties.Test_config_template,
+		benchmark.Properties.Test_suites,
+		nil,
+		benchmark.Properties.Auto_gen_config)
+
+	// default relative install path is module name
+	if !Bool(benchmark.Properties.No_named_install_directory) {
+		benchmark.baseCompiler.relative = ctx.ModuleName()
+	} else if String(benchmark.baseCompiler.Properties.Relative_install_path) == "" {
+		ctx.PropertyErrorf("no_named_install_directory", "Module install directory may only be disabled if relative_install_path is set")
+	}
+
+	benchmark.binaryDecorator.install(ctx)
+}
diff --git a/tradefed/autogen.go b/tradefed/autogen.go
index 27d71e8..3d96c84 100644
--- a/tradefed/autogen.go
+++ b/tradefed/autogen.go
@@ -245,6 +245,25 @@
 	return path
 }
 
+func AutoGenRustBenchmarkConfig(ctx android.ModuleContext, testConfigProp *string,
+	testConfigTemplateProp *string, testSuites []string, config []Config, autoGenConfig *bool) android.Path {
+	path, autogenPath := testConfigPath(ctx, testConfigProp, testSuites, autoGenConfig, testConfigTemplateProp)
+	if autogenPath != nil {
+		templatePath := getTestConfigTemplate(ctx, testConfigTemplateProp)
+		if templatePath.Valid() {
+			autogenTemplate(ctx, autogenPath, templatePath.String(), config, "")
+		} else {
+			if ctx.Device() {
+				autogenTemplate(ctx, autogenPath, "${RustDeviceBenchmarkConfigTemplate}", config, "")
+			} else {
+				autogenTemplate(ctx, autogenPath, "${RustHostBenchmarkConfigTemplate}", config, "")
+			}
+		}
+		return autogenPath
+	}
+	return path
+}
+
 func AutoGenRobolectricTestConfig(ctx android.ModuleContext, testConfigProp *string, testConfigTemplateProp *string,
 	testSuites []string, autoGenConfig *bool) android.Path {
 	path, autogenPath := testConfigPath(ctx, testConfigProp, testSuites, autoGenConfig, testConfigTemplateProp)
diff --git a/tradefed/config.go b/tradefed/config.go
index f3566a8..999424c 100644
--- a/tradefed/config.go
+++ b/tradefed/config.go
@@ -34,6 +34,8 @@
 	pctx.SourcePathVariable("PythonBinaryHostTestConfigTemplate", "build/make/core/python_binary_host_test_config_template.xml")
 	pctx.SourcePathVariable("RustDeviceTestConfigTemplate", "build/make/core/rust_device_test_config_template.xml")
 	pctx.SourcePathVariable("RustHostTestConfigTemplate", "build/make/core/rust_host_test_config_template.xml")
+	pctx.SourcePathVariable("RustDeviceBenchmarkConfigTemplate", "build/make/core/rust_device_benchmark_config_template.xml")
+	pctx.SourcePathVariable("RustHostBenchmarkConfigTemplate", "build/make/core/rust_host_benchmark_config_template.xml")
 	pctx.SourcePathVariable("RobolectricTestConfigTemplate", "build/make/core/robolectric_test_config_template.xml")
 	pctx.SourcePathVariable("ShellTestConfigTemplate", "build/make/core/shell_test_config_template.xml")
 
diff --git a/tradefed/makevars.go b/tradefed/makevars.go
index f9682e4..9b5a20f 100644
--- a/tradefed/makevars.go
+++ b/tradefed/makevars.go
@@ -33,6 +33,8 @@
 	ctx.Strict("PYTHON_BINARY_HOST_TEST_CONFIG_TEMPLATE", "${PythonBinaryHostTestConfigTemplate}")
 	ctx.Strict("RUST_DEVICE_TEST_CONFIG_TEMPLATE", "${RustDeviceTestConfigTemplate}")
 	ctx.Strict("RUST_HOST_TEST_CONFIG_TEMPLATE", "${RustHostTestConfigTemplate}")
+	ctx.Strict("RUST_DEVICE_BENCHMARK_CONFIG_TEMPLATE", "${RustDeviceBenchmarkConfigTemplate}")
+	ctx.Strict("RUST_HOST_BENCHMARK_CONFIG_TEMPLATE", "${RustHostBenchmarkConfigTemplate}")
 	ctx.Strict("SHELL_TEST_CONFIG_TEMPLATE", "${ShellTestConfigTemplate}")
 
 	ctx.Strict("EMPTY_TEST_CONFIG", "${EmptyTestConfig}")