icrc 0.0.1

Collection of ICRC Packages
Documentation
use std::collections::BTreeMap;
use candid::{CandidType, Deserialize, Int, Nat, Principal};
use candid::types::principal;
use icrc_ledger_types::icrc1::account::{Account};
use num_bigint::BigUint;
use num_traits::ToPrimitive;
use serde::Serialize;

use crate::icrc7::consts::*;
use crate::icrc7::types::{TransferArgs, TransferError, TransferResult};

#[derive(CandidType, Serialize, Deserialize, Clone, Debug, PartialEq, Eq)]
pub enum Value {
    Blob(Vec<u8>),
    Text(String),
    Nat(Nat),
    Int(Int),
    Array(Vec<Value>),
    Map(Vec<(String, Value)>),
}

pub struct Token {
    owner: Account,
    metadata: Vec<(String, Value)>,
}

impl Token {
    pub fn new() -> Self {
        Self {
            owner: Account::from(principal::Principal::management_canister()),
            metadata: Vec::new(),
        }
    }

    pub fn with_owner(mut self, owner: Account) -> Self {
        self.owner = owner;
        self
    }

    pub fn with_metadata(mut self, metadata: Vec<(String, Value)>) -> Self {
        self.metadata = metadata;
        self
    }
}

pub struct Ledger {
    pub symbol: String,
    pub name: String,
    pub description: Option<String>,
    pub logo: Option<String>,
    pub tokens: BTreeMap<Nat, Token>,
    pub owners: BTreeMap<Account, Vec<Nat>>,
    pub supply_cap: Option<Nat>,
}

impl Ledger {
    pub fn new(symbol: String, name: String) -> Self {
        Self {
            symbol,
            name,
            description: None,
            logo: None,
            tokens: BTreeMap::new(),
            owners: BTreeMap::new(),
            supply_cap: None,
        }
    }

    pub fn collection_metadata(&self) -> Vec<(String, Value)> {
        let mut vec = Vec::new();
        vec.push(("icrc7:symbol".to_string(), Value::Text(self.symbol.clone())));
        vec.push(("icrc7:name".to_string(), Value::Text(self.name.clone())));
        if let Some(description) = &self.description {
            vec.push(("icrc7:description".to_string(), Value::Text(description.clone())));
        }
        if let Some(logo) = &self.logo {
            vec.push(("icrc7:logo".to_string(), Value::Text(logo.clone())));
        }
        vec.push(("icrc7:total_supply".to_string(), Value::Nat(self.tokens.len().into())));
        if let Some(supply_cap) = &self.supply_cap {
            vec.push(("icrc7:supply_cap".to_string(), Value::Nat(supply_cap.clone())));
        }
        vec
    }

    pub fn token_metadata(&self, token_ids: Vec<Nat>) -> Vec<Option<Vec<(String, Value)>>> {
        token_ids
            .iter()
            .map(|token_id| {
                self.tokens
                    .get(token_id)
                    .map(|token| token.metadata.clone())
            })
            .collect()
    }

    pub fn owner_of(&self, token_ids: Vec<Nat>) -> Vec<Option<Account>> {
        token_ids
            .iter()
            .map(|token_id| {
                self.tokens
                    .get(token_id)
                    .map(|token| token.owner.clone())
            })
            .collect()
    }

    pub fn balance_of(&self, owners: Vec<Account>) -> Vec<Nat> {
        owners
            .iter()
            .map(|owner| {
                self.owners
                    .get(owner)
                    .map(|tokens| Nat::from(tokens.len()))
                    .unwrap_or(Nat::from(0u64))
            })
            .collect()
    }

    pub fn create_tokens(&mut self, token_ids: Vec<Nat>, owner: Account) {
        self.tokens.extend(token_ids.iter().map(|token_id| {
            (token_id.clone(), Token::new().with_owner(owner.clone()))
        }));
        self.owners.entry(owner).or_insert(Vec::new()).extend(token_ids);
        self.owners.get_mut(&owner).unwrap().sort(); // Should be fast with nearly sorted data.
    }

