1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
use std::collections::BTreeMap;

use cosmwasm_std::{Addr, Deps, StdError, StdResult};

use cw_asset::AssetInfo;

use crate::memory::state::{ASSET_ADDRESSES, CONTRACT_ADDRESSES};

use super::{asset_entry::AssetEntry, contract_entry::ContractEntry};

/// Struct that provides easy in-contract memory querying.
#[cosmwasm_schema::cw_serde]
pub struct Memory {
    /// Address of the memory contract
    pub address: Addr,
}

impl Memory {
    /// Raw Query to Memory contract
    pub fn query_contracts(
        &self,
        deps: Deps,
        contracts: Vec<ContractEntry>,
    ) -> StdResult<BTreeMap<ContractEntry, Addr>> {
        let mut resolved_contracts: BTreeMap<ContractEntry, Addr> = BTreeMap::new();

        // Query over keys
        for key in contracts.into_iter() {
            let result: Addr = CONTRACT_ADDRESSES
                .query(&deps.querier, self.address.clone(), key.clone())?
                .ok_or_else(|| {
                    StdError::generic_err(format!("contract {} not found in memory", key))
                })?;
            resolved_contracts.insert(key, result);
        }
        Ok(resolved_contracts)
    }

    /// Raw query of a single contract Addr
    pub fn query_contract(&self, deps: Deps, contract: &ContractEntry) -> StdResult<Addr> {
        let result: Addr = CONTRACT_ADDRESSES
            .query(&deps.querier, self.address.clone(), contract.clone())?
            .ok_or_else(|| {
                StdError::generic_err(format!("contract {} not found in memory", contract))
            })?;
        // Addresses are checked when stored.
        Ok(Addr::unchecked(result))
    }

    /// Raw Query to Memory contract
    pub fn query_assets(
        &self,
        deps: Deps,
        assets: Vec<AssetEntry>,
    ) -> StdResult<BTreeMap<AssetEntry, AssetInfo>> {
        let mut resolved_assets: BTreeMap<AssetEntry, AssetInfo> = BTreeMap::new();

        for asset in assets.into_iter() {
            let result = ASSET_ADDRESSES
                .query(&deps.querier, self.address.clone(), asset.clone())?
                .ok_or_else(|| {
                    StdError::generic_err(format!("asset {} not found in memory", &asset))
                })?;
            resolved_assets.insert(asset, result);
        }
        Ok(resolved_assets)
    }

    /// Raw query of a single AssetInfo
    pub fn query_asset(&self, deps: Deps, asset: &AssetEntry) -> StdResult<AssetInfo> {
        let result = ASSET_ADDRESSES
            .query(&deps.querier, self.address.clone(), asset.clone())?
            .ok_or_else(|| {
                StdError::generic_err(format!("asset {} not found in memory", &asset))
            })?;
        Ok(result)
    }

    // Query single pair address from mem
    // pub fn query_pair_address(
    //     &self,
    //     deps: Deps,
    //     asset_names: [String; 2],
    //     dex: &str,
    // ) -> StdResult<Addr> {
    //     let mut lowercase = asset_names.map(|s| s.to_ascii_lowercase());
    //     lowercase.sort();
    //     let key = format!("{}_{}", lowercase[0], lowercase[1]);
    //     query_contract_from_mem(deps, &self.address, &ContractEntry::new(dex, &key))
    // }
}