1 files changed, 15 insertions, 18 deletions

diff --git a/tools/aconfig/aconfig/src/commands.rs b/tools/aconfig/aconfig/src/commands.rs
index 7736ce75ee..6945fd4649 100644
--- a/tools/aconfig/aconfig/src/commands.rs
+++ b/tools/aconfig/aconfig/src/commands.rs
@@ -202,7 +202,11 @@ pub fn create_java_lib(mut input: Input, codegen_mode: CodegenMode) -> Result<Ve
     generate_java_code(&package, modified_parsed_flags.into_iter(), codegen_mode)
 }
 
-pub fn create_cpp_lib(mut input: Input, codegen_mode: CodegenMode) -> Result<Vec<OutputFile>> {
+pub fn create_cpp_lib(
+    mut input: Input,
+    codegen_mode: CodegenMode,
+    allow_instrumentation: bool,
+) -> Result<Vec<OutputFile>> {
     // TODO(327420679): Enable export mode for native flag library
     ensure!(
         codegen_mode != CodegenMode::Exported,
@@ -214,8 +218,14 @@ pub fn create_cpp_lib(mut input: Input, codegen_mode: CodegenMode) -> Result<Vec
         bail!("no parsed flags, or the parsed flags use different packages");
     };
     let package = package.to_string();
-    let _flag_ids = assign_flag_ids(&package, modified_parsed_flags.iter())?;
-    generate_cpp_code(&package, modified_parsed_flags.into_iter(), codegen_mode)
+    let flag_ids = assign_flag_ids(&package, modified_parsed_flags.iter())?;
+    generate_cpp_code(
+        &package,
+        modified_parsed_flags.into_iter(),
+        codegen_mode,
+        flag_ids,
+        allow_instrumentation,
+    )
 }
 
 pub fn create_rust_lib(mut input: Input, codegen_mode: CodegenMode) -> Result<OutputFile> {
@@ -239,13 +249,8 @@ pub fn create_storage(
     container: &str,
     file: &StorageFileType,
 ) -> Result<Vec<u8>> {
-    let parsed_flags_vec: Vec<ProtoParsedFlags> = caches
-        .into_iter()
-        .map(|mut input| input.try_parse_flags())
-        .collect::<Result<Vec<_>>>()?
-        .into_iter()
-        .filter(|pfs| find_unique_container(pfs) == Some(container))
-        .collect();
+    let parsed_flags_vec: Vec<ProtoParsedFlags> =
+        caches.into_iter().map(|mut input| input.try_parse_flags()).collect::<Result<Vec<_>>>()?;
     generate_storage_file(container, parsed_flags_vec.iter(), file)
 }
 
@@ -324,14 +329,6 @@ fn find_unique_package(parsed_flags: &[ProtoParsedFlag]) -> Option<&str> {
     Some(package)
 }
 
-fn find_unique_container(parsed_flags: &ProtoParsedFlags) -> Option<&str> {
-    let container = parsed_flags.parsed_flag.first().map(|pf| pf.container())?;
-    if parsed_flags.parsed_flag.iter().any(|pf| pf.container() != container) {
-        return None;
-    }
-    Some(container)
-}
-
 pub fn modify_parsed_flags_based_on_mode(
     parsed_flags: ProtoParsedFlags,
     codegen_mode: CodegenMode,