From b80b67e2cf8f3ce726d41e69749aea35e1682ef5 Mon Sep 17 00:00:00 2001
From: Julien Cretin <cretin@google.com>
Date: Fri, 4 Mar 2022 10:23:44 +0100
Subject: [PATCH] Factorize tock storage in assert_once

---
 src/env/tock/mod.rs | 12 ++++++++----
 1 file changed, 8 insertions(+), 4 deletions(-)

diff --git a/src/env/tock/mod.rs b/src/env/tock/mod.rs
index e5973c0..1facc85 100644
--- a/src/env/tock/mod.rs
+++ b/src/env/tock/mod.rs
@@ -77,18 +77,22 @@ impl Env for TockEnv {
     }
 
     fn storage(&mut self) -> StorageResult<Self::Storage> {
-        assert!(!self.storage);
-        self.storage = true;
+        assert_once(&mut self.storage);
         unsafe { steal_storage() }
     }
 
     fn upgrade_storage(&mut self) -> StorageResult<Self::UpgradeStorage> {
-        assert!(!self.upgrade_storage);
-        self.upgrade_storage = true;
+        assert_once(&mut self.upgrade_storage);
         SyscallUpgradeStorage::new()
     }
 }
 
+/// Asserts a boolean is false and sets it to true.
+fn assert_once(b: &mut bool) {
+    assert!(!*b);
+    *b = true;
+}
+
 // Returns whether the keepalive was sent, or false if cancelled.
 fn send_keepalive_up_needed(
     cid: ChannelID,