Add tests to verify more use cases of afdo in cc and rust

Test: go test
Bug: 267229065
Change-Id: I4aa9538e92aa27e0f842bdce3725bc82739196cb
diff --git a/cc/afdo_test.go b/cc/afdo_test.go
index 335910c..40f705b 100644
--- a/cc/afdo_test.go
+++ b/cc/afdo_test.go
@@ -180,3 +180,100 @@
 		}
 	}
 }
+
+func TestAfdoEnabledWithMultiArchs(t *testing.T) {
+	t.Parallel()
+	bp := `
+	cc_library_shared {
+		name: "foo",
+		srcs: ["test.c"],
+		afdo: true,
+		compile_multilib: "both",
+	}
+`
+	result := android.GroupFixturePreparers(
+		prepareForCcTest,
+		android.FixtureAddTextFile("toolchain/pgo-profiles/sampling/foo_arm.afdo", "TEST"),
+		android.FixtureAddTextFile("toolchain/pgo-profiles/sampling/foo_arm64.afdo", "TEST"),
+	).RunTestWithBp(t, bp)
+
+	fooArm := result.ModuleForTests("foo", "android_arm_armv7-a-neon_shared")
+	fooArmCFlags := fooArm.Rule("cc").Args["cFlags"]
+	if w := "-fprofile-sample-use=toolchain/pgo-profiles/sampling/foo_arm.afdo"; !strings.Contains(fooArmCFlags, w) {
+		t.Errorf("Expected 'foo' to enable afdo, but did not find %q in cflags %q", w, fooArmCFlags)
+	}
+
+	fooArm64 := result.ModuleForTests("foo", "android_arm64_armv8-a_shared")
+	fooArm64CFlags := fooArm64.Rule("cc").Args["cFlags"]
+	if w := "-fprofile-sample-use=toolchain/pgo-profiles/sampling/foo_arm64.afdo"; !strings.Contains(fooArm64CFlags, w) {
+		t.Errorf("Expected 'foo' to enable afdo, but did not find %q in cflags %q", w, fooArm64CFlags)
+	}
+}
+
+func TestMultipleAfdoRDeps(t *testing.T) {
+	t.Parallel()
+	bp := `
+	cc_library_shared {
+		name: "libTest",
+		srcs: ["test.c"],
+		static_libs: ["libFoo"],
+		afdo: true,
+	}
+
+	cc_library_shared {
+		name: "libBar",
+		srcs: ["bar.c"],
+		static_libs: ["libFoo"],
+		afdo: true,
+	}
+
+	cc_library_static {
+		name: "libFoo",
+		srcs: ["foo.c"],
+	}
+	`
+
+	result := android.GroupFixturePreparers(
+		prepareForCcTest,
+		android.FixtureAddTextFile("toolchain/pgo-profiles/sampling/libTest.afdo", "TEST"),
+		android.FixtureAddTextFile("toolchain/pgo-profiles/sampling/libBar.afdo", "TEST"),
+	).RunTestWithBp(t, bp)
+
+	expectedCFlagLibTest := "-fprofile-sample-use=toolchain/pgo-profiles/sampling/libTest.afdo"
+	expectedCFlagLibBar := "-fprofile-sample-use=toolchain/pgo-profiles/sampling/libBar.afdo"
+
+	libTest := result.ModuleForTests("libTest", "android_arm64_armv8-a_shared")
+	libTestAfdoVariantOfLibFoo := result.ModuleForTests("libFoo", "android_arm64_armv8-a_static_afdo-libTest")
+
+	libBar := result.ModuleForTests("libBar", "android_arm64_armv8-a_shared")
+	libBarAfdoVariantOfLibFoo := result.ModuleForTests("libFoo", "android_arm64_armv8-a_static_afdo-libBar")
+
+	// Check cFlags of afdo-enabled modules and the afdo-variant of their static deps
+	cFlags := libTest.Rule("cc").Args["cFlags"]
+	if !strings.Contains(cFlags, expectedCFlagLibTest) {
+		t.Errorf("Expected 'libTest' to enable afdo, but did not find %q in cflags %q", expectedCFlagLibTest, cFlags)
+	}
+	cFlags = libBar.Rule("cc").Args["cFlags"]
+	if !strings.Contains(cFlags, expectedCFlagLibBar) {
+		t.Errorf("Expected 'libBar' to enable afdo, but did not find %q in cflags %q", expectedCFlagLibBar, cFlags)
+	}
+
+	cFlags = libTestAfdoVariantOfLibFoo.Rule("cc").Args["cFlags"]
+	if !strings.Contains(cFlags, expectedCFlagLibTest) {
+		t.Errorf("Expected 'libTestAfdoVariantOfLibFoo' to enable afdo, but did not find %q in cflags %q", expectedCFlagLibTest, cFlags)
+	}
+
+	cFlags = libBarAfdoVariantOfLibFoo.Rule("cc").Args["cFlags"]
+	if !strings.Contains(cFlags, expectedCFlagLibBar) {
+		t.Errorf("Expected 'libBarAfdoVariantOfLibFoo' to enable afdo, but did not find %q in cflags %q", expectedCFlagLibBar, cFlags)
+	}
+
+	// Check dependency edges of static deps
+	if !hasDirectDep(result, libTest.Module(), libTestAfdoVariantOfLibFoo.Module()) {
+		t.Errorf("libTest missing dependency on afdo variant of libFoo")
+	}
+
+	if !hasDirectDep(result, libBar.Module(), libBarAfdoVariantOfLibFoo.Module()) {
+		t.Errorf("libBar missing dependency on afdo variant of libFoo")
+	}
+}
diff --git a/rust/Android.bp b/rust/Android.bp
index 3fd68e5..b01a94a 100644
--- a/rust/Android.bp
+++ b/rust/Android.bp
@@ -42,6 +42,7 @@
         "toolchain_library.go",
     ],
     testSrcs: [
+        "afdo_test.go",
         "benchmark_test.go",
         "binary_test.go",
         "bindgen_test.go",
diff --git a/rust/afdo_test.go b/rust/afdo_test.go
new file mode 100644
index 0000000..fa20eef
--- /dev/null
+++ b/rust/afdo_test.go
@@ -0,0 +1,76 @@
+// Copyright 2023 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 rust
+
+import (
+	"android/soong/android"
+	"fmt"
+	"strings"
+	"testing"
+)
+
+func TestAfdoEnabled(t *testing.T) {
+	bp := `
+	rust_binary {
+		name: "foo",
+		srcs: ["foo.rs"],
+		afdo: true,
+	}
+`
+	result := android.GroupFixturePreparers(
+		prepareForRustTest,
+		android.FixtureAddTextFile("toolchain/pgo-profiles/sampling/foo.afdo", ""),
+		rustMockedFiles.AddToFixture(),
+	).RunTestWithBp(t, bp)
+
+	foo := result.ModuleForTests("foo", "android_arm64_armv8-a").Rule("rustc")
+
+	expectedCFlag := fmt.Sprintf(afdoFlagFormat, "toolchain/pgo-profiles/sampling/foo.afdo")
+
+	if !strings.Contains(foo.Args["rustcFlags"], expectedCFlag) {
+		t.Errorf("Expected 'foo' to enable afdo, but did not find %q in cflags %q", expectedCFlag, foo.Args["rustcFlags"])
+	}
+}
+
+func TestAfdoEnabledWithMultiArchs(t *testing.T) {
+	bp := `
+	rust_binary {
+		name: "foo",
+		srcs: ["foo.rs"],
+		afdo: true,
+		compile_multilib: "both",
+	}
+`
+	result := android.GroupFixturePreparers(
+		prepareForRustTest,
+		android.FixtureAddTextFile("toolchain/pgo-profiles/sampling/foo_arm.afdo", ""),
+		android.FixtureAddTextFile("toolchain/pgo-profiles/sampling/foo_arm64.afdo", ""),
+		rustMockedFiles.AddToFixture(),
+	).RunTestWithBp(t, bp)
+
+	fooArm := result.ModuleForTests("foo", "android_arm_armv7-a-neon").Rule("rustc")
+	fooArm64 := result.ModuleForTests("foo", "android_arm64_armv8-a").Rule("rustc")
+
+	expectedCFlagArm := fmt.Sprintf(afdoFlagFormat, "toolchain/pgo-profiles/sampling/foo_arm.afdo")
+	expectedCFlagArm64 := fmt.Sprintf(afdoFlagFormat, "toolchain/pgo-profiles/sampling/foo_arm64.afdo")
+
+	if !strings.Contains(fooArm.Args["rustcFlags"], expectedCFlagArm) {
+		t.Errorf("Expected 'fooArm' to enable afdo, but did not find %q in cflags %q", expectedCFlagArm, fooArm.Args["rustcFlags"])
+	}
+
+	if !strings.Contains(fooArm64.Args["rustcFlags"], expectedCFlagArm64) {
+		t.Errorf("Expected 'fooArm64' to enable afdo, but did not find %q in cflags %q", expectedCFlagArm64, fooArm64.Args["rustcFlags"])
+	}
+}