Add filter by container to aflags

Test: m -j120 && acloud create --local-image && adb shell aflags list --c system
Test: cargo t
Bug: 340840507
Change-Id: I5db7f204673accdbd3c4ad62e88b213028a8d5ab
diff --git a/tools/aconfig/aflags/src/main.rs b/tools/aconfig/aflags/src/main.rs
index 516b773..05c15bb 100644
--- a/tools/aconfig/aflags/src/main.rs
+++ b/tools/aconfig/aflags/src/main.rs
@@ -35,10 +35,14 @@
 
 impl std::fmt::Display for FlagPermission {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        write!(f, "{}", match &self {
-            Self::ReadOnly => "read-only",
-            Self::ReadWrite => "read-write",
-        })
+        write!(
+            f,
+            "{}",
+            match &self {
+                Self::ReadOnly => "read-only",
+                Self::ReadWrite => "read-write",
+            }
+        )
     }
 }
 
@@ -50,10 +54,14 @@
 
 impl std::fmt::Display for ValuePickedFrom {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        write!(f, "{}", match &self {
-            Self::Default => "default",
-            Self::Server => "server",
-        })
+        write!(
+            f,
+            "{}",
+            match &self {
+                Self::Default => "default",
+                Self::Server => "server",
+            }
+        )
     }
 }
 
@@ -77,10 +85,14 @@
 
 impl std::fmt::Display for FlagValue {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        write!(f, "{}", match &self {
-            Self::Enabled => "enabled",
-            Self::Disabled => "disabled",
-        })
+        write!(
+            f,
+            "{}",
+            match &self {
+                Self::Enabled => "enabled",
+                Self::Disabled => "disabled",
+            }
+        )
     }
 }
 
@@ -153,6 +165,10 @@
         /// Read from the new flag storage.
         #[clap(long)]
         use_new_storage: bool,
+
+        /// Optionally filter by container name.
+        #[clap(short = 'c', long = "container")]
+        container: Option<String>,
     },
 
     /// Enable an aconfig flag on this device, on the next boot.
@@ -176,6 +192,23 @@
     longest_permission_col: usize,
 }
 
+struct Filter {
+    container: Option<String>,
+}
+
+impl Filter {
+    fn apply(&self, flags: &[Flag]) -> Vec<Flag> {
+        flags
+            .iter()
+            .filter(|flag| match &self.container {
+                Some(c) => flag.container == *c,
+                None => true,
+            })
+            .cloned()
+            .collect()
+    }
+}
+
 fn format_flag_row(flag: &Flag, info: &PaddingInfo) -> String {
     let full_name = flag.qualified_name();
     let p0 = info.longest_flag_col + 1;
@@ -215,11 +248,12 @@
     Ok(())
 }
 
-fn list(source_type: FlagSourceType) -> Result<String> {
-    let flags = match source_type {
+fn list(source_type: FlagSourceType, container: Option<String>) -> Result<String> {
+    let flags_unfiltered = match source_type {
         FlagSourceType::DeviceConfig => DeviceConfigSource::list_flags()?,
         FlagSourceType::AconfigStorage => AconfigStorageSource::list_flags()?,
     };
+    let flags = (Filter { container }).apply(&flags_unfiltered);
     let padding_info = PaddingInfo {
         longest_flag_col: flags.iter().map(|f| f.qualified_name().len()).max().unwrap_or(0),
         longest_val_col: flags.iter().map(|f| f.value.to_string().len()).max().unwrap_or(0),
@@ -251,8 +285,12 @@
 fn main() {
     let cli = Cli::parse();
     let output = match cli.command {
-        Command::List { use_new_storage: true } => list(FlagSourceType::AconfigStorage).map(Some),
-        Command::List { use_new_storage: false } => list(FlagSourceType::DeviceConfig).map(Some),
+        Command::List { use_new_storage: true, container } => {
+            list(FlagSourceType::AconfigStorage, container).map(Some)
+        }
+        Command::List { use_new_storage: false, container } => {
+            list(FlagSourceType::DeviceConfig, container).map(Some)
+        }
         Command::Enable { qualified_name } => set_flag(&qualified_name, "true").map(|_| None),
         Command::Disable { qualified_name } => set_flag(&qualified_name, "false").map(|_| None),
     };
@@ -262,3 +300,84 @@
         Err(message) => println!("Error: {message}"),
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_filter_container() {
+        let flags = vec![
+            Flag {
+                namespace: "namespace".to_string(),
+                name: "test1".to_string(),
+                package: "package".to_string(),
+                value: FlagValue::Disabled,
+                staged_value: None,
+                permission: FlagPermission::ReadWrite,
+                value_picked_from: ValuePickedFrom::Default,
+                container: "system".to_string(),
+            },
+            Flag {
+                namespace: "namespace".to_string(),
+                name: "test2".to_string(),
+                package: "package".to_string(),
+                value: FlagValue::Disabled,
+                staged_value: None,
+                permission: FlagPermission::ReadWrite,
+                value_picked_from: ValuePickedFrom::Default,
+                container: "not_system".to_string(),
+            },
+            Flag {
+                namespace: "namespace".to_string(),
+                name: "test3".to_string(),
+                package: "package".to_string(),
+                value: FlagValue::Disabled,
+                staged_value: None,
+                permission: FlagPermission::ReadWrite,
+                value_picked_from: ValuePickedFrom::Default,
+                container: "system".to_string(),
+            },
+        ];
+
+        assert_eq!((Filter { container: Some("system".to_string()) }).apply(&flags).len(), 2);
+    }
+
+    #[test]
+    fn test_filter_no_container() {
+        let flags = vec![
+            Flag {
+                namespace: "namespace".to_string(),
+                name: "test1".to_string(),
+                package: "package".to_string(),
+                value: FlagValue::Disabled,
+                staged_value: None,
+                permission: FlagPermission::ReadWrite,
+                value_picked_from: ValuePickedFrom::Default,
+                container: "system".to_string(),
+            },
+            Flag {
+                namespace: "namespace".to_string(),
+                name: "test2".to_string(),
+                package: "package".to_string(),
+                value: FlagValue::Disabled,
+                staged_value: None,
+                permission: FlagPermission::ReadWrite,
+                value_picked_from: ValuePickedFrom::Default,
+                container: "not_system".to_string(),
+            },
+            Flag {
+                namespace: "namespace".to_string(),
+                name: "test3".to_string(),
+                package: "package".to_string(),
+                value: FlagValue::Disabled,
+                staged_value: None,
+                permission: FlagPermission::ReadWrite,
+                value_picked_from: ValuePickedFrom::Default,
+                container: "system".to_string(),
+            },
+        ];
+
+        assert_eq!((Filter { container: None }).apply(&flags).len(), 3);
+    }
+}