    pub fn tokens(&self, prev: Option<Nat>, take: Option<Nat>) -> Vec<Nat> {
        let n = take
            .unwrap_or(Nat::from(DEFAULT_TAKE_VALUE)).0 // Default to DEFAULT_TAKE_VALUE.
            .min(BigUint::from(MAX_TAKE_VALUE)) // Limit to MAX_TAKE_VALUE.
            .to_usize().unwrap(); // Convert to usize.
        return if let Some(start) = prev {
            self.tokens.iter()
                .skip_while(|(token_id, _)| token_id.0.le(&start.0))
                .map(|(token_id, _)| token_id.clone())
                .take(n)
                .collect()
        } else {
            // If prev is None, return the first n tokens.
            self.tokens.keys().take(n).cloned().collect()
        };
    }

    pub fn tokens_of(&self, account: Account, prev: Option<Nat>, take: Option<Nat>) -> Vec<Nat> {
        let n = take
            .unwrap_or(Nat::from(DEFAULT_TAKE_VALUE)).0 // Default to DEFAULT_TAKE_VALUE.
            .min(BigUint::from(MAX_TAKE_VALUE)) // Limit to MAX_TAKE_VALUE.
            .to_usize().unwrap(); // Convert to usize.

        return if let Some(start) = prev {
            self.owners.get(&account).unwrap_or(&Vec::new())
                .iter()
                .skip_while(|token_id| token_id.0.le(&start.0))
                .take(n)
                .cloned()
                .collect()
        } else {
            // If prev is None, return the first n tokens.
            self.owners.get(&account).unwrap_or(&Vec::new())
                .iter()
                .take(n)
                .cloned()
                .collect()
        };
    }

    pub fn transfer(&self, caller: Principal, args: Vec<TransferArgs>, atomic: bool) -> Vec<Option<TransferResult>> {
        if args.len() > MAX_UPDATE_BATCH_SIZE as usize {
            return vec![Some(Err(TransferError::GenericBatchError {
                error_code: Nat::from(1u64), // TODO: Define error codes.
                message: "Batch size exceeds the maximum allowed".to_string(),
            })); args.len()];
        }
        let mut results = vec![None; args.len()];
        for (i, args) in args.iter().enumerate() {
            let result = self.transfer_single(caller, args);
            if result.is_err() && atomic {
                results = vec![None; i + 1];
                results[i] = Some(result);
                return results;
            }
            results[i] = Some(result);
        }
        results
    }

    fn transfer_single(&self, caller: Principal, args: &TransferArgs) -> TransferResult {
        let token = self.tokens.get(&args.token_id)
            .ok_or(TransferError::NonExistingTokenId)?; // Non existing token.
        if token.owner.owner != caller {
            return Err(TransferError::Unauthorized);
        }
        Ok(Nat::from(0u64))
    }
}

#[test]
fn test_token_metadata() {
    let mut ledger = Ledger::new("ICRC7".to_string(), "icrc7".to_string());

    let metadata = ledger.token_metadata(vec![Nat::from(0u64)]);
    assert_eq!(metadata.len(), 1);
    assert_eq!(metadata[0], None);

    ledger.tokens.insert(Nat::from(0u64), Token::new());
    let metadata = ledger.token_metadata(vec![Nat::from(0u64)]);
    assert_eq!(metadata.len(), 1);
    assert_eq!(metadata[0].is_some(), true);
}

#[test]
fn test_collection_metadata() {
    let ledger = Ledger::new("ICRC7".to_string(), "icrc7".to_string());
    let metadata = ledger.collection_metadata();
    assert_eq!(metadata.len(), 3);
    assert_eq!(metadata[0].0, "icrc7:symbol");
    assert_eq!(metadata[0].1, Value::Text("ICRC7".to_string()));
    assert_eq!(metadata[1].0, "icrc7:name");
    assert_eq!(metadata[1].1, Value::Text("icrc7".to_string()));
    assert_eq!(metadata[2].0, "icrc7:total_supply");
    assert_eq!(metadata[2].1, Value::Nat(Nat::from(0u64)));
}

