Reimplement afdo support for rust

Ignore-AOSP-First: The parent CL is internal
Bug: 267229065
Test: go test
Change-Id: Ia14679285b92f3f14ff269392a61f978c71311b2
Merged-In: Ia14679285b92f3f14ff269392a61f978c71311b2
diff --git a/rust/rust.go b/rust/rust.go
index f85babc..ea7f9c6 100644
--- a/rust/rust.go
+++ b/rust/rust.go
@@ -39,7 +39,6 @@
 		ctx.BottomUp("rust_libraries", LibraryMutator).Parallel()
 		ctx.BottomUp("rust_stdlinkage", LibstdMutator).Parallel()
 		ctx.BottomUp("rust_begin", BeginMutator).Parallel()
-
 	})
 	android.PostDepsMutators(func(ctx android.RegisterMutatorsContext) {
 		ctx.BottomUp("rust_sanitizers", rustSanitizerRuntimeMutator).Parallel()
@@ -920,7 +919,7 @@
 
 	// Calculate rustc flags
 	if mod.afdo != nil {
-		flags, deps = mod.afdo.flags(ctx, flags, deps)
+		flags, deps = mod.afdo.flags(actx, flags, deps)
 	}
 	if mod.compiler != nil {
 		flags = mod.compiler.compilerFlags(ctx, flags)
@@ -1613,6 +1612,8 @@
 
 	// proc_macros are compiler plugins, and so we need the host arch variant as a dependendcy.
 	actx.AddFarVariationDependencies(ctx.Config().BuildOSTarget.Variations(), procMacroDepTag, deps.ProcMacros...)
+
+	mod.afdo.addDep(ctx, actx)
 }
 
 // addRlibDependency will add an rlib dependency, rewriting to the snapshot library if available.