address feedback

gear-tech · Nov 16, 2023 · b2e6e74 · b2e6e74
1 parent fca257c
commit b2e6e74
Show file tree

Hide file tree

Showing 4 changed files with 22 additions and 12 deletions.
diff --git a/core-processor/src/ext.rs b/core-processor/src/ext.rs
@@ -721,6 +721,13 @@ impl Externalities for Ext {
             .map_err(Into::into)
     }
 
+    fn free(&mut self, page: WasmPage) -> Result<(), Self::AllocError> {
+        self.context
+            .allocations_context
+            .free(page)
+            .map_err(Into::into)
+    }
+
     fn free_range(&mut self, start: WasmPage, end: WasmPage) -> Result<(), Self::AllocError> {
         if start > end {
             return Err(AllocExtError::Alloc(AllocError::InvalidRange(
@@ -1314,7 +1321,7 @@ mod tests {
         assert!(ext.free(existing_page).is_ok());
         assert_eq!(ext.gas_left(), gas_left);
 
-        assert!(ext.free(non_existing_page).is_ok());
+        assert!(ext.free(non_existing_page).is_err());
         assert_eq!(ext.gas_left(), gas_left);
     }
 

diff --git a/core/src/env.rs b/core/src/env.rs
@@ -192,10 +192,8 @@ pub trait Externalities {
         mem: &mut impl Memory,
     ) -> Result<WasmPage, Self::AllocError>;
 
-    /// Free specific memory page
-    fn free(&mut self, page: WasmPage) -> Result<(), Self::AllocError> {
-        self.free_range(page, page)
-    }
+    /// Free specific page
+    fn free(&mut self, page: WasmPage) -> Result<(), Self::AllocError>;
 
     /// Free specific memory range
     fn free_range(&mut self, start: WasmPage, end: WasmPage) -> Result<(), Self::AllocError>;

diff --git a/core/src/memory.rs b/core/src/memory.rs
@@ -381,20 +381,23 @@ impl AllocationsContext {
 
     /// Free specific memory page.
     pub fn free(&mut self, page: WasmPage) -> Result<(), AllocError> {
-        self.free_range(page..=page)
+        if page < self.static_pages || page >= self.max_pages {
+            return Err(AllocError::InvalidFree(page.0));
+        }
+
+        if !self.allocations.remove(&page) {
+            return Err(AllocError::InvalidFree(page.0));
+        }
+
+        Ok(())
     }
 
     /// Try to free pages in range. Will only return error if range is invalid.
     ///
     /// Currently running program should own this pages.
     pub fn free_range(&mut self, range: RangeInclusive<WasmPage>) -> Result<(), AllocError> {
         if *range.start() < self.static_pages || *range.end() >= self.max_pages {
-            let page = if *range.start() < self.static_pages {
-                range.start().0
-            } else {
-                range.end().0
-            };
-            return Err(AllocError::InvalidFree(page));
+            return Err(AllocError::InvalidRange(range.start().0, range.end().0));
         }
 
         self.allocations.retain(|p| !range.contains(p));
@@ -600,6 +603,7 @@ mod tests {
         fn assert_free_error(err: AllocError) {
             match err {
                 AllocError::InvalidFree(_) => {}
+                AllocError::InvalidRange(_, _) => {}
                 err => panic!("{err:?}"),
             }
         }

diff --git a/pallets/gear/src/tests.rs b/pallets/gear/src/tests.rs
@@ -13810,6 +13810,7 @@ fn free_range_success() {
         run_to_next_block(None);
 
         assert_succeed(mid);
+        assert!(Gear::is_initialized(pid));
         assert!(Gear::is_active(pid));
     });
 }