From 50aeacd6001f224044d058275e9eb7709593aa6c Mon Sep 17 00:00:00 2001
From: Corto <adrena.corto@gmail.com>
Date: Fri, 1 Nov 2024 10:17:43 +0800
Subject: [PATCH] Add helpers around StakingType

---
 src/types.rs | 28 ++++++++++++++++++++++++++++
 1 file changed, 28 insertions(+)

diff --git a/src/types.rs b/src/types.rs
index 62d44de..b74ff19 100644
--- a/src/types.rs
+++ b/src/types.rs
@@ -524,6 +524,28 @@ pub enum StakingType {
     LP = 2,
 }
 
+impl From<StakingType> for u8 {
+    fn from(val: StakingType) -> Self {
+        match val {
+            StakingType::LM => 1,
+            StakingType::LP => 2,
+        }
+    }
+}
+
+impl TryFrom<u8> for StakingType {
+    type Error = anyhow::Error;
+
+    fn try_from(value: u8) -> std::result::Result<Self, Self::Error> {
+        Ok(match value {
+            1 => StakingType::LM,
+            2 => StakingType::LP,
+            // Return an error if unknown value
+            _ => anyhow::bail!("Invalid staking type"),
+        })
+    }
+}
+
 #[account(zero_copy)]
 #[derive(Default, Debug, PartialEq, AnchorSerialize, AnchorDeserialize)]
 #[repr(C)]
@@ -537,6 +559,12 @@ pub struct UserStaking {
     pub locked_stakes: [LockedStake; MAX_LOCKED_STAKE_COUNT],
 }
 
+impl UserStaking {
+    pub fn get_staking_type(&self) -> StakingType {
+        StakingType::try_from(self.staking_type).unwrap()
+    }
+}
+
 #[derive(
     Copy, Clone, PartialEq, AnchorSerialize, AnchorDeserialize, Default, Debug, Zeroable, Pod,
 )]