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
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
use std::collections::HashMap;

use cosmwasm_std::Addr;
use cosmwasm_std::{Env, MessageInfo, Reply};

use cw_utils::NativeBalance;

use crate::prefixed_storage::get_full_contract_storage_namespace;
use crate::wasm::{CodeData, ContractData};

use super::contract::WasmContract;

#[derive(Debug, Clone, Default)]
pub struct WasmStorage {
    pub contracts: HashMap<String, ContractData>,
    pub codes: HashMap<usize, WasmContract>,
    pub code_data: HashMap<usize, CodeData>,
    pub storage: Vec<(Vec<u8>, Vec<u8>)>,
}

impl WasmStorage {
    pub fn get_contract_storage(&self, contract_addr: &Addr) -> Vec<(Vec<u8>, Vec<u8>)> {
        let namespace =
            get_full_contract_storage_namespace(&Addr::unchecked(contract_addr)).to_vec();
        let namespace_len = namespace.len();
        let keys: Vec<(Vec<u8>, Vec<u8>)> = self
            .storage
            .iter()
            // We filter only value in this namespace
            .filter(|(k, _)| k.len() >= namespace_len && k[..namespace_len] == namespace)
            .cloned()
            // We remove the namespace prefix from the key
            .map(|(k, value)| (k[namespace_len..].to_vec(), value))
            .collect();

        keys
    }
}

#[derive(Debug, Clone, Default)]
pub struct BankStorage {
    pub storage: Vec<(Addr, NativeBalance)>,
}

#[derive(Clone, Default)]
pub struct QuerierStorage {
    pub wasm: WasmStorage,
    pub bank: BankStorage,
}

#[derive(Debug)]
pub struct InstanceArguments {
    pub function: WasmFunction,
    pub init_storage: Vec<(Vec<u8>, Vec<u8>)>,
}

#[derive(Debug)]
pub enum WasmFunction {
    Execute(ExecuteArgs),
    Instantiate(InstantiateArgs),
    Query(QueryArgs),
    Sudo(SudoArgs),
    Reply(ReplyArgs),
    Migrate(MigrateArgs),
}

#[derive(Debug)]
pub struct ExecuteArgs {
    pub env: Env,
    pub info: MessageInfo,
    pub msg: Vec<u8>,
}

#[derive(Debug)]
pub struct InstantiateArgs {
    pub env: Env,
    pub info: MessageInfo,
    pub msg: Vec<u8>,
}

#[derive(Debug)]
pub struct QueryArgs {
    pub env: Env,
    pub msg: Vec<u8>,
}

#[derive(Debug)]
pub struct SudoArgs {
    pub env: Env,
    pub msg: Vec<u8>,
}

#[derive(Debug)]
pub struct ReplyArgs {
    pub env: Env,
    pub reply: Reply,
}

#[derive(Debug)]
pub struct MigrateArgs {
    pub env: Env,
    pub msg: Vec<u8>,
}

impl WasmFunction {
    pub fn get_address(&self) -> Addr {
        match self {
            WasmFunction::Execute(ExecuteArgs { env, .. }) => env.contract.address.clone(),
            WasmFunction::Instantiate(InstantiateArgs { env, .. }) => env.contract.address.clone(),
            WasmFunction::Query(QueryArgs { env, .. }) => env.contract.address.clone(),
            WasmFunction::Reply(ReplyArgs { env, .. }) => env.contract.address.clone(),
            WasmFunction::Sudo(SudoArgs { env, .. }) => env.contract.address.clone(),
            WasmFunction::Migrate(MigrateArgs { env, .. }) => env.contract.address.clone(),
        }
    }
}