From 4b2d3c9551bb62276e0713beb8e940d04049bdd2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jo=C3=A3o=20M=2E=20Bezerra?= <marcospb19@hotmail.com>
Date: Wed, 28 Jun 2023 12:44:52 -0300
Subject: [PATCH] toml_edit: add `retain` to arrays and tables

added to: `Table`, `InlineTable`, `Array` and `ArrayOfTables`
---
 crates/toml_edit/src/array.rs           | 14 ++++++++++++++
 crates/toml_edit/src/array_of_tables.rs | 14 ++++++++++++++
 crates/toml_edit/src/inline_table.rs    | 18 ++++++++++++++++++
 crates/toml_edit/src/table.rs           | 14 ++++++++++++++
 4 files changed, 60 insertions(+)

diff --git a/crates/toml_edit/src/array.rs b/crates/toml_edit/src/array.rs
index 68cb4467..045b451e 100644
--- a/crates/toml_edit/src/array.rs
+++ b/crates/toml_edit/src/array.rs
@@ -303,6 +303,20 @@ impl Array {
         }
     }
 
+    /// Retains only the values specified by the `keep` predicate.
+    ///
+    /// In other words, remove all values for which `keep(&value)` returns `false`.
+    ///
+    /// This method operates in place, visiting each element exactly once in the
+    /// original order, and preserves the order of the retained elements.
+    pub fn retain<F>(&mut self, mut keep: F)
+    where
+        F: FnMut(&Value) -> bool,
+    {
+        self.values
+            .retain(|item| item.as_value().map(&mut keep).unwrap_or(false));
+    }
+
     fn value_op<T>(
         &mut self,
         v: Value,
diff --git a/crates/toml_edit/src/array_of_tables.rs b/crates/toml_edit/src/array_of_tables.rs
index 461a716c..c4d71948 100644
--- a/crates/toml_edit/src/array_of_tables.rs
+++ b/crates/toml_edit/src/array_of_tables.rs
@@ -91,6 +91,20 @@ impl ArrayOfTables {
     pub fn remove(&mut self, index: usize) {
         self.values.remove(index);
     }
+
+    /// Retains only the elements specified by the `keep` predicate.
+    ///
+    /// In other words, remove all tables for which `keep(&table)` returns `false`.
+    ///
+    /// This method operates in place, visiting each element exactly once in the
+    /// original order, and preserves the order of the retained elements.
+    pub fn retain<F>(&mut self, mut keep: F)
+    where
+        F: FnMut(&Table) -> bool,
+    {
+        self.values
+            .retain(|item| item.as_table().map(&mut keep).unwrap_or(false));
+    }
 }
 
 /// An iterator type over `ArrayOfTables`'s values.
diff --git a/crates/toml_edit/src/inline_table.rs b/crates/toml_edit/src/inline_table.rs
index 9327818e..3dc6c0c0 100644
--- a/crates/toml_edit/src/inline_table.rs
+++ b/crates/toml_edit/src/inline_table.rs
@@ -363,6 +363,24 @@ impl InlineTable {
             kv.value.into_value().ok().map(|value| (key, value))
         })
     }
+
+    /// Retains only the elements specified by the `keep` predicate.
+    ///
+    /// In other words, remove all pairs `(key, value)` for which
+    /// `keep(&key, &mut value)` returns `false`.
+    ///
+    /// The elements are visited in iteration order.
+    pub fn retain<F>(&mut self, mut keep: F)
+    where
+        F: FnMut(&str, &mut Value) -> bool,
+    {
+        self.items.retain(|key, item| {
+            item.value
+                .as_value_mut()
+                .map(|value| keep(key, value))
+                .unwrap_or(false)
+        });
+    }
 }
 
 impl std::fmt::Display for InlineTable {
diff --git a/crates/toml_edit/src/table.rs b/crates/toml_edit/src/table.rs
index 2f61abf7..45d6d61b 100644
--- a/crates/toml_edit/src/table.rs
+++ b/crates/toml_edit/src/table.rs
@@ -397,6 +397,20 @@ impl Table {
     pub fn remove_entry(&mut self, key: &str) -> Option<(Key, Item)> {
         self.items.shift_remove(key).map(|kv| (kv.key, kv.value))
     }
+
+    /// Retains only the elements specified by the `keep` predicate.
+    ///
+    /// In other words, remove all pairs `(key, item)` for which
+    /// `keep(&key, &mut item)` returns `false`.
+    ///
+    /// The elements are visited in iteration order.
+    pub fn retain<F>(&mut self, mut keep: F)
+    where
+        F: FnMut(&str, &mut Item) -> bool,
+    {
+        self.items
+            .retain(|key, key_value| keep(key, &mut key_value.value));
+    }
 }
 
 impl std::fmt::Display for Table {