Merge pull request #36 from invariant-labs/add-claim-fee

Added claim_fee entrypoint
invariant-labs · Dec 12, 2023 · a54c33e · a54c33e
2 parents ef793f8 + aa2d617
commit a54c33e
Show file tree

Hide file tree

Showing 3 changed files with 59 additions and 0 deletions.
diff --git a/resources/invariant_schema.json b/resources/invariant_schema.json
@@ -216,6 +216,27 @@
         }
       }
     },
+    {
+      "name": "claim_fee",
+      "is_mutable": true,
+      "args": [
+        {
+          "name": "index",
+          "ty": "U32"
+        }
+      ],
+      "return_ty": {
+        "Result": {
+          "ok": {
+            "Tuple2": [
+              "Any",
+              "Any"
+            ]
+          },
+          "err": "U32"
+        }
+      }
+    },
     {
       "name": "create_position",
       "is_mutable": true,

diff --git a/src/contracts/entrypoints.rs b/src/contracts/entrypoints.rs
@@ -41,6 +41,8 @@ pub trait Entrypoints {
     fn is_tick_initialized(&self, key: PoolKey, index: i32) -> bool;
     fn get_tick(&self, key: PoolKey, index: i32) -> Result<Tick, InvariantError>;
 
+    fn claim_fee(&mut self, index: u32) -> Result<(TokenAmount, TokenAmount), InvariantError>;
+
     fn create_position(
         &mut self,
         pool_key: PoolKey,

diff --git a/src/lib.rs b/src/lib.rs
@@ -333,6 +333,42 @@ impl Entrypoints for Invariant {
         self.ticks.get(key, index)
     }
 
+    pub fn claim_fee(&mut self, index: u32) -> Result<(TokenAmount, TokenAmount), InvariantError> {
+        let caller = odra::contract_env::caller();
+        let current_timestamp = odra::contract_env::get_block_time();
+        let mut position = self.positions.get(caller, index)?;
+        let mut lower_tick = self
+            .ticks
+            .get(position.pool_key, position.lower_tick_index)?;
+        let mut upper_tick = self
+            .ticks
+            .get(position.pool_key, position.upper_tick_index)?;
+        let mut pool = self.pools.get(position.pool_key)?;
+
+        let (x, y) = position.claim_fee(
+            &mut pool,
+            &mut upper_tick,
+            &mut lower_tick,
+            current_timestamp,
+        );
+
+        self.positions.update(caller, index, &position)?;
+        self.pools.update(position.pool_key, &pool)?;
+        self.ticks
+            .update(position.pool_key, position.lower_tick_index, &lower_tick)?;
+        self.ticks
+            .update(position.pool_key, position.upper_tick_index, &upper_tick)?;
+
+        if x.get().is_zero() {
+            Erc20Ref::at(&position.pool_key.token_x).transfer(&caller, &x.get());
+        }
+
+        if y.get().is_zero() {
+            Erc20Ref::at(&position.pool_key.token_y).transfer(&caller, &y.get());
+        }
+
+        Ok((x, y))
+    }
     pub fn create_position(
         &mut self,
         pool_key: PoolKey,