From 9904a29ccc8db5a3b6dca6dd90d393facc4f64dc Mon Sep 17 00:00:00 2001
From: Nika Layzell <nika@thelayzells.com>
Date: Tue, 3 Sep 2024 15:44:15 -0400
Subject: [PATCH] More consistently tidy files before serializing

---
 src/format.rs                                 | 30 ++++++++++++++++++-
 src/main.rs                                   |  4 +++
 src/serialization.rs                          | 19 ++++++++++++
 src/storage.rs                                | 21 +++++--------
 ...__aggregate__merge_audits_files_basic.snap |  8 ++---
 5 files changed, 63 insertions(+), 19 deletions(-)

diff --git a/src/format.rs b/src/format.rs
index 40ccb34a..d1e3c87c 100644
--- a/src/format.rs
+++ b/src/format.rs
@@ -2,7 +2,7 @@
 
 use crate::errors::{StoreVersionParseError, VersionParseError};
 use crate::resolver::{DiffRecommendation, ViolationConflict};
-use crate::serialization::spanned::Spanned;
+use crate::serialization::{spanned::Spanned, Tidyable};
 use crate::{flock::Filesystem, serialization};
 use core::{cmp, fmt};
 use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet};
@@ -235,6 +235,14 @@ pub struct AuditsFile {
     pub trusted: TrustedPackages,
 }
 
+impl Tidyable for AuditsFile {
+    fn tidy(&mut self) {
+        self.audits.tidy();
+        self.wildcard_audits.tidy();
+        self.trusted.tidy();
+    }
+}
+
 /// Foreign audits.toml with unparsed entries and audits. Should have the same
 /// structure as `AuditsFile`, but with individual audits and criteria unparsed.
 #[derive(serde::Deserialize, Clone, Debug)]
@@ -583,6 +591,12 @@ pub struct ConfigFile {
     pub exemptions: SortedMap<PackageName, Vec<ExemptedDependency>>,
 }
 
+impl Tidyable for ConfigFile {
+    fn tidy(&mut self) {
+        self.exemptions.tidy();
+    }
+}
+
 pub static SAFE_TO_DEPLOY: CriteriaStr = "safe-to-deploy";
 pub static SAFE_TO_RUN: CriteriaStr = "safe-to-run";
 pub static DEFAULT_CRITERIA: CriteriaStr = SAFE_TO_DEPLOY;
@@ -1011,6 +1025,16 @@ pub struct ImportsFile {
     pub audits: SortedMap<ImportName, AuditsFile>,
 }
 
+impl Tidyable for ImportsFile {
+    fn tidy(&mut self) {
+        self.unpublished.tidy();
+        self.publisher.tidy();
+        for audits_file in self.audits.values_mut() {
+            audits_file.tidy();
+        }
+    }
+}
+
 /// Information about who published a specific version of a crate to be cached
 /// in imports.lock.
 #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
@@ -1073,6 +1097,10 @@ pub enum DiffCache {
     },
 }
 
+impl Tidyable for DiffCache {
+    fn tidy(&mut self) {}
+}
+
 impl Default for DiffCache {
     fn default() -> Self {
         DiffCache::V2 {
diff --git a/src/main.rs b/src/main.rs
index 57608150..9eeb9f52 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -28,6 +28,7 @@ use out::{progress_bar, IncProgressOnDrop};
 use reqwest::Url;
 use serde::de::Deserialize;
 use serialization::spanned::Spanned;
+use serialization::Tidyable;
 use storage::fetch_registry;
 use thiserror::Error;
 use tracing::{error, info, trace, warn};
@@ -2553,6 +2554,9 @@ fn do_aggregate_audits(sources: Vec<(String, AuditsFile)>) -> Result<AuditsFile,
                 }));
         }
     }
+
+    aggregate.tidy();
+
     if errors.is_empty() {
         Ok(aggregate)
     } else {
diff --git a/src/serialization.rs b/src/serialization.rs
index 32c5f015..6c1942bb 100644
--- a/src/serialization.rs
+++ b/src/serialization.rs
@@ -372,6 +372,25 @@ pub mod audit {
     }
 }
 
