aflags: show staged values for each flag.

Lines now have a new column, and look like this:

com.android.example.flag1 enabled -            server read-write system
com.android.example.flag2 enabled (->disabled) server read-write system

The dash represents no change on boot, the (->value) represents a change
to value on the next boot.

Test: adb shell device_config put staged multitasking*com.android.wm.shell.enable_taskbar_navbar_unification true && adb shell aflags list | grep navbar_unification
Bug: 324436145
Change-Id: I022460bc69fcb2ccd9c6db8f060fcbd0337d1ea6
diff --git a/tools/aconfig/aflags/src/device_config_source.rs b/tools/aconfig/aflags/src/device_config_source.rs
index 2589f3d..882120f 100644
--- a/tools/aconfig/aflags/src/device_config_source.rs
+++ b/tools/aconfig/aflags/src/device_config_source.rs
@@ -55,6 +55,7 @@
         name,
         container,
         value,
+        staged_value: None,
         permission,
         value_picked_from: ValuePickedFrom::Default,
     }
@@ -118,24 +119,56 @@
     parse_device_config(&list_output)
 }
 
-fn reconcile(pb_flags: &[Flag], dc_flags: HashMap<String, FlagValue>) -> Vec<Flag> {
+/// Parse the list of newline-separated staged flags.
+///
+/// The output is a newline-sepaarated list of entries which follow this format:
+///   `namespace*flagname=value`
+///
+/// The resulting map maps from `namespace/flagname` to `value`, if a staged flag exists for
+/// `namespace/flagname`.
+fn parse_staged_flags(raw: &str) -> Result<HashMap<String, FlagValue>> {
+    let mut flags = HashMap::new();
+    for line in raw.split('\n') {
+        match (line.find('*'), line.find('=')) {
+            (Some(star_index), Some(equal_index)) => {
+                let namespace = &line[..star_index];
+                let flag = &line[star_index + 1..equal_index];
+                if let Ok(value) = FlagValue::try_from(&line[equal_index + 1..]) {
+                    flags.insert(namespace.to_owned() + "/" + flag, value);
+                }
+            }
+            _ => continue,
+        };
+    }
+    Ok(flags)
+}
+
+fn read_staged_flags() -> Result<HashMap<String, FlagValue>> {
+    let staged_flags_output = read_device_config_output(&["list", "staged"])?;
+    parse_staged_flags(&staged_flags_output)
+}
+
+fn reconcile(
+    pb_flags: &[Flag],
+    dc_flags: HashMap<String, FlagValue>,
+    staged_flags: HashMap<String, FlagValue>,
+) -> Vec<Flag> {
     pb_flags
         .iter()
         .map(|f| {
-            dc_flags
-                .get(&format!("{}/{}.{}", f.namespace, f.package, f.name))
-                .map(|value| {
-                    if *value == f.value {
-                        Flag { value_picked_from: ValuePickedFrom::Default, ..f.clone() }
-                    } else {
-                        Flag {
-                            value_picked_from: ValuePickedFrom::Server,
-                            value: *value,
-                            ..f.clone()
-                        }
-                    }
-                })
-                .unwrap_or(f.clone())
+            let server_override = dc_flags.get(&format!("{}/{}", f.namespace, f.qualified_name()));
+            let (value_picked_from, selected_value) = match server_override {
+                Some(value) if *value != f.value => (ValuePickedFrom::Server, *value),
+                _ => (ValuePickedFrom::Default, f.value),
+            };
+            Flag { value_picked_from, value: selected_value, ..f.clone() }
+        })
+        .map(|f| {
+            let staged_value = staged_flags
+                .get(&format!("{}/{}", f.namespace, f.qualified_name()))
+                .map(|value| if *value != f.value { Some(*value) } else { None })
+                .unwrap_or(None);
+            Flag { staged_value, ..f }
         })
         .collect()
 }
@@ -144,8 +177,9 @@
     fn list_flags() -> Result<Vec<Flag>> {
         let pb_flags = read_pb_files()?;
         let dc_flags = read_device_config_flags()?;
+        let staged_flags = read_staged_flags()?;
 
-        let flags = reconcile(&pb_flags, dc_flags);
+        let flags = reconcile(&pb_flags, dc_flags, staged_flags);
         Ok(flags)
     }
 
diff --git a/tools/aconfig/aflags/src/main.rs b/tools/aconfig/aflags/src/main.rs
index febd567..808ffa0 100644
--- a/tools/aconfig/aflags/src/main.rs
+++ b/tools/aconfig/aflags/src/main.rs
@@ -22,7 +22,7 @@
 mod device_config_source;
 use device_config_source::DeviceConfigSource;
 
-#[derive(Clone, PartialEq)]
+#[derive(Clone, PartialEq, Debug)]
 enum FlagPermission {
     ReadOnly,
     ReadWrite,
@@ -37,7 +37,7 @@
     }
 }
 
