Redirect memtag libraries to snapshot

memtag sanitizer libraries are vendor available and can be captured as
snapshots. This change adds a redirection logic for memtag libraries.

This is just a workaround, just like other SnapshotInfoProvider calls.
In the future we need to refactor these codes. So TODO is added to
remind refactoring.

Bug: 178470649
Test: soong test
Change-Id: Id77f1ce94255b56a68f3e1d7446a68189c45ac54
diff --git a/cc/sanitize.go b/cc/sanitize.go
index e1ac9f0..8f17e4e 100644
--- a/cc/sanitize.go
+++ b/cc/sanitize.go
@@ -1080,6 +1080,12 @@
 			if Bool(c.sanitize.Properties.Sanitize.Diag.Memtag_heap) {
 				noteDep = "note_memtag_heap_sync"
 			}
+			// If we're using snapshots, redirect to snapshot whenever possible
+			// TODO(b/178470649): clean manual snapshot redirections
+			snapshot := mctx.Provider(SnapshotInfoProvider).(SnapshotInfo)
+			if lib, ok := snapshot.StaticLibs[noteDep]; ok {
+				noteDep = lib
+			}
 			depTag := libraryDependencyTag{Kind: staticLibraryDependency, wholeStatic: true}
 			variations := append(mctx.Target().Variations(),
 				blueprint.Variation{Mutator: "link", Variation: "static"})
diff --git a/cc/vendor_snapshot_test.go b/cc/vendor_snapshot_test.go
index 8f77c28..340b317 100644
--- a/cc/vendor_snapshot_test.go
+++ b/cc/vendor_snapshot_test.go
@@ -819,6 +819,18 @@
 
 func TestVendorSnapshotSanitizer(t *testing.T) {
 	bp := `
+	vendor_snapshot {
+		name: "vendor_snapshot",
+		version: "28",
+		arch: {
+			arm64: {
+				static_libs: [
+					"libsnapshot",
+					"note_memtag_heap_sync",
+				],
+			},
+		},
+	}
 	vendor_snapshot_static {
 		name: "libsnapshot",
 		vendor: true,
@@ -833,8 +845,41 @@
 			},
 		},
 	}
+
+	vendor_snapshot_static {
+		name: "note_memtag_heap_sync",
+		vendor: true,
+		target_arch: "arm64",
+		version: "28",
+		arch: {
+			arm64: {
+				src: "note_memtag_heap_sync.a",
+			},
+		},
+	}
+
+	cc_test {
+		name: "vstest",
+		gtest: false,
+		vendor: true,
+		compile_multilib: "64",
+		nocrt: true,
+		no_libcrt: true,
+		stl: "none",
+		static_libs: ["libsnapshot"],
+		system_shared_libs: [],
+	}
 `
-	config := TestConfig(t.TempDir(), android.Android, nil, bp, nil)
+
+	mockFS := map[string][]byte{
+		"vendor/Android.bp":              []byte(bp),
+		"vendor/libc++demangle.a":        nil,
+		"vendor/libsnapshot.a":           nil,
+		"vendor/libsnapshot.cfi.a":       nil,
+		"vendor/note_memtag_heap_sync.a": nil,
+	}
+
+	config := TestConfig(t.TempDir(), android.Android, nil, "", mockFS)
 	config.TestProductVariables.DeviceVndkVersion = StringPtr("28")
 	config.TestProductVariables.Platform_vndk_version = StringPtr("29")
 	ctx := testCcWithConfig(t, config)