#[test]
fn test_owner_of() {
    let mut ledger = Ledger::new("ICRC7".to_string(), "icrc7".to_string());

    let owners = ledger.owner_of(vec![Nat::from(0u64)]);
    assert_eq!(owners.len(), 1);
    assert_eq!(owners[0], None);

    ledger.tokens.insert(Nat::from(0u64), Token::new());
    let owners = ledger.owner_of(vec![Nat::from(0u64)]);
    assert_eq!(owners.len(), 1);
    assert_eq!(owners[0].unwrap().owner, Principal::management_canister());
}

#[test]
fn test_tokens() {
    let account = Account::from(Principal::management_canister());
    let mut ledger = Ledger::new("ICRC7".to_string(), "icrc7".to_string());

    assert_eq!(ledger.tokens(None, None).len(), 0);

    ledger.create_tokens(vec![Nat::from(0u64)], account.clone());
    assert_eq!(ledger.tokens(None, None).len(), 1);
    assert_eq!(ledger.tokens(Some(Nat::from(0u64)), None).len(), 0);

    ledger.create_tokens((1..=MAX_TAKE_VALUE + 10).map(Nat::from).collect(), account.clone());
    assert_eq!(ledger.tokens(None, None).len(), DEFAULT_TAKE_VALUE as usize);
    assert_eq!(ledger.tokens(None, Some(Nat::from(MAX_TAKE_VALUE))).len(), MAX_TAKE_VALUE as usize);

    let tokens = ledger.tokens(Some(Nat::from(DEFAULT_TAKE_VALUE)), None);
    assert_eq!(tokens.len(), DEFAULT_TAKE_VALUE as usize);
    assert_eq!(tokens[0], Nat::from(DEFAULT_TAKE_VALUE + 1));
}

#[test]
fn test_tokens_of() {
    let account = Account::from(Principal::management_canister());
    let mut ledger = Ledger::new("ICRC7".to_string(), "icrc7".to_string());

    assert_eq!(ledger.tokens_of(account, None, None).len(), 0);

    ledger.create_tokens(vec![Nat::from(0u64)], account.clone());
    assert_eq!(ledger.tokens_of(account, None, None).len(), 1);
    assert_eq!(ledger.tokens_of(account, Some(Nat::from(0u64)), None).len(), 0);

    ledger.create_tokens((1..=MAX_TAKE_VALUE + 10).rev().map(Nat::from).collect(), account.clone());
    assert_eq!(ledger.tokens_of(account, None, None).len(), DEFAULT_TAKE_VALUE as usize);
    assert_eq!(ledger.tokens_of(account, None, Some(Nat::from(MAX_TAKE_VALUE))).len(), MAX_TAKE_VALUE as usize);

    let tokens = ledger.tokens_of(account, Some(Nat::from(DEFAULT_TAKE_VALUE)), None);
    assert_eq!(tokens.len(), DEFAULT_TAKE_VALUE as usize);
    assert_eq!(tokens[0], Nat::from(DEFAULT_TAKE_VALUE + 1));

    // Check that tokens are sorted.
    assert_eq!(tokens, (DEFAULT_TAKE_VALUE + 1..=DEFAULT_TAKE_VALUE + DEFAULT_TAKE_VALUE).map(Nat::from).collect::<Vec<Nat>>());
}

#[test]
fn test_transfer() {
    let account = Account::from(Principal::management_canister());
    let mut ledger = Ledger::new("ICRC7".to_string(), "icrc7".to_string());
    ledger.create_tokens(vec![Nat::from(123u64)], account.clone());

    // Non existing token.
    assert_eq!(
        ledger.transfer(Principal::anonymous(), vec![TransferArgs::new(account, Nat::from(0u64))], false),
        vec![Some(Err(TransferError::NonExistingTokenId))]
    );

    // Not the owner.
    assert_eq!(
        ledger.transfer(Principal::anonymous(), vec![TransferArgs::new(account, Nat::from(123u64))], false),
        vec![Some(Err(TransferError::Unauthorized))]
    );
}