-#[derive(Clone)]
+#[derive(Clone, Debug)]
 enum ValuePickedFrom {
     Default,
     Server,
@@ -79,13 +79,14 @@
     }
 }
 
-#[derive(Clone)]
+#[derive(Clone, Debug)]
 struct Flag {
     namespace: String,
     name: String,
     package: String,
     container: String,
     value: FlagValue,
+    staged_value: Option<FlagValue>,
     permission: FlagPermission,
     value_picked_from: ValuePickedFrom,
 }
@@ -94,6 +95,13 @@
     fn qualified_name(&self) -> String {
         format!("{}.{}", self.package, self.name)
     }
+
+    fn display_staged_value(&self) -> String {
+        match self.staged_value {
+            Some(v) => format!("(->{})", v.to_string()),
+            None => "-".to_string(),
+        }
+    }
 }
 
 trait FlagSource {
@@ -110,6 +118,10 @@
   * `package`: package set for this flag in its .aconfig definition.
   * `flag_name`: flag name, also set in definition.
   * `value`: the value read from the flag.
+  * `staged_value`: the value on next boot:
+    + `-`: same as current value
+    + `(->enabled) flipped to enabled on boot.
+    + `(->disabled) flipped to disabled on boot.
   * `provenance`: one of:
     + `default`: the flag value comes from its build-time default.
     + `server`: the flag value comes from a server override.
@@ -145,6 +157,7 @@
 struct PaddingInfo {
     longest_flag_col: usize,
     longest_val_col: usize,
+    longest_staged_val_col: usize,
     longest_value_picked_from_col: usize,
     longest_permission_col: usize,
 }
@@ -156,15 +169,20 @@
     let val = flag.value.to_string();
     let p1 = info.longest_val_col + 1;
 
+    let staged_val = flag.display_staged_value();
+    let p2 = info.longest_staged_val_col + 1;
+
     let value_picked_from = flag.value_picked_from.to_string();
-    let p2 = info.longest_value_picked_from_col + 1;
+    let p3 = info.longest_value_picked_from_col + 1;
 
     let perm = flag.permission.to_string();
-    let p3 = info.longest_permission_col + 1;
+    let p4 = info.longest_permission_col + 1;
 
     let container = &flag.container;
 
-    format!("{full_name:p0$}{val:p1$}{value_picked_from:p2$}{perm:p3$}{container}\n")
+    format!(
+        "{full_name:p0$}{val:p1$}{staged_val:p2$}{value_picked_from:p3$}{perm:p4$}{container}\n"
+    )
 }
 
 fn set_flag(qualified_name: &str, value: &str) -> Result<()> {
@@ -188,6 +206,11 @@
     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),
+        longest_staged_val_col: flags
+            .iter()
+            .map(|f| f.display_staged_value().len())
+            .max()
+            .unwrap_or(0),
         longest_value_picked_from_col: flags
             .iter()
             .map(|f| f.value_picked_from.to_string().len())