From dc5b40a36021d6d6261f88e7d905a8ff6ff0e240 Mon Sep 17 00:00:00 2001
From: JakeSilverman <JakeRossSilverman@gmail.com>
Date: Thu, 9 Oct 2025 16:01:46 -0700
Subject: [PATCH] [Specs] ordered_map remove_or_none

---
 .../aptos-framework/doc/ordered_map.md        |  8 ++++
 .../sources/datastructures/ordered_map.move   | 43 +++++++++++++++++++
 .../datastructures/ordered_map.spec.move      | 19 +++++---
 3 files changed, 64 insertions(+), 6 deletions(-)

diff --git a/aptos-move/framework/aptos-framework/doc/ordered_map.md b/aptos-move/framework/aptos-framework/doc/ordered_map.md
index 842cddd962aa1..6158b381b9dd1 100644
--- a/aptos-move/framework/aptos-framework/doc/ordered_map.md
+++ b/aptos-move/framework/aptos-framework/doc/ordered_map.md
@@ -2129,6 +2129,14 @@ Apply the function to a mutable reference of each key-value pair in the map.
 
 <pre><code><b>pragma</b> opaque;
 <b>pragma</b> verify = <b>false</b>;
+<b>aborts_if</b> [abstract] <b>false</b>;
+<b>ensures</b> [abstract] <a href="ordered_map.md#0x1_ordered_map_spec_contains_key">spec_contains_key</a>(<b>old</b>(self), key) ==&gt; result == <a href="../../aptos-stdlib/../move-stdlib/doc/option.md#0x1_option_spec_some">option::spec_some</a>(<a href="ordered_map.md#0x1_ordered_map_spec_get">spec_get</a>(<b>old</b>(self), key));
+<b>ensures</b> [abstract] !<a href="ordered_map.md#0x1_ordered_map_spec_contains_key">spec_contains_key</a>(<b>old</b>(self), key) ==&gt; result == <a href="../../aptos-stdlib/../move-stdlib/doc/option.md#0x1_option_spec_none">option::spec_none</a>();
+<b>ensures</b> [abstract] !<a href="ordered_map.md#0x1_ordered_map_spec_contains_key">spec_contains_key</a>(self, key);
+<b>ensures</b> [abstract] <a href="../../aptos-stdlib/../move-stdlib/doc/option.md#0x1_option_spec_is_none">option::spec_is_none</a>(result) ==&gt; self == <b>old</b>(self);
+<b>ensures</b> [abstract] <a href="../../aptos-stdlib/../move-stdlib/doc/option.md#0x1_option_spec_is_some">option::spec_is_some</a>(result) ==&gt; <a href="ordered_map.md#0x1_ordered_map_spec_len">spec_len</a>(<b>old</b>(self)) == <a href="ordered_map.md#0x1_ordered_map_spec_len">spec_len</a>(self) + 1;
+<b>ensures</b> [abstract] <b>forall</b> k: K <b>where</b> k != key: <a href="ordered_map.md#0x1_ordered_map_spec_contains_key">spec_contains_key</a>(self, k) ==&gt; <a href="ordered_map.md#0x1_ordered_map_spec_get">spec_get</a>(self, k) == <a href="ordered_map.md#0x1_ordered_map_spec_get">spec_get</a>(<b>old</b>(self), k);
+<b>ensures</b> [abstract] <b>forall</b> k: K <b>where</b> k != key: <a href="ordered_map.md#0x1_ordered_map_spec_contains_key">spec_contains_key</a>(<b>old</b>(self), k) == <a href="ordered_map.md#0x1_ordered_map_spec_contains_key">spec_contains_key</a>(self, k);
 </code></pre>
 
 
diff --git a/aptos-move/framework/aptos-framework/sources/datastructures/ordered_map.move b/aptos-move/framework/aptos-framework/sources/datastructures/ordered_map.move
index f5c27551f0c4a..db35394c18ac1 100644
--- a/aptos-move/framework/aptos-framework/sources/datastructures/ordered_map.move
+++ b/aptos-move/framework/aptos-framework/sources/datastructures/ordered_map.move
@@ -1449,4 +1449,47 @@ module aptos_std::ordered_map {
         aborts_if !spec_contains_key(map, 1);
      }
 
+    #[verify_only]
+    fun test_verify_remove_or_none() {
+        let keys: vector<u64> = vector[1, 2, 3];
+        let values: vector<u64> = vector[4, 5, 6];
+        let map = new_from(keys, values);
+        spec {
+            assert spec_len(map) == 3;
+        };
+        let (_key, _value) = map.borrow_back();
+        spec{
+            assert keys[0] == 1;
+            assert keys[1] == 2;
+            assert spec_contains_key(map, 1);
+            assert spec_contains_key(map, 2);
+        };
+        let result_1 = map.remove_or_none(&1);
+        spec {
+            assert spec_contains_key(map, 2);
+            assert spec_get(map, 2) == 5;
+            assert option::spec_is_some(result_1);
+            assert option::spec_borrow(result_1) == 4;
+            assert spec_len(map) == 2;
+            assert !spec_contains_key(map, 1);
+            assert !spec_contains_key(map, 4);
+        };
+        let result_2 = map.remove_or_none(&4);
+        spec {
+            assert spec_contains_key(map, 2);
+            assert spec_get(map, 2) == 5;
+            assert option::spec_is_none(result_2);
+            assert spec_len(map) == 2;
+            assert !spec_contains_key(map, 4);
+        };
+        map.remove(&2);
+        map.remove(&3);
+        spec {
+            assert !spec_contains_key(map, 1);
+            assert !spec_contains_key(map, 2);
+            assert !spec_contains_key(map, 3);
+            assert spec_len(map) == 0;
+        };
+        map.destroy_empty();
+    }
 }
diff --git a/aptos-move/framework/aptos-framework/sources/datastructures/ordered_map.spec.move b/aptos-move/framework/aptos-framework/sources/datastructures/ordered_map.spec.move
index 8e7d64ab75f92..f36e6f6a6250e 100644
--- a/aptos-move/framework/aptos-framework/sources/datastructures/ordered_map.spec.move
+++ b/aptos-move/framework/aptos-framework/sources/datastructures/ordered_map.spec.move
@@ -62,6 +62,19 @@ spec aptos_std::ordered_map {
         ensures [abstract] forall k: K where k != key: spec_contains_key(old(self), k) == spec_contains_key(self, k);
     }
 
+    spec remove_or_none {
+        pragma opaque;
+        pragma verify = false;
+        aborts_if [abstract] false;
+        ensures [abstract] spec_contains_key(old(self), key) ==> result == option::spec_some(spec_get(old(self), key));
+        ensures [abstract] !spec_contains_key(old(self), key) ==> result == option::spec_none();
+        ensures [abstract] !spec_contains_key(self, key);
+        ensures [abstract] option::spec_is_none(result) ==> self == old(self);
+        ensures [abstract] option::spec_is_some(result) ==> spec_len(old(self)) == spec_len(self) + 1;
+        ensures [abstract] forall k: K where k != key: spec_contains_key(self, k) ==> spec_get(self, k) == spec_get(old(self), k);
+        ensures [abstract] forall k: K where k != key: spec_contains_key(old(self), k) == spec_contains_key(self, k);
+    }
+
     spec is_empty {
         pragma intrinsic;
     }
@@ -281,10 +294,4 @@ spec aptos_std::ordered_map {
         pragma opaque;
         pragma verify = false;
     }
-
-    spec remove_or_none {
-        pragma opaque;
-        pragma verify = false;
-    }
-
 }