+/// Trait implemented by format data types which may want to be cleaned up
+/// before they are serialized.
+pub trait Tidyable {
+    /// Ensure that the data structure is tidy and ready to be serialized.
+    /// This may remove empty entries from maps, ensure lists are sorted, etc.
+    fn tidy(&mut self);
+}
+
+/// Helper for tidying the common audit data structure, removing empty entries,
+/// and sorting audit lists.
+impl<K: Ord, E: Ord> Tidyable for SortedMap<K, Vec<E>> {
+    fn tidy(&mut self) {
+        self.retain(|_, entries| {
+            entries.sort();
+            !entries.is_empty()
+        });
+    }
+}
+
 /// Inline arrays which have a representation longer than this will be rendered
 /// over multiple lines.
 const ARRAY_WRAP_THRESHOLD: usize = 80;
diff --git a/src/storage.rs b/src/storage.rs
index e9442026..a4ec287d 100644
--- a/src/storage.rs
+++ b/src/storage.rs
@@ -41,7 +41,7 @@ use crate::{
     },
     network::Network,
     out::{progress_bar, IncProgressOnDrop},
-    serialization::{parse_from_value, spanned::Spanned, to_formatted_toml},
+    serialization::{parse_from_value, spanned::Spanned, to_formatted_toml, Tidyable},
     Config, PackageExt, PartialConfig, CARGO_ENV,
 };
 
@@ -2992,12 +2992,14 @@ where
 }
 fn store_toml<T>(
     heading: &str,
-    val: T,
+    mut val: T,
     user_info: Option<&FastMap<CratesUserId, CratesCacheUser>>,
 ) -> Result<String, StoreTomlError>
 where
-    T: Serialize,
+    T: Serialize + Tidyable,
 {
+    val.tidy();
+
     let toml_document = to_formatted_toml(val, user_info)?;
     Ok(format!("{heading}{toml_document}"))
 }
@@ -3019,25 +3021,16 @@ where
     Ok(json_string)
 }
 fn store_audits(
-    mut audits: AuditsFile,
+    audits: AuditsFile,
     user_info: &FastMap<CratesUserId, CratesCacheUser>,
 ) -> Result<String, StoreTomlError> {
     let heading = r###"
 # cargo-vet audits file
 "###;
-    audits
-        .audits
-        .values_mut()
-        .for_each(|entries| entries.sort());
 
     store_toml(heading, audits, Some(user_info))
 }
-fn store_config(mut config: ConfigFile) -> Result<String, StoreTomlError> {
-    config
-        .exemptions
-        .values_mut()
-        .for_each(|entries| entries.sort());
-
+fn store_config(config: ConfigFile) -> Result<String, StoreTomlError> {
     let heading = r###"
 # cargo-vet config file
 "###;
diff --git a/src/tests/snapshots/cargo_vet__tests__aggregate__merge_audits_files_basic.snap b/src/tests/snapshots/cargo_vet__tests__aggregate__merge_audits_files_basic.snap
index 2ceadf2f..79bd607e 100644
--- a/src/tests/snapshots/cargo_vet__tests__aggregate__merge_audits_files_basic.snap
+++ b/src/tests/snapshots/cargo_vet__tests__aggregate__merge_audits_files_basic.snap
@@ -26,13 +26,13 @@ aggregated-from = "https://source2.example.com/supply_chain/audits.toml"
 
 [[audits.package1]]
 criteria = "safe-to-deploy"
-version = "10.0.0"
-aggregated-from = "https://source1.example.com/supply_chain/audits.toml"
+version = "5.0.0"
+aggregated-from = "https://source2.example.com/supply_chain/audits.toml"
 
 [[audits.package1]]
 criteria = "safe-to-deploy"
-version = "5.0.0"
-aggregated-from = "https://source2.example.com/supply_chain/audits.toml"
+version = "10.0.0"
+aggregated-from = "https://source1.example.com/supply_chain/audits.toml"
 
 [[audits.package1]]
 criteria = "safe-to-deploy"