From a8af6b449ac8ad607595649efde08e2d2b46d668 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?S=C3=A9bastien=20d=27Herbais=20de=20Thun?=
 <sebastien.d.herbais@gmail.com>
Date: Tue, 10 Oct 2023 11:44:59 +0200
Subject: [PATCH] Adds a default value to `.remove()` on `dict` and `array`
 (#2346)

---
 crates/typst/src/eval/array.rs   |  9 +++++++--
 crates/typst/src/eval/dict.rs    | 16 +++++++++++-----
 crates/typst/src/eval/methods.rs | 11 +++++++++--
 tests/typ/compiler/array.typ     | 13 +++++++++++++
 tests/typ/compiler/dict.typ      | 12 ++++++++++++
 5 files changed, 52 insertions(+), 9 deletions(-)

diff --git a/crates/typst/src/eval/array.rs b/crates/typst/src/eval/array.rs
index 926427027..bee3fae75 100644
--- a/crates/typst/src/eval/array.rs
+++ b/crates/typst/src/eval/array.rs
@@ -247,9 +247,14 @@ impl Array {
         /// The index at which to remove the item. If negative, indexes from
         /// the back.
         index: i64,
+        /// A default value to return if the index is out of bounds.
+        #[named]
+        default: Option<Value>,
     ) -> StrResult<Value> {
-        let i = self.locate(index, false)?;
-        Ok(self.0.remove(i))
+        self.locate_opt(index, false)
+            .map(|i| self.0.remove(i))
+            .or(default)
+            .ok_or_else(|| out_of_bounds_no_default(index, self.len()))
     }
 
     /// Extracts a subslice of the array. Fails with an error if the start or
diff --git a/crates/typst/src/eval/dict.rs b/crates/typst/src/eval/dict.rs
index c5da7c158..b77686fee 100644
--- a/crates/typst/src/eval/dict.rs
+++ b/crates/typst/src/eval/dict.rs
@@ -169,11 +169,17 @@ impl Dict {
 
     /// Removes a pair from the dictionary by key and return the value.
     #[func]
-    pub fn remove(&mut self, key: Str) -> StrResult<Value> {
-        match Arc::make_mut(&mut self.0).shift_remove(&key) {
-            Some(value) => Ok(value),
-            None => Err(missing_key(&key)),
-        }
+    pub fn remove(
+        &mut self,
+        key: Str,
+        /// A default value to return if the key does not exist.
+        #[named]
+        default: Option<Value>,
+    ) -> StrResult<Value> {
+        Arc::make_mut(&mut self.0)
+            .shift_remove(&key)
+            .or(default)
+            .ok_or_else(|| missing_key(&key))
     }
 
     /// Returns the keys of the dictionary as an array in insertion order.
diff --git a/crates/typst/src/eval/methods.rs b/crates/typst/src/eval/methods.rs
index b8d71c76d..843549e65 100644
--- a/crates/typst/src/eval/methods.rs
+++ b/crates/typst/src/eval/methods.rs
@@ -51,13 +51,20 @@ pub fn call_mut(
             "insert" => {
                 array.insert(args.expect("index")?, args.expect("value")?).at(span)?
             }
-            "remove" => output = array.remove(args.expect("index")?).at(span)?,
+            "remove" => {
+                output = array
+                    .remove(args.expect("index")?, args.named("default")?)
+                    .at(span)?
+            }
             _ => return missing(),
         },
 
         Value::Dict(dict) => match method {
             "insert" => dict.insert(args.expect::<Str>("key")?, args.expect("value")?),
-            "remove" => output = dict.remove(args.expect::<Str>("key")?).at(span)?,
+            "remove" => {
+                output =
+                    dict.remove(args.expect("key")?, args.named("default")?).at(span)?
+            }
             _ => return missing(),
         },
 
diff --git a/tests/typ/compiler/array.typ b/tests/typ/compiler/array.typ
index 92d926a07..12c3c3a01 100644
--- a/tests/typ/compiler/array.typ
+++ b/tests/typ/compiler/array.typ
@@ -63,6 +63,19 @@
 #test((1, 2, 3).at(2, default: 5), 3)
 #test((1, 2, 3).at(3, default: 5), 5)
 
+---
+// Test remove with default value.
+
+#{
+  let array = (1, 2, 3)
+  test(array.remove(2, default: 5), 3)
+}
+
+#{
+  let array = (1, 2, 3)
+  test(array.remove(3, default: 5), 5)
+}
+
 ---
 // Test bad lvalue.
 // Error: 2:3-2:14 cannot mutate a temporary value
diff --git a/tests/typ/compiler/dict.typ b/tests/typ/compiler/dict.typ
index 2668f347d..f3a70c2b6 100644
--- a/tests/typ/compiler/dict.typ
+++ b/tests/typ/compiler/dict.typ
@@ -40,6 +40,18 @@
 #test((a: 1, b: 2).at("b", default: 3), 2)
 #test((a: 1, b: 2).at("c", default: 3), 3)
 
+---
+// Test remove with default value.
+#{
+  let dict = (a: 1, b: 2)
+  test(dict.remove("b", default: 3), 2)
+}
+
+#{
+  let dict = (a: 1, b: 2)
+  test(dict.remove("c", default: 3), 3)
+}
+
 ---
 // Missing lvalue is not automatically none-initialized.
 #{