├── examples ├── fib │ ├── candid.did │ ├── src │ │ ├── main.rs │ │ ├── lib.rs │ │ └── canister.rs │ └── Cargo.toml ├── counter │ ├── src │ │ ├── main.rs │ │ ├── lib.rs │ │ └── canister.rs │ ├── candid.did │ └── Cargo.toml ├── factory_counter │ ├── src │ │ ├── main.rs │ │ ├── lib.rs │ │ └── canister.rs │ ├── candid.did │ └── Cargo.toml ├── multi_counter │ ├── src │ │ ├── main.rs │ │ ├── lib.rs │ │ └── canister.rs │ ├── candid.did │ └── Cargo.toml ├── naming_system │ ├── src │ │ ├── main.rs │ │ ├── lib.rs │ │ └── canister.rs │ ├── candid.did │ └── Cargo.toml └── README.md ├── ic-kit-stable ├── src │ ├── lib.rs │ └── core │ │ ├── copy.rs │ │ ├── mod.rs │ │ ├── utils.rs │ │ ├── memory.rs │ │ ├── global.rs │ │ ├── checksum.rs │ │ ├── allocator.rs │ │ └── pointer.rs ├── README.md └── Cargo.toml ├── ic-kit-macros ├── README.md ├── Cargo.toml └── src │ ├── test.rs │ ├── metadata.rs │ ├── lib.rs │ └── export_service.rs ├── ic-kit-certified ├── README.md ├── src │ ├── collections │ │ ├── mod.rs │ │ ├── group │ │ │ └── builder.rs │ │ ├── paged.rs │ │ ├── seq.rs │ │ └── map.rs │ ├── lib.rs │ ├── rbtree │ │ ├── debug_alloc.rs │ │ ├── iterator.rs │ │ ├── entry.rs │ │ └── test.rs │ ├── label.rs │ ├── as_hash_tree.rs │ └── hashtree.rs └── Cargo.toml ├── rust-toolchain.toml ├── .gitignore ├── ic-kit ├── README.md ├── src │ ├── ic │ │ ├── mod.rs │ │ ├── spawn.rs │ │ ├── cycles.rs │ │ ├── canister.rs │ │ ├── stable.rs │ │ ├── storage.rs │ │ └── call.rs │ ├── canister.rs │ ├── setup.rs │ ├── lib.rs │ ├── utils.rs │ ├── stable.rs │ └── storage.rs └── Cargo.toml ├── ic-kit-sys ├── README.md ├── src │ ├── lib.rs │ └── types.rs └── Cargo.toml ├── ic-kit-management ├── Cargo.toml └── src │ └── lib.rs ├── ic-kit-runtime ├── README.md ├── src │ ├── users.rs │ ├── lib.rs │ ├── handle.rs │ ├── stable.rs │ ├── call.rs │ └── types.rs └── Cargo.toml ├── e2e ├── canisters │ ├── reverse.rs │ ├── simple_kv_store.rs │ └── async.rs ├── Cargo.toml ├── src │ └── lib.rs └── tests │ └── e2e.rs ├── Cargo.toml ├── dfx.json └── README.md /examples/fib/candid.did: -------------------------------------------------------------------------------- 1 | service : { fib : (nat64) -> (nat64) } -------------------------------------------------------------------------------- /examples/fib/src/main.rs: -------------------------------------------------------------------------------- 1 | mod canister; 2 | fn main() {} 3 | -------------------------------------------------------------------------------- /examples/counter/src/main.rs: -------------------------------------------------------------------------------- 1 | mod canister; 2 | 3 | fn main() {} 4 | -------------------------------------------------------------------------------- /examples/factory_counter/src/main.rs: -------------------------------------------------------------------------------- 1 | mod canister; 2 | fn main() {} 3 | -------------------------------------------------------------------------------- /examples/multi_counter/src/main.rs: -------------------------------------------------------------------------------- 1 | mod canister; 2 | fn main() {} 3 | -------------------------------------------------------------------------------- /examples/naming_system/src/main.rs: -------------------------------------------------------------------------------- 1 | mod canister; 2 | fn main() {} 3 | -------------------------------------------------------------------------------- /ic-kit-stable/src/lib.rs: -------------------------------------------------------------------------------- 1 | mod core; 2 | 3 | pub use crate::core::*; 4 | -------------------------------------------------------------------------------- /examples/factory_counter/candid.did: -------------------------------------------------------------------------------- 1 | service : { deploy_counter : () -> (principal) } -------------------------------------------------------------------------------- /examples/fib/src/lib.rs: -------------------------------------------------------------------------------- 1 | pub mod canister; 2 | pub use canister::FibCanister; 3 | -------------------------------------------------------------------------------- /examples/counter/src/lib.rs: -------------------------------------------------------------------------------- 1 | pub mod canister; 2 | pub use canister::CounterCanister; 3 | -------------------------------------------------------------------------------- /examples/multi_counter/candid.did: -------------------------------------------------------------------------------- 1 | service : { add_counter : (principal) -> (); increment : () -> () } -------------------------------------------------------------------------------- /examples/naming_system/src/lib.rs: -------------------------------------------------------------------------------- 1 | pub mod canister; 2 | pub use canister::NamingSystemCanister; 3 | -------------------------------------------------------------------------------- /ic-kit-macros/README.md: -------------------------------------------------------------------------------- 1 | # Ic Kit Macros 2 | 3 | A collection of macros for headaches canister development with ic-kit stack. -------------------------------------------------------------------------------- /examples/naming_system/candid.did: -------------------------------------------------------------------------------- 1 | service : { 2 | get_name : (principal) -> (opt text) query; 3 | register : (text) -> (); 4 | } -------------------------------------------------------------------------------- /ic-kit-certified/README.md: -------------------------------------------------------------------------------- 1 | # Ic Kit Certified 2 | 3 | Data structures implementing AsHashTree suitable for IC's certified variables. 4 | -------------------------------------------------------------------------------- /rust-toolchain.toml: -------------------------------------------------------------------------------- 1 | [toolchain] 2 | channel = "1.60.0" 3 | targets = ["wasm32-unknown-unknown"] 4 | components = ["rustfmt", "clippy"] 5 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | /target 2 | 3 | 4 | # Added by cargo 5 | # 6 | # already existing elements were commented out 7 | 8 | #/target 9 | Cargo.lock 10 | .dfx -------------------------------------------------------------------------------- /examples/counter/candid.did: -------------------------------------------------------------------------------- 1 | service : { 2 | get_counter : () -> (nat64) query; 3 | increment : () -> (nat64); 4 | increment_by : (nat8) -> (nat64); 5 | } -------------------------------------------------------------------------------- /ic-kit-certified/src/collections/mod.rs: -------------------------------------------------------------------------------- 1 | //! Useful collections that implement [`crate::AsHashTree`] 2 | 3 | pub mod group; 4 | pub mod map; 5 | pub mod paged; 6 | pub mod seq; 7 | -------------------------------------------------------------------------------- /ic-kit/README.md: -------------------------------------------------------------------------------- 1 | # Ic Kit 2 | 3 | Psychedelic's Canister Development Kit which provides a full stack of tools for creating 4 | testing and build canisters on the Internet Computer. 5 | -------------------------------------------------------------------------------- /ic-kit-stable/src/core/copy.rs: -------------------------------------------------------------------------------- 1 | /// A marker for any type that its content can just be copied to stable storage as is. 2 | pub trait StableCopy {} 3 | 4 | impl StableCopy for T {} 5 | -------------------------------------------------------------------------------- /examples/multi_counter/src/lib.rs: -------------------------------------------------------------------------------- 1 | //! A simple example of a counter service that can proxy an increment call to a list of counter 2 | //! canister. 3 | 4 | pub mod canister; 5 | pub use canister::MultiCounterCanister; 6 | -------------------------------------------------------------------------------- /examples/factory_counter/src/lib.rs: -------------------------------------------------------------------------------- 1 | //! A simple example of a counter service that can proxy an increment call to a list of counter 2 | //! canister. 3 | 4 | pub mod canister; 5 | pub use canister::FactoryCounterCanister; 6 | -------------------------------------------------------------------------------- /ic-kit-stable/README.md: -------------------------------------------------------------------------------- 1 | # IC-Kit Stable 2 | 3 | Implementation of data structures that live on the stable storage for the Intent Computer. 4 | Using this structures allows your canister to not have to deal with pre- / post-upgrade. -------------------------------------------------------------------------------- /ic-kit/src/ic/mod.rs: -------------------------------------------------------------------------------- 1 | mod call; 2 | mod canister; 3 | mod cycles; 4 | mod spawn; 5 | mod stable; 6 | mod storage; 7 | 8 | pub use call::*; 9 | pub use canister::*; 10 | pub use cycles::*; 11 | pub use spawn::*; 12 | pub use stable::*; 13 | pub use storage::*; 14 | -------------------------------------------------------------------------------- /ic-kit-sys/README.md: -------------------------------------------------------------------------------- 1 | # IC Kit's Bindings 2 | 3 | This crate provides bindings to the Internet Computer's wasm runtime in both 4 | wasm and non-wasm runtimes. 5 | 6 | This is low-level APIs, and we encourage you to look at using [IC Kit](https://crates.io/crates/ic-kit) 7 | directly. 8 | -------------------------------------------------------------------------------- /ic-kit-stable/src/core/mod.rs: -------------------------------------------------------------------------------- 1 | mod allocator; 2 | mod checksum; 3 | mod copy; 4 | mod global; 5 | mod hole; 6 | mod lru; 7 | mod memory; 8 | mod pointer; 9 | mod utils; 10 | 11 | pub use copy::StableCopy; 12 | 13 | pub use allocator::*; 14 | pub use global::*; 15 | pub use lru::*; 16 | pub use pointer::*; 17 | -------------------------------------------------------------------------------- /ic-kit/src/ic/spawn.rs: -------------------------------------------------------------------------------- 1 | use crate::futures; 2 | 3 | /// Execute a future without blocking the current call. The given future is polled once initially 4 | /// to kickstart the async calls. 5 | #[inline(always)] 6 | pub fn spawn>(future: F) { 7 | futures::spawn(future) 8 | } 9 | -------------------------------------------------------------------------------- /examples/README.md: -------------------------------------------------------------------------------- 1 | # IC Kit Examples 2 | 3 | This directory contains some example canister's implemented using the IC Kit. Each example tries 4 | to be as simple as possible to demonstrate one aspect of the IC Kit and a possible design pattern 5 | you can also use to develop canisters. 6 | 7 | Simple State Manipulation: 8 | - Counter -------------------------------------------------------------------------------- /ic-kit-management/Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "ic-kit-management" 3 | version = "0.1.0" 4 | edition = "2021" 5 | 6 | # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html 7 | 8 | [dependencies] 9 | ic-kit = {path="../ic-kit", version="0.5.0-alpha.4"} 10 | candid="0.7" 11 | serde="1.0" 12 | -------------------------------------------------------------------------------- /ic-kit-runtime/README.md: -------------------------------------------------------------------------------- 1 | # IC Kit Runtime 2 | 3 | This crate implements a canister runtime by implementing `ic-kit-sys` mock 4 | handlers. This provides the ability to compile a canister as a binary 5 | application and run it natively and yet have access to the Internet Computer's 6 | WASM runtime APIs as described in the Interface Spec. 7 | -------------------------------------------------------------------------------- /examples/fib/Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "ic_kit_example_fib" 3 | version = "0.1.0" 4 | edition = "2021" 5 | 6 | # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html 7 | 8 | [dependencies] 9 | ic-kit = {path="../../ic-kit"} 10 | 11 | [[bin]] 12 | name = "ic_kit_example_fib" 13 | path = "src/main.rs" 14 | -------------------------------------------------------------------------------- /examples/counter/Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "ic_kit_example_counter" 3 | version = "0.1.0" 4 | edition = "2021" 5 | 6 | # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html 7 | 8 | [dependencies] 9 | ic-kit = {path="../../ic-kit"} 10 | 11 | [[bin]] 12 | name = "ic_kit_example_counter" 13 | path = "src/main.rs" 14 | -------------------------------------------------------------------------------- /examples/naming_system/Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "ic_kit_example_naming_system" 3 | version = "0.1.0" 4 | edition = "2021" 5 | 6 | # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html 7 | 8 | [dependencies] 9 | ic-kit = {path="../../ic-kit"} 10 | 11 | [[bin]] 12 | name = "ic_kit_example_naming_system" 13 | path = "src/main.rs" 14 | -------------------------------------------------------------------------------- /ic-kit-certified/src/lib.rs: -------------------------------------------------------------------------------- 1 | pub mod as_hash_tree; 2 | pub mod collections; 3 | pub mod hashtree; 4 | pub mod label; 5 | pub mod rbtree; 6 | 7 | pub use as_hash_tree::AsHashTree; 8 | pub use collections::group::builder::GroupBuilder; 9 | pub use collections::group::Group; 10 | pub use collections::map::Map; 11 | pub use collections::paged::Paged; 12 | pub use collections::seq::Seq; 13 | pub use hashtree::{Hash, HashTree}; 14 | -------------------------------------------------------------------------------- /examples/multi_counter/Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "ic_kit_example_multi_counter" 3 | version = "0.1.0" 4 | edition = "2021" 5 | 6 | # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html 7 | 8 | [dependencies] 9 | ic-kit = {path="../../ic-kit"} 10 | ic_kit_example_counter = {path="../counter"} 11 | 12 | [[bin]] 13 | name = "ic_kit_example_multi_counter" 14 | path = "src/main.rs" 15 | -------------------------------------------------------------------------------- /examples/factory_counter/Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "ic_kit_example_factory_counter" 3 | version = "0.1.0" 4 | edition = "2021" 5 | 6 | # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html 7 | 8 | [dependencies] 9 | ic-kit = {path="../../ic-kit"} 10 | ic_kit_example_counter = {path="../counter"} 11 | 12 | [[bin]] 13 | name = "ic_kit_example_factory_counter" 14 | path = "src/main.rs" 15 | -------------------------------------------------------------------------------- /e2e/canisters/reverse.rs: -------------------------------------------------------------------------------- 1 | use ic_kit::utils::{arg_data_raw, arg_data_size, reply}; 2 | 3 | #[export_name = "canister_query reverse"] 4 | fn reverse() { 5 | let arg_bytes: Vec = arg_data_raw(); 6 | assert_eq!(arg_bytes.len(), arg_data_size()); 7 | reply(arg_bytes.into_iter().rev().collect::>().as_ref()); 8 | } 9 | 10 | #[export_name = "canister_update empty_call"] 11 | fn empty_call() { 12 | reply(&[]); 13 | } 14 | 15 | fn main() {} 16 | -------------------------------------------------------------------------------- /examples/factory_counter/src/canister.rs: -------------------------------------------------------------------------------- 1 | use ic_kit::prelude::*; 2 | use ic_kit_example_counter::CounterCanister; 3 | 4 | #[update] 5 | async fn deploy_counter() -> Principal { 6 | // TODO(qti3e) we should make it possible to deploy an instance of CounterCanister. 7 | // It should work in the IC and the ic-kit-runtime environment. 8 | // CounterCanister::install_code(CANISTER_ID); 9 | todo!() 10 | } 11 | 12 | #[derive(KitCanister)] 13 | #[candid_path("candid.did")] 14 | pub struct FactoryCounterCanister; 15 | -------------------------------------------------------------------------------- /Cargo.toml: -------------------------------------------------------------------------------- 1 | [workspace] 2 | members = [ 3 | "e2e", 4 | "examples/counter", 5 | "examples/factory_counter", 6 | "examples/fib", 7 | "examples/multi_counter", 8 | "examples/naming_system", 9 | "ic-kit", 10 | "ic-kit-certified", 11 | "ic-kit-macros", 12 | "ic-kit-management", 13 | "ic-kit-runtime", 14 | "ic-kit-stable", 15 | "ic-kit-sys", 16 | ] 17 | 18 | [profile.canister-release] 19 | inherits = "release" 20 | debug = false 21 | panic = "abort" 22 | lto = true 23 | opt-level = 'z' 24 | -------------------------------------------------------------------------------- /ic-kit-runtime/src/users.rs: -------------------------------------------------------------------------------- 1 | //! A set of mock principal ids. 2 | 3 | use candid::Principal; 4 | use lazy_static::lazy_static; 5 | 6 | lazy_static! { 7 | pub static ref ALICE: Principal = Principal::self_authenticating("ALICE"); 8 | pub static ref BOB: Principal = Principal::self_authenticating("BOB"); 9 | pub static ref JOHN: Principal = Principal::self_authenticating("JOHN"); 10 | pub static ref PARSA: Principal = Principal::self_authenticating("PARSA"); 11 | pub static ref OZ: Principal = Principal::self_authenticating("OZ"); 12 | } 13 | -------------------------------------------------------------------------------- /ic-kit-management/src/lib.rs: -------------------------------------------------------------------------------- 1 | // TODO(qti3e Implement management interface and types. 2 | 3 | use ic_kit::prelude::*; 4 | 5 | #[derive(Deserialize, Debug, Clone, PartialOrd, PartialEq, CandidType)] 6 | pub struct CreateCanisterArgument { 7 | pub settings: Option, 8 | } 9 | 10 | #[derive(Deserialize, Debug, Clone, PartialOrd, PartialEq, CandidType)] 11 | pub struct CanisterSettings { 12 | pub controllers: Option>, 13 | pub compute_allocation: Option, 14 | pub memory_allocation: Option, 15 | pub freezing_threshold: Option, 16 | } 17 | -------------------------------------------------------------------------------- /ic-kit/src/canister.rs: -------------------------------------------------------------------------------- 1 | /// A canister. 2 | pub trait KitCanister { 3 | /// Create a new instance of this canister using the provided canister id. 4 | #[cfg(not(target_family = "wasm"))] 5 | fn build(canister_id: candid::Principal) -> ic_kit_runtime::Canister; 6 | 7 | /// Create a new instance of this canister with the anonymous principal id. 8 | #[cfg(not(target_family = "wasm"))] 9 | fn anonymous() -> ic_kit_runtime::Canister { 10 | Self::build(candid::Principal::anonymous()) 11 | } 12 | 13 | /// The candid description of the canister. 14 | fn candid() -> String; 15 | } 16 | -------------------------------------------------------------------------------- /ic-kit-sys/src/lib.rs: -------------------------------------------------------------------------------- 1 | //! This crate provides the bindings to the Internet Computer runtime that can work in wasm 2 | //! environments, and also provides ability to mock this methods for non-wasm envs, this is 3 | //! part of the Psychedelic's Canister Development kit, [IC-Kit](https://github.com/psychedelic/ic-kit). 4 | //! 5 | //! This is a low level crate, and we don't recommend you to use this directly, and encourage you 6 | //! to look at ic-kit itself. 7 | 8 | /// System APIs exposed by the Internet Computer's WASM runtime. 9 | pub mod ic0; 10 | 11 | /// The common types related to the system API. 12 | pub mod types; 13 | -------------------------------------------------------------------------------- /dfx.json: -------------------------------------------------------------------------------- 1 | { 2 | "version": 1, 3 | "canisters": { 4 | "counter": { 5 | "candid": "examples/counter/candid.did", 6 | "package": "ic_kit_example_counter", 7 | "type": "rust" 8 | }, 9 | "fib": { 10 | "candid": "examples/fib/candid.did", 11 | "package": "ic_kit_example_fib", 12 | "type": "rust" 13 | }, 14 | "naming_system": { 15 | "candid": "examples/naming_system/candid.did", 16 | "package": "naming_system", 17 | "type": "rust" 18 | } 19 | }, 20 | "networks": { 21 | "local": { 22 | "bind": "127.0.0.1:8000", 23 | "type": "ephemeral" 24 | } 25 | } 26 | } 27 | -------------------------------------------------------------------------------- /ic-kit-stable/src/core/utils.rs: -------------------------------------------------------------------------------- 1 | use crate::core::memory::Memory; 2 | 3 | // Reads a struct from memory. 4 | pub fn read_struct(addr: u64) -> T { 5 | let mut t: T = unsafe { core::mem::zeroed() }; 6 | let t_slice = unsafe { 7 | core::slice::from_raw_parts_mut(&mut t as *mut _ as *mut u8, core::mem::size_of::()) 8 | }; 9 | M::stable_read(addr, t_slice); 10 | t 11 | } 12 | 13 | // Writes a struct to memory. 14 | pub fn write_struct(addr: u64, t: &T) { 15 | let slice = unsafe { 16 | core::slice::from_raw_parts(t as *const _ as *const u8, core::mem::size_of::()) 17 | }; 18 | M::stable_write(addr, slice); 19 | } 20 | -------------------------------------------------------------------------------- /e2e/canisters/simple_kv_store.rs: -------------------------------------------------------------------------------- 1 | use ic_kit::macros::{post_upgrade, pre_upgrade, query, update}; 2 | use serde_bytes::ByteBuf; 3 | use std::collections::BTreeMap; 4 | 5 | type Store = BTreeMap; 6 | 7 | #[update] 8 | fn insert(store: &mut Store, key: String, value: ByteBuf) { 9 | store.insert(key, value); 10 | } 11 | 12 | #[query] 13 | fn lookup(store: &Store, key: String) -> Option<&ByteBuf> { 14 | store.get(&key) 15 | } 16 | 17 | #[pre_upgrade] 18 | fn pre_upgrade(store: &Store) { 19 | ic_kit::stable::stable_store((store,)).unwrap(); 20 | } 21 | 22 | #[post_upgrade] 23 | fn post_upgrade() { 24 | let (persisted_store,): (Store,) = ic_kit::stable::stable_restore().unwrap(); 25 | ic_kit::ic::swap(persisted_store); 26 | } 27 | 28 | fn main() {} 29 | -------------------------------------------------------------------------------- /ic-kit-sys/Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "ic-kit-sys" 3 | version = "0.1.3" 4 | edition = "2021" 5 | authors = ["Parsa Ghadimi ", "Ossian Mapes "] 6 | description = "IC-Kit's API bindings to the Internet Computer's WASM runtime." 7 | license = "MIT" 8 | readme = "README.md" 9 | repository = "https://github.com/Psychedelic/ic-kit" 10 | documentation = "https://docs.rs/ic-kit-sys" 11 | homepage = "https://sly.ooo" 12 | categories = ["api-bindings", "development-tools::testing"] 13 | keywords = ["internet-computer", "canister", "fleek", "psychedelic"] 14 | include = ["src", "Cargo.toml", "README.md"] 15 | 16 | # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html 17 | 18 | [target.'cfg(not(target_family = "wasm"))'.dependencies] 19 | tokio = {version="1.20", features=["sync"]} 20 | futures = {version="0.3"} -------------------------------------------------------------------------------- /ic-kit-stable/Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "ic-kit-stable" 3 | version = "0.1.0-alpha.0" 4 | edition = "2021" 5 | authors = ["Parsa Ghadimi ", "Ossian Mapes "] 6 | description = "IC-Kit's data structures living on the stable storage." 7 | license = "MIT" 8 | readme = "README.md" 9 | repository = "https://github.com/Psychedelic/ic-kit" 10 | documentation = "https://docs.rs/ic-kit-stable" 11 | homepage = "https://sly.ooo" 12 | categories = ["data-structures"] 13 | keywords = ["internet-computer", "canister", "fleek", "psychedelic"] 14 | include = ["src", "Cargo.toml", "README.md"] 15 | 16 | # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html 17 | 18 | [dependencies] 19 | ic-kit = {path="../ic-kit", version="0.5.0-alpha.3"} 20 | 21 | [features] 22 | experimental-stable64 = ["ic-kit/experimental-stable64"] 23 | -------------------------------------------------------------------------------- /e2e/Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "ic-kit-e2e-tests" 3 | version = "0.1.0" 4 | authors = ["DFINITY Stiftung "] 5 | edition = "2021" 6 | description = "End-to-end tests for the IC-Kit" 7 | license = "Apache-2.0" 8 | repository = "https://github.com/psychedelic/ic-kit" 9 | 10 | [dependencies] 11 | ic-kit = { path = "../ic-kit" } 12 | candid = "0.7.4" 13 | cargo_metadata = "0.14.2" 14 | escargot = { version = "0.5.7", features = ["print"] } 15 | lazy_static = "1.4.0" 16 | serde_bytes = "0.11" 17 | 18 | [[bin]] 19 | name = "simple-kv-store" 20 | path = "canisters/simple_kv_store.rs" 21 | 22 | [[bin]] 23 | name = "async" 24 | path = "canisters/async.rs" 25 | 26 | [[bin]] 27 | name = "reverse" 28 | path = "canisters/reverse.rs" 29 | 30 | [dev-dependencies] 31 | ic-state-machine-tests = { git = "https://github.com/dfinity/ic", rev = "02a4a828f2f4d3b1dcb93a84e60672a3f3fdb400" } 32 | -------------------------------------------------------------------------------- /ic-kit-certified/Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "ic-kit-certified" 3 | version = "0.1.0-alpha.0" 4 | edition = "2018" 5 | description = "Certified variable friendly data structures for the Internet Computer." 6 | authors = ["Parsa Ghadimi ", "Ossian Mapes "] 7 | license = "MIT" 8 | readme = "README.md" 9 | repository = "https://github.com/Psychedelic/ic-kit" 10 | documentation = "https://docs.rs/ic-kit-certified" 11 | homepage = "https://sly.ooo" 12 | categories = ["data-structures"] 13 | keywords = ["internet-computer", "canister", "fleek", "psychedelic"] 14 | include = ["src", "Cargo.toml", "README.md"] 15 | 16 | [dependencies] 17 | ic-kit-stable = {path="../ic-kit-stable"} 18 | ic-types = "0.4.1" 19 | candid = "0.7" 20 | sha2 = "0.10.2" 21 | serde = { version="1.0.116", features = ["derive"] } 22 | serde_bytes = "0.11.5" 23 | serde_cbor = "0.11.2" 24 | hex = "0.4.3" 25 | -------------------------------------------------------------------------------- /ic-kit-macros/Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "ic-kit-macros" 3 | version = "0.1.1-alpha.0" 4 | edition = "2021" 5 | authors = ["Parsa Ghadimi ", "Ossian Mapes "] 6 | description = "IC-Kit's macros for canister development" 7 | license = "MIT" 8 | readme = "README.md" 9 | repository = "https://github.com/Psychedelic/ic-kit" 10 | documentation = "https://docs.rs/ic-kit-macros" 11 | homepage = "https://sly.ooo" 12 | categories = ["api-bindings"] 13 | keywords = ["internet-computer", "canister", "fleek", "psychedelic"] 14 | include = ["src", "Cargo.toml", "README.md"] 15 | 16 | # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html 17 | 18 | [dependencies] 19 | quote = "1.0" 20 | proc-macro2 = "1.0" 21 | syn = "1.0" 22 | serde = "1.0" 23 | serde_tokenstream = "0.1" 24 | lazy_static = "1.4" 25 | compile-time-run = "0.2.12" 26 | 27 | [lib] 28 | proc-macro = true 29 | -------------------------------------------------------------------------------- /ic-kit-runtime/src/lib.rs: -------------------------------------------------------------------------------- 1 | // We normally wouldn't have to do this, but since most of ic-kit users will build for wasm, we 2 | // should handle this and print a nice compiler error to not confuse the users with 177 errors 3 | // printed on their screen. 4 | cfg_if::cfg_if! { 5 | if #[cfg(target_family = "wasm")] { 6 | compile_error!("IC-Kit runtime does not support builds for WASM."); 7 | } else { 8 | pub mod call; 9 | pub mod canister; 10 | pub mod replica; 11 | pub mod stable; 12 | pub mod types; 13 | pub mod users; 14 | pub mod handle; 15 | 16 | pub use canister::{Canister, CanisterMethod}; 17 | pub use replica::Replica; 18 | pub use tokio::runtime::Builder as TokioRuntimeBuilder; 19 | 20 | pub mod prelude { 21 | pub use crate::replica::Replica; 22 | pub use crate::users; 23 | } 24 | } 25 | } 26 | -------------------------------------------------------------------------------- /ic-kit/Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "ic-kit" 3 | version = "0.5.0-alpha.4" 4 | description = "Testable Canister Developer Kit for the Internet Computer." 5 | authors = ["Parsa Ghadimi ", "Ossian Mapes "] 6 | edition = "2018" 7 | license = "MIT" 8 | readme = "README.md" 9 | repository = "https://github.com/Psychedelic/ic-kit" 10 | documentation = "https://docs.rs/ic-kit" 11 | homepage = "https://sly.ooo" 12 | categories = ["api-bindings", "development-tools::testing"] 13 | keywords = ["internet-computer", "canister", "fleek", "psychedelic"] 14 | include = ["src", "Cargo.toml", "README.md"] 15 | 16 | [dependencies] 17 | ic-kit-sys = { path = "../ic-kit-sys", version = "0.1.3" } 18 | ic-kit-macros = { path = "../ic-kit-macros", version = "0.1.1-alpha.0" } 19 | candid = "0.8" 20 | serde = "1.0" 21 | 22 | [target.'cfg(not(target_family = "wasm"))'.dependencies] 23 | ic-kit-runtime = { path = "../ic-kit-runtime", version = "0.1.0-alpha.1" } 24 | 25 | [features] 26 | experimental-stable64 = [] 27 | experimental-cycles128 = [] 28 | -------------------------------------------------------------------------------- /ic-kit-runtime/Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "ic-kit-runtime" 3 | version = "0.1.0-alpha.1" 4 | edition = "2021" 5 | authors = ["Parsa Ghadimi ", "Ossian Mapes "] 6 | description = "IC-Kit's Canister Simulator" 7 | license = "MIT" 8 | readme = "README.md" 9 | repository = "https://github.com/Psychedelic/ic-kit" 10 | documentation = "https://docs.rs/ic-kit-runtime" 11 | homepage = "https://sly.ooo" 12 | categories = ["api-bindings", "development-tools::testing"] 13 | keywords = ["internet-computer", "canister", "fleek", "psychedelic"] 14 | include = ["src", "Cargo.toml", "README.md"] 15 | 16 | # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html 17 | 18 | [dependencies] 19 | cfg-if = "1.0.0" 20 | 21 | [target.'cfg(not(target_family = "wasm"))'.dependencies] 22 | ic-kit-sys = { path = "../ic-kit-sys", version = "0.1.3" } 23 | ic-types = "0.6" 24 | tokio = { version = "1.20", features = ["sync", "macros", "rt"] } 25 | thread-local-panic-hook = "0.1.0" 26 | lazy_static = "1.4" 27 | memmap = "0.7.0" 28 | futures = "0.3" 29 | actix = "0.13" 30 | candid = "0.8" 31 | serde = "1.0" 32 | -------------------------------------------------------------------------------- /ic-kit/src/setup.rs: -------------------------------------------------------------------------------- 1 | use crate::ic; 2 | use std::panic; 3 | 4 | #[cfg(target_family = "wasm")] 5 | static mut DONE: bool = false; 6 | 7 | #[cfg(target_family = "wasm")] 8 | pub fn setup_hooks() { 9 | unsafe { 10 | if DONE { 11 | return; 12 | } 13 | DONE = true; 14 | } 15 | 16 | set_panic_hook(); 17 | } 18 | 19 | #[cfg(not(target_family = "wasm"))] 20 | pub fn setup_hooks() { 21 | set_panic_hook(); 22 | } 23 | 24 | /// Sets a custom panic hook, uses debug.trace 25 | fn set_panic_hook() { 26 | panic::set_hook(Box::new(|info| { 27 | let file = info.location().unwrap().file(); 28 | let line = info.location().unwrap().line(); 29 | let col = info.location().unwrap().column(); 30 | 31 | let msg = match info.payload().downcast_ref::<&'static str>() { 32 | Some(s) => *s, 33 | None => match info.payload().downcast_ref::() { 34 | Some(s) => &s[..], 35 | None => "Box", 36 | }, 37 | }; 38 | 39 | let err_info = format!("Panicked at '{}', {}:{}:{}", msg, file, line, col); 40 | ic::print(&err_info); 41 | ic::trap(&err_info); 42 | })); 43 | } 44 | -------------------------------------------------------------------------------- /examples/fib/src/canister.rs: -------------------------------------------------------------------------------- 1 | use ic_kit::prelude::*; 2 | 3 | #[update] 4 | async fn fib(n: u64) -> u64 { 5 | println!("fib({})", n); 6 | if n > 20 { 7 | ic::trap("Let's not kill IC."); 8 | } 9 | 10 | if n <= 1 { 11 | return n; 12 | } 13 | 14 | let a = CallBuilder::new(id(), "fib") 15 | .with_arg(n - 1) 16 | .perform_one::() 17 | .await 18 | .unwrap(); 19 | 20 | let b = CallBuilder::new(id(), "fib") 21 | .with_arg(n - 2) 22 | .perform_one::() 23 | .await 24 | .unwrap(); 25 | 26 | a + b 27 | } 28 | 29 | #[derive(KitCanister)] 30 | #[candid_path("candid.did")] 31 | pub struct FibCanister; 32 | 33 | #[cfg(test)] 34 | mod test { 35 | use super::*; 36 | 37 | #[kit_test] 38 | async fn fib_6(replica: Replica) { 39 | let canister = replica.add_canister(FibCanister::anonymous()); 40 | 41 | let fib_6 = canister 42 | .new_call("fib") 43 | .with_caller(*users::ALICE) 44 | .with_arg(6u64) 45 | .perform() 46 | .await 47 | .decode_one::() 48 | .unwrap(); 49 | 50 | assert_eq!(fib_6, 8); 51 | } 52 | } 53 | -------------------------------------------------------------------------------- /ic-kit-certified/src/rbtree/debug_alloc.rs: -------------------------------------------------------------------------------- 1 | use std::cell::RefCell; 2 | use std::collections::HashSet; 3 | 4 | thread_local! { 5 | static ALLOCATED_POINTERS: RefCell> = RefCell::new(HashSet::new()); 6 | } 7 | 8 | /// Marks the specified pointer as being reachable. 9 | /// Should be invoked right after the memory is allocated. 10 | pub fn mark_pointer_allocated(ptr: *const T) { 11 | ALLOCATED_POINTERS.with(move |ptrs| { 12 | assert!(ptrs.borrow_mut().insert(ptr as *const ())); 13 | }) 14 | } 15 | 16 | /// Marks the specified pointer as deleted. 17 | /// Should be invoked right after the pointer is freed. 18 | pub fn mark_pointer_deleted(ptr: *const T) { 19 | ALLOCATED_POINTERS.with(move |ptrs| { 20 | assert!( 21 | ptrs.borrow_mut().remove(&(ptr as *const ())), 22 | "DOUBLE FREE: deleted pointer {:?} that is not allocated", 23 | ptr 24 | ); 25 | }) 26 | } 27 | 28 | /// Returns the number of pointer that were allocated and not yet 29 | /// deleted. 30 | pub fn count_allocated_pointers() -> usize { 31 | ALLOCATED_POINTERS.with(|ptrs| ptrs.borrow().len()) 32 | } 33 | 34 | /// Returns true if the specified pointer was allocated and not yet 35 | /// deleted. 36 | pub fn is_live(ptr: *const T) -> bool { 37 | ALLOCATED_POINTERS.with(move |ptrs| ptrs.borrow().contains(&(ptr as *const ()))) 38 | } 39 | -------------------------------------------------------------------------------- /ic-kit-macros/src/test.rs: -------------------------------------------------------------------------------- 1 | use proc_macro2::{Span, TokenStream}; 2 | use quote::quote; 3 | use syn::spanned::Spanned; 4 | use syn::{parse2, Error, ItemFn}; 5 | 6 | pub fn gen_test_code(_: TokenStream, item: TokenStream) -> Result { 7 | let fun: ItemFn = parse2::(item.clone()).map_err(|e| { 8 | Error::new( 9 | item.span(), 10 | format!( 11 | "The #[kit_test] macro must be on top of a function.\n {}", 12 | e 13 | ), 14 | ) 15 | })?; 16 | 17 | let signature = &fun.sig; 18 | let visibility = &fun.vis; 19 | let is_async = signature.asyncness.is_some(); 20 | let name = &signature.ident; 21 | 22 | if !is_async { 23 | return Err(Error::new( 24 | Span::call_site(), 25 | format!("The #[kit_test] can only be used on top of an async function.",), 26 | )); 27 | } 28 | 29 | Ok(quote! { 30 | #[test] 31 | #visibility fn #name() { 32 | #item 33 | 34 | let rt = ic_kit::rt::TokioRuntimeBuilder::new_current_thread() 35 | .build() 36 | .expect("ic-kit: Could not build tokio runtime."); 37 | 38 | rt.block_on(async { 39 | let replica = ic_kit::rt::replica::Replica::default(); 40 | #name(replica).await; 41 | }); 42 | } 43 | }) 44 | } 45 | -------------------------------------------------------------------------------- /e2e/canisters/async.rs: -------------------------------------------------------------------------------- 1 | use ic_kit::prelude::*; 2 | 3 | #[derive(Default)] 4 | struct Resource(u64); 5 | 6 | #[derive(Default)] 7 | struct NotificationsReceived(u64); 8 | 9 | #[query] 10 | fn inc(n: u64) -> u64 { 11 | n + 1 12 | } 13 | 14 | #[query] 15 | fn invocation_count(r: &Resource) -> u64 { 16 | r.0 17 | } 18 | 19 | #[update] 20 | async fn panic_after_async() { 21 | let x = with_mut(|r: &mut Resource| { 22 | r.0 += 1; 23 | r.0 24 | }); 25 | 26 | CallBuilder::new(id(), "inc") 27 | .with_arg(x) 28 | .perform_rejection() 29 | .await 30 | .expect("failed to call self"); 31 | 32 | ic::trap("Goodbye, cruel world.") 33 | } 34 | 35 | #[query] 36 | fn notifications_received(notifications: &NotificationsReceived) -> u64 { 37 | notifications.0 38 | } 39 | 40 | #[update] 41 | fn on_notify(notifications: &mut NotificationsReceived) { 42 | notifications.0 += 1; 43 | } 44 | 45 | #[update] 46 | fn notify(whom: Principal, method: String) { 47 | CallBuilder::new(whom, method.as_str()) 48 | .perform_one_way() 49 | .unwrap_or_else(|reject| { 50 | ic::trap(&format!( 51 | "failed to notify (callee={}, method={}): {:?}", 52 | whom, method, reject 53 | )) 54 | }); 55 | } 56 | 57 | #[query] 58 | fn greet(name: String) -> String { 59 | format!("Hello, {}", name) 60 | } 61 | 62 | fn main() {} 63 | -------------------------------------------------------------------------------- /e2e/src/lib.rs: -------------------------------------------------------------------------------- 1 | use cargo_metadata::MetadataCommand; 2 | use escargot::CargoBuild; 3 | use std::path::PathBuf; 4 | 5 | /// Builds a canister with the specified name from the current 6 | /// package and returns the WebAssembly module. 7 | pub fn cargo_build_canister(bin_name: &str) -> Vec { 8 | let dir = PathBuf::from(std::env::var("CARGO_MANIFEST_DIR").unwrap()); 9 | 10 | let cargo_toml_path = dir.join("Cargo.toml"); 11 | 12 | let target_dir = MetadataCommand::new() 13 | .manifest_path(&cargo_toml_path) 14 | .no_deps() 15 | .exec() 16 | .expect("failed to run cargo metadata") 17 | .target_directory; 18 | 19 | // We use a different target path to stop the native cargo build 20 | // cache being invalidated every time we run this function 21 | let wasm_target_dir = target_dir.join("canister-build"); 22 | 23 | let cargo_build = CargoBuild::new() 24 | .target("wasm32-unknown-unknown") 25 | .bin(bin_name) 26 | .args(["--profile", "canister-release"]) 27 | .manifest_path(&cargo_toml_path) 28 | .target_dir(&wasm_target_dir); 29 | 30 | let binary = cargo_build 31 | .run() 32 | .expect("Cargo failed to compile the wasm binary"); 33 | 34 | std::fs::read(binary.path()).unwrap_or_else(|e| { 35 | panic!( 36 | "failed to read compiled Wasm file from {}: {}", 37 | binary.path().display(), 38 | e 39 | ) 40 | }) 41 | } 42 | -------------------------------------------------------------------------------- /ic-kit/src/lib.rs: -------------------------------------------------------------------------------- 1 | mod canister; 2 | mod futures; 3 | mod setup; 4 | mod storage; 5 | 6 | /// System APIs for the Internet Computer. 7 | pub mod ic; 8 | 9 | /// Helper methods around the stable storage. 10 | pub mod stable; 11 | 12 | /// Internal utility methods to deal with reading data. 13 | pub mod utils; 14 | 15 | // re-exports. 16 | pub use candid::{self, CandidType, Nat, Principal}; 17 | pub use ic_kit_macros as macros; 18 | pub use setup::setup_hooks; 19 | 20 | // The KitCanister derive macro. 21 | pub use canister::KitCanister; 22 | pub use ic_kit_macros::KitCanister; 23 | 24 | /// The IC-kit runtime, which can be used for testing the canister in non-wasm environments. 25 | #[cfg(not(target_family = "wasm"))] 26 | pub use ic_kit_runtime as rt; 27 | 28 | /// The famous prelude module which re exports the most useful methods. 29 | pub mod prelude { 30 | pub use super::canister::KitCanister; 31 | pub use super::ic; 32 | pub use super::ic::CallBuilder; 33 | pub use super::ic::{balance, caller, id, spawn}; 34 | pub use super::ic::{maybe_with, maybe_with_mut, swap, take, with, with_mut}; 35 | pub use super::ic::{Cycles, StableSize}; 36 | pub use candid::{CandidType, Nat, Principal}; 37 | pub use serde::{Deserialize, Serialize}; 38 | 39 | pub use ic_kit_macros::*; 40 | 41 | #[cfg(not(target_family = "wasm"))] 42 | pub use ic_kit_runtime as rt; 43 | 44 | #[cfg(not(target_family = "wasm"))] 45 | pub use ic_kit_runtime::prelude::*; 46 | } 47 | -------------------------------------------------------------------------------- /ic-kit-macros/src/metadata.rs: -------------------------------------------------------------------------------- 1 | use compile_time_run::run_command_str; 2 | use proc_macro2::{Ident, Literal, Span, TokenStream}; 3 | use quote::quote; 4 | 5 | pub fn generate_static_string(key: T, val: T) -> TokenStream { 6 | let val = val.to_string(); 7 | let key = Ident::new(&key.to_string(), Span::call_site()); 8 | let val_code = Literal::byte_string(val.as_bytes()); 9 | let val_len = val.len(); 10 | quote! { pub static #key: [u8; #val_len] = *#val_code; } 11 | } 12 | 13 | pub fn generate_metadata() -> TokenStream { 14 | // TODO(oz): Gracefully handle errors if the project is not a git repository 15 | let git_commit = 16 | generate_static_string("GIT_COMMIT", run_command_str!("git", "rev-parse", "HEAD")); 17 | 18 | let git_url = generate_static_string( 19 | "GIT_URL", 20 | run_command_str!("git", "config", "--get", "remote.origin.url"), 21 | ); 22 | 23 | let cdk = generate_static_string( 24 | "CDK_VERSION", 25 | run_command_str!("cargo", "tree", "-i", "ic-kit", "-e", "build"), 26 | ); 27 | 28 | let compiler = generate_static_string( 29 | "COMPILER", 30 | run_command_str!("rustc", "--version", "--verbose"), 31 | ); 32 | 33 | let dfx = generate_static_string("DFX_VERSION", run_command_str!("dfx", "--version")); 34 | 35 | quote!( 36 | #[link_section = "icp:public env:git_commit"] 37 | #git_commit 38 | #[link_section = "icp:public env:git_url"] 39 | #git_url 40 | #[link_section = "icp:public env:cdk"] 41 | #cdk 42 | #[link_section = "icp:public env:compiler"] 43 | #compiler 44 | #[link_section = "icp:public env:dfx"] 45 | #dfx 46 | ) 47 | } 48 | -------------------------------------------------------------------------------- /examples/naming_system/src/canister.rs: -------------------------------------------------------------------------------- 1 | use ic_kit::prelude::*; 2 | use std::collections::HashMap; 3 | 4 | #[derive(Default)] 5 | struct Registry { 6 | names: HashMap, 7 | } 8 | 9 | #[derive(Default)] 10 | struct Stats { 11 | called_register: u64, 12 | } 13 | 14 | #[update] 15 | fn register(registry: &mut Registry, stats: &mut Stats, name: String) { 16 | stats.called_register += 1; 17 | registry.names.insert(caller(), name); 18 | } 19 | 20 | #[query] 21 | fn get_name(registry: &Registry, user: Principal) -> Option<&String> { 22 | registry.names.get(&user) 23 | } 24 | 25 | #[derive(KitCanister)] 26 | #[candid_path("candid.did")] 27 | pub struct NamingSystemCanister; 28 | 29 | #[cfg(test)] 30 | mod tests { 31 | use super::*; 32 | 33 | #[kit_test] 34 | async fn test(replica: Replica) { 35 | let ns = replica.add_canister(NamingSystemCanister::anonymous()); 36 | 37 | ns.new_call("register") 38 | .with_caller(*users::ALICE) 39 | .with_arg("Alice") 40 | .perform() 41 | .await 42 | .assert_ok(); 43 | 44 | ns.new_call("register") 45 | .with_caller(*users::BOB) 46 | .with_arg("Bob") 47 | .perform() 48 | .await 49 | .assert_ok(); 50 | 51 | let alice_name = ns 52 | .new_call("get_name") 53 | .with_arg(*users::ALICE) 54 | .perform() 55 | .await 56 | .decode_one::>() 57 | .unwrap(); 58 | 59 | assert_eq!(alice_name, Some("Alice".to_string())); 60 | 61 | let bob_name = ns 62 | .new_call("get_name") 63 | .with_arg(*users::BOB) 64 | .perform() 65 | .await 66 | .decode_one::>() 67 | .unwrap(); 68 | 69 | assert_eq!(bob_name, Some("Bob".to_string())); 70 | } 71 | } 72 | -------------------------------------------------------------------------------- /ic-kit/src/utils.rs: -------------------------------------------------------------------------------- 1 | use ic_kit_sys::ic0; 2 | 3 | /// Return the size of the raw argument to this entry point. 4 | pub fn arg_data_size() -> usize { 5 | unsafe { ic0::msg_arg_data_size() as usize } 6 | } 7 | 8 | /// Return the raw argument data to this entry point. 9 | pub fn arg_data_raw() -> Vec { 10 | unsafe { 11 | let len: usize = ic0::msg_arg_data_size() as usize; 12 | if len == 0 { 13 | return Vec::new(); 14 | } 15 | 16 | let mut bytes = Vec::with_capacity(len); 17 | ic0::msg_arg_data_copy(bytes.as_mut_ptr() as isize, 0, len as isize); 18 | bytes.set_len(len); 19 | bytes 20 | } 21 | } 22 | 23 | /// Reply to the current call using the provided buffer. 24 | pub fn reply(buf: &[u8]) { 25 | unsafe { 26 | if !buf.is_empty() { 27 | ic0::msg_reply_data_append(buf.as_ptr() as isize, buf.len() as isize) 28 | } 29 | ic0::msg_reply() 30 | } 31 | } 32 | 33 | /// Reject the current call. 34 | pub fn reject(message: &str) { 35 | unsafe { ic0::msg_reject(message.as_ptr() as isize, message.len() as isize) } 36 | } 37 | 38 | /// Accept the incoming message. 39 | pub fn accept() { 40 | unsafe { 41 | ic0::accept_message(); 42 | } 43 | } 44 | 45 | /// Return the name of the current canister method. 46 | pub fn method_name() -> String { 47 | let len = unsafe { ic0::msg_method_name_size() as usize }; 48 | let mut bytes = vec![0u8; len]; 49 | unsafe { 50 | ic0::msg_method_name_copy(bytes.as_mut_ptr() as isize, 0, len as isize); 51 | } 52 | String::from_utf8_lossy(&bytes).to_string() 53 | } 54 | 55 | /// Get the value of specified performance counter. 56 | /// 57 | /// Supported counter type: 58 | /// 59 | /// 0 : Instruction counter. 60 | /// The number of WebAssembly instructions the system has determined that the canister has executed. 61 | pub fn performance_counter(counter_type: u32) -> u64 { 62 | unsafe { ic0::performance_counter(counter_type as i32) as u64 } 63 | } 64 | -------------------------------------------------------------------------------- /ic-kit-stable/src/core/memory.rs: -------------------------------------------------------------------------------- 1 | /// The memory interface. temp remove this once ic-kit-runtime has stable* support. 2 | pub trait Memory { 3 | fn stable_size() -> u64; 4 | fn stable_grow(new_pages: u64) -> i64; 5 | fn stable_read(offset: u64, buf: &mut [u8]); 6 | fn stable_write(offset: u64, buf: &[u8]); 7 | } 8 | 9 | #[cfg(not(target_family = "wasm"))] 10 | pub mod mock { 11 | use super::Memory; 12 | use ic_kit::rt::stable::{HeapStableMemory, StableMemoryBackend}; 13 | use std::cell::RefCell; 14 | 15 | thread_local! { 16 | static MEMORY: RefCell = RefCell::new(HeapStableMemory::default()); 17 | } 18 | 19 | // A memory interface that uses ic-kit's HeapStableMemory. 20 | pub struct MockMemory; 21 | 22 | impl Memory for MockMemory { 23 | fn stable_size() -> u64 { 24 | MEMORY.with(|c| c.borrow_mut().stable_size()) 25 | } 26 | 27 | fn stable_grow(new_pages: u64) -> i64 { 28 | MEMORY.with(|c| c.borrow_mut().stable_grow(new_pages)) 29 | } 30 | 31 | fn stable_read(offset: u64, buf: &mut [u8]) { 32 | MEMORY.with(|c| c.borrow_mut().stable_read(offset, buf)) 33 | } 34 | 35 | fn stable_write(offset: u64, buf: &[u8]) { 36 | MEMORY.with(|c| c.borrow_mut().stable_write(offset, buf)) 37 | } 38 | } 39 | } 40 | 41 | /// A memory backend using the IC. 42 | pub struct IcMemory; 43 | 44 | impl Memory for IcMemory { 45 | fn stable_size() -> u64 { 46 | ic_kit::ic::stable_size() as u64 47 | } 48 | 49 | fn stable_grow(new_pages: u64) -> i64 { 50 | match ic_kit::ic::stable_grow(new_pages as ic_kit::ic::StableSize) { 51 | Ok(s) => s as i64, 52 | Err(_) => -1, 53 | } 54 | } 55 | 56 | fn stable_read(offset: u64, buf: &mut [u8]) { 57 | ic_kit::ic::stable_read(offset as ic_kit::ic::StableSize, buf) 58 | } 59 | 60 | fn stable_write(offset: u64, buf: &[u8]) { 61 | ic_kit::ic::stable_write(offset as ic_kit::ic::StableSize, buf) 62 | } 63 | } 64 | 65 | #[cfg(test)] 66 | pub type DefaultMemory = mock::MockMemory; 67 | 68 | #[cfg(not(test))] 69 | pub type DefaultMemory = IcMemory; 70 | -------------------------------------------------------------------------------- /ic-kit-certified/src/rbtree/iterator.rs: -------------------------------------------------------------------------------- 1 | use super::{Node, RbTree}; 2 | use crate::label::Label; 3 | use crate::AsHashTree; 4 | use std::marker::PhantomData; 5 | 6 | /// An iterator over key-values in a RbTree. 7 | pub struct RbTreeIterator<'tree, K: 'static + Label, V: AsHashTree + 'static> { 8 | visit: *mut Node, 9 | stack: Vec<*mut Node>, 10 | remaining_elements: usize, 11 | lifetime: PhantomData<&'tree RbTree>, 12 | } 13 | 14 | impl<'tree, K: 'static + Label, V: AsHashTree + 'static> RbTreeIterator<'tree, K, V> { 15 | pub fn new(tree: &'tree RbTree) -> Self { 16 | Self { 17 | visit: tree.root, 18 | stack: Vec::with_capacity(8), 19 | remaining_elements: tree.len(), 20 | lifetime: PhantomData::default(), 21 | } 22 | } 23 | } 24 | 25 | impl<'tree, K: 'static + Label, V: AsHashTree + 'static> Iterator for RbTreeIterator<'tree, K, V> { 26 | type Item = (&'tree K, &'tree V); 27 | 28 | #[inline] 29 | fn next(&mut self) -> Option { 30 | unsafe { 31 | while !self.visit.is_null() { 32 | self.stack.push(self.visit); 33 | self.visit = (*self.visit).left; 34 | } 35 | 36 | if let Some(node) = self.stack.pop() { 37 | self.visit = (*node).right; 38 | self.remaining_elements -= 1; 39 | return Some((&(*node).key, &(*node).value)); 40 | } 41 | 42 | None 43 | } 44 | } 45 | 46 | #[inline] 47 | fn size_hint(&self) -> (usize, Option) { 48 | (self.remaining_elements, Some(self.remaining_elements)) 49 | } 50 | } 51 | 52 | #[cfg(test)] 53 | mod tests { 54 | use super::*; 55 | 56 | #[test] 57 | fn should_visit_all() { 58 | let mut tree = RbTree::<[u8; 1], u8>::new(); 59 | 60 | for i in 0..250u8 { 61 | tree.insert([i], i); 62 | } 63 | 64 | let iter = RbTreeIterator::new(&tree); 65 | 66 | let mut expected_v = 0u8; 67 | 68 | for (_, v) in iter { 69 | assert_eq!(v, &expected_v); 70 | expected_v += 1; 71 | } 72 | 73 | assert_eq!(expected_v, 250); 74 | } 75 | } 76 | -------------------------------------------------------------------------------- /ic-kit/src/ic/cycles.rs: -------------------------------------------------------------------------------- 1 | use ic_kit_sys::ic0; 2 | 3 | /// The type used to represent the cycles amount, which is u128 when the `experimental-cycles128` 4 | /// feature is on and a u64 otherwise. 5 | #[cfg(feature = "experimental-cycles128")] 6 | pub type Cycles = u128; 7 | 8 | /// The type used to represent the cycles amount, which is u128 when the `experimental-cycles128` 9 | /// feature is on and a u64 otherwise. 10 | #[cfg(not(feature = "experimental-cycles128"))] 11 | pub type Cycles = u64; 12 | 13 | /// Return the number of available cycles that is sent by the caller. 14 | pub fn msg_cycles_available() -> Cycles { 15 | #[cfg(not(feature = "experimental-cycles128"))] 16 | unsafe { 17 | ic0::msg_cycles_available() as u64 18 | } 19 | 20 | #[cfg(feature = "experimental-cycles128")] 21 | { 22 | let mut recv = 0u128; 23 | unsafe { ic0::msg_cycles_available128(&mut recv as *mut u128 as isize) } 24 | u128::from_le(recv) 25 | } 26 | } 27 | 28 | /// Accept the given amount of cycles, returns the actual amount of accepted cycles. 29 | #[inline(always)] 30 | pub fn msg_cycles_accept(max_amount: Cycles) -> Cycles { 31 | #[cfg(not(feature = "experimental-cycles128"))] 32 | unsafe { 33 | ic0::msg_cycles_accept(max_amount as i64) as u64 34 | } 35 | 36 | #[cfg(feature = "experimental-cycles128")] 37 | { 38 | if max_amount < (u64::MAX as u128) { 39 | return unsafe { ic0::msg_cycles_accept(max_amount as i64) as u128 }; 40 | } 41 | 42 | let high = (max_amount >> 64) as u64 as i64; 43 | let low = (max_amount & (1 << 64)) as u64 as i64; 44 | let mut recv = 0u128; 45 | unsafe { 46 | ic0::msg_cycles_accept128(high, low, &mut recv as *mut u128 as isize); 47 | } 48 | u128::from_le(recv) 49 | } 50 | } 51 | 52 | /// Return the cycles that were sent back by the canister that was just called. 53 | /// This method should only be called right after an inter-canister call. 54 | #[inline(always)] 55 | pub fn msg_cycles_refunded() -> Cycles { 56 | #[cfg(not(feature = "experimental-cycles128"))] 57 | unsafe { 58 | ic0::msg_cycles_refunded() as u64 59 | } 60 | 61 | #[cfg(feature = "experimental-cycles128")] 62 | { 63 | let mut recv = 0u128; 64 | unsafe { ic0::msg_cycles_refunded128(&mut recv as *mut u128 as isize) } 65 | u128::from_le(recv) 66 | } 67 | } 68 | -------------------------------------------------------------------------------- /ic-kit-stable/src/core/global.rs: -------------------------------------------------------------------------------- 1 | use crate::core::allocator::{BlockAddress, BlockSize, StableAllocator}; 2 | use crate::core::lru::LruCache; 3 | use ic_kit::stable::StableMemoryError; 4 | use std::borrow::BorrowMut; 5 | use std::cell::RefCell; 6 | 7 | thread_local! { 8 | static ALLOCATOR: RefCell> = RefCell::new(None); 9 | static LRU: RefCell> = RefCell::new(None); 10 | } 11 | 12 | /// Set the stable storage allocator instance to be used for the canister. 13 | /// 14 | /// # Panics 15 | /// 16 | /// If called more than once throughout the canister's lifetime. 17 | pub fn set_global_allocator(allocator: StableAllocator) { 18 | ALLOCATOR.with(|cell| { 19 | let mut option = cell.borrow_mut(); 20 | 21 | if option.is_some() { 22 | panic!("set_global_allocator is only supposed to be called once."); 23 | } 24 | 25 | option.replace(allocator); 26 | }); 27 | } 28 | 29 | /// Set a custom LRU cache for the canister. 30 | /// 31 | /// # Panics 32 | /// 33 | /// This method must only be called once during the initialization of the canister, either during 34 | /// `init` or `post_upgrade`. 35 | pub fn set_global_lru(lru: LruCache) { 36 | LRU.with(|cell| { 37 | let mut option = cell.borrow_mut(); 38 | 39 | if option.is_some() { 40 | panic!("set_global_lru is only supposed to be called once during initialization."); 41 | } 42 | 43 | option.replace(lru); 44 | }); 45 | } 46 | 47 | /// Allocate a block with the given size from the global stable storage allocator. 48 | pub fn allocate(size: BlockSize) -> Result { 49 | ALLOCATOR.with(|cell| { 50 | cell.borrow_mut() 51 | .as_mut() 52 | .expect("ic_kit_stable::set_global_allocator must have been called.") 53 | .allocate(size) 54 | }) 55 | } 56 | 57 | /// Free the block at the given address, the address must be the one you've retrieved earlier using 58 | /// a call to [`allocate`]. 59 | pub fn free(address: BlockSize) { 60 | ALLOCATOR.with(|cell| { 61 | cell.borrow_mut() 62 | .as_mut() 63 | .expect("ic_kit_stable::set_global_allocator must have been called.") 64 | .free(address) 65 | }) 66 | } 67 | 68 | /// Use the LRU cache instance to the callback. 69 | #[inline] 70 | pub(crate) fn with_lru U>(f: F) -> U { 71 | LRU.with(|cell| { 72 | let mut lru = cell.borrow_mut(); 73 | let lru_mut = lru.get_or_insert_with(|| LruCache::default()); 74 | f(lru_mut) 75 | }) 76 | } 77 | -------------------------------------------------------------------------------- /examples/counter/src/canister.rs: -------------------------------------------------------------------------------- 1 | use ic_kit::prelude::*; 2 | 3 | #[derive(Default)] 4 | pub struct Counter { 5 | number: u64, 6 | } 7 | 8 | impl Counter { 9 | /// Increment the counter by one. 10 | pub fn increment(&mut self) -> u64 { 11 | self.number += 1; 12 | self.number 13 | } 14 | 15 | /// Increment the counter by the provided value. 16 | pub fn increment_by(&mut self, n: u8) -> u64 { 17 | self.number += n as u64; 18 | self.number 19 | } 20 | } 21 | 22 | #[update] 23 | pub fn increment(counter: &mut Counter) -> u64 { 24 | println!("Counter Increment!"); 25 | counter.increment() 26 | } 27 | 28 | #[update] 29 | pub fn increment_by(counter: &mut Counter, n: u8) -> u64 { 30 | counter.increment_by(n) 31 | } 32 | 33 | #[query] 34 | pub fn get_counter(counter: &Counter) -> u64 { 35 | counter.number 36 | } 37 | 38 | #[derive(KitCanister)] 39 | #[candid_path("candid.did")] 40 | pub struct CounterCanister; 41 | 42 | #[cfg(test)] 43 | mod tests { 44 | use super::*; 45 | 46 | #[kit_test] 47 | async fn test_increment(replica: Replica) { 48 | let c = replica.add_canister(CounterCanister::anonymous()); 49 | 50 | let r = c 51 | .new_call("increment") 52 | .perform() 53 | .await 54 | .decode_one::() 55 | .unwrap(); 56 | 57 | assert_eq!(r, 1); 58 | 59 | assert_eq!( 60 | c.new_call("get_counter") 61 | .perform() 62 | .await 63 | .decode_one::() 64 | .unwrap(), 65 | 1 66 | ); 67 | 68 | let r = c 69 | .new_call("increment") 70 | .perform() 71 | .await 72 | .decode_one::() 73 | .unwrap(); 74 | 75 | assert_eq!(r, 2); 76 | 77 | assert_eq!( 78 | c.new_call("get_counter") 79 | .perform() 80 | .await 81 | .decode_one::() 82 | .unwrap(), 83 | 2 84 | ); 85 | } 86 | 87 | #[kit_test] 88 | async fn test_increment_by(replica: Replica) { 89 | let c = replica.add_canister(CounterCanister::anonymous()); 90 | assert_eq!( 91 | c.new_call("increment_by") 92 | .with_arg(2u8) 93 | .perform() 94 | .await 95 | .decode_one::() 96 | .unwrap(), 97 | 2 98 | ); 99 | } 100 | } 101 | -------------------------------------------------------------------------------- /examples/multi_counter/src/canister.rs: -------------------------------------------------------------------------------- 1 | use ic_kit::prelude::*; 2 | use std::collections::HashSet; 3 | 4 | #[derive(Default)] 5 | struct MultiCounter { 6 | canister_ids: HashSet, 7 | } 8 | 9 | #[update] 10 | fn increment(counters: &MultiCounter) { 11 | println!("MultiCounter Increment!"); 12 | 13 | for &canister_id in counters.canister_ids.iter() { 14 | println!("Increment on {}", canister_id); 15 | 16 | CallBuilder::new(canister_id, "increment") 17 | .perform_one_way() 18 | .expect("Expected the one way call to succeed."); 19 | } 20 | } 21 | 22 | #[update] 23 | fn add_counter(counters: &mut MultiCounter, canister_id: Principal) { 24 | println!("Add counter: {}", canister_id); 25 | counters.canister_ids.insert(canister_id); 26 | } 27 | 28 | #[derive(KitCanister)] 29 | #[candid_path("candid.did")] 30 | pub struct MultiCounterCanister; 31 | 32 | #[cfg(test)] 33 | mod tests { 34 | use super::*; 35 | use ic_kit_example_counter::CounterCanister; 36 | 37 | #[kit_test] 38 | async fn test_multi_canister(replica: Replica) { 39 | let counter1_id = Principal::from_text("whq4n-xiaaa-aaaam-qaazq-cai").unwrap(); 40 | let counter2_id = Principal::from_text("lj532-6iaaa-aaaah-qcc7a-cai").unwrap(); 41 | 42 | let canister = replica.add_canister(MultiCounterCanister::anonymous()); 43 | let counter1 = replica.add_canister(CounterCanister::build(counter1_id)); 44 | let counter2 = replica.add_canister(CounterCanister::build(counter2_id)); 45 | 46 | // Test individual counters. 47 | 48 | let r = counter1 49 | .new_call("increment") 50 | .perform() 51 | .await 52 | .decode_one::() 53 | .unwrap(); 54 | 55 | assert_eq!(r, 1); 56 | 57 | let r = counter2 58 | .new_call("increment") 59 | .perform() 60 | .await 61 | .decode_one::() 62 | .unwrap(); 63 | 64 | assert_eq!(r, 1); 65 | 66 | // Add the counters to the multi-counter. 67 | 68 | canister 69 | .new_call("add_counter") 70 | .with_arg(&counter1_id) 71 | .perform() 72 | .await; 73 | 74 | canister 75 | .new_call("add_counter") 76 | .with_arg(&counter2_id) 77 | .perform() 78 | .await; 79 | 80 | // Do a proxy increment call. 81 | let x = canister.new_call("increment").perform().await; 82 | 83 | // TODO(qti3e) replica.idle.await 84 | 85 | println!("{:#?}", x); 86 | } 87 | } 88 | -------------------------------------------------------------------------------- /ic-kit-runtime/src/handle.rs: -------------------------------------------------------------------------------- 1 | use std::panic::{RefUnwindSafe, UnwindSafe}; 2 | 3 | use candid::Principal; 4 | use tokio::sync::oneshot; 5 | 6 | use crate::call::{CallBuilder, CallReply}; 7 | use crate::types::{Env, Message, RequestId}; 8 | use crate::Replica; 9 | 10 | pub struct CanisterHandle<'a> { 11 | pub(crate) replica: &'a Replica, 12 | pub(crate) canister_id: Principal, 13 | } 14 | 15 | impl<'a> CanisterHandle<'a> { 16 | /// Create a new call builder to call this canister. 17 | pub fn new_call>(&self, method_name: S) -> CallBuilder { 18 | CallBuilder::new(self.replica, self.canister_id, method_name.into()) 19 | } 20 | 21 | /// Run the given custom function in the execution thread of the canister. 22 | pub async fn custom( 23 | &self, 24 | f: F, 25 | env: Env, 26 | ) -> CallReply { 27 | let (tx, rx) = oneshot::channel(); 28 | 29 | self.replica.enqueue_request( 30 | self.canister_id, 31 | Message::CustomTask { 32 | request_id: RequestId::new(), 33 | task: Box::new(f), 34 | env, 35 | }, 36 | Some(tx), 37 | ); 38 | 39 | rx.await.unwrap() 40 | } 41 | 42 | /// Run the given raw message in the canister's execution thread. 43 | pub async fn run_env(&self, env: Env) -> CallReply { 44 | let (tx, rx) = oneshot::channel(); 45 | 46 | self.replica.enqueue_request( 47 | self.canister_id, 48 | Message::Request { 49 | request_id: RequestId::new(), 50 | env, 51 | }, 52 | Some(tx), 53 | ); 54 | 55 | rx.await.unwrap() 56 | } 57 | 58 | /// Runs the init hook of the canister. For more customization use [`CanisterHandle::run_env`] 59 | /// with [`Env::init()`]. 60 | pub async fn init(&self) -> CallReply { 61 | self.run_env(Env::init()).await 62 | } 63 | 64 | /// Runs the pre_upgrade hook of the canister. For more customization use 65 | /// [`CanisterHandle::run_env`] with [`Env::pre_upgrade()`]. 66 | pub async fn pre_upgrade(&self) -> CallReply { 67 | self.run_env(Env::pre_upgrade()).await 68 | } 69 | 70 | /// Runs the post_upgrade hook of the canister. For more customization use 71 | /// [`CanisterHandle::run_env`] with [`Env::post_upgrade()`]. 72 | pub async fn post_upgrade(&self) -> CallReply { 73 | self.run_env(Env::post_upgrade()).await 74 | } 75 | 76 | /// Runs the post_upgrade hook of the canister. For more customization use 77 | /// [`CanisterHandle::run_env`] with [`Env::heartbeat()`]. 78 | pub async fn heartbeat(&self) -> CallReply { 79 | self.run_env(Env::heartbeat()).await 80 | } 81 | } 82 | -------------------------------------------------------------------------------- /ic-kit-sys/src/types.rs: -------------------------------------------------------------------------------- 1 | use std::error; 2 | use std::fmt; 3 | 4 | /// The result of `candid::encode_args(())` which is used as the default argument. 5 | pub const CANDID_EMPTY_ARG: &[u8] = &[68, 73, 68, 76, 0, 0]; 6 | 7 | /// Rejection code from calling another canister. 8 | #[allow(missing_docs)] 9 | #[repr(i32)] 10 | #[derive(Debug, Clone, Copy)] 11 | pub enum RejectionCode { 12 | NoError = 0, 13 | SysFatal = 1, 14 | SysTransient = 2, 15 | DestinationInvalid = 3, 16 | CanisterReject = 4, 17 | CanisterError = 5, 18 | Unknown, 19 | } 20 | 21 | impl From for RejectionCode { 22 | fn from(code: i32) -> Self { 23 | match code { 24 | 0 => RejectionCode::NoError, 25 | 1 => RejectionCode::SysFatal, 26 | 2 => RejectionCode::SysTransient, 27 | 3 => RejectionCode::DestinationInvalid, 28 | 4 => RejectionCode::CanisterReject, 29 | 5 => RejectionCode::CanisterError, 30 | _ => RejectionCode::Unknown, 31 | } 32 | } 33 | } 34 | 35 | impl From for RejectionCode { 36 | fn from(code: u32) -> Self { 37 | RejectionCode::from(code as i32) 38 | } 39 | } 40 | 41 | #[derive(Debug)] 42 | pub enum CallError { 43 | /// Indicates that the `ic0::call_perform` failed and the call is not queued. 44 | CouldNotSend, 45 | /// The rejection callback wsa called from the IC, the call failed with the given rejection 46 | /// code and message. 47 | Rejected(RejectionCode, String), 48 | /// The call happened successfully, but there was an error during deserialization of the 49 | /// response. 50 | /// The raw response is captured here. 51 | ResponseDeserializationError(Vec), 52 | } 53 | 54 | impl fmt::Display for CallError { 55 | fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { 56 | match self { 57 | CallError::CouldNotSend => f.write_str("Could not send message"), 58 | CallError::Rejected(c, m) => write!(f, "Call rejected (code={:?}): '{}'", c, m), 59 | CallError::ResponseDeserializationError(..) => { 60 | f.write_str("Could not deserialize the response.") 61 | } 62 | } 63 | } 64 | } 65 | 66 | impl error::Error for CallError {} 67 | 68 | /// A possible error value when dealing with stable memory. 69 | #[derive(Debug, Eq, PartialEq)] 70 | pub enum StableMemoryError { 71 | /// No more stable memory could be allocated. 72 | OutOfMemory, 73 | /// Attempted to read more stable memory than had been allocated. 74 | OutOfBounds, 75 | } 76 | 77 | impl fmt::Display for StableMemoryError { 78 | fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { 79 | match self { 80 | Self::OutOfMemory => f.write_str("Out of memory"), 81 | Self::OutOfBounds => f.write_str("Read exceeds allocated memory"), 82 | } 83 | } 84 | } 85 | 86 | impl error::Error for StableMemoryError {} 87 | -------------------------------------------------------------------------------- /ic-kit-runtime/src/stable.rs: -------------------------------------------------------------------------------- 1 | use memmap::MmapMut; 2 | 3 | /// A dynamic backend that can be used to handle stable storage. An implementation can decide 4 | /// where to store the data as long as it provides the given functionalities. 5 | pub trait StableMemoryBackend { 6 | fn stable_size(&mut self) -> u64; 7 | fn stable_grow(&mut self, new_pages: u64) -> i64; 8 | fn stable_read(&mut self, offset: u64, buf: &mut [u8]); 9 | fn stable_write(&mut self, offset: u64, buf: &[u8]); 10 | } 11 | 12 | /// An stable storage backend that uses a mapped file under the hood to provide the storage space. 13 | pub struct FileSystemStableMemory { 14 | _file: MmapMut, 15 | } 16 | 17 | impl StableMemoryBackend for FileSystemStableMemory { 18 | fn stable_size(&mut self) -> u64 { 19 | todo!() 20 | } 21 | fn stable_grow(&mut self, _new_pages: u64) -> i64 { 22 | todo!() 23 | } 24 | fn stable_read(&mut self, _offset: u64, _buf: &mut [u8]) { 25 | todo!() 26 | } 27 | fn stable_write(&mut self, _offset: u64, _buf: &[u8]) { 28 | todo!() 29 | } 30 | } 31 | 32 | /// An stable storage backend that stores everything in the heap. By default it has a 128MB limit. 33 | pub struct HeapStableMemory { 34 | pages: Vec<[u8; 1 << 16]>, 35 | max_pages: u64, 36 | } 37 | 38 | impl Default for HeapStableMemory { 39 | fn default() -> Self { 40 | Self { 41 | pages: Vec::new(), 42 | max_pages: 128 << 20 >> 16, 43 | } 44 | } 45 | } 46 | 47 | impl HeapStableMemory { 48 | /// Create a stable storage backend with the provided max page. 49 | pub fn new(max_pages: u64) -> Self { 50 | Self { 51 | pages: Vec::new(), 52 | max_pages, 53 | } 54 | } 55 | } 56 | 57 | impl StableMemoryBackend for HeapStableMemory { 58 | fn stable_size(&mut self) -> u64 { 59 | self.pages.len() as u64 60 | } 61 | 62 | fn stable_grow(&mut self, new_pages: u64) -> i64 { 63 | let size = self.pages.len() as u64; 64 | if new_pages + size > self.max_pages { 65 | -1 66 | } else { 67 | for _ in 0..new_pages { 68 | self.pages.push([0; 1 << 16]); 69 | } 70 | size as i64 71 | } 72 | } 73 | 74 | fn stable_read(&mut self, offset: u64, buf: &mut [u8]) { 75 | // TODO(qti3e) This can be optimized. 76 | for i in 0..buf.len() { 77 | let offset = offset + i as u64; 78 | let page = offset >> 16; 79 | let byte = offset - (page << 16); 80 | buf[i] = self.pages[page as usize][byte as usize]; 81 | } 82 | } 83 | 84 | fn stable_write(&mut self, offset: u64, buf: &[u8]) { 85 | // TODO(qti3e) This can be optimized. 86 | for i in 0..buf.len() { 87 | let offset = offset + i as u64; 88 | let page = offset >> 16; 89 | let byte = offset - (page << 16); 90 | self.pages[page as usize][byte as usize] = buf[i]; 91 | } 92 | } 93 | } 94 | -------------------------------------------------------------------------------- /ic-kit-macros/src/lib.rs: -------------------------------------------------------------------------------- 1 | use proc_macro::TokenStream; 2 | 3 | use syn::parse_macro_input; 4 | 5 | use entry::{gen_entry_point_code, EntryPoint}; 6 | use test::gen_test_code; 7 | 8 | mod entry; 9 | mod export_service; 10 | mod metadata; 11 | mod test; 12 | 13 | fn process_entry_point( 14 | entry_point: EntryPoint, 15 | attr: TokenStream, 16 | item: TokenStream, 17 | ) -> TokenStream { 18 | gen_entry_point_code(entry_point, attr.into(), item.into()) 19 | .unwrap_or_else(|error| error.to_compile_error()) 20 | .into() 21 | } 22 | 23 | /// Export the function as the init hook of the canister. 24 | #[proc_macro_attribute] 25 | pub fn init(attr: TokenStream, item: TokenStream) -> TokenStream { 26 | process_entry_point(EntryPoint::Init, attr, item) 27 | } 28 | 29 | /// Export the function as the pre_upgrade hook of the canister. 30 | #[proc_macro_attribute] 31 | pub fn pre_upgrade(attr: TokenStream, item: TokenStream) -> TokenStream { 32 | process_entry_point(EntryPoint::PreUpgrade, attr, item) 33 | } 34 | 35 | /// Export the function as the post_upgrade hook of the canister. 36 | #[proc_macro_attribute] 37 | pub fn post_upgrade(attr: TokenStream, item: TokenStream) -> TokenStream { 38 | process_entry_point(EntryPoint::PostUpgrade, attr, item) 39 | } 40 | 41 | /// Export the function as the inspect_message hook of the canister. 42 | #[proc_macro_attribute] 43 | pub fn inspect_message(attr: TokenStream, item: TokenStream) -> TokenStream { 44 | process_entry_point(EntryPoint::InspectMessage, attr, item) 45 | } 46 | 47 | /// Export the function as the heartbeat hook of the canister. 48 | #[proc_macro_attribute] 49 | pub fn heartbeat(attr: TokenStream, item: TokenStream) -> TokenStream { 50 | process_entry_point(EntryPoint::Heartbeat, attr, item) 51 | } 52 | 53 | /// Export an update method for the canister. 54 | #[proc_macro_attribute] 55 | pub fn update(attr: TokenStream, item: TokenStream) -> TokenStream { 56 | process_entry_point(EntryPoint::Update, attr, item) 57 | } 58 | 59 | /// Export a query method for the canister. 60 | #[proc_macro_attribute] 61 | pub fn query(attr: TokenStream, item: TokenStream) -> TokenStream { 62 | process_entry_point(EntryPoint::Query, attr, item) 63 | } 64 | 65 | /// A macro to generate IC-Kit tests. 66 | #[proc_macro_attribute] 67 | pub fn kit_test(attr: TokenStream, item: TokenStream) -> TokenStream { 68 | gen_test_code(attr.into(), item.into()) 69 | .unwrap_or_else(|error| error.to_compile_error()) 70 | .into() 71 | } 72 | 73 | #[proc_macro_derive(KitCanister, attributes(candid_path))] 74 | pub fn kit_export(input: TokenStream) -> TokenStream { 75 | let input = parse_macro_input!(input as syn::DeriveInput); 76 | let save_candid_path_result = get_save_candid_path(&input); 77 | 78 | match save_candid_path_result { 79 | Ok(save_candid_path) => export_service::export_service(input, save_candid_path).into(), 80 | Err(e) => e.to_compile_error().into(), 81 | } 82 | } 83 | 84 | fn get_save_candid_path(input: &syn::DeriveInput) -> syn::Result> { 85 | let candid_path_helper_attribute_option = input 86 | .attrs 87 | .iter() 88 | .find(|attr| attr.path.is_ident("candid_path")); 89 | 90 | match candid_path_helper_attribute_option { 91 | Some(candid_path_helper_attribute) => { 92 | let custom_candid_path_lit: syn::LitStr = candid_path_helper_attribute.parse_args()?; 93 | Ok(Some(custom_candid_path_lit)) 94 | } 95 | None => Ok(None), 96 | } 97 | } 98 | -------------------------------------------------------------------------------- /ic-kit/src/ic/canister.rs: -------------------------------------------------------------------------------- 1 | use crate::ic::{with, Cycles}; 2 | use candid::Principal; 3 | use ic_kit_sys::ic0; 4 | use std::convert::TryFrom; 5 | 6 | /// A type wrapper for the current canister's Principal ID. 7 | #[derive(Clone)] 8 | struct CanisterPrincipalId(Principal); 9 | 10 | impl Default for CanisterPrincipalId { 11 | fn default() -> Self { 12 | let len: usize = unsafe { ic0::canister_self_size() as usize }; 13 | let mut bytes = vec![0u8; len]; 14 | unsafe { 15 | ic0::canister_self_copy(bytes.as_mut_ptr() as isize, 0, len as isize); 16 | } 17 | Self(Principal::try_from(&bytes).unwrap()) 18 | } 19 | } 20 | 21 | /// Trap the canister with the provided message. This will rollback the canister state at the 22 | /// beginning of the current entry point, consider rejecting the message gracefully. 23 | #[inline(always)] 24 | pub fn trap(message: &str) -> ! { 25 | unsafe { 26 | ic0::trap(message.as_ptr() as isize, message.len() as isize); 27 | } 28 | unreachable!() 29 | } 30 | 31 | /// Print a debug message from the canister that can be viewed during local development. 32 | #[inline(always)] 33 | pub fn print>(s: S) { 34 | let s = s.as_ref(); 35 | unsafe { 36 | ic0::debug_print(s.as_ptr() as isize, s.len() as isize); 37 | } 38 | } 39 | 40 | /// ID of the current canister. 41 | #[inline(always)] 42 | pub fn id() -> Principal { 43 | // Hmm... so you've made it this far to this codebase, we simple only ask the runtime once 44 | // to return the principal id of the current canister, and then we store it. 45 | with(CanisterPrincipalId::clone).0 46 | } 47 | 48 | /// The time in nanoseconds. 49 | #[inline(always)] 50 | pub fn time() -> u64 { 51 | unsafe { ic0::time() as u64 } 52 | } 53 | 54 | /// The balance of the canister. 55 | #[inline(always)] 56 | pub fn balance() -> Cycles { 57 | #[cfg(not(feature = "experimental-cycles128"))] 58 | unsafe { 59 | ic0::canister_cycle_balance() as u64 60 | } 61 | 62 | #[cfg(feature = "experimental-cycles128")] 63 | unsafe { 64 | let mut recv = 0u128; 65 | ic0::canister_cycle_balance128(&mut recv as *mut u128 as isize); 66 | u128::from_le(recv) 67 | } 68 | } 69 | 70 | /// The caller who has invoked this method on the canister. 71 | /// 72 | /// # Panics 73 | /// 74 | /// If called after a reply/reject callback. 75 | #[inline(always)] 76 | pub fn caller() -> Principal { 77 | let len = unsafe { ic0::msg_caller_size() as usize }; 78 | let mut bytes = vec![0u8; len]; 79 | unsafe { 80 | ic0::msg_caller_copy(bytes.as_mut_ptr() as isize, 0, len as isize); 81 | } 82 | Principal::try_from(&bytes).unwrap() 83 | } 84 | 85 | /// Set the certified data of the canister, this method traps if data.len > 32. 86 | #[inline(always)] 87 | pub fn set_certified_data(data: &[u8]) { 88 | unsafe { ic0::certified_data_set(data.as_ptr() as isize, data.len() as isize) } 89 | } 90 | 91 | /// Returns the data certificate authenticating certified_data set by this canister. 92 | #[inline(always)] 93 | pub fn data_certificate() -> Option> { 94 | if unsafe { ic0::data_certificate_present() } == 0 { 95 | return None; 96 | } 97 | 98 | let n = unsafe { ic0::data_certificate_size() }; 99 | let mut buf = vec![0u8; n as usize]; 100 | unsafe { 101 | ic0::data_certificate_copy(buf.as_mut_ptr() as isize, 0, n); 102 | } 103 | Some(buf) 104 | } 105 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # IC Kit 2 | 3 | [![Docs](https://docs.rs/ic-kit/badge.svg)](https://docs.rs/ic-kit) 4 | 5 | This library provides an alternative to `ic-cdk` that can help developers write canisters 6 | and unit test them in their Rust code. 7 | 8 | Blog post: [IC Kit: Rust Canister Testing Kit we Used to Improve Code Coverage at Fleek](https://blog.fleek.co/posts/ickit-rust-canister-testing) 9 | 10 | # THIS IS ALPHA SOFTWARE 11 | 12 | IC-Kit v0.5 is still alpha and a developer preview. we're working towards making sure it's feature complete, and advise 13 | you to not use it in any sensitive canister before it comes out of the alpha. 14 | 15 | ## Install 16 | 17 | Add this to your `Cargo.toml` 18 | 19 | ```toml 20 | [dependencies] 21 | ic-kit = "0.5.0-alpha.6" 22 | candid = "0.8" 23 | ``` 24 | 25 | ## Example Usage 26 | 27 | See [the examples](./examples) directory. 28 | 29 | ## What's different? 30 | 31 | IC-Kit 0.5 is breaking drop-in replacement guarantee of the CDK, which allows us to go one step further in improving 32 | canister development experience. 33 | 34 | > Note: The kit includes a breaking change to the Principal type. The kit is upgraded to candid 0.8 and ic_types 0.6, where `Principal` lives directly in the candid crate. 35 | 36 | ### Fully Simulated Replica 37 | 38 | Now we have a `#[kit_test]` macro which gives you access to a replica simulator that you can use for testing your 39 | canister. 40 | 41 | ```rust 42 | use ic_kit::prelude::*; 43 | 44 | #[kit_test] 45 | async fn test(replica: Replica) { 46 | // let handle = replica.add_canister(Canister::anonymous()); 47 | } 48 | ``` 49 | 50 | ### Inspect Message 51 | 52 | It makes it easier to use the `inspect_message` feature of the Interest Computer, your function only 53 | needs to return a `bool` and we take care of the rest. 54 | 55 | ```rust 56 | use ic_kit::prelude::*; 57 | 58 | #[inspect_message] 59 | fn inspect_message() -> bool { 60 | // Only accept ingress message calls which have a payload 61 | // smaller than 1000 bytes. 62 | ic_kit::arg_data_size::arg_data_size() <= 1000 63 | } 64 | ``` 65 | 66 | ### Secure Memory Helpers 67 | 68 | No need to use `thread_local!` to get rid of the `get_mut` anymore, we have deprecated `get[/_mut` method 69 | and now have `with` variation. 70 | 71 | ```rust 72 | let count = with(|counter: &Counter| { 73 | *counter.get() 74 | }); 75 | 76 | let count = with_mut(|counter: &mut Counter| { 77 | *counter.get() 78 | }); 79 | ``` 80 | 81 | ### Dependency Injection 82 | 83 | Now your sync (non-async) methods can be simplified, we wrap them in the appropriate `with` and `with_mut` methods for you 84 | so you don't have to think about it. 85 | 86 | ```rust 87 | use ic_kit::prelude::*; 88 | use std::collections::HashMap; 89 | 90 | #[derive(Default)] 91 | struct Registry { 92 | names: HashMap, 93 | } 94 | 95 | #[update] 96 | fn register(registry: &mut Registry, name: String) { 97 | registry.names.insert(caller(), name); 98 | } 99 | 100 | #[query] 101 | fn get_name(registry: &Registry, user: Principal) -> Option<&String> { 102 | registry.names.get(&user) 103 | } 104 | ``` 105 | 106 | ### Native Macros 107 | 108 | Now we no longer rely on the `ic-cdk-macros` allowing us to host our version of macros and innovate even more. 109 | 110 | ### Importable Canisters 111 | 112 | The `KitCanister` derive macro allows you to export a canister to be imported from other canisters. 113 | 114 | ### Easy Auto Candid Generation 115 | 116 | ```rust 117 | #[derive(KitCanister)] 118 | #[candid_path("candid.did")] 119 | pub struct NamingSystemCanister; 120 | ``` 121 | -------------------------------------------------------------------------------- /ic-kit-stable/src/core/checksum.rs: -------------------------------------------------------------------------------- 1 | /// A u40 value that uses the extra padding of u64 for a checksum. 2 | #[repr(packed)] 3 | pub struct CheckedU40(u64); 4 | 5 | impl CheckedU40 { 6 | /// Create a new checked u40 value. 7 | /// 8 | /// # Panics 9 | /// 10 | /// If the provided value is larger than 40bits, or if it's zero. 11 | pub fn new(value: u64) -> Self { 12 | if value & 0xffffff0000000000 > 0 { 13 | panic!("only 40bit integers are supported."); 14 | } 15 | 16 | if value == 0 { 17 | panic!("zero for checksum is not supported.") 18 | } 19 | 20 | let a = (value & 0xff00000000) >> 32; 21 | let b = (value & 0x00ff000000) >> 24; 22 | let c = (value & 0x0000ff0000) >> 16; 23 | let d = (value & 0x000000ff00) >> 8; 24 | let e = value & 0x00000000ff; 25 | let x = a ^ b ^ c; 26 | let y = c ^ d ^ e; 27 | let z = x ^ y; 28 | let s = (x << 16) | (y << 8) | z; 29 | let r = (s << 40) | value; 30 | 31 | CheckedU40(r) 32 | } 33 | 34 | /// Verify and unpack the wrapped value. 35 | pub fn verify(&self) -> Option { 36 | let value = self.0; 37 | 38 | // we don't support zero. 39 | if value == 0 { 40 | return None; 41 | } 42 | 43 | let x = (value & 0xff00000000000000) >> 56; 44 | let y = (value & 0x00ff000000000000) >> 48; 45 | let z = (value & 0x0000ff0000000000) >> 40; 46 | let a = (value & 0xff00000000) >> 32; 47 | let b = (value & 0x00ff000000) >> 24; 48 | let c = (value & 0x0000ff0000) >> 16; 49 | let d = (value & 0x000000ff00) >> 8; 50 | let e = value & 0x00000000ff; 51 | 52 | let xx = a ^ b ^ c; 53 | let yy = c ^ d ^ e; 54 | let zz = x ^ y; 55 | 56 | if xx == x && yy == y && zz == z { 57 | Some(value & 0x000000ffffffffff) 58 | } else { 59 | None 60 | } 61 | } 62 | 63 | /// Return the protected number without any checks. 64 | pub fn unchecked(&self) -> u64 { 65 | self.0 & 0x000000ffffffffff 66 | } 67 | } 68 | 69 | #[cfg(test)] 70 | mod tests { 71 | use super::*; 72 | 73 | #[test] 74 | fn test_verify() { 75 | assert_eq!(CheckedU40::new(0x34653443ab).verify(), Some(0x34653443ab)); 76 | assert_eq!(CheckedU40(0x34653443ab).verify(), None); 77 | 78 | for i in 0..1_000 { 79 | let n = ((0x121212 * i) << 40) | (0x1010101010 * (i + 0x12345)); 80 | assert_eq!(CheckedU40(n).verify(), None); 81 | 82 | let n = n & 0x000000ffffffffff; 83 | assert_eq!(CheckedU40::new(n).verify(), Some(n)); 84 | } 85 | 86 | assert_eq!( 87 | CheckedU40::new(0x000000ffffffffff).verify(), 88 | Some(0x000000ffffffffff) 89 | ); 90 | } 91 | 92 | #[test] 93 | #[should_panic] 94 | fn test_zero() { 95 | CheckedU40::new(0); 96 | } 97 | 98 | #[test] 99 | fn test_unchecked() { 100 | for i in 0..1_000 { 101 | let n = ((0x121212 * i) << 40) | (0x1010101010 * (i + 0x12345)); 102 | let n = n & 0x000000ffffffffff; 103 | assert_eq!(CheckedU40::new(n).unchecked(), n); 104 | } 105 | } 106 | 107 | #[test] 108 | #[should_panic] 109 | fn test_non_40bit() { 110 | CheckedU40::new(0x000000ffffffffff + 1); 111 | } 112 | 113 | #[test] 114 | fn test_invalid() { 115 | assert_eq!(CheckedU40(0xffffffffffffffff).verify(), None); 116 | assert_eq!(CheckedU40(0).verify(), None); 117 | } 118 | } 119 | -------------------------------------------------------------------------------- /ic-kit/src/ic/stable.rs: -------------------------------------------------------------------------------- 1 | use ic_kit_sys::ic0; 2 | use ic_kit_sys::types::StableMemoryError; 3 | 4 | /// A type which represents either a page count or an offset in the stable storage, it's a u64 5 | /// when the `experimental-stable64` feature is enabled, otherwise a u32. 6 | #[cfg(feature = "experimental-stable64")] 7 | pub type StableSize = u64; 8 | 9 | /// A type which represents either a page count or an offset in the stable storage, it's a u64 10 | /// when the `experimental-stable64` feature is enabled, otherwise a u32. 11 | #[cfg(not(feature = "experimental-stable64"))] 12 | pub type StableSize = u32; 13 | 14 | /// Returns the current size of the stable memory in WebAssembly pages. 15 | /// Note: One WebAssembly page is 64KiB 16 | #[inline(always)] 17 | pub fn stable_size() -> StableSize { 18 | #[cfg(not(feature = "experimental-stable64"))] 19 | unsafe { 20 | ic0::stable_size() as u32 21 | } 22 | 23 | #[cfg(feature = "experimental-stable64")] 24 | unsafe { 25 | ic0::stable64_size() as u64 26 | } 27 | } 28 | 29 | /// Attempts to grow the stable memory by `new_pages` (added pages). 30 | /// 31 | /// Returns an error if it wasn't possible. Otherwise, returns the previous 32 | /// size that was reserved. 33 | /// 34 | /// Note: One WebAssembly page is 64KiB 35 | #[inline(always)] 36 | pub fn stable_grow(new_pages: StableSize) -> Result { 37 | #[cfg(not(feature = "experimental-stable64"))] 38 | unsafe { 39 | match ic0::stable_grow(new_pages as i32) { 40 | -1 => Err(StableMemoryError::OutOfMemory), 41 | x => Ok(x as u32), 42 | } 43 | } 44 | 45 | #[cfg(feature = "experimental-stable64")] 46 | unsafe { 47 | match if new_pages < (u32::MAX as u64 - 1) { 48 | ic0::stable_grow(new_pages as i32) as i64 49 | } else { 50 | ic0::stable64_grow(new_pages as i64) as i64 51 | } { 52 | -1 => Err(StableMemoryError::OutOfMemory), 53 | x => Ok(x as u64), 54 | } 55 | } 56 | } 57 | 58 | /// Writes data to the stable memory location specified by an offset. 59 | #[inline(always)] 60 | pub fn stable_write(offset: StableSize, buf: &[u8]) { 61 | #[cfg(not(feature = "experimental-stable64"))] 62 | unsafe { 63 | ic0::stable_write(offset as i32, buf.as_ptr() as isize, buf.len() as isize) 64 | } 65 | 66 | #[cfg(feature = "experimental-stable64")] 67 | unsafe { 68 | if offset < (u32::MAX as u64 - 1) { 69 | ic0::stable_write(offset as i32, buf.as_ptr() as isize, buf.len() as isize) 70 | } else { 71 | ic0::stable64_write(offset as i64, buf.as_ptr() as i64, buf.len() as i64) 72 | } 73 | } 74 | } 75 | 76 | /// Reads data from the stable memory location specified by an offset. 77 | #[inline(always)] 78 | pub fn stable_read(offset: StableSize, buf: &mut [u8]) { 79 | #[cfg(not(feature = "experimental-stable64"))] 80 | unsafe { 81 | ic0::stable_read(buf.as_ptr() as isize, offset as i32, buf.len() as isize); 82 | }; 83 | 84 | #[cfg(feature = "experimental-stable64")] 85 | unsafe { 86 | if offset < (u32::MAX as u64 - 1) { 87 | ic0::stable_read(buf.as_ptr() as isize, offset as i32, buf.len() as isize); 88 | } else { 89 | ic0::stable64_read(buf.as_ptr() as i64, offset as i64, buf.len() as i64); 90 | } 91 | } 92 | } 93 | 94 | pub(crate) fn stable_bytes() -> Vec { 95 | let size = (stable_size() as usize) << 16; 96 | let mut vec = Vec::with_capacity(size); 97 | unsafe { 98 | ic0::stable_read(vec.as_ptr() as isize, 0, size as isize); 99 | vec.set_len(size); 100 | } 101 | vec 102 | } 103 | -------------------------------------------------------------------------------- /ic-kit-stable/src/core/allocator.rs: -------------------------------------------------------------------------------- 1 | use crate::core::checksum::CheckedU40; 2 | use crate::core::hole::HoleList; 3 | use crate::core::memory::{DefaultMemory, IcMemory, Memory}; 4 | use crate::core::utils::read_struct; 5 | use ic_kit::stable::StableMemoryError; 6 | 7 | /// An address to a block. 8 | pub type BlockAddress = u64; 9 | 10 | /// Size of a block. 11 | pub type BlockSize = u64; 12 | 13 | /// The internal minimum allocation size (includes size header) 14 | /// size : u64 = 8 bytes 15 | /// next : u64 = 8 bytes 16 | /// If the node is used then next is overwritten by content. 17 | pub const MIN_ALLOCATION_SIZE: BlockSize = 16; 18 | 19 | // TODO(qti3e) next steps: 20 | // write the HoleList root to stable storage at the first block. 21 | // load the HoleList from stable storage if present. 22 | 23 | /// An allocator over the stable storage. This allocator assumes that it owns the entire stable 24 | /// storage if there are already data in the stable storage. 25 | pub struct StableAllocator { 26 | hole_list: HoleList, 27 | } 28 | 29 | impl StableAllocator { 30 | pub fn new() -> Self { 31 | Self { 32 | hole_list: HoleList::new(), 33 | } 34 | } 35 | 36 | /// Allocate a stable storage block with the given size. 37 | pub fn allocate(&mut self, size: BlockSize) -> Result { 38 | // we need 8 more bytes to store the CheckedU40 for the block size. 39 | let size = size + 8; 40 | 41 | if let Some((addr, _)) = self.hole_list.find(size) { 42 | // skip the block's size which is inserted into the first 8 bytes of the block. 43 | return Ok(addr + 8); 44 | } 45 | 46 | // number of pages we need to grow in order to fit this size. this is a ceiling division. 47 | // by 1 WebAssembly page. 48 | let new_pages = (size + (1 << 16) - 1) >> 16; 49 | let start = M::stable_grow(new_pages); 50 | 51 | // we couldn't allocate anymore. 52 | if start == -1 { 53 | return Err(StableMemoryError::OutOfMemory); 54 | } 55 | 56 | let addr = (start as u64) << 16; 57 | self.hole_list.insert(addr, new_pages << 16); 58 | 59 | let addr = self 60 | .hole_list 61 | .find(size) 62 | .expect("unreachable allocation condition.") 63 | .0; 64 | 65 | Ok(addr + 8) 66 | } 67 | 68 | /// Free the stable storage block at the given address. The address must be an address returned 69 | /// by a previous invocation to the [`allocate`] method. 70 | pub fn free(&mut self, addr: BlockAddress) { 71 | if addr < 8 { 72 | return; 73 | } 74 | 75 | let addr = addr - 8; 76 | 77 | // guard the api misuse by checking the checksum. 78 | if let Some(size) = read_struct::(addr).verify() { 79 | self.hole_list.insert(addr, size); 80 | } else { 81 | #[cfg(test)] 82 | panic!("Invalid pointer passed to free().") 83 | } 84 | } 85 | } 86 | 87 | #[cfg(test)] 88 | mod tests { 89 | use super::*; 90 | 91 | #[test] 92 | #[should_panic] 93 | fn free_misuse() { 94 | let mut allocator = StableAllocator::::new(); 95 | assert_eq!(allocator.allocate(100), Ok(8)); 96 | assert_eq!(allocator.allocate(100), Ok(116)); 97 | allocator.free(100); 98 | } 99 | 100 | #[test] 101 | fn allocate_after_free() { 102 | let mut allocator = StableAllocator::::new(); 103 | assert_eq!(allocator.allocate(100), Ok(8)); 104 | assert_eq!(allocator.allocate(100), Ok(116)); 105 | allocator.free(8); 106 | assert_eq!(allocator.allocate(100), Ok(8)); 107 | } 108 | } 109 | -------------------------------------------------------------------------------- /ic-kit/src/ic/storage.rs: -------------------------------------------------------------------------------- 1 | use crate::storage::{BorrowMany, BorrowMutMany, Storage}; 2 | 3 | thread_local! { 4 | static STORAGE: Storage = Storage::default(); 5 | } 6 | 7 | /// Pass an immutable reference to the value associated with the given type to the closure. 8 | /// 9 | /// If no value is currently associated to the type `T`, this method will insert the default 10 | /// value in its place before invoking the callback. Use `maybe_with` if you don't want the 11 | /// default value to be inserted or if your type does not implement the [`Default`] trait. 12 | /// 13 | /// This is a safe replacement for the previously known `ic_kit::ic::get` API, and you can use it 14 | /// instead of `lazy_static` or `local_thread`. 15 | pub fn with U>(callback: F) -> U { 16 | STORAGE.with(|storage| storage.with(callback)) 17 | } 18 | 19 | /// Like [`with`], but does not initialize the data with the default value and simply returns None, 20 | /// if there is no value associated with the type. 21 | pub fn maybe_with U>(callback: F) -> Option { 22 | STORAGE.with(|storage| storage.maybe_with(callback)) 23 | } 24 | 25 | /// Pass a mutable reference to the value associated with the given type to the closure. 26 | /// 27 | /// If no value is currently associated to the type `T`, this method will insert the default 28 | /// value in its place before invoking the callback. Use `maybe_with_mut` if you don't want the 29 | /// default value to be inserted or if your type does not implement the [`Default`] trait. 30 | /// 31 | /// This is a safe replacement for the previously known `ic_kit::ic::get` API, and you can use it 32 | /// instead of `lazy_static` or `local_thread`. 33 | pub fn with_mut U>(callback: F) -> U { 34 | STORAGE.with(|storage| storage.with_mut(callback)) 35 | } 36 | 37 | /// Like [`with_mut`], but does not initialize the data with the default value and simply returns 38 | /// None, if there is no value associated with the type. 39 | pub fn maybe_with_mut U>(callback: F) -> Option { 40 | STORAGE.with(|storage| storage.maybe_with_mut(callback)) 41 | } 42 | 43 | /// Remove the current value associated with the type and return it. 44 | pub fn take() -> Option { 45 | STORAGE.with(|storage| storage.take::()) 46 | } 47 | 48 | /// Swaps the value associated with type `T` with the given value, returns the old one. 49 | pub fn swap(value: T) -> Option { 50 | STORAGE.with(|storage| storage.swap(value)) 51 | } 52 | 53 | /// Like [`crate::ic::with`] but passes the immutable reference of multiple variables to the 54 | /// closure as a tuple. 55 | /// 56 | /// # Example 57 | /// ``` 58 | /// use ic_kit::ic; 59 | /// 60 | /// #[derive(Default)] 61 | /// struct S1 { 62 | /// a: u64, 63 | /// } 64 | /// 65 | /// #[derive(Default)] 66 | /// struct S2 { 67 | /// a: u64, 68 | /// } 69 | /// 70 | /// ic::with_many(|(a, b): (&S1, &S2)| { 71 | /// // Now we have access to both S1 and S2. 72 | /// println!("S1: {}, S2: {}", a.a, b.a); 73 | /// }); 74 | /// ``` 75 | pub fn with_many U>(callback: F) -> U { 76 | STORAGE.with(|storage| storage.with_many(callback)) 77 | } 78 | 79 | /// Like [`crate::ic::with_mut`] but passes the mutable reference of multiple variables to the 80 | /// closure as a tuple. 81 | /// 82 | /// # Example 83 | /// ``` 84 | /// use ic_kit::ic; 85 | /// 86 | /// #[derive(Default)] 87 | /// struct S1 { 88 | /// a: u64, 89 | /// } 90 | /// 91 | /// #[derive(Default)] 92 | /// struct S2 { 93 | /// a: u64, 94 | /// } 95 | /// 96 | /// ic::with_many_mut(|(a, b): (&mut S1, &mut S2)| { 97 | /// // Now we have access to both S1 and S2 and can mutate them. 98 | /// a.a += 1; 99 | /// b.a += 1; 100 | /// }); 101 | /// ``` 102 | pub fn with_many_mut U>(callback: F) -> U { 103 | STORAGE.with(|storage| storage.with_many_mut(callback)) 104 | } 105 | -------------------------------------------------------------------------------- /ic-kit-certified/src/label.rs: -------------------------------------------------------------------------------- 1 | use candid::Principal; 2 | use std::borrow::{Borrow, Cow}; 3 | use std::ptr::NonNull; 4 | use std::rc::Rc; 5 | use std::sync::Arc; 6 | 7 | /// Any value that can be used as a label in the [`HashTree`] and can be a key 8 | /// in the [`RbTree`]. 9 | /// 10 | /// [`HashTree`]: crate::HashTree 11 | /// [`RbTree`]: crate::rbtree::RbTree 12 | pub trait Label: Ord { 13 | fn as_label(&self) -> Cow<[u8]>; 14 | } 15 | 16 | /// A type `T` can be defined as prefix of type `U`, if they follow the same 17 | /// representation and any valid value of `T` is also a valid head for a value 18 | /// of type `U`. 19 | /// 20 | /// The implementation should guarantee that the ordering is preserved which 21 | /// implies: 22 | /// For any `u: U = x . y` where `x: T`: 23 | /// 1. `x0 < x1 => u0 < u1` 24 | /// 2. `x0 > x1 => u0 > u1` 25 | /// 3. `u0 == u1 => x0 == x1` 26 | /// 27 | /// To implement this type, the Self (i.e `U`) should be borrowable as a `T`. 28 | pub trait Prefix: Label + Borrow { 29 | /// Check if the provided value is the prefix of self. The default 30 | /// implementation only extracts the prefix from Self and checks 31 | /// for their equality which might not be true for some cases 32 | /// where we're dealing with slices of variable length for example. 33 | fn is_prefix(&self, prefix: &T) -> bool { 34 | self.borrow() == prefix 35 | } 36 | } 37 | 38 | impl Label for Vec { 39 | fn as_label(&self) -> Cow<[u8]> { 40 | Cow::Borrowed(self) 41 | } 42 | } 43 | 44 | impl Prefix<[u8]> for Vec { 45 | fn is_prefix(&self, prefix: &[u8]) -> bool { 46 | self.starts_with(prefix) 47 | } 48 | } 49 | 50 | impl Label for Box<[u8]> { 51 | fn as_label(&self) -> Cow<[u8]> { 52 | Cow::Borrowed(self) 53 | } 54 | } 55 | 56 | impl Prefix<[u8]> for Box<[u8]> { 57 | fn is_prefix(&self, prefix: &[u8]) -> bool { 58 | self.starts_with(prefix) 59 | } 60 | } 61 | 62 | impl Label for Principal { 63 | fn as_label(&self) -> Cow<[u8]> { 64 | Cow::Borrowed(self.as_slice()) 65 | } 66 | } 67 | 68 | impl Label for String { 69 | fn as_label(&self) -> Cow<[u8]> { 70 | Cow::Borrowed(self.as_bytes()) 71 | } 72 | } 73 | 74 | impl Prefix for String { 75 | fn is_prefix(&self, prefix: &str) -> bool { 76 | self.as_bytes().starts_with(prefix.as_bytes()) 77 | } 78 | } 79 | 80 | impl Label for bool { 81 | fn as_label(&self) -> Cow<[u8]> { 82 | if *self { 83 | Cow::Owned(vec![1]) 84 | } else { 85 | Cow::Owned(vec![0]) 86 | } 87 | } 88 | } 89 | 90 | impl Label for Box 91 | where 92 | T: Label, 93 | { 94 | #[inline] 95 | fn as_label(&self) -> Cow<[u8]> { 96 | self.as_ref().as_label() 97 | } 98 | } 99 | 100 | impl Label for Rc 101 | where 102 | T: Label, 103 | { 104 | #[inline] 105 | fn as_label(&self) -> Cow<[u8]> { 106 | self.as_ref().as_label() 107 | } 108 | } 109 | 110 | impl Label for Arc 111 | where 112 | T: Label, 113 | { 114 | #[inline] 115 | fn as_label(&self) -> Cow<[u8]> { 116 | self.as_ref().as_label() 117 | } 118 | } 119 | 120 | impl Label for NonNull 121 | where 122 | T: Label, 123 | { 124 | #[inline] 125 | fn as_label(&self) -> Cow<[u8]> { 126 | unsafe { self.as_ref().as_label() } 127 | } 128 | } 129 | 130 | macro_rules! impl_fixed_size { 131 | ( $($size:expr),* ) => { 132 | $( 133 | impl Label for [u8; $size] { 134 | #[inline] 135 | fn as_label(&self) -> Cow<[u8]> { 136 | Cow::Borrowed(self) 137 | } 138 | } 139 | 140 | impl Prefix<[u8]> for [u8; $size] { 141 | #[inline] 142 | fn is_prefix(&self, prefix: &[u8]) -> bool { 143 | self.starts_with(prefix) 144 | } 145 | } 146 | )* 147 | } 148 | } 149 | 150 | impl_fixed_size!( 151 | 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, 152 | 27, 28, 29, 30, 31, 32 153 | ); 154 | 155 | macro_rules! impl_num { 156 | ( $($name:ty),* ) => { 157 | $( 158 | impl Label for $name { 159 | fn as_label(&self) -> Cow<[u8]> { 160 | Cow::Owned(self.to_be_bytes().into()) 161 | } 162 | } 163 | )* 164 | } 165 | } 166 | 167 | impl_num!(u8, i8, u16, i16, u32, i32, u64, i64, u128, i128, usize, isize); 168 | -------------------------------------------------------------------------------- /ic-kit-certified/src/collections/group/builder.rs: -------------------------------------------------------------------------------- 1 | use super::{Group, GroupLeaf, GroupNode, GroupNodeInner}; 2 | use std::any::{type_name, TypeId}; 3 | use std::collections::{BTreeMap, HashMap, VecDeque}; 4 | 5 | pub struct GroupBuilder { 6 | root: GroupBuilderNode, 7 | data: HashMap>, 8 | } 9 | 10 | enum GroupBuilderNode { 11 | Directory { 12 | children: BTreeMap>, 13 | }, 14 | Leaf { 15 | tid: TypeId, 16 | }, 17 | } 18 | 19 | impl GroupBuilder { 20 | pub fn new() -> Self { 21 | Self { 22 | root: GroupBuilderNode::Directory { 23 | children: BTreeMap::new(), 24 | }, 25 | data: HashMap::new(), 26 | } 27 | } 28 | 29 | pub fn insert, P: IntoIterator>( 30 | mut self, 31 | path: P, 32 | data: T, 33 | ) -> Self { 34 | let path = path 35 | .into_iter() 36 | .map(|x| x.into()) 37 | .collect::>(); 38 | 39 | let tid = TypeId::of::(); 40 | 41 | if self.data.insert(tid, Box::new(data)).is_some() { 42 | panic!("Type '{}' is already used in the group.", type_name::()) 43 | } 44 | 45 | self.root.insert(path, tid); 46 | 47 | self 48 | } 49 | 50 | #[must_use = "The constructed group must be used."] 51 | pub fn build(self) -> Group { 52 | let mut group = Group { 53 | root: self.root.build(), 54 | data: self.data, 55 | dependencies: Default::default(), 56 | }; 57 | 58 | group.init(); 59 | 60 | group 61 | } 62 | } 63 | 64 | impl Default for GroupBuilder { 65 | fn default() -> Self { 66 | Self::new() 67 | } 68 | } 69 | 70 | impl GroupBuilderNode { 71 | pub fn insert(&mut self, mut path: VecDeque, tid: TypeId) { 72 | if let GroupBuilderNode::Directory { children } = self { 73 | if path.len() == 1 { 74 | let name = path.pop_back().unwrap(); 75 | 76 | let leaf = GroupBuilderNode::Leaf { tid }; 77 | 78 | children 79 | .entry(name) 80 | .and_modify(|_| panic!("Path is already used.")) 81 | .or_insert_with(|| Box::new(leaf)); 82 | 83 | return; 84 | } 85 | let dir_name = path.pop_front().unwrap(); 86 | 87 | children 88 | .entry(dir_name) 89 | .or_insert_with(|| { 90 | Box::new(GroupBuilderNode::Directory { 91 | children: BTreeMap::new(), 92 | }) 93 | }) 94 | .insert(path, tid); 95 | return; 96 | } 97 | 98 | panic!("Can not insert to a leaf node."); 99 | } 100 | 101 | pub fn build(self) -> GroupNode { 102 | match self { 103 | GroupBuilderNode::Directory { children } => { 104 | let mut children = children 105 | .into_iter() 106 | .map(|(k, v)| GroupNode { 107 | id: 0, 108 | data: GroupNodeInner::Labeled(k, Box::new(v.build())), 109 | }) 110 | .collect::>(); 111 | 112 | // Create a semi-balanced binary tree out of the child nodes. 113 | // (Using `reduce` will not generate a balanced tree.) 114 | while children.len() > 1 { 115 | let mut new_children = VecDeque::with_capacity(children.len() / 2); 116 | 117 | while children.len() > 1 { 118 | let a = Box::new(children.pop_front().unwrap()); 119 | let b = Box::new(children.pop_front().unwrap()); 120 | new_children.push_back(GroupNode { 121 | id: 0, 122 | data: GroupNodeInner::Fork(a, b), 123 | }) 124 | } 125 | 126 | if let Some(last) = children.pop_front() { 127 | new_children.push_back(last); 128 | } 129 | 130 | children = new_children; 131 | } 132 | 133 | children.pop_front().unwrap() 134 | } 135 | GroupBuilderNode::Leaf { tid } => GroupNode { 136 | id: 0, 137 | data: GroupNodeInner::Leaf(tid), 138 | }, 139 | } 140 | } 141 | } 142 | -------------------------------------------------------------------------------- /ic-kit/src/stable.rs: -------------------------------------------------------------------------------- 1 | /// Provides utility methods to deal with stable storage on your canister. 2 | // This file is copied from ic_cdk, but changed so that it works with IC-Kit. 3 | use crate::ic::{stable_bytes, stable_grow, stable_read, stable_size, stable_write, StableSize}; 4 | use candid::utils::{ArgumentDecoder, ArgumentEncoder}; 5 | use std::io; 6 | 7 | pub use ic_kit_sys::types::StableMemoryError; 8 | 9 | /// A writer to the stable memory. 10 | /// 11 | /// Will attempt to grow the memory as it writes, 12 | /// and keep offsets and total capacity. 13 | pub struct StableWriter { 14 | /// The offset of the next write. 15 | offset: StableSize, 16 | /// The capacity, in pages. 17 | capacity: StableSize, 18 | } 19 | 20 | impl Default for StableWriter { 21 | fn default() -> Self { 22 | let capacity = stable_size(); 23 | 24 | Self { 25 | offset: 0, 26 | capacity, 27 | } 28 | } 29 | } 30 | 31 | impl StableWriter { 32 | /// Create a new stable writer that writes from the given offset forward. 33 | pub fn new(offset: StableSize) -> Self { 34 | StableWriter { 35 | offset, 36 | capacity: stable_size(), 37 | } 38 | } 39 | 40 | /// Returns the current offset of the writer. 41 | pub fn offset(&self) -> StableSize { 42 | self.offset 43 | } 44 | 45 | /// Attempts to grow the memory by adding new pages. 46 | pub fn grow(&mut self, added_pages: StableSize) -> Result<(), StableMemoryError> { 47 | let old_page_count = stable_grow(added_pages)?; 48 | self.capacity = old_page_count + added_pages; 49 | Ok(()) 50 | } 51 | 52 | /// Writes a byte slice to the buffer. 53 | /// 54 | /// The only condition where this will error out is if it cannot grow the memory. 55 | pub fn write(&mut self, buf: &[u8]) -> Result { 56 | if self.offset + (buf.len() as StableSize) > (self.capacity << 16) { 57 | self.grow((buf.len() >> 16) as StableSize + 1)?; 58 | } 59 | 60 | stable_write(self.offset, buf); 61 | self.offset += buf.len() as StableSize; 62 | Ok(buf.len()) 63 | } 64 | } 65 | 66 | impl io::Write for StableWriter { 67 | fn write(&mut self, buf: &[u8]) -> Result { 68 | self.write(buf) 69 | .map_err(|_| io::Error::new(io::ErrorKind::Other, "Out Of Memory")) 70 | } 71 | 72 | fn flush(&mut self) -> Result<(), io::Error> { 73 | // Noop. 74 | Ok(()) 75 | } 76 | } 77 | 78 | /// A reader to the stable memory. 79 | /// 80 | /// Keeps an offset and reads off stable memory consecutively. 81 | pub struct StableReader { 82 | /// The offset of the next write. 83 | offset: StableSize, 84 | } 85 | 86 | impl Default for StableReader { 87 | fn default() -> Self { 88 | Self { offset: 0 } 89 | } 90 | } 91 | 92 | impl StableReader { 93 | /// Create a new stable reader that reads from the given offset forward. 94 | pub fn new(offset: StableSize) -> Self { 95 | StableReader { offset } 96 | } 97 | 98 | /// Reads data from the stable memory location specified by an offset. 99 | pub fn read(&mut self, buf: &mut [u8]) -> Result { 100 | stable_read(self.offset, buf); 101 | self.offset += buf.len() as StableSize; 102 | Ok(buf.len()) 103 | } 104 | } 105 | 106 | impl io::Read for StableReader { 107 | fn read(&mut self, buf: &mut [u8]) -> Result { 108 | self.read(buf) 109 | .map_err(|_| io::Error::new(io::ErrorKind::Other, "Unexpected error.")) 110 | } 111 | } 112 | 113 | /// Store the given data to the stable storage. 114 | #[deprecated( 115 | since = "0.5.0", 116 | note = "This is a non-performant legacy from IC-CDK for us to deal with." 117 | )] 118 | pub fn stable_store(data: T) -> Result<(), candid::Error> 119 | where 120 | T: ArgumentEncoder, 121 | { 122 | candid::write_args(&mut StableWriter::default(), data) 123 | } 124 | 125 | /// Restore the data from the stable storage. If the data is not already stored the None value 126 | /// is returned. 127 | #[deprecated( 128 | since = "0.5.0", 129 | note = "This is a non-performant legacy from IC-CDK for us to deal with." 130 | )] 131 | pub fn stable_restore() -> Result 132 | where 133 | T: for<'de> ArgumentDecoder<'de>, 134 | { 135 | let bytes = stable_bytes(); 136 | let mut de = 137 | candid::de::IDLDeserialize::new(bytes.as_slice()).map_err(|e| format!("{:?}", e))?; 138 | let res = ArgumentDecoder::decode(&mut de).map_err(|e| format!("{:?}", e))?; 139 | Ok(res) 140 | } 141 | -------------------------------------------------------------------------------- /ic-kit-stable/src/core/pointer.rs: -------------------------------------------------------------------------------- 1 | use crate::core::allocator::{BlockAddress, BlockSize}; 2 | use crate::core::copy::StableCopy; 3 | use crate::core::global::{allocate, with_lru}; 4 | use crate::core::memory::DefaultMemory; 5 | use crate::core::utils::write_struct; 6 | use ic_kit::stable::StableMemoryError; 7 | use std::marker::PhantomData; 8 | use std::mem::ManuallyDrop; 9 | use std::ops::{Deref, DerefMut}; 10 | 11 | /// A smart pointer for data that lives on the stable storage, this uses the LRU layer to GC the 12 | /// data from heap and prevent multiple reads of the same block by keeping the most recently used 13 | /// addresses in the heap. 14 | #[repr(packed)] 15 | pub struct StablePtr(BlockAddress, PhantomData); 16 | 17 | impl Clone for StablePtr { 18 | fn clone(&self) -> Self { 19 | *self 20 | } 21 | } 22 | 23 | impl Copy for StablePtr {} 24 | 25 | impl StablePtr 26 | where 27 | T: StableCopy, 28 | { 29 | /// Allocate space for the given data on the stable storage and return a stable pointer. 30 | pub fn new(data: T) -> Result { 31 | let data = ManuallyDrop::new(data); 32 | let addr = allocate(std::mem::size_of::() as BlockSize)?; 33 | write_struct::(addr, &data); 34 | Ok(Self::from_address(addr)) 35 | } 36 | 37 | /// Create a new pointer at the given address. 38 | pub fn from_address(address: BlockAddress) -> Self { 39 | StablePtr(address, PhantomData::default()) 40 | } 41 | 42 | /// Creates an returns a null reference. 43 | pub fn null() -> Self { 44 | Self::from_address(BlockAddress::MAX) 45 | } 46 | 47 | /// Returns true if the pointer is not null. 48 | pub fn is_null(&self) -> bool { 49 | self.0 == BlockAddress::MAX 50 | } 51 | 52 | /// Returns an immutable reference to the data. 53 | pub unsafe fn as_ref(&self) -> Option> { 54 | if self.is_null() { 55 | None 56 | } else { 57 | let data = with_lru(|lru| { 58 | lru.pin(self.0); 59 | lru.get(self.0) 60 | }); 61 | 62 | Some(StableRef { 63 | data: unsafe { data as *mut T }, 64 | ptr: &self, 65 | }) 66 | } 67 | } 68 | 69 | /// Returns a mutable reference to the data. Calling this method marks the block as modified. 70 | pub unsafe fn as_mut(&self) -> Option> { 71 | if self.is_null() { 72 | None 73 | } else { 74 | let data = with_lru(|lru| { 75 | lru.pin(self.0); 76 | let data = lru.get(self.0); 77 | lru.mark_modified(self.0); 78 | data 79 | }); 80 | 81 | Some(StableRefMut { 82 | data: unsafe { data as *mut T }, 83 | ptr: &self, 84 | }) 85 | } 86 | } 87 | } 88 | 89 | /// An immutable reference to the data at the given block. 90 | pub struct StableRef<'b, T> { 91 | data: *mut T, 92 | ptr: &'b StablePtr, 93 | } 94 | 95 | /// A mutable reference to the data at the given block. 96 | pub struct StableRefMut<'b, T> { 97 | data: *mut T, 98 | ptr: &'b StablePtr, 99 | } 100 | 101 | impl Deref for StableRef<'_, T> { 102 | type Target = T; 103 | 104 | fn deref(&self) -> &Self::Target { 105 | unsafe { self.data.as_ref().unwrap() } 106 | } 107 | } 108 | 109 | impl Deref for StableRefMut<'_, T> { 110 | type Target = T; 111 | 112 | fn deref(&self) -> &Self::Target { 113 | unsafe { self.data.as_ref().unwrap() } 114 | } 115 | } 116 | 117 | impl DerefMut for StableRefMut<'_, T> { 118 | fn deref_mut(&mut self) -> &mut Self::Target { 119 | unsafe { self.data.as_mut().unwrap() } 120 | } 121 | } 122 | 123 | impl Drop for StableRef<'_, T> { 124 | fn drop(&mut self) { 125 | with_lru(|lru| { 126 | lru.unpin(self.ptr.0); 127 | }); 128 | } 129 | } 130 | 131 | impl Drop for StableRefMut<'_, T> { 132 | fn drop(&mut self) { 133 | with_lru(|lru| { 134 | lru.unpin(self.ptr.0); 135 | }); 136 | } 137 | } 138 | 139 | #[cfg(test)] 140 | mod tests { 141 | use crate::core::allocator::StableAllocator; 142 | use crate::core::copy::StableCopy; 143 | use crate::core::global::{set_global_allocator, with_lru}; 144 | use crate::core::pointer::StablePtr; 145 | 146 | struct Counter { 147 | count: u128, 148 | } 149 | 150 | impl StableCopy for Counter {} 151 | 152 | #[test] 153 | fn test_ref() { 154 | let counter = Counter { 155 | count: 0xaabbccddeeff, 156 | }; 157 | 158 | // Setup the env. 159 | set_global_allocator(StableAllocator::new()); 160 | 161 | // Create a pointer from the address. 162 | let ptr = StablePtr::new(counter).unwrap(); 163 | 164 | { 165 | let counter_ref = unsafe { ptr.as_ref().unwrap() }; 166 | assert_eq!(counter_ref.count, 0xaabbccddeeff); 167 | 168 | let mut mut_ref = unsafe { ptr.as_mut().unwrap() }; 169 | mut_ref.count = 0x1234; 170 | assert_eq!(counter_ref.count, 0x1234); 171 | } 172 | 173 | with_lru(|lru| lru.clear()); 174 | 175 | // Trying to use the ptr should still have the data. 176 | let counter_ref = unsafe { ptr.as_ref().unwrap() }; 177 | assert_eq!(counter_ref.count, 0x1234); 178 | } 179 | } 180 | -------------------------------------------------------------------------------- /ic-kit-certified/src/as_hash_tree.rs: -------------------------------------------------------------------------------- 1 | use crate::hashtree::leaf_hash; 2 | use crate::{Hash, HashTree}; 3 | use candid::{Nat, Principal}; 4 | use std::borrow::Cow; 5 | use std::ptr::NonNull; 6 | use std::rc::Rc; 7 | use std::sync::Arc; 8 | 9 | /// Defines any type that can be converted to a [`HashTree`]. 10 | /// 11 | /// For number types this trait is implemented to use the BigEndian byte ordering to represent 12 | /// the number as a [`[u8]`]. 13 | pub trait AsHashTree { 14 | /// This method should return the root hash of this hash tree. 15 | /// Must be equivalent to `as_hash_tree().reconstruct()`. 16 | /// 17 | /// Only change the default implementation if you have a better implementation 18 | /// that does not rely on reconstructing a tree. 19 | fn root_hash(&self) -> Hash { 20 | self.as_hash_tree().reconstruct() 21 | } 22 | 23 | /// Constructs a hash tree corresponding to the data. 24 | fn as_hash_tree(&self) -> HashTree<'_>; 25 | } 26 | 27 | impl AsHashTree for Vec { 28 | #[inline] 29 | fn root_hash(&self) -> Hash { 30 | leaf_hash(self.as_slice()) 31 | } 32 | 33 | #[inline] 34 | fn as_hash_tree(&self) -> HashTree<'_> { 35 | HashTree::Leaf(Cow::from(self.as_slice())) 36 | } 37 | } 38 | 39 | impl AsHashTree for &[u8] { 40 | #[inline] 41 | fn root_hash(&self) -> Hash { 42 | leaf_hash(self) 43 | } 44 | 45 | #[inline] 46 | fn as_hash_tree(&self) -> HashTree<'_> { 47 | HashTree::Leaf(Cow::from(*self)) 48 | } 49 | } 50 | 51 | impl AsHashTree for bool { 52 | #[inline] 53 | fn root_hash(&self) -> Hash { 54 | if *self { 55 | leaf_hash(&[1]) 56 | } else { 57 | leaf_hash(&[0]) 58 | } 59 | } 60 | 61 | #[inline] 62 | fn as_hash_tree(&self) -> HashTree<'_> { 63 | let value = if *self { vec![1] } else { vec![0] }; 64 | HashTree::Leaf(Cow::Owned(value)) 65 | } 66 | } 67 | 68 | impl AsHashTree for String { 69 | #[inline] 70 | fn root_hash(&self) -> Hash { 71 | leaf_hash(self.as_bytes()) 72 | } 73 | 74 | #[inline] 75 | fn as_hash_tree(&self) -> HashTree<'_> { 76 | HashTree::Leaf(Cow::from(self.as_bytes())) 77 | } 78 | } 79 | 80 | impl AsHashTree for &str { 81 | #[inline] 82 | fn root_hash(&self) -> Hash { 83 | leaf_hash(self.as_bytes()) 84 | } 85 | 86 | #[inline] 87 | fn as_hash_tree(&self) -> HashTree<'_> { 88 | HashTree::Leaf(Cow::from(self.as_bytes())) 89 | } 90 | } 91 | 92 | impl AsHashTree for Principal { 93 | #[inline] 94 | fn root_hash(&self) -> Hash { 95 | leaf_hash(self.as_slice()) 96 | } 97 | 98 | #[inline] 99 | fn as_hash_tree(&self) -> HashTree<'_> { 100 | HashTree::Leaf(Cow::from(self.as_slice())) 101 | } 102 | } 103 | 104 | impl AsHashTree for Nat { 105 | #[inline] 106 | fn root_hash(&self) -> Hash { 107 | leaf_hash(&self.0.to_bytes_be()) 108 | } 109 | 110 | #[inline] 111 | fn as_hash_tree(&self) -> HashTree<'_> { 112 | HashTree::Leaf(Cow::Owned(self.0.to_bytes_be())) 113 | } 114 | } 115 | 116 | impl AsHashTree for Box 117 | where 118 | T: AsHashTree, 119 | { 120 | #[inline] 121 | fn root_hash(&self) -> Hash { 122 | self.as_ref().root_hash() 123 | } 124 | 125 | #[inline] 126 | fn as_hash_tree(&self) -> HashTree<'_> { 127 | self.as_ref().as_hash_tree() 128 | } 129 | } 130 | 131 | impl AsHashTree for Rc 132 | where 133 | T: AsHashTree, 134 | { 135 | #[inline] 136 | fn root_hash(&self) -> Hash { 137 | self.as_ref().root_hash() 138 | } 139 | 140 | #[inline] 141 | fn as_hash_tree(&self) -> HashTree<'_> { 142 | self.as_ref().as_hash_tree() 143 | } 144 | } 145 | 146 | impl AsHashTree for Arc 147 | where 148 | T: AsHashTree, 149 | { 150 | #[inline] 151 | fn root_hash(&self) -> Hash { 152 | self.as_ref().root_hash() 153 | } 154 | 155 | #[inline] 156 | fn as_hash_tree(&self) -> HashTree<'_> { 157 | self.as_ref().as_hash_tree() 158 | } 159 | } 160 | 161 | impl AsHashTree for NonNull 162 | where 163 | T: AsHashTree, 164 | { 165 | #[inline] 166 | fn root_hash(&self) -> Hash { 167 | unsafe { self.as_ref().root_hash() } 168 | } 169 | 170 | #[inline] 171 | fn as_hash_tree(&self) -> HashTree<'_> { 172 | unsafe { self.as_ref().as_hash_tree() } 173 | } 174 | } 175 | 176 | macro_rules! impl_fixed_size { 177 | ( $($size:expr),* ) => { 178 | $( 179 | impl AsHashTree for [u8; $size] { 180 | #[inline] 181 | fn root_hash(&self) -> Hash { 182 | leaf_hash(self) 183 | } 184 | 185 | #[inline] 186 | fn as_hash_tree(&self) -> HashTree<'_> { 187 | HashTree::Leaf(Cow::from(self as &[u8])) 188 | } 189 | } 190 | )* 191 | } 192 | } 193 | 194 | macro_rules! impl_num { 195 | ( $($name:ty),* ) => { 196 | $( 197 | impl AsHashTree for $name { 198 | #[inline] 199 | fn root_hash(&self) -> Hash { 200 | leaf_hash(&self.to_be_bytes()) 201 | } 202 | 203 | #[inline] 204 | fn as_hash_tree(&self) -> HashTree<'_> { 205 | let bytes = self.to_be_bytes(); 206 | HashTree::Leaf(Cow::Owned(bytes.into())) 207 | } 208 | } 209 | )* 210 | } 211 | } 212 | 213 | impl_num!(u8, i8, u16, i16, u32, i32, u64, i64, u128, i128, usize, isize, f32, f64); 214 | impl_fixed_size!( 215 | 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, 216 | 27, 28, 29, 30, 31, 32 217 | ); 218 | -------------------------------------------------------------------------------- /e2e/tests/e2e.rs: -------------------------------------------------------------------------------- 1 | use candid::utils::{decode_args, encode_args, ArgumentDecoder, ArgumentEncoder}; 2 | use ic_kit_e2e_tests::cargo_build_canister; 3 | use ic_state_machine_tests::{CanisterId, ErrorCode, StateMachine, UserError, WasmResult}; 4 | use serde_bytes::ByteBuf; 5 | 6 | #[derive(Debug)] 7 | enum CallError { 8 | Reject(String), 9 | UserError(UserError), 10 | } 11 | 12 | /// A helper function that we use to implement both [`call_candid`] and [`query_candid`]. 13 | fn with_candid( 14 | input: Input, 15 | f: impl FnOnce(Vec) -> Result, 16 | ) -> Result 17 | where 18 | Input: ArgumentEncoder, 19 | Output: for<'a> ArgumentDecoder<'a>, 20 | { 21 | let in_bytes = encode_args(input).expect("failed to encode args"); 22 | match f(in_bytes) { 23 | Ok(WasmResult::Reply(out_bytes)) => Ok(decode_args(&out_bytes).unwrap_or_else(|e| { 24 | panic!( 25 | "Failed to decode bytes {:?} as candid type: {:?}", 26 | std::any::type_name::(), 27 | e 28 | ) 29 | })), 30 | Ok(WasmResult::Reject(message)) => Err(CallError::Reject(message)), 31 | Err(user_error) => Err(CallError::UserError(user_error)), 32 | } 33 | } 34 | 35 | /// Call a canister candid method. 36 | fn call_candid( 37 | env: &StateMachine, 38 | canister_id: CanisterId, 39 | method: &str, 40 | input: Input, 41 | ) -> Result 42 | where 43 | Input: ArgumentEncoder, 44 | Output: for<'a> ArgumentDecoder<'a>, 45 | { 46 | with_candid(input, |bytes| { 47 | env.execute_ingress(canister_id, method, bytes) 48 | }) 49 | } 50 | 51 | /// Query a canister candid method. 52 | fn query_candid( 53 | env: &StateMachine, 54 | canister_id: CanisterId, 55 | method: &str, 56 | input: Input, 57 | ) -> Result 58 | where 59 | Input: ArgumentEncoder, 60 | Output: for<'a> ArgumentDecoder<'a>, 61 | { 62 | with_candid(input, |bytes| env.query(canister_id, method, bytes)) 63 | } 64 | 65 | /// Checks that a canister that uses [`ic_cdk::storage::stable_store`] 66 | /// and [`ic_cdk::storage::stable_restore`] functions can keep its data 67 | /// across upgrades. 68 | #[test] 69 | fn test_storage_roundtrip() { 70 | let env = StateMachine::new(); 71 | let kv_store_wasm = cargo_build_canister("simple-kv-store"); 72 | let canister_id = env 73 | .install_canister(kv_store_wasm.clone(), vec![], None) 74 | .unwrap(); 75 | 76 | let () = call_candid(&env, canister_id, "insert", (&"candid", &b"did")) 77 | .expect("failed to insert 'candid'"); 78 | 79 | env.upgrade_canister(canister_id, kv_store_wasm, vec![]) 80 | .expect("failed to upgrade the simple-kv-store canister"); 81 | 82 | let (result,): (Option,) = 83 | query_candid(&env, canister_id, "lookup", (&"candid",)).expect("failed to lookup 'candid'"); 84 | assert_eq!(result, Some(ByteBuf::from(b"did".to_vec()))); 85 | } 86 | 87 | #[test] 88 | fn test_panic_after_async_frees_resources() { 89 | let env = StateMachine::new(); 90 | let wasm = cargo_build_canister("async"); 91 | let canister_id = env 92 | .install_canister(wasm, vec![], None) 93 | .expect("failed to install a canister"); 94 | 95 | for i in 1..3 { 96 | match call_candid(&env, canister_id, "panic_after_async", ()) { 97 | Ok(()) => (), 98 | Err(CallError::Reject(msg)) => panic!("unexpected reject: {}", msg), 99 | Err(CallError::UserError(e)) => { 100 | println!("Got a user error as expected: {}", e); 101 | 102 | assert_eq!(e.code(), ErrorCode::CanisterCalledTrap); 103 | let expected_message = "Goodbye, cruel world."; 104 | assert!( 105 | e.description().contains(expected_message), 106 | "Expected the user error to contain '{}', got: {}", 107 | expected_message, 108 | e.description() 109 | ); 110 | } 111 | } 112 | 113 | let (n,): (u64,) = call_candid(&env, canister_id, "invocation_count", ()) 114 | .expect("failed to call invocation_count"); 115 | 116 | assert_eq!(i, n, "expected the invocation count to be {}, got {}", i, n); 117 | } 118 | } 119 | 120 | #[test] 121 | fn test_raw_api() { 122 | let env = StateMachine::new(); 123 | let rev = cargo_build_canister("reverse"); 124 | let canister_id = env.install_canister(rev, vec![], None).unwrap(); 125 | 126 | let result = env.query(canister_id, "reverse", vec![1, 2, 3, 4]).unwrap(); 127 | assert_eq!(result, WasmResult::Reply(vec![4, 3, 2, 1])); 128 | 129 | let result = env 130 | .execute_ingress(canister_id, "empty_call", Default::default()) 131 | .unwrap(); 132 | assert_eq!(result, WasmResult::Reply(Default::default())); 133 | } 134 | 135 | #[test] 136 | fn test_notify_calls() { 137 | let env = StateMachine::new(); 138 | let wasm = cargo_build_canister("async"); 139 | let sender_id = env 140 | .install_canister(wasm.clone(), vec![], None) 141 | .expect("failed to install a canister"); 142 | 143 | let receiver_id = env 144 | .install_canister(wasm, vec![], None) 145 | .expect("failed to install a canister"); 146 | 147 | let (n,): (u64,) = query_candid(&env, receiver_id, "notifications_received", ()) 148 | .expect("failed to query 'notifications_received'"); 149 | assert_eq!(n, 0); 150 | 151 | let () = call_candid(&env, sender_id, "notify", (receiver_id, "on_notify")) 152 | .expect("failed to call 'notify'"); 153 | 154 | let (n,): (u64,) = query_candid(&env, receiver_id, "notifications_received", ()) 155 | .expect("failed to query 'notifications_received'"); 156 | assert_eq!(n, 1); 157 | } 158 | -------------------------------------------------------------------------------- /ic-kit-certified/src/rbtree/entry.rs: -------------------------------------------------------------------------------- 1 | use super::{Node, RbTree}; 2 | use crate::label::Label; 3 | use crate::AsHashTree; 4 | use std::fmt::{self, Debug}; 5 | use Entry::{Occupied, Vacant}; 6 | 7 | /// A view into a single entry in a map, which may either be vacant or occupied. 8 | /// 9 | /// This `enum` is constructed from the [`entry`] method on [`RbTree`]. 10 | /// 11 | /// [`entry`]: RbTree::entry 12 | pub enum Entry<'a, K: 'static + Label, V: AsHashTree + 'static> { 13 | Vacant(VacantEntry<'a, K, V>), 14 | Occupied(OccupiedEntry<'a, K, V>), 15 | } 16 | 17 | /// A view into a vacant entry in a [`RbTree`]. It is part of the [`Entry`] enum. 18 | pub struct VacantEntry<'a, K: 'static + Label, V: AsHashTree + 'static> { 19 | pub(super) map: &'a mut RbTree, 20 | pub(super) key: K, 21 | } 22 | 23 | pub struct OccupiedEntry<'a, K: 'static + Label, V: AsHashTree + 'static> { 24 | pub(super) map: &'a mut RbTree, 25 | pub(super) key: K, 26 | pub(super) node: *mut Node, 27 | } 28 | 29 | impl<'a, K: 'static + Label, V: AsHashTree + 'static> VacantEntry<'a, K, V> { 30 | /// Sets the value of the entry with the VacantEntry’s key, and returns a mutable 31 | /// reference to it. 32 | #[inline] 33 | pub fn insert(self, value: V) -> &'a mut V { 34 | self.map.insert(self.key, value).1 35 | } 36 | 37 | /// Take ownership of the key. 38 | #[inline] 39 | pub fn into_key(self) -> K { 40 | self.key 41 | } 42 | 43 | /// Gets a reference to the key that would be used when inserting a value through 44 | /// the [`VacantEntry`]. 45 | #[inline] 46 | pub fn key(&self) -> &K { 47 | &self.key 48 | } 49 | } 50 | 51 | impl<'a, K: 'static + Label, V: AsHashTree + 'static> OccupiedEntry<'a, K, V> { 52 | /// Gets a reference to the value in the entry. 53 | #[inline] 54 | pub fn get(&self) -> &V { 55 | unsafe { &(*self.node).value } 56 | } 57 | 58 | /// Gets a mutable reference to the value in the entry. 59 | /// 60 | /// If you need a reference to the `OccupiedEntry` that may outlive the destruction of 61 | /// the `Entry` value, see [`into_mut`]. 62 | /// 63 | /// [`into_mut`]: OccupiedEntry::into_mut 64 | #[inline] 65 | pub fn get_mut(&mut self) -> &mut V { 66 | unsafe { &mut (*self.node).value } 67 | } 68 | 69 | /// Converts the entry into a mutable reference to its value. 70 | /// 71 | /// If you need multiple references to the OccupiedEntry, see [`get_mut`]. 72 | /// 73 | /// [`get_mut`]: OccupiedEntry::get_mut 74 | #[inline] 75 | pub fn into_mut(self) -> &'a mut V { 76 | unsafe { &mut (*self.node).value } 77 | } 78 | 79 | /// Gets a reference to the key in the entry. 80 | #[inline] 81 | pub fn key(&self) -> &K { 82 | &self.key 83 | } 84 | 85 | /// Takes the value of the entry out of the map, and returns it. 86 | #[inline] 87 | pub fn remove(self) -> V { 88 | self.map.delete(&self.key).unwrap().1 89 | } 90 | 91 | /// Take ownership of the key and value from the map. 92 | #[inline] 93 | pub fn remove_entry(self) -> (K, V) { 94 | self.map.delete(&self.key).unwrap() 95 | } 96 | } 97 | 98 | impl<'a, K: 'static + Label, V: AsHashTree + 'static> Entry<'a, K, V> { 99 | /// Provides in-place mutable access to an occupied entry before any 100 | /// potential inserts into the map. 101 | #[inline] 102 | pub fn and_modify(self, f: F) -> Self 103 | where 104 | F: FnOnce(&mut V), 105 | { 106 | match self { 107 | Occupied(mut entry) => { 108 | f(entry.get_mut()); 109 | Occupied(entry) 110 | } 111 | Vacant(entry) => Vacant(entry), 112 | } 113 | } 114 | 115 | /// Ensures a value is in the entry by inserting the default value if empty, 116 | /// and returns a mutable reference to the value in the entry. 117 | #[inline] 118 | pub fn or_default(self) -> &'a mut V 119 | where 120 | V: Default, 121 | { 122 | match self { 123 | Occupied(entry) => entry.into_mut(), 124 | Vacant(entry) => entry.insert(Default::default()), 125 | } 126 | } 127 | 128 | /// Ensures a value is in the entry by inserting the default if empty, and returns 129 | /// a mutable reference to the value in the entry. 130 | #[inline] 131 | pub fn or_insert(self, default: V) -> &'a mut V { 132 | match self { 133 | Occupied(entry) => entry.into_mut(), 134 | Vacant(entry) => entry.insert(default), 135 | } 136 | } 137 | 138 | /// Ensures a value is in the entry by inserting the result of the default function if empty, 139 | /// and returns a mutable reference to the value in the entry. 140 | #[inline] 141 | pub fn or_insert_with V>(self, default: F) -> &'a mut V { 142 | match self { 143 | Occupied(entry) => entry.into_mut(), 144 | Vacant(entry) => entry.insert(default()), 145 | } 146 | } 147 | 148 | /// Ensures a value is in the entry by inserting, if empty, the result of the default function. 149 | /// This method allows for generating key-derived values for insertion by providing the default 150 | /// function a reference to the key that was moved during the `.entry(key)` method call. 151 | /// 152 | /// The reference to the moved key is provided so that cloning or copying the key is 153 | /// unnecessary, unlike with `.or_insert_with(|| ... )`. 154 | #[inline] 155 | pub fn or_insert_with_key V>(self, default: F) -> &'a mut V { 156 | match self { 157 | Occupied(entry) => entry.into_mut(), 158 | Vacant(entry) => { 159 | let value = default(entry.key()); 160 | entry.insert(value) 161 | } 162 | } 163 | } 164 | 165 | /// Returns a reference to this entry’s key. 166 | #[inline] 167 | pub fn key(&self) -> &K { 168 | match self { 169 | Occupied(entry) => entry.key(), 170 | Vacant(entry) => entry.key(), 171 | } 172 | } 173 | } 174 | 175 | impl<'a, K: 'static + Label, V: AsHashTree + 'static> Debug for Entry<'a, K, V> 176 | where 177 | K: Debug, 178 | V: Debug, 179 | { 180 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { 181 | match *self { 182 | Vacant(ref v) => f.debug_tuple("Entry").field(v).finish(), 183 | Occupied(ref o) => f.debug_tuple("Entry").field(o).finish(), 184 | } 185 | } 186 | } 187 | 188 | impl<'a, K: 'static + Label, V: AsHashTree + 'static> Debug for VacantEntry<'a, K, V> 189 | where 190 | K: Debug, 191 | V: Debug, 192 | { 193 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { 194 | f.debug_tuple("VacantEntry").field(self.key()).finish() 195 | } 196 | } 197 | 198 | impl<'a, K: 'static + Label, V: AsHashTree + 'static> Debug for OccupiedEntry<'a, K, V> 199 | where 200 | K: Debug, 201 | V: Debug, 202 | { 203 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { 204 | f.debug_struct("OccupiedEntry") 205 | .field("key", self.key()) 206 | .field("value", self.get()) 207 | .finish() 208 | } 209 | } 210 | -------------------------------------------------------------------------------- /ic-kit-certified/src/collections/paged.rs: -------------------------------------------------------------------------------- 1 | use crate::label::{Label, Prefix}; 2 | use crate::{AsHashTree, Hash, HashTree, Map, Seq}; 3 | use candid::CandidType; 4 | use serde::{Deserialize, Serialize}; 5 | use std::borrow::{Borrow, Cow}; 6 | 7 | #[derive(CandidType, Serialize, Deserialize, Debug)] 8 | pub struct Paged { 9 | data: Map, Seq>, 10 | } 11 | 12 | #[derive(Ord, CandidType, Serialize, Deserialize, PartialOrd, Eq, PartialEq, Debug)] 13 | struct PagedKey { 14 | key: K, 15 | page: u32, 16 | } 17 | 18 | impl Label for PagedKey { 19 | #[inline] 20 | fn as_label(&self) -> Cow<[u8]> { 21 | let mut data = self.key.as_label().to_vec(); 22 | data.extend_from_slice(&self.page.to_be_bytes()); 23 | Cow::Owned(data) 24 | } 25 | } 26 | 27 | impl Borrow for PagedKey { 28 | #[inline] 29 | fn borrow(&self) -> &K { 30 | &self.key 31 | } 32 | } 33 | 34 | impl Prefix for PagedKey {} 35 | 36 | impl Default for Paged { 37 | fn default() -> Self { 38 | Self::new() 39 | } 40 | } 41 | 42 | impl Paged { 43 | pub fn new() -> Self { 44 | Self { data: Map::new() } 45 | } 46 | 47 | pub fn insert(&mut self, key: K, item: V) { 48 | let tree = &mut self.data.inner; 49 | let mut item = Some(item); 50 | 51 | let page = tree 52 | .modify_max_with_prefix(&key, |key, seq| { 53 | if seq.len() == S { 54 | return Some(key.page + 1); 55 | } 56 | seq.append(item.take().unwrap()); 57 | None 58 | }) 59 | .unwrap_or(Some(0)); 60 | 61 | if let Some(page) = page { 62 | let key = PagedKey { key, page }; 63 | let mut value = Seq::new(); 64 | value.append(item.take().unwrap()); 65 | tree.insert(key, value); 66 | } 67 | } 68 | 69 | pub fn get_last_page_number(&self, key: &K) -> Option { 70 | self.data 71 | .inner 72 | .max_entry_with_prefix(key) 73 | .map(|(k, _)| k.page as usize) 74 | } 75 | 76 | // TODO(qti3e) Remove the Clone. 77 | pub fn witness_last_page_number(&self, key: &K) -> HashTree<'_> 78 | where 79 | K: Clone, 80 | { 81 | let page = self 82 | .data 83 | .inner 84 | .max_entry_with_prefix(key) 85 | .map(|(k, _)| k.page + 1) 86 | .unwrap_or(0); 87 | let key = PagedKey { 88 | key: key.clone(), 89 | page, 90 | }; 91 | self.data.witness(&key) 92 | } 93 | 94 | pub fn get(&self, key: &K, page: usize) -> Option<&Seq> { 95 | let page = page as u32; 96 | let key = (key, page); 97 | self.data.inner.get_with(|k| key.cmp(&(&k.key, k.page))) 98 | } 99 | 100 | // TODO(qti3e) Remove the Clone in future. 101 | pub fn witness(&self, key: &K, page: usize) -> HashTree<'_> 102 | where 103 | K: Clone, 104 | { 105 | let key = PagedKey { 106 | key: key.clone(), 107 | page: page as u32, 108 | }; 109 | self.data.witness(&key) 110 | } 111 | } 112 | 113 | impl AsHashTree 114 | for Paged 115 | { 116 | fn root_hash(&self) -> Hash { 117 | self.data.root_hash() 118 | } 119 | 120 | fn as_hash_tree(&self) -> HashTree<'_> { 121 | self.data.as_hash_tree() 122 | } 123 | } 124 | 125 | #[cfg(test)] 126 | mod tests { 127 | use super::*; 128 | 129 | #[test] 130 | fn modify_max_with_prefix() { 131 | let mut paged = Paged::::new(); 132 | paged.data.append_deep(PagedKey { key: 1, page: 0 }, 0); 133 | paged.data.append_deep(PagedKey { key: 1, page: 0 }, 1); 134 | paged.data.append_deep(PagedKey { key: 1, page: 0 }, 2); 135 | paged.data.append_deep(PagedKey { key: 1, page: 1 }, 3); 136 | paged.data.append_deep(PagedKey { key: 1, page: 1 }, 4); 137 | paged.data.append_deep(PagedKey { key: 1, page: 1 }, 5); 138 | paged.data.append_deep(PagedKey { key: 1, page: 2 }, 18); 139 | 140 | paged.data.append_deep(PagedKey { key: 3, page: 0 }, 6); 141 | paged.data.append_deep(PagedKey { key: 3, page: 0 }, 7); 142 | paged.data.append_deep(PagedKey { key: 3, page: 0 }, 8); 143 | paged.data.append_deep(PagedKey { key: 3, page: 1 }, 9); 144 | paged.data.append_deep(PagedKey { key: 3, page: 1 }, 10); 145 | paged.data.append_deep(PagedKey { key: 3, page: 1 }, 11); 146 | 147 | paged.data.append_deep(PagedKey { key: 5, page: 0 }, 12); 148 | paged.data.append_deep(PagedKey { key: 5, page: 0 }, 13); 149 | paged.data.append_deep(PagedKey { key: 5, page: 0 }, 14); 150 | paged.data.append_deep(PagedKey { key: 5, page: 1 }, 15); 151 | paged.data.append_deep(PagedKey { key: 5, page: 1 }, 16); 152 | paged.data.append_deep(PagedKey { key: 5, page: 1 }, 17); 153 | 154 | assert_eq!(paged.data.inner.modify_max_with_prefix(&0, |k, _| k), None); 155 | 156 | assert_eq!( 157 | paged.data.inner.modify_max_with_prefix(&1, |k, _| k), 158 | Some(&PagedKey { key: 1, page: 2 }) 159 | ); 160 | 161 | assert_eq!(paged.data.inner.modify_max_with_prefix(&2, |k, _| k), None); 162 | 163 | assert_eq!( 164 | paged.data.inner.modify_max_with_prefix(&3, |k, _| k), 165 | Some(&PagedKey { key: 3, page: 1 }) 166 | ); 167 | 168 | assert_eq!(paged.data.inner.modify_max_with_prefix(&4, |k, _| k), None); 169 | 170 | assert_eq!( 171 | paged.data.inner.modify_max_with_prefix(&5, |k, _| k), 172 | Some(&PagedKey { key: 5, page: 1 }) 173 | ); 174 | 175 | assert_eq!(paged.data.inner.modify_max_with_prefix(&6, |k, _| k), None); 176 | } 177 | 178 | #[test] 179 | fn get() { 180 | let mut paged = Paged::::new(); 181 | 182 | // For each key from 0 to 4, insert 10 items, which is 3 full page + one last page 183 | // with one item. 184 | // 0: [0 5 10] [15 20 25] [30 35 40] [45] -> 15p + 5i + 0 185 | // 1: [1 6 11] [16 21 26] [31 36 41] [46] -> 15p + 5i + 1 186 | // 2: [2 7 12] [17 22 27] [32 37 42] [47] -> 15p + 5i + 2 187 | // 3: [3 8 13] [18 23 28] [33 38 43] [48] -> 15p + 5i + 3 188 | // 4: [4 9 14] [19 24 29] [34 39 44] [49] -> 15p + 5i + 4 189 | for i in 0..50 { 190 | paged.insert(i % 5, i); 191 | } 192 | 193 | for k in 0..5 { 194 | for p in 0..3 { 195 | let seq = (0..3).map(|i| 15 * p + 5 * i + k).collect::>(); 196 | assert_eq!(paged.get(&k, p as usize), Some(&seq)); 197 | } 198 | } 199 | 200 | for k in 0..5 { 201 | let seq = vec![45 + k].into_iter().collect::>(); 202 | assert_eq!(paged.get(&k, 3), Some(&seq)); 203 | assert_eq!(paged.get(&k, 4), None); 204 | } 205 | } 206 | } 207 | -------------------------------------------------------------------------------- /ic-kit-certified/src/hashtree.rs: -------------------------------------------------------------------------------- 1 | // This file is copied from ic-certified-map which was released under Apache V2. 2 | // Some modifications are made to improve the code quality. 3 | use serde::{ser::SerializeSeq, Serialize, Serializer}; 4 | use serde_bytes::Bytes; 5 | use sha2::{Digest, Sha256}; 6 | use std::borrow::Cow; 7 | 8 | /// SHA-256 hash bytes. 9 | pub type Hash = [u8; 32]; 10 | 11 | #[derive(Debug, Eq, PartialEq)] 12 | pub struct ForkInner<'a>(pub HashTree<'a>, pub HashTree<'a>); 13 | 14 | impl<'a> ForkInner<'a> { 15 | pub fn left(&self) -> &HashTree<'a> { 16 | &self.0 17 | } 18 | 19 | pub fn right(&self) -> &HashTree<'a> { 20 | &self.1 21 | } 22 | } 23 | 24 | /// HashTree as defined in the interfaces spec. 25 | /// https://sdk.dfinity.org/docs/interface-spec/index.html#_certificate 26 | #[derive(Debug, Eq, PartialEq)] 27 | pub enum HashTree<'a> { 28 | Empty, 29 | Fork(Box>), 30 | Labeled(Cow<'a, [u8]>, Box>), 31 | Leaf(Cow<'a, [u8]>), 32 | Pruned(Hash), 33 | } 34 | 35 | pub fn fork<'a>(l: HashTree<'a>, r: HashTree<'a>) -> HashTree<'a> { 36 | HashTree::Fork(Box::new(ForkInner(l, r))) 37 | } 38 | 39 | pub fn labeled<'a>(l: &'a [u8], t: HashTree<'a>) -> HashTree<'a> { 40 | HashTree::Labeled(Cow::Borrowed(l), Box::new(t)) 41 | } 42 | 43 | pub fn fork_hash(l: &Hash, r: &Hash) -> Hash { 44 | let mut h = domain_sep("ic-hashtree-fork"); 45 | h.update(&l[..]); 46 | h.update(&r[..]); 47 | h.finalize().into() 48 | } 49 | 50 | pub fn leaf_hash(data: &[u8]) -> Hash { 51 | let mut h = domain_sep("ic-hashtree-leaf"); 52 | h.update(data); 53 | h.finalize().into() 54 | } 55 | 56 | pub fn labeled_hash(label: &[u8], content_hash: &Hash) -> Hash { 57 | let mut h = domain_sep("ic-hashtree-labeled"); 58 | h.update(label); 59 | h.update(&content_hash[..]); 60 | h.finalize().into() 61 | } 62 | 63 | impl<'a> HashTree<'a> { 64 | pub fn reconstruct(&self) -> Hash { 65 | match self { 66 | Self::Empty => domain_sep("ic-hashtree-empty").finalize().into(), 67 | Self::Fork(f) => fork_hash(&f.0.reconstruct(), &f.1.reconstruct()), 68 | Self::Labeled(l, t) => { 69 | let thash = t.reconstruct(); 70 | labeled_hash(l, &thash) 71 | } 72 | Self::Leaf(data) => leaf_hash(data), 73 | Self::Pruned(h) => *h, 74 | } 75 | } 76 | 77 | /// Collect and return all of the labels in this HashTree. 78 | /// 79 | /// This method is intended for testing purposes. 80 | pub fn get_labels<'b: 'a>(&'b self) -> Vec<&'b [u8]> { 81 | fn go<'a>(keys: &mut Vec<&'a [u8]>, tree: &'a HashTree<'a>) { 82 | match tree { 83 | HashTree::Labeled(key, value) => { 84 | keys.push(key); 85 | go(keys, value); 86 | } 87 | HashTree::Fork(lr) => { 88 | go(keys, lr.left()); 89 | go(keys, lr.right()); 90 | } 91 | _ => (), 92 | } 93 | } 94 | 95 | let mut keys = Vec::new(); 96 | go(&mut keys, self); 97 | keys 98 | } 99 | 100 | /// Collect and return all of the values in this HashTree. 101 | /// 102 | /// This method is intended for testing purposes. 103 | pub fn get_leaf_values<'b: 'a>(&'b self) -> Vec<&'b [u8]> { 104 | fn go<'a>(values: &mut Vec<&'a [u8]>, tree: &'a HashTree<'a>) { 105 | match tree { 106 | HashTree::Leaf(value) => { 107 | values.push(value); 108 | } 109 | HashTree::Fork(lr) => { 110 | go(values, lr.left()); 111 | go(values, lr.right()); 112 | } 113 | HashTree::Labeled(_, t) => { 114 | go(values, &*t); 115 | } 116 | _ => (), 117 | } 118 | } 119 | 120 | let mut values = Vec::new(); 121 | go(&mut values, self); 122 | values 123 | } 124 | } 125 | 126 | impl Serialize for HashTree<'_> { 127 | fn serialize(&self, serializer: S) -> Result<::Ok, ::Error> 128 | where 129 | S: Serializer, 130 | { 131 | match self { 132 | HashTree::Empty => { 133 | let mut seq = serializer.serialize_seq(Some(1))?; 134 | seq.serialize_element(&0u8)?; 135 | seq.end() 136 | } 137 | HashTree::Fork(p) => { 138 | let mut seq = serializer.serialize_seq(Some(3))?; 139 | seq.serialize_element(&1u8)?; 140 | seq.serialize_element(&p.0)?; 141 | seq.serialize_element(&p.1)?; 142 | seq.end() 143 | } 144 | HashTree::Labeled(label, tree) => { 145 | let mut seq = serializer.serialize_seq(Some(3))?; 146 | seq.serialize_element(&2u8)?; 147 | seq.serialize_element(Bytes::new(label))?; 148 | seq.serialize_element(&tree)?; 149 | seq.end() 150 | } 151 | HashTree::Leaf(leaf_bytes) => { 152 | let mut seq = serializer.serialize_seq(Some(2))?; 153 | seq.serialize_element(&3u8)?; 154 | seq.serialize_element(Bytes::new(leaf_bytes))?; 155 | seq.end() 156 | } 157 | HashTree::Pruned(digest) => { 158 | let mut seq = serializer.serialize_seq(Some(2))?; 159 | seq.serialize_element(&4u8)?; 160 | seq.serialize_element(Bytes::new(&digest[..]))?; 161 | seq.end() 162 | } 163 | } 164 | } 165 | } 166 | 167 | fn domain_sep(s: &str) -> sha2::Sha256 { 168 | let buf: [u8; 1] = [s.len() as u8]; 169 | let mut h = Sha256::new(); 170 | h.update(&buf[..]); 171 | h.update(s.as_bytes()); 172 | h 173 | } 174 | 175 | #[cfg(test)] 176 | mod tests { 177 | use super::{ 178 | fork, labeled, 179 | HashTree::{Empty, Leaf}, 180 | }; 181 | use std::borrow::Cow; 182 | 183 | //─┬─┬╴"a" ─┬─┬╴"x" ─╴"hello" 184 | // │ │ │ └╴Empty 185 | // │ │ └╴ "y" ─╴"world" 186 | // │ └╴"b" ──╴"good" 187 | // └─┬╴"c" ──╴Empty 188 | // └╴"d" ──╴"morning" 189 | #[test] 190 | fn test_public_spec_example() { 191 | let t = fork( 192 | fork( 193 | labeled( 194 | b"a", 195 | fork( 196 | fork(labeled(b"x", Leaf(Cow::Borrowed(b"hello"))), Empty), 197 | labeled(b"y", Leaf(Cow::Borrowed(b"world"))), 198 | ), 199 | ), 200 | labeled(b"b", Leaf(Cow::Borrowed(b"good"))), 201 | ), 202 | fork( 203 | labeled(b"c", Empty), 204 | labeled(b"d", Leaf(Cow::Borrowed(b"morning"))), 205 | ), 206 | ); 207 | 208 | assert_eq!( 209 | hex::encode(&t.reconstruct()[..]), 210 | "eb5c5b2195e62d996b84c9bcc8259d19a83786a2f59e0878cec84c811f669aa0".to_string() 211 | ); 212 | 213 | assert_eq!( 214 | hex::encode(serde_cbor::to_vec(&t).unwrap()), 215 | "8301830183024161830183018302417882034568656c6c6f810083024179820345776f726c6483024162820344676f6f648301830241638100830241648203476d6f726e696e67".to_string()); 216 | } 217 | } 218 | -------------------------------------------------------------------------------- /ic-kit/src/storage.rs: -------------------------------------------------------------------------------- 1 | #![allow(non_snake_case)] 2 | use std::any::{Any, TypeId}; 3 | use std::borrow::{Borrow, BorrowMut}; 4 | use std::cell::RefCell; 5 | use std::collections::hash_map::Entry; 6 | use std::collections::HashMap; 7 | use std::ops::DerefMut; 8 | 9 | type StorageMap = HashMap>>; 10 | 11 | /// An storage implementation for singleton design pattern, where we only have one value 12 | /// associated with each types. 13 | #[derive(Default)] 14 | pub struct Storage { 15 | storage: RefCell, 16 | } 17 | 18 | impl Storage { 19 | /// Ensure the default value exists on the map. 20 | #[inline(always)] 21 | fn ensure_default(&self, tid: TypeId) { 22 | self.storage 23 | .borrow_mut() 24 | .entry(tid) 25 | .or_insert_with(|| RefCell::new(Box::new(T::default()))); 26 | } 27 | 28 | /// Pass an immutable reference to the stored data of the type `T` to the closure, 29 | /// if there is no data associated with the type, store the `Default` and then perform the 30 | /// operation. 31 | #[inline] 32 | pub fn with U>(&self, callback: F) -> U { 33 | let tid = TypeId::of::(); 34 | self.ensure_default::(tid); 35 | let cell = unsafe { self.storage.try_borrow_unguarded() } 36 | .unwrap() 37 | .get(&tid) 38 | .unwrap() 39 | .borrow(); 40 | let borrow = cell.downcast_ref::().unwrap(); 41 | callback(borrow) 42 | } 43 | 44 | /// Pass an immutable reference to the stored data of the type `T` to the closure, 45 | /// if there is no data associated with the type, just return None. 46 | #[inline] 47 | pub fn maybe_with U>(&self, callback: F) -> Option { 48 | let tid = TypeId::of::(); 49 | unsafe { self.storage.try_borrow_unguarded() } 50 | .unwrap() 51 | .get(&tid) 52 | .map(|c| c.borrow()) 53 | .map(|c| callback(c.borrow().downcast_ref::().unwrap())) 54 | } 55 | 56 | /// Like [`Self::with`] but passes a mutable reference. 57 | #[inline] 58 | pub fn with_mut U>(&self, callback: F) -> U { 59 | let tid = TypeId::of::(); 60 | self.ensure_default::(tid); 61 | let mut cell = unsafe { self.storage.try_borrow_unguarded() } 62 | .unwrap() 63 | .get(&tid) 64 | .unwrap() 65 | .borrow_mut(); 66 | let borrow = cell.downcast_mut::().unwrap(); 67 | callback(borrow) 68 | } 69 | 70 | /// Like [`Self::maybe_with`] but passes a mutable reference. 71 | #[inline] 72 | pub fn maybe_with_mut U>(&self, callback: F) -> Option { 73 | let tid = TypeId::of::(); 74 | unsafe { self.storage.try_borrow_unguarded() } 75 | .unwrap() 76 | .get(&tid) 77 | .map(|c| c.borrow_mut()) 78 | .map(|mut c| callback(c.borrow_mut().downcast_mut::().unwrap())) 79 | } 80 | 81 | /// Remove the data associated with the type `T`, and returns it if any. 82 | #[inline] 83 | pub fn take(&self) -> Option { 84 | let tid = TypeId::of::(); 85 | self.storage 86 | .borrow_mut() 87 | .remove(&tid) 88 | .map(|cell| *cell.into_inner().downcast::().unwrap()) 89 | } 90 | 91 | /// Store the given value for type `T`, returns the previously stored value if any. 92 | #[inline] 93 | pub fn swap(&self, value: T) -> Option { 94 | let tid = TypeId::of::(); 95 | match self.storage.borrow_mut().entry(tid) { 96 | Entry::Occupied(mut o) => Some( 97 | *o.get_mut() 98 | .replace(Box::new(value)) 99 | .downcast::() 100 | .unwrap(), 101 | ), 102 | Entry::Vacant(v) => { 103 | v.insert(RefCell::new(Box::new(value))); 104 | None 105 | } 106 | } 107 | } 108 | 109 | /// Just like `.with` but can pass the immutable reference to many items in one closure. 110 | #[inline] 111 | pub fn with_many U>(&self, callback: F) -> U { 112 | { 113 | A::ensure_default(self.storage.borrow_mut().deref_mut()) 114 | } 115 | 116 | let storage = unsafe { self.storage.try_borrow_unguarded() }.unwrap(); 117 | A::with(storage, callback) 118 | } 119 | 120 | /// Just like `.with_mut` but can pass the mutable reference to many items in one closure. 121 | #[inline] 122 | pub fn with_many_mut U>(&self, callback: F) -> U { 123 | { 124 | A::ensure_default(self.storage.borrow_mut().deref_mut()) 125 | } 126 | 127 | let storage = unsafe { self.storage.try_borrow_unguarded() }.unwrap(); 128 | A::with_mut(storage, callback) 129 | } 130 | } 131 | 132 | pub trait BorrowMany: Sized { 133 | fn ensure_default(storage: &mut StorageMap); 134 | 135 | fn with U>(storage: &StorageMap, callback: F) -> U; 136 | } 137 | 138 | pub trait BorrowMutMany: Sized { 139 | fn ensure_default(storage: &mut StorageMap); 140 | 141 | fn with_mut U>(storage: &StorageMap, callback: F) -> U; 142 | } 143 | 144 | macro_rules! implement_borrow_many { 145 | ($(($($name: ident)+))+) => { 146 | $( 147 | impl<'a, $($name: 'static,)+> BorrowMany for ($(&'a $name,)+) 148 | where 149 | $($name: Default,)+ 150 | { 151 | #[inline(always)] 152 | fn ensure_default(storage: &mut StorageMap) { 153 | $( 154 | storage 155 | .entry(TypeId::of::<$name>()) 156 | .or_insert_with(|| RefCell::new(Box::new($name::default()))); 157 | )+ 158 | } 159 | 160 | #[inline(always)] 161 | fn with U>(storage: &StorageMap, callback: F) -> U { 162 | $( 163 | let $name = storage.get(&TypeId::of::<$name>()).unwrap().borrow(); 164 | )+ 165 | 166 | callback(( 167 | $( 168 | unsafe { &*(&*$name as *const Box) } 169 | .downcast_ref::<$name>() 170 | .unwrap(), 171 | )+ 172 | )) 173 | } 174 | } 175 | 176 | impl<'a, $($name: 'static,)+> BorrowMutMany for ($(&'a mut $name,)+) 177 | where 178 | $($name: Default,)+ 179 | { 180 | #[inline(always)] 181 | fn ensure_default(storage: &mut StorageMap) { 182 | $( 183 | storage 184 | .entry(TypeId::of::<$name>()) 185 | .or_insert_with(|| RefCell::new(Box::new($name::default()))); 186 | )+ 187 | } 188 | 189 | #[inline(always)] 190 | fn with_mut U>(storage: &StorageMap, callback: F) -> U { 191 | $( 192 | let mut $name = storage.get(&TypeId::of::<$name>()).unwrap().borrow_mut(); 193 | )+ 194 | 195 | callback(( 196 | $( 197 | unsafe { &mut *(&mut *$name as *mut Box) } 198 | .downcast_mut::<$name>() 199 | .unwrap(), 200 | )+ 201 | )) 202 | } 203 | } 204 | )+ 205 | }; 206 | } 207 | 208 | implement_borrow_many!( 209 | (A0) 210 | (A0 A1) 211 | (A0 A1 A3) 212 | (A0 A1 A3 A4) 213 | (A0 A1 A3 A4 A5) 214 | (A0 A1 A3 A4 A5 A6) 215 | (A0 A1 A3 A4 A5 A6 A7) 216 | (A0 A1 A3 A4 A5 A6 A7 A8) 217 | (A0 A1 A3 A4 A5 A6 A7 A8 A9) 218 | ); 219 | -------------------------------------------------------------------------------- /ic-kit-runtime/src/call.rs: -------------------------------------------------------------------------------- 1 | use candid::utils::{ArgumentDecoder, ArgumentEncoder}; 2 | use candid::{decode_args, decode_one, encode_args, encode_one, CandidType, Principal}; 3 | use serde::de::DeserializeOwned; 4 | 5 | use ic_kit_sys::types::{CallError, RejectionCode, CANDID_EMPTY_ARG}; 6 | 7 | use crate::types::*; 8 | use crate::Replica; 9 | 10 | /// A CallBuilder for a replica. 11 | #[derive(Clone)] 12 | pub struct CallBuilder<'a> { 13 | replica: &'a Replica, 14 | canister_id: Principal, 15 | method_name: String, 16 | sender: Principal, 17 | payment: u128, 18 | arg: Option>, 19 | } 20 | 21 | /// A reply by the canister. 22 | #[derive(Debug)] 23 | pub enum CallReply { 24 | Reply { 25 | data: Vec, 26 | cycles_refunded: u128, 27 | }, 28 | Reject { 29 | rejection_code: RejectionCode, 30 | rejection_message: String, 31 | cycles_refunded: u128, 32 | }, 33 | } 34 | 35 | impl<'a> CallBuilder<'a> { 36 | /// Create a new call builder for the given type. 37 | pub fn new(replica: &'a Replica, canister_id: Principal, method_name: String) -> Self { 38 | Self { 39 | replica, 40 | canister_id, 41 | sender: Principal::anonymous(), 42 | method_name, 43 | payment: 0, 44 | arg: None, 45 | } 46 | } 47 | 48 | /// Use the given candid tuple value as the argument for this mock call. 49 | /// 50 | /// # Panics 51 | /// 52 | /// This method panics if the argument for this call is already set via a prior 53 | /// call to any of the `with_args`, `with_arg` or `with_arg_raw`. 54 | pub fn with_args(mut self, arguments: T) -> Self { 55 | assert!(self.arg.is_none(), "Arguments may only be set once."); 56 | self.arg = Some(encode_args(arguments).unwrap()); 57 | self 58 | } 59 | 60 | /// Shorthand for `with_args((argument, ))` to pass tuples with only one element to the call. 61 | /// 62 | /// # Panics 63 | /// 64 | /// This method panics if the argument for this call is already set via a prior 65 | /// call to any of the `with_args`, `with_arg` or `with_arg_raw`. 66 | pub fn with_arg(mut self, argument: T) -> Self { 67 | assert!(self.arg.is_none(), "Arguments may only be set once."); 68 | self.arg = Some(encode_one(argument).unwrap()); 69 | self 70 | } 71 | 72 | /// Pass the given raw buffer as the call argument, this does not perform any serialization on 73 | /// the data. 74 | /// 75 | /// # Panics 76 | /// 77 | /// This method panics if the argument for this call is already set via a prior 78 | /// call to any of the `with_args`, `with_arg` or `with_arg_raw`. 79 | pub fn with_arg_raw>>(mut self, argument: A) -> Self { 80 | assert!(self.arg.is_none(), "Arguments may only be set once."); 81 | self.arg = Some(argument.into()); 82 | self 83 | } 84 | 85 | /// Use the given amount of cycles for this mock call. 86 | pub fn with_payment(mut self, cycles: u128) -> Self { 87 | self.payment = cycles; 88 | self 89 | } 90 | 91 | /// Make the call from this sender. 92 | pub fn with_caller>(mut self, caller: I) -> Self { 93 | self.sender = caller.into(); 94 | self 95 | } 96 | 97 | /// Perform the call and returns the reply from the canister. 98 | pub async fn perform(&self) -> CallReply { 99 | self.replica.perform_call(self.into()).await 100 | } 101 | } 102 | 103 | impl CallReply { 104 | /// Convert the reply to a message that can be delivered to a canister. 105 | pub(crate) fn to_message(self, reply_to: OutgoingRequestId) -> Message { 106 | match self { 107 | CallReply::Reply { 108 | data, 109 | cycles_refunded, 110 | } => Message::Reply { 111 | reply_to, 112 | env: Env::default() 113 | .with_entry_mode(EntryMode::ReplyCallback) 114 | .with_raw_args(data) 115 | .with_cycles_refunded(cycles_refunded), 116 | }, 117 | CallReply::Reject { 118 | rejection_code, 119 | rejection_message, 120 | cycles_refunded, 121 | } => Message::Reply { 122 | reply_to, 123 | env: Env::default() 124 | .with_entry_mode(EntryMode::RejectCallback) 125 | .with_cycles_refunded(cycles_refunded) 126 | .with_rejection_code(rejection_code) 127 | .with_rejection_message(rejection_message), 128 | }, 129 | } 130 | } 131 | 132 | /// Return the raw response bytes from this call. 133 | pub fn bytes(&self) -> Result<&[u8], CallError> { 134 | self.into() 135 | } 136 | 137 | /// Try to decode the response to the provided candid tuple. 138 | pub fn decode ArgumentDecoder<'a>>(&self) -> Result { 139 | let bytes = self.bytes()?; 140 | match decode_args(bytes) { 141 | Err(_) => Err(CallError::ResponseDeserializationError(bytes.to_vec())), 142 | Ok(r) => Ok(r), 143 | } 144 | } 145 | 146 | /// Tries to decode a single argument. 147 | pub fn decode_one(&self) -> Result 148 | where 149 | T: DeserializeOwned + CandidType, 150 | { 151 | let bytes = self.bytes()?; 152 | match decode_one(bytes) { 153 | Err(_) => Err(CallError::ResponseDeserializationError(bytes.to_vec())), 154 | Ok(r) => Ok(r), 155 | } 156 | } 157 | 158 | /// Return the rejection code from this call, returns `RejectionCode::NoError` when the call 159 | /// succeed. 160 | pub fn rejection_code(&self) -> RejectionCode { 161 | match &self { 162 | CallReply::Reply { .. } => RejectionCode::NoError, 163 | CallReply::Reject { rejection_code, .. } => *rejection_code, 164 | } 165 | } 166 | 167 | /// Returns the possible rejection message. 168 | pub fn rejection_message(&self) -> Option<&str> { 169 | match &self { 170 | CallReply::Reply { .. } => None, 171 | CallReply::Reject { 172 | rejection_message, .. 173 | } => Some(rejection_message.as_str()), 174 | } 175 | } 176 | 177 | //// Returns the number of cycles refunded from this canister. 178 | pub fn cycles_refunded(&self) -> u128 { 179 | match &self { 180 | CallReply::Reply { 181 | cycles_refunded, .. 182 | } => *cycles_refunded, 183 | CallReply::Reject { 184 | cycles_refunded, .. 185 | } => *cycles_refunded, 186 | } 187 | } 188 | 189 | /// Returns true if the call was okay. 190 | pub fn is_ok(&self) -> bool { 191 | match &self { 192 | CallReply::Reply { .. } => true, 193 | CallReply::Reject { .. } => false, 194 | } 195 | } 196 | 197 | /// Returns true if the call was rejected. 198 | pub fn is_error(&self) -> bool { 199 | match &self { 200 | CallReply::Reply { .. } => false, 201 | CallReply::Reject { .. } => true, 202 | } 203 | } 204 | 205 | /// Assert the response is okay. 206 | pub fn assert_ok(&self) { 207 | assert!(self.is_ok(), "The call was rejected."); 208 | } 209 | 210 | /// Assert the response is a rejection. 211 | pub fn assert_error(&self) { 212 | assert!(self.is_error(), "Expected a rejection, but got a reply."); 213 | } 214 | } 215 | 216 | impl<'a> From<&'a CallReply> for Result<&'a [u8], CallError> { 217 | fn from(reply: &'a CallReply) -> Self { 218 | match reply { 219 | CallReply::Reply { data, .. } => Ok(data.as_slice()), 220 | CallReply::Reject { 221 | rejection_code, 222 | rejection_message, 223 | .. 224 | } => Err(CallError::Rejected( 225 | *rejection_code, 226 | rejection_message.clone(), 227 | )), 228 | } 229 | } 230 | } 231 | 232 | impl<'a> From<&'a CallBuilder<'a>> for CanisterCall { 233 | fn from(builder: &'a CallBuilder) -> Self { 234 | CanisterCall { 235 | sender: builder.sender, 236 | request_id: RequestId::new(), 237 | callee: builder.canister_id, 238 | method: builder.method_name.clone(), 239 | payment: builder.payment, 240 | arg: builder 241 | .arg 242 | .clone() 243 | .unwrap_or_else(|| CANDID_EMPTY_ARG.to_vec()), 244 | } 245 | } 246 | } 247 | -------------------------------------------------------------------------------- /ic-kit-macros/src/export_service.rs: -------------------------------------------------------------------------------- 1 | use std::collections::BTreeMap; 2 | use std::sync::Mutex; 3 | 4 | use lazy_static::lazy_static; 5 | use proc_macro2::{Ident, Span, TokenStream}; 6 | use quote::{quote, ToTokens}; 7 | use syn::{DeriveInput, Error}; 8 | 9 | use crate::metadata::generate_metadata; 10 | use crate::EntryPoint; 11 | 12 | struct Method { 13 | hidden: bool, 14 | mode: EntryPoint, 15 | rust_name: String, 16 | _arg_names: Vec, 17 | arg_types: Vec, 18 | rets: Vec, 19 | } 20 | 21 | lazy_static! { 22 | static ref METHODS: Mutex> = Mutex::new(Default::default()); 23 | static ref LIFE_CYCLES: Mutex> = Mutex::new(Default::default()); 24 | } 25 | 26 | pub(crate) fn declare( 27 | entry_point: EntryPoint, 28 | rust_name: Ident, 29 | name: String, 30 | hidden: bool, 31 | can_args: Vec, 32 | can_types: Vec, 33 | rt: &syn::ReturnType, 34 | ) -> Result<(), Error> { 35 | let rets = match rt { 36 | syn::ReturnType::Default => Vec::new(), 37 | syn::ReturnType::Type(_, ty) => match ty.as_ref() { 38 | syn::Type::Tuple(tuple) => tuple 39 | .elems 40 | .iter() 41 | .cloned() 42 | .map(remove_reference_recursive) 43 | .collect(), 44 | _ => vec![remove_reference_recursive(ty.as_ref().clone())], 45 | }, 46 | }; 47 | 48 | let method = Method { 49 | hidden, 50 | mode: entry_point, 51 | rust_name: rust_name.to_string(), 52 | _arg_names: can_args.iter().map(|i| i.to_string()).collect(), 53 | arg_types: can_types 54 | .iter() 55 | .map(|t| format!("{}", t.to_token_stream())) 56 | .collect(), 57 | rets: rets 58 | .into_iter() 59 | .map(|c| format!("{}", c.to_token_stream())) 60 | .collect(), 61 | }; 62 | 63 | if entry_point.is_lifecycle() { 64 | if LIFE_CYCLES 65 | .lock() 66 | .unwrap() 67 | .insert(entry_point, method) 68 | .is_some() 69 | { 70 | return Err(Error::new( 71 | rust_name.span(), 72 | format!("Canister's '{}' method already defined.", entry_point), 73 | )); 74 | } 75 | } else if METHODS 76 | .lock() 77 | .unwrap() 78 | .insert(name.clone(), method) 79 | .is_some() 80 | { 81 | return Err(Error::new( 82 | rust_name.span(), 83 | format!("Method '{}' is already defined.", name), 84 | )); 85 | }; 86 | 87 | Ok(()) 88 | } 89 | 90 | pub fn export_service(input: DeriveInput, save_candid_path: Option) -> TokenStream { 91 | let methods = { 92 | let mut map = METHODS.lock().unwrap(); 93 | std::mem::replace(&mut *map, BTreeMap::new()) 94 | }; 95 | 96 | let mut life_cycles = { 97 | let mut map = LIFE_CYCLES.lock().unwrap(); 98 | std::mem::replace(&mut *map, BTreeMap::new()) 99 | }; 100 | 101 | let mut rust_methods = Vec::new(); 102 | rust_methods.extend( 103 | life_cycles 104 | .values() 105 | .map(|m| Ident::new(m.rust_name.as_str(), Span::call_site())), 106 | ); 107 | rust_methods.extend( 108 | methods 109 | .values() 110 | .map(|m| Ident::new(m.rust_name.as_str(), Span::call_site())), 111 | ); 112 | 113 | let gen_tys = methods.iter().map( 114 | |( 115 | name, 116 | Method { 117 | arg_types, 118 | rets, 119 | mode, 120 | hidden, 121 | .. 122 | }, 123 | )| { 124 | let args = arg_types 125 | .iter() 126 | .map(|t| generate_arg(quote! { args }, t)) 127 | .collect::>(); 128 | 129 | let rets = rets 130 | .iter() 131 | .map(|t| generate_arg(quote! { rets }, t)) 132 | .collect::>(); 133 | 134 | let modes = match mode { 135 | EntryPoint::Update => quote! { vec![] }, 136 | EntryPoint::Query => { 137 | quote! { vec![ic_kit::candid::parser::types::FuncMode::Query] } 138 | } 139 | _ => unreachable!(), 140 | }; 141 | 142 | if !hidden { 143 | quote! { 144 | { 145 | let mut args = Vec::new(); 146 | #(#args)* 147 | let mut rets = Vec::new(); 148 | #(#rets)* 149 | let func = Function { args, rets, modes: #modes }; 150 | service.push((#name.to_string(), Type::Func(func))); 151 | } 152 | } 153 | } else { 154 | quote! {} 155 | } 156 | }, 157 | ); 158 | 159 | let service = quote! { 160 | use ic_kit::candid::types::{CandidType, Function, Type}; 161 | let mut service = Vec::<(String, Type)>::new(); 162 | let mut env = ic_kit::candid::types::internal::TypeContainer::new(); 163 | #(#gen_tys)* 164 | service.sort_unstable_by_key(|(name, _)| name.clone()); 165 | let ty = Type::Service(service); 166 | }; 167 | 168 | let actor = if let Some(init) = life_cycles.remove(&EntryPoint::Init) { 169 | let args = init 170 | .arg_types 171 | .iter() 172 | .map(|t| generate_arg(quote! { init_args }, t)) 173 | .collect::>(); 174 | 175 | quote! { 176 | let mut init_args = Vec::new(); 177 | #(#args)* 178 | let actor = Some(Type::Class(init_args, Box::new(ty))); 179 | } 180 | } else { 181 | quote! { let actor = Some(ty); } 182 | }; 183 | 184 | let name = input.ident; 185 | 186 | let save_candid = if let Some(path) = save_candid_path { 187 | quote! { 188 | #[cfg(test)] 189 | #[test] 190 | fn ic_kit_save_candid() { 191 | use ic_kit::KitCanister; 192 | use std::env; 193 | use std::fs; 194 | use std::path::PathBuf; 195 | 196 | let candid = #name::candid(); 197 | let mut path = PathBuf::from(env::var("CARGO_MANIFEST_DIR").unwrap()); 198 | path.push(#path); 199 | let dir = path.parent().unwrap(); 200 | 201 | fs::create_dir_all(dir).unwrap_or_else(|e| { 202 | panic!( 203 | "Failed to create the directory '{}': {}", 204 | dir.as_os_str().to_string_lossy(), 205 | e 206 | ) 207 | }); 208 | 209 | fs::write(&path, candid).unwrap_or_else(|e| { 210 | panic!( 211 | "Failed to write to the file '{}': {}", 212 | path.as_os_str().to_string_lossy(), 213 | e 214 | ) 215 | }); 216 | 217 | println!("Saved candid to: {}", path.as_os_str().to_string_lossy()); 218 | } 219 | } 220 | } else { 221 | quote! {} 222 | }; 223 | 224 | let metadata = generate_metadata(); 225 | 226 | quote! { 227 | #metadata 228 | 229 | impl ic_kit::KitCanister for #name { 230 | #[cfg(not(target_family = "wasm"))] 231 | fn build(canister_id: ic_kit::Principal) -> ic_kit::rt::Canister { 232 | ic_kit::rt::Canister::new(canister_id) 233 | #( 234 | .with_method::<#rust_methods>() 235 | )* 236 | } 237 | 238 | fn candid() -> String { 239 | #service 240 | #actor 241 | let result = ic_kit::candid::bindings::candid::compile(&env.env, &actor); 242 | format!("{}", result) 243 | } 244 | } 245 | 246 | #[cfg(target_family = "wasm")] 247 | #[doc(hidden)] 248 | #[export_name = "canister_query __get_candid_interface_tmp_hack"] 249 | fn _ic_kit_canister_query___get_candid_interface_tmp_hack() { 250 | let candid = #name::candid(); 251 | let bytes = ic_kit::candid::encode_one(candid) 252 | .expect("Could not encode canister's response."); 253 | ic_kit::utils::reply(&bytes); 254 | } 255 | 256 | #save_candid 257 | } 258 | } 259 | 260 | fn generate_arg(name: TokenStream, ty: &str) -> TokenStream { 261 | let ty = syn::parse_str::(ty).unwrap(); 262 | quote! { 263 | #name.push(env.add::<#ty>()); 264 | } 265 | } 266 | 267 | /// Remove the references in a type and makes it an owned type, this is used to parse the return 268 | /// type when it's using Kit's DI. 269 | fn remove_reference_recursive(ty: syn::Type) -> syn::Type { 270 | match ty { 271 | syn::Type::Reference(r) => *r.elem, 272 | syn::Type::Tuple(tuple) => syn::Type::Tuple(syn::TypeTuple { 273 | paren_token: tuple.paren_token, 274 | elems: syn::punctuated::Punctuated::from_iter( 275 | tuple.elems.into_iter().map(remove_reference_recursive), 276 | ), 277 | }), 278 | syn::Type::Group(group) => syn::Type::Group(syn::TypeGroup { 279 | group_token: group.group_token, 280 | elem: Box::new(remove_reference_recursive(*group.elem)), 281 | }), 282 | syn::Type::Paren(paren) => syn::Type::Paren(syn::TypeParen { 283 | paren_token: paren.paren_token, 284 | elem: Box::new(remove_reference_recursive(*paren.elem)), 285 | }), 286 | syn::Type::Slice(slice) => syn::Type::Slice(syn::TypeSlice { 287 | bracket_token: slice.bracket_token, 288 | elem: Box::new(remove_reference_recursive(*slice.elem)), 289 | }), 290 | t => t, 291 | } 292 | } 293 | -------------------------------------------------------------------------------- /ic-kit-certified/src/rbtree/test.rs: -------------------------------------------------------------------------------- 1 | use super::{KeyBound, RbTree}; 2 | use crate::{AsHashTree, HashTree}; 3 | use std::convert::AsRef; 4 | 5 | #[test] 6 | fn test_witness() { 7 | let mut t = RbTree::, Vec>::new(); 8 | 9 | for i in 0u64..10 { 10 | let key = (1 + 2 * i).to_be_bytes(); 11 | let val = (1 + 2 * i).to_le_bytes(); 12 | t.insert(key.into(), val.into()); 13 | assert_eq!(t.get(&key[..]).map(|v| &v[..]), Some(&val[..])); 14 | } 15 | 16 | for i in 0u64..10 { 17 | let key = (1 + 2 * i).to_be_bytes(); 18 | let key = key.as_ref(); 19 | 20 | let ht = t.witness(key); 21 | assert_eq!( 22 | ht.reconstruct(), 23 | t.root_hash(), 24 | "key: {}, witness {:?}", 25 | hex::encode(key), 26 | ht 27 | ); 28 | 29 | let ht = t.keys_with_prefix(key); 30 | assert_eq!( 31 | ht.reconstruct(), 32 | t.root_hash(), 33 | "key: {}, lower bound: {:?}, upper_bound: {:?}, witness {:?}", 34 | hex::encode(key), 35 | t.lower_bound(key).map(hex::encode), 36 | t.right_prefix_neighbor(key).map(hex::encode), 37 | ht 38 | ); 39 | } 40 | 41 | for i in 0u64..10 { 42 | for j in i..10 { 43 | let start = (2 * i).to_be_bytes(); 44 | let end = (2 * j).to_be_bytes(); 45 | let ht = t.key_range(&start[..], &end[..]); 46 | assert_eq!( 47 | ht.reconstruct(), 48 | t.root_hash(), 49 | "key range: [{}, {}], witness {:?}", 50 | hex::encode(&start[..]), 51 | hex::encode(&end[..]), 52 | ht 53 | ); 54 | 55 | let ht = t.value_range(&start[..], &end[..]); 56 | assert_eq!( 57 | ht.reconstruct(), 58 | t.root_hash(), 59 | "key range: [{}, {}], witness {:?}", 60 | hex::encode(&start[..]), 61 | hex::encode(&end[..]), 62 | ht 63 | ); 64 | } 65 | } 66 | 67 | for i in 0u64..11 { 68 | let key = (2 * i).to_be_bytes(); 69 | let ht = t.witness(&key[..]); 70 | assert_eq!( 71 | ht.reconstruct(), 72 | t.root_hash(), 73 | "key: {}, witness {:?}", 74 | hex::encode(&key[..]), 75 | ht 76 | ); 77 | } 78 | 79 | for i in 0u64..10 { 80 | let key = (1 + 2 * i).to_be_bytes(); 81 | let val = (1 + 2 * i).to_le_bytes(); 82 | 83 | assert_eq!(t.get(&key[..]).map(|v| &v[..]), Some(&val[..])); 84 | 85 | t.delete(&key[..]); 86 | for j in 0u64..10 { 87 | let witness_key = (1 + 2 * j).to_be_bytes(); 88 | let ht = t.witness(&witness_key[..]); 89 | assert_eq!( 90 | ht.reconstruct(), 91 | t.root_hash(), 92 | "key: {}, witness {:?}", 93 | hex::encode(&key[..]), 94 | ht 95 | ); 96 | } 97 | assert_eq!(t.get(&key[..]), None); 98 | } 99 | } 100 | 101 | #[test] 102 | fn test_key_bounds() { 103 | let mut t = RbTree::, Vec>::new(); 104 | t.insert(vec![1], vec![10]); 105 | t.insert(vec![3], vec![30]); 106 | 107 | assert_eq!(t.lower_bound(&[0u8][..]), None); 108 | assert_eq!(t.lower_bound(&[1u8][..]), Some(KeyBound::Exact(&vec![1u8]))); 109 | assert_eq!( 110 | t.lower_bound(&[2u8][..]), 111 | Some(KeyBound::Neighbor(&vec![1u8])) 112 | ); 113 | assert_eq!(t.lower_bound(&[3u8][..]), Some(KeyBound::Exact(&vec![3u8]))); 114 | assert_eq!( 115 | t.lower_bound(&[4u8][..]), 116 | Some(KeyBound::Neighbor(&vec![3u8])) 117 | ); 118 | 119 | assert_eq!( 120 | t.upper_bound(&[0u8][..]), 121 | Some(KeyBound::Neighbor(&vec![1u8])) 122 | ); 123 | assert_eq!(t.upper_bound(&[1u8][..]), Some(KeyBound::Exact(&vec![1u8]))); 124 | assert_eq!( 125 | t.upper_bound(&[2u8][..]), 126 | Some(KeyBound::Neighbor(&vec![3u8])) 127 | ); 128 | assert_eq!(t.upper_bound(&[3u8][..]), Some(KeyBound::Exact(&vec![3u8]))); 129 | assert_eq!(t.upper_bound(&[4u8][..]), None); 130 | } 131 | 132 | #[test] 133 | fn test_prefix_neighbor() { 134 | let mut t = RbTree::>::new(); 135 | t.insert("a/b".into(), vec![0]); 136 | t.insert("a/b".into(), vec![0]); 137 | t.insert("a/b/c".into(), vec![1]); 138 | t.insert("a/b/d".into(), vec![2]); 139 | t.insert("a/c/d".into(), vec![3]); 140 | 141 | assert_eq!( 142 | t.right_prefix_neighbor("a/b/c"), 143 | Some(KeyBound::Neighbor(&"a/b/d".into())) 144 | ); 145 | assert_eq!( 146 | t.right_prefix_neighbor("a/b"), 147 | Some(KeyBound::Neighbor(&"a/c/d".into())) 148 | ); 149 | assert_eq!(t.right_prefix_neighbor("a/c/d"), None); 150 | assert_eq!(t.right_prefix_neighbor("a"), None); 151 | } 152 | 153 | #[test] 154 | fn simple_delete_test() { 155 | let mut t = RbTree::::new(); 156 | t.insert("x".into(), "a".into()); 157 | t.insert("y".into(), "b".into()); 158 | t.insert("z".into(), "c".into()); 159 | 160 | t.delete("x"); 161 | assert_eq!(t.get("x"), None); 162 | assert_eq!(t.get("y"), Some(&"b".into())); 163 | assert_eq!(t.get("z"), Some(&"c".into())); 164 | 165 | t.delete("y"); 166 | assert_eq!(t.get("y"), None); 167 | assert_eq!(t.get("z"), Some(&"c".into())); 168 | 169 | t.delete("z"); 170 | assert_eq!(t.get("z"), None); 171 | } 172 | 173 | #[test] 174 | fn simple_delete_test_2() { 175 | let mut t = RbTree::::new(); 176 | t.insert("x".into(), "y".into()); 177 | t.insert("z".into(), "w".into()); 178 | 179 | t.delete("z"); 180 | assert_eq!(t.get("z"), None); 181 | assert_eq!(t.get("x"), Some(&"y".into())); 182 | } 183 | 184 | #[test] 185 | fn map_model_test() { 186 | use std::collections::HashMap; 187 | 188 | let mut hm: HashMap, Vec> = HashMap::new(); 189 | let mut rb = RbTree::, Vec>::new(); 190 | 191 | for i in 0..100u64 { 192 | hm.insert(i.to_be_bytes().to_vec(), i.to_be_bytes().to_vec()); 193 | rb.insert(i.to_be_bytes().to_vec(), i.to_be_bytes().to_vec()); 194 | 195 | for k in hm.keys() { 196 | assert_eq!(hm.get(k), rb.get(k)); 197 | } 198 | } 199 | let keys: Vec<_> = hm.keys().cloned().collect(); 200 | 201 | for k in keys { 202 | hm.remove(&k); 203 | 204 | assert!(rb.get(&k).is_some()); 205 | rb.delete(&k); 206 | assert!(rb.get(&k).is_none()); 207 | 208 | for k in hm.keys() { 209 | assert_eq!(hm.get(k), rb.get(k)); 210 | } 211 | } 212 | assert_eq!(super::debug_alloc::count_allocated_pointers(), 0); 213 | } 214 | 215 | #[test] 216 | fn test_nested_witness() { 217 | let mut rb: RbTree> = RbTree::new(); 218 | let mut nested = RbTree::new(); 219 | nested.insert("bottom".into(), "data".into()); 220 | rb.insert("top".into(), nested); 221 | 222 | let ht = rb.nested_witness("top", |v| v.witness("bottom")); 223 | 224 | assert_eq!(ht.reconstruct(), rb.root_hash()); 225 | match ht { 226 | HashTree::Labeled(lt, tt) => { 227 | assert_eq!(lt.as_ref(), b"top"); 228 | match &(*tt) { 229 | HashTree::Labeled(lb, _) => { 230 | assert_eq!(lb.as_ref(), b"bottom"); 231 | } 232 | other => panic!("unexpected nested tree: {:?}", other), 233 | } 234 | } 235 | other => panic!("expected a labeled tree, got {:?}", other), 236 | } 237 | 238 | rb.modify("top", |m| { 239 | m.delete("bottom"); 240 | }); 241 | let ht = rb.nested_witness("top", |v| v.witness("bottom")); 242 | assert_eq!(ht.reconstruct(), rb.root_hash()); 243 | } 244 | 245 | #[test] 246 | fn test_witness_key_range() { 247 | let mut t = RbTree::::new(); 248 | t.insert("b".into(), "x".into()); 249 | t.insert("d".into(), "y".into()); 250 | t.insert("f".into(), "z".into()); 251 | 252 | assert_eq!(t.key_range("a", "a").get_labels(), vec![b"b"]); 253 | assert_eq!(t.key_range("a", "b").get_labels(), vec![b"b"]); 254 | assert_eq!(t.key_range("a", "c").get_labels(), vec![b"b", b"d"]); 255 | assert_eq!(t.key_range("a", "f").get_labels(), vec![b"b", b"d", b"f"]); 256 | assert_eq!(t.key_range("a", "z").get_labels(), vec![b"b", b"d", b"f"]); 257 | 258 | assert_eq!(t.key_range("b", "b").get_labels(), vec![b"b"]); 259 | assert_eq!(t.key_range("b", "c").get_labels(), vec![b"b", b"d"]); 260 | assert_eq!(t.key_range("b", "f").get_labels(), vec![b"b", b"d", b"f"]); 261 | assert_eq!(t.key_range("b", "z").get_labels(), vec![b"b", b"d", b"f"]); 262 | 263 | assert_eq!(t.key_range("d", "e").get_labels(), vec![b"d", b"f"]); 264 | assert_eq!(t.key_range("d", "f").get_labels(), vec![b"d", b"f"]); 265 | assert_eq!(t.key_range("d", "z").get_labels(), vec![b"d", b"f"]); 266 | assert_eq!(t.key_range("y", "z").get_labels(), vec![b"f"]); 267 | 268 | assert!(t.key_range("a", "z").get_leaf_values().is_empty()); 269 | } 270 | 271 | #[test] 272 | fn test_witness_value_range() { 273 | let mut t = RbTree::::new(); 274 | t.insert("b".into(), "x".into()); 275 | t.insert("d".into(), "y".into()); 276 | t.insert("f".into(), "z".into()); 277 | 278 | assert_eq!(t.key_range("a", "a").get_labels(), vec![b"b"]); 279 | assert!(t.value_range("a", "a").get_leaf_values().is_empty()); 280 | 281 | assert_eq!(t.value_range("a", "b").get_labels(), vec![b"b"]); 282 | assert_eq!(t.value_range("a", "b").get_leaf_values(), vec![b"x"]); 283 | 284 | assert_eq!(t.value_range("f", "z").get_labels(), vec![b"f"]); 285 | assert_eq!(t.value_range("f", "z").get_leaf_values(), vec![b"z"]); 286 | 287 | assert_eq!(t.value_range("g", "z").get_labels(), vec![b"f"]); 288 | assert!(t.value_range("g", "z").get_leaf_values().is_empty()); 289 | 290 | assert_eq!(t.value_range("a", "z").get_labels(), vec![b"b", b"d", b"f"]); 291 | assert_eq!( 292 | t.value_range("a", "z").get_leaf_values(), 293 | vec![b"x", b"y", b"z"] 294 | ); 295 | } 296 | -------------------------------------------------------------------------------- /ic-kit/src/ic/call.rs: -------------------------------------------------------------------------------- 1 | use crate::futures; 2 | use crate::futures::CallFuture; 3 | use crate::ic::Cycles; 4 | use crate::utils::arg_data_raw; 5 | use candid::utils::{ArgumentDecoder, ArgumentEncoder}; 6 | use candid::{decode_args, decode_one, encode_args, encode_one, CandidType, Principal}; 7 | use ic_kit_sys::ic0; 8 | use serde::de::DeserializeOwned; 9 | 10 | use ic_kit_sys::types::RejectionCode; 11 | pub use ic_kit_sys::types::{CallError, CANDID_EMPTY_ARG}; 12 | 13 | /// A call builder that let's you create an inter-canister call which can be then sent to the 14 | /// destination. 15 | pub struct CallBuilder { 16 | canister_id: Principal, 17 | method_name: String, 18 | payment: Cycles, 19 | arg: Option>, 20 | } 21 | 22 | impl CallBuilder { 23 | /// Create a new call constructor, calling this method does nothing unless one of the perform 24 | /// methods are called. 25 | pub fn new>(canister_id: Principal, method_name: S) -> Self { 26 | Self { 27 | canister_id, 28 | method_name: method_name.into(), 29 | payment: 0, 30 | arg: None, 31 | } 32 | } 33 | 34 | /// Use the given candid tuple value as the argument. 35 | /// 36 | /// # Panics 37 | /// 38 | /// This method panics if the argument for this call is already set via a prior 39 | /// call to any of the `with_args`, `with_arg` or `with_arg_raw`. 40 | /// 41 | /// Use `clear_args` if you want to reset the arguments. 42 | pub fn with_args(mut self, arguments: T) -> Self { 43 | assert!(self.arg.is_none(), "Call arguments can only be set once."); 44 | self.arg = Some(encode_args(arguments).unwrap()); 45 | self 46 | } 47 | 48 | /// Shorthand for `with_args((argument, ))`. 49 | /// 50 | /// # Panics 51 | /// 52 | /// This method panics if the argument for this call is already set via a prior 53 | /// call to any of the `with_args`, `with_arg` or `with_arg_raw`. 54 | /// 55 | /// Use `clear_args` if you want to reset the arguments. 56 | pub fn with_arg(mut self, argument: T) -> Self { 57 | assert!(self.arg.is_none(), "Call arguments can only be set once."); 58 | self.arg = Some(encode_one(argument).unwrap()); 59 | self 60 | } 61 | 62 | /// Set the raw argument that can be used for this call, this does not use candid to serialize 63 | /// the call argument and uses the provided raw buffer as the argument. 64 | /// 65 | /// Be sure that you know what you're doing when using this method. 66 | /// 67 | /// # Panics 68 | /// 69 | /// This method panics if the argument for this call is already set via a prior 70 | /// call to any of the `with_args`, `with_arg` or `with_arg_raw`. 71 | /// 72 | /// Use `clear_args` if you want to reset the arguments. 73 | pub fn with_arg_raw>>(mut self, argument: A) -> Self { 74 | assert!(self.arg.is_none(), "Call arguments can only be set once."); 75 | self.arg = Some(argument.into()); 76 | self 77 | } 78 | 79 | /// Clear any arguments set for this call. After calling this method you can call with_arg* 80 | /// methods again without the panic. 81 | pub fn clear_args(&mut self) { 82 | self.arg = None; 83 | } 84 | 85 | /// Set the payment amount for the canister. This will overwrite any previously added cycles 86 | /// to this call, use `add_payment` if you want to increment the amount of used cycles in 87 | /// this call. 88 | /// 89 | /// # Safety 90 | /// 91 | /// Be sure that your canister has the provided amount of cycles upon performing the call, 92 | /// since any of the perform methods will just trap the canister if the provided payment 93 | /// amount is larger than the amount of canister's balance. 94 | pub fn with_payment(mut self, payment: Cycles) -> Self { 95 | self.payment = payment; 96 | self 97 | } 98 | 99 | /// Add the given provided amount of cycles to the cycles already provided to this call. 100 | pub fn add_payment(mut self, payment: Cycles) -> Self { 101 | self.payment += payment; 102 | self 103 | } 104 | 105 | /// Should be called after the `ic0::call_new` to set the call arguments. 106 | #[inline(always)] 107 | unsafe fn ic0_internal_call_perform(&self) -> i32 { 108 | #[cfg(not(feature = "experimental-cycles128"))] 109 | ic0::call_cycles_add(self.payment as i64); 110 | 111 | #[cfg(feature = "experimental-cycles128")] 112 | if self.payment > 0 && self.payment < (u64::MAX as u128) { 113 | ic0::call_cycles_add(self.payment as i64); 114 | } else if self.payment > 0 { 115 | let high = (self.payment >> 64) as u64 as i64; 116 | let low = (self.payment & (1 << 64)) as u64 as i64; 117 | ic0::call_cycles_add128(high, low); 118 | } 119 | 120 | let args_raw = self.arg.as_deref().unwrap_or(CANDID_EMPTY_ARG); 121 | 122 | if !args_raw.is_empty() { 123 | ic0::call_data_append(args_raw.as_ptr() as isize, args_raw.len() as isize); 124 | } 125 | 126 | ic0::call_perform() 127 | } 128 | 129 | /// Perform a call when you do not care about the response in anyway. We advise you to use this 130 | /// method when you can since it is probably cheaper. 131 | /// 132 | /// # Traps 133 | /// 134 | /// This method traps if the amount determined in the `payment` is larger than the canister's 135 | /// balance at the time of invocation. 136 | pub fn perform_one_way(self) -> Result<(), RejectionCode> { 137 | let callee = self.canister_id.as_slice(); 138 | let method = self.method_name.as_str(); 139 | 140 | let e_code = unsafe { 141 | ic0::call_new( 142 | callee.as_ptr() as isize, 143 | callee.len() as isize, 144 | method.as_ptr() as isize, 145 | method.len() as isize, 146 | -1, 147 | -1, 148 | -1, 149 | -1, 150 | ); 151 | 152 | self.ic0_internal_call_perform() 153 | }; 154 | 155 | if e_code != 0 { 156 | Err(e_code.into()) 157 | } else { 158 | Ok(()) 159 | } 160 | } 161 | 162 | /// Perform the call and return a future that can will be resolved in any of the callbacks. 163 | /// 164 | /// # Traps 165 | /// 166 | /// This method traps if the amount determined in the `payment` is larger than the canister's 167 | /// balance at the time of invocation. 168 | #[must_use] 169 | fn perform_internal(&self) -> CallFuture { 170 | let future = unsafe { 171 | let future = futures::call_new(self.canister_id, self.method_name.as_str()); 172 | let e_code = self.ic0_internal_call_perform(); 173 | 174 | if e_code != 0 { 175 | future.mark_ready() 176 | } else { 177 | future 178 | } 179 | }; 180 | 181 | future 182 | } 183 | 184 | /// Use this method when you want to perform a call and only care about the delivery status 185 | /// of the call and don't need the returned buffer in anyway. 186 | /// 187 | /// # Traps 188 | /// 189 | /// This method traps if the amount determined in the `payment` is larger than the canister's 190 | /// balance at the time of invocation. 191 | pub async fn perform_rejection(&self) -> Result<(), CallError> { 192 | let future = self.perform_internal(); 193 | 194 | // if the future is already ready, it indicates a `ic0::call_perform` non-zero response. 195 | if future.is_ready() { 196 | return Err(CallError::CouldNotSend); 197 | } 198 | 199 | // await for the call to comeback. 200 | future.await; 201 | 202 | let rejection_code = unsafe { ic0::msg_reject_code() }; 203 | if rejection_code == 0 { 204 | return Ok(()); 205 | } 206 | 207 | let rejection_message_size = unsafe { ic0::msg_reject_msg_size() } as usize; 208 | let mut bytes = vec![0u8; rejection_message_size]; 209 | unsafe { 210 | ic0::msg_reject_msg_copy( 211 | bytes.as_mut_ptr() as isize, 212 | 0, 213 | rejection_message_size as isize, 214 | ); 215 | } 216 | 217 | Err(CallError::Rejected( 218 | rejection_code.into(), 219 | String::from_utf8_lossy(&bytes).to_string(), 220 | )) 221 | } 222 | 223 | /// Perform the call and return the raw response buffer without decoding it. 224 | /// 225 | /// # Traps 226 | /// 227 | /// This method traps if the amount determined in the `payment` is larger than the canister's 228 | /// balance at the time of invocation. 229 | pub async fn perform_raw(&self) -> Result, CallError> { 230 | self.perform_rejection().await?; 231 | Ok(arg_data_raw()) 232 | } 233 | 234 | /// Perform the call and return a future which will resolve to the candid decoded response. Or 235 | /// any of the errors that might happen, consider looking at other alternatives of this method 236 | /// as well if you don't care about the response or want the raw/non-decoded response. 237 | /// 238 | /// # Traps 239 | /// 240 | /// This method traps if the amount determined in the `payment` is larger than the canister's 241 | /// balance at the time of invocation. 242 | pub async fn perform ArgumentDecoder<'a>>(&self) -> Result { 243 | let bytes = self.perform_raw().await?; 244 | 245 | match decode_args(&bytes) { 246 | Err(_) => Err(CallError::ResponseDeserializationError(bytes)), 247 | Ok(r) => Ok(r), 248 | } 249 | } 250 | 251 | /// Perform the call and return a future which will resolve to the candid decoded response. 252 | /// Unlink perform, this method only expects a result with one argument from the canister, 253 | /// and decodes the arguments using the candid's decode_one. 254 | /// 255 | /// 256 | /// # Traps 257 | /// 258 | /// This method traps if the amount determined in the `payment` is larger than the canister's 259 | /// balance at the time of invocation. 260 | pub async fn perform_one(&self) -> Result 261 | where 262 | T: DeserializeOwned + CandidType, 263 | { 264 | let bytes = self.perform_raw().await?; 265 | 266 | match decode_one(&bytes) { 267 | Err(_) => Err(CallError::ResponseDeserializationError(bytes)), 268 | Ok(r) => Ok(r), 269 | } 270 | } 271 | } 272 | -------------------------------------------------------------------------------- /ic-kit-certified/src/collections/seq.rs: -------------------------------------------------------------------------------- 1 | use crate::{AsHashTree, Hash, HashTree}; 2 | use candid::types::Type; 3 | use candid::CandidType; 4 | use serde::{Deserialize, Deserializer, Serialize, Serializer}; 5 | use sha2::{Digest, Sha256}; 6 | use std::borrow::Borrow; 7 | use std::iter::FromIterator; 8 | use std::ops::Index; 9 | use std::slice::{Iter, SliceIndex}; 10 | 11 | /// An append only list of `T`. 12 | /// 13 | /// # Example 14 | /// 15 | /// ``` 16 | /// use ic_kit_certified::Seq; 17 | /// 18 | /// let mut seq = Seq::::new(); 19 | /// 20 | /// seq.append(0); 21 | /// seq.append(1); 22 | /// 23 | /// assert_eq!(seq.len(), 2); 24 | /// ``` 25 | #[derive(Default, Eq, PartialEq, Clone, Debug)] 26 | pub struct Seq { 27 | hash: Hash, 28 | items: Vec, 29 | } 30 | 31 | impl Seq { 32 | /// Create a new, empty Stack 33 | #[inline] 34 | pub const fn new() -> Self { 35 | Self { 36 | items: Vec::new(), 37 | hash: [0; 32], 38 | } 39 | } 40 | 41 | /// Construct a new, empty Stack with the specified capacity. 42 | #[inline] 43 | pub fn with_capacity(capacity: usize) -> Self { 44 | Self { 45 | items: Vec::with_capacity(capacity), 46 | hash: [0; 32], 47 | } 48 | } 49 | } 50 | 51 | impl Seq { 52 | /// Append a new item to the sequence and update the hash. 53 | pub fn append(&mut self, item: T) { 54 | let mut h = Sha256::new(); 55 | h.update(&self.hash); 56 | h.update(item.root_hash()); 57 | 58 | self.hash = h.finalize().into(); 59 | self.items.push(item); 60 | } 61 | 62 | /// Clear the sequence by removing all of the items. This method does not have 63 | /// any effects on the allocated memory. 64 | #[inline] 65 | pub fn clear(&mut self) { 66 | self.hash = [0; 32]; 67 | self.items.clear(); 68 | } 69 | 70 | /// Shrinks the capacity of the seq as much as possible. 71 | #[inline] 72 | pub fn shrink_to_fit(&mut self) { 73 | self.items.shrink_to_fit(); 74 | } 75 | 76 | /// Reserve space for at least `additional` more elements. 77 | #[inline] 78 | pub fn reserve(&mut self, additional: usize) { 79 | self.items.reserve(additional) 80 | } 81 | 82 | /// Reserve space for exactly `additional` more elements. 83 | #[inline] 84 | pub fn reserve_exact(&mut self, additional: usize) { 85 | self.items.reserve_exact(additional) 86 | } 87 | 88 | /// Return the underlying vector containing the items. 89 | #[inline] 90 | pub fn as_vec(&self) -> &Vec { 91 | &self.items 92 | } 93 | 94 | /// Returns `true` if the sequence does not have any elements. 95 | #[inline] 96 | pub fn is_empty(&self) -> bool { 97 | self.items.len() == 0 98 | } 99 | 100 | /// Returns the number of elements in the sequence, also referred to as its ‘length’. 101 | #[inline] 102 | pub fn len(&self) -> usize { 103 | self.items.len() 104 | } 105 | 106 | /// Returns the number of elements the sequence can hold without reallocating. 107 | #[inline] 108 | pub fn capacity(&self) -> usize { 109 | self.items.capacity() 110 | } 111 | 112 | /// Returns an iterator over the data. 113 | #[inline] 114 | pub fn iter(&self) -> Iter<'_, T> { 115 | self.items.iter() 116 | } 117 | 118 | /// Recompute the hash of the sequence. 119 | #[inline] 120 | fn recompute_hash(&mut self, prev_len: usize) { 121 | let mut hash = self.hash; 122 | 123 | for item in &self.items[prev_len..] { 124 | let mut h = Sha256::new(); 125 | h.update(&hash); 126 | h.update(item.root_hash()); 127 | hash = h.finalize().into(); 128 | } 129 | 130 | self.hash = hash; 131 | } 132 | } 133 | 134 | impl AsHashTree for Seq { 135 | #[inline] 136 | fn root_hash(&self) -> Hash { 137 | self.hash 138 | } 139 | 140 | #[inline] 141 | fn as_hash_tree(&self) -> HashTree<'_> { 142 | HashTree::Pruned(self.hash) 143 | } 144 | } 145 | 146 | impl From> for Seq { 147 | #[inline] 148 | fn from(items: Vec) -> Self { 149 | let mut seq = Seq { 150 | items, 151 | hash: [0; 32], 152 | }; 153 | 154 | seq.recompute_hash(0); 155 | 156 | seq 157 | } 158 | } 159 | 160 | impl FromIterator for Seq { 161 | #[inline] 162 | fn from_iter>(iter: I) -> Self { 163 | let mut seq = Seq { 164 | items: iter.into_iter().collect(), 165 | hash: [0; 32], 166 | }; 167 | 168 | seq.recompute_hash(0); 169 | 170 | seq 171 | } 172 | } 173 | 174 | impl From> for Vec { 175 | #[inline] 176 | fn from(seq: Seq) -> Self { 177 | seq.items 178 | } 179 | } 180 | 181 | impl AsRef<[T]> for Seq { 182 | #[inline] 183 | fn as_ref(&self) -> &[T] { 184 | self.items.as_ref() 185 | } 186 | } 187 | 188 | impl Borrow<[T]> for Seq { 189 | #[inline] 190 | fn borrow(&self) -> &[T] { 191 | self.items.borrow() 192 | } 193 | } 194 | 195 | impl<'a, T: AsHashTree + Copy + 'a> Extend<&'a T> for Seq { 196 | #[inline] 197 | fn extend>(&mut self, iter: I) { 198 | let prev_len = self.items.len(); 199 | self.items.extend(iter); 200 | self.recompute_hash(prev_len) 201 | } 202 | } 203 | 204 | impl Extend for Seq { 205 | #[inline] 206 | fn extend>(&mut self, iter: I) { 207 | let prev_len = self.items.len(); 208 | self.items.extend(iter); 209 | self.recompute_hash(prev_len) 210 | } 211 | } 212 | 213 | impl<'a, T: AsHashTree + Clone> From<&'a [T]> for Seq { 214 | #[inline] 215 | fn from(items: &'a [T]) -> Self { 216 | let mut seq = Seq { 217 | items: items.into(), 218 | hash: [0; 32], 219 | }; 220 | seq.recompute_hash(0); 221 | seq 222 | } 223 | } 224 | 225 | impl<'a, T: AsHashTree + Clone> From<&'a mut [T]> for Seq { 226 | #[inline] 227 | fn from(items: &'a mut [T]) -> Self { 228 | let mut seq = Seq { 229 | items: items.into(), 230 | hash: [0; 32], 231 | }; 232 | seq.recompute_hash(0); 233 | seq 234 | } 235 | } 236 | 237 | impl> Index for Seq { 238 | type Output = I::Output; 239 | 240 | #[inline] 241 | fn index(&self, index: I) -> &Self::Output { 242 | self.items.index(index) 243 | } 244 | } 245 | 246 | impl Serialize for Seq { 247 | fn serialize(&self, serializer: S) -> Result 248 | where 249 | S: Serializer, 250 | { 251 | self.items.serialize(serializer) 252 | } 253 | } 254 | 255 | impl<'de, T: AsHashTree + Deserialize<'de>> Deserialize<'de> for Seq { 256 | fn deserialize(deserializer: D) -> Result 257 | where 258 | D: Deserializer<'de>, 259 | { 260 | let mut seq = Seq { 261 | items: >::deserialize(deserializer)?, 262 | hash: [0; 32], 263 | }; 264 | 265 | seq.recompute_hash(0); 266 | 267 | Ok(seq) 268 | } 269 | } 270 | 271 | impl CandidType for Seq { 272 | fn _ty() -> Type { 273 | >::_ty() 274 | } 275 | 276 | fn idl_serialize(&self, serializer: S) -> Result<(), S::Error> 277 | where 278 | S: candid::types::Serializer, 279 | { 280 | self.items.idl_serialize(serializer) 281 | } 282 | } 283 | 284 | #[cfg(test)] 285 | mod tests { 286 | use super::*; 287 | use candid::{decode_one, encode_one}; 288 | 289 | #[test] 290 | fn append() { 291 | let mut seq = Seq::::with_capacity(1000); 292 | let mut hash = seq.root_hash(); 293 | assert_eq!(seq.is_empty(), true); 294 | 295 | for i in 0..1000 { 296 | seq.append(i); 297 | assert_eq!(seq.len(), i + 1); 298 | let new_hash = seq.root_hash(); 299 | assert_ne!(hash, new_hash); 300 | hash = new_hash; 301 | } 302 | 303 | assert_eq!(seq.is_empty(), false); 304 | 305 | seq.clear(); 306 | assert_eq!(seq.len(), 0); 307 | assert_eq!(seq.is_empty(), true); 308 | assert_eq!(seq.root_hash(), Seq::::new().root_hash()); 309 | 310 | for i in 0..1000 { 311 | seq.append(i); 312 | } 313 | 314 | assert_eq!(hash, seq.root_hash()); 315 | } 316 | 317 | #[test] 318 | fn extend() { 319 | let manual = { 320 | let mut seq = Seq::::with_capacity(100); 321 | for i in 0..100 { 322 | seq.append(i); 323 | } 324 | seq 325 | }; 326 | 327 | { 328 | let mut seq = Seq::::new(); 329 | seq.extend(0..100); 330 | 331 | assert_eq!(manual.len(), seq.len()); 332 | assert_eq!(manual.root_hash(), seq.root_hash()); 333 | } 334 | 335 | { 336 | let mut seq = Seq::::new(); 337 | seq.extend(0..50); 338 | seq.extend(50..100); 339 | assert_eq!(manual.len(), seq.len()); 340 | assert_eq!(manual.root_hash(), seq.root_hash()); 341 | } 342 | 343 | { 344 | let mut seq = Seq::::new(); 345 | seq.extend(0..50); 346 | seq.append(50); 347 | seq.extend(51..100); 348 | assert_eq!(manual.len(), seq.len()); 349 | assert_eq!(manual.root_hash(), seq.root_hash()); 350 | } 351 | } 352 | 353 | #[test] 354 | fn index() { 355 | let seq = (0..100).collect::>(); 356 | 357 | for i in 0..100 { 358 | assert_eq!(seq[i], i); 359 | } 360 | } 361 | 362 | #[test] 363 | #[should_panic] 364 | fn index_out_of_range() { 365 | let seq = Seq::::new(); 366 | seq[0]; 367 | } 368 | 369 | #[test] 370 | fn serde_cbor() { 371 | let seq = (0..10).collect::>(); 372 | let serialized = serde_cbor::to_vec(&seq).unwrap(); 373 | let actual: Seq = serde_cbor::from_slice(&serialized).unwrap(); 374 | assert_eq!(actual.len(), 10); 375 | assert_eq!(actual.hash, seq.hash); 376 | assert_eq!(actual, seq); 377 | let expected = (0..10).collect::>(); 378 | let deserialized_as_vec: Vec = serde_cbor::from_slice(&serialized).unwrap(); 379 | assert_eq!(deserialized_as_vec, expected); 380 | } 381 | 382 | #[test] 383 | fn candid() { 384 | let seq = (0..10).collect::>(); 385 | let encoded = encode_one(&seq).unwrap(); 386 | let decoded: Seq = decode_one(&encoded).unwrap(); 387 | assert_eq!(seq, decoded); 388 | } 389 | } 390 | -------------------------------------------------------------------------------- /ic-kit-runtime/src/types.rs: -------------------------------------------------------------------------------- 1 | use std::panic::{RefUnwindSafe, UnwindSafe}; 2 | use std::sync::atomic::{AtomicU64, Ordering}; 3 | use std::time::{SystemTime, UNIX_EPOCH}; 4 | 5 | use candid::utils::ArgumentEncoder; 6 | use candid::Principal; 7 | use candid::{encode_args, encode_one, CandidType}; 8 | 9 | use ic_kit_sys::types::{RejectionCode, CANDID_EMPTY_ARG}; 10 | 11 | static REQUEST_ID: AtomicU64 = AtomicU64::new(0); 12 | 13 | /// A request ID for a request that is coming to this canister from the outside. 14 | pub type IncomingRequestId = RequestId; 15 | /// A request ID for a request that this canister has submitted. 16 | pub type OutgoingRequestId = RequestId; 17 | 18 | /// An opaque request id. 19 | #[derive(Hash, Clone, Copy, Ord, PartialOrd, Eq, PartialEq, Debug)] 20 | pub struct RequestId(u64); 21 | 22 | impl RequestId { 23 | /// Create a new request id and return it. 24 | pub fn new() -> Self { 25 | Self(REQUEST_ID.fetch_add(1, Ordering::SeqCst)) 26 | } 27 | } 28 | 29 | /// The entry method for a request. 30 | #[derive(Debug, PartialEq, Copy, Clone)] 31 | pub enum EntryMode { 32 | Init, 33 | PreUpgrade, 34 | PostUpgrade, 35 | Heartbeat, 36 | InspectMessage, 37 | Update, 38 | Query, 39 | ReplyCallback, 40 | RejectCallback, 41 | CleanupCallback, 42 | CustomTask, 43 | } 44 | 45 | /// The canister's environment that should be used during a message. 46 | pub struct Env { 47 | /// Determines the canister' balance. 48 | pub balance: u128, 49 | /// The type of the entry point that should be simulated, this enables trapping when a the 50 | /// method is calling a system api call that it should not be able to call during the 51 | /// execution of that entry point. 52 | pub entry_mode: EntryMode, 53 | /// The principal id of the sender. 54 | pub sender: Principal, 55 | /// The method to call. Only applies to update/query calls. 56 | pub method_name: Option, 57 | /// The cycles provided to the canister during this call. 58 | pub cycles_available: u128, 59 | /// The amount of refunded cycles. 60 | pub cycles_refunded: u128, 61 | /// The arguments provided to the canister during this call. 62 | pub args: Vec, 63 | /// The reply rejection code. Default to `0` 64 | pub rejection_code: RejectionCode, 65 | /// The rejection message. Only applicable when `rejection_code != 0` 66 | pub rejection_message: String, 67 | /// The current time in nanoseconds. 68 | pub time: u64, 69 | } 70 | 71 | pub type TaskFn = Box; 72 | 73 | /// A message sent to a canister that trigger execution of a task on the canister's execution thread 74 | /// based on the type of the message. 75 | pub enum Message { 76 | /// A custom function that you want to be executed in the canister's execution thread. 77 | CustomTask { 78 | /// The request id of this incoming message. 79 | request_id: IncomingRequestId, 80 | /// the task handler that should be executed in the canister's execution thread. 81 | task: TaskFn, 82 | /// The env to use for this custom execution. 83 | env: Env, 84 | }, 85 | /// A top-level request to the canister. 86 | Request { 87 | /// The request id of the incoming message. Must be None if the reply_to is set. 88 | request_id: IncomingRequestId, 89 | /// The env to use during the execution of this task. 90 | env: Env, 91 | }, 92 | // Either a reply_callback or reject_callbacks. 93 | Reply { 94 | /// Which request is this reply for. 95 | reply_to: OutgoingRequestId, 96 | /// The env to use for this, assert: 97 | /// env.entry_mode == ReplyCallback 98 | /// env.entry_mode == RejectCallback 99 | env: Env, 100 | }, 101 | } 102 | 103 | /// A call that has made to another canister. 104 | #[derive(Debug)] 105 | pub struct CanisterCall { 106 | pub sender: Principal, 107 | pub request_id: RequestId, 108 | pub callee: Principal, 109 | pub method: String, 110 | pub payment: u128, 111 | pub arg: Vec, 112 | } 113 | 114 | impl From for Message { 115 | fn from(call: CanisterCall) -> Self { 116 | Message::Request { 117 | request_id: call.request_id, 118 | env: Env::default() 119 | .with_entry_mode(EntryMode::Update) 120 | .with_sender(call.sender) 121 | .with_method_name(call.method) 122 | .with_cycles_available(call.payment) 123 | .with_raw_args(call.arg), 124 | } 125 | } 126 | } 127 | 128 | impl Default for Env { 129 | fn default() -> Self { 130 | Env { 131 | balance: 100_000_000_000_000, 132 | entry_mode: EntryMode::CustomTask, 133 | sender: Principal::anonymous(), 134 | method_name: None, 135 | cycles_available: 0, 136 | cycles_refunded: 0, 137 | args: CANDID_EMPTY_ARG.to_vec(), 138 | rejection_code: RejectionCode::NoError, 139 | rejection_message: String::new(), 140 | time: now(), 141 | } 142 | } 143 | } 144 | 145 | impl Env { 146 | /// Create a new env for an update call. 147 | pub fn update>(method_name: S) -> Self { 148 | Self::default() 149 | .with_entry_mode(EntryMode::Update) 150 | .with_method_name(method_name) 151 | } 152 | 153 | /// Create a new env for a query call. 154 | pub fn query>(method_name: S) -> Self { 155 | Self::default() 156 | .with_entry_mode(EntryMode::Query) 157 | .with_method_name(method_name) 158 | } 159 | 160 | /// Create a new env for a call to the init function. 161 | pub fn init() -> Self { 162 | Self::default().with_entry_mode(EntryMode::Init) 163 | } 164 | 165 | /// Create a new env for a call to the pre_upgrade function. 166 | pub fn pre_upgrade() -> Self { 167 | Self::default().with_entry_mode(EntryMode::PreUpgrade) 168 | } 169 | 170 | /// Create a new env for a call to the post_upgrade function. 171 | pub fn post_upgrade() -> Self { 172 | Self::default().with_entry_mode(EntryMode::PostUpgrade) 173 | } 174 | 175 | /// Create a new env for a call to the heartbeat function. 176 | pub fn heartbeat() -> Self { 177 | Self::default().with_entry_mode(EntryMode::Heartbeat) 178 | } 179 | 180 | /// Determines the canister's cycle balance for this call. 181 | pub fn with_balance(mut self, balance: u128) -> Self { 182 | self.balance = balance; 183 | self 184 | } 185 | 186 | /// Use the provided time for this env. 187 | pub fn with_time(mut self, time: u64) -> Self { 188 | self.time = time; 189 | self 190 | } 191 | 192 | /// Use the given entry mode in this env. 193 | pub fn with_entry_mode(mut self, mode: EntryMode) -> Self { 194 | self.entry_mode = mode; 195 | self 196 | } 197 | 198 | /// Provide this environment with the given principal id as the caller. 199 | pub fn with_sender(mut self, sender: Principal) -> Self { 200 | self.sender = sender; 201 | self 202 | } 203 | 204 | /// Provide the given env with the given method name to execute. 205 | pub fn with_method_name>(mut self, method_name: S) -> Self { 206 | self.method_name = Some(method_name.into()); 207 | self 208 | } 209 | 210 | /// Provide the current env with the given amount of cycles to execute. 211 | pub fn with_cycles_available(mut self, cycles: u128) -> Self { 212 | self.cycles_available = cycles; 213 | self 214 | } 215 | 216 | /// Provide the current env with the given amount of refunded cycles, only applicable 217 | /// if this is reply/reject callback. 218 | pub fn with_cycles_refunded(mut self, cycles: u128) -> Self { 219 | self.cycles_refunded = cycles; 220 | self 221 | } 222 | 223 | /// The arguments in this environment, in a reply mode this is the data returned to the 224 | /// canister. 225 | pub fn with_raw_args>>(mut self, argument: A) -> Self { 226 | self.args = argument.into(); 227 | self 228 | } 229 | 230 | /// Encode the provided tuple using candid and use it as arguments during this execution. 231 | pub fn with_args(mut self, arguments: T) -> Self { 232 | self.args = encode_args(arguments).unwrap(); 233 | self 234 | } 235 | 236 | /// Shorthand for `with_args((argument, ))` to pass tuples with only one element to the call. 237 | pub fn with_arg(mut self, argument: T) -> Self { 238 | self.args = encode_one(argument).unwrap(); 239 | self 240 | } 241 | 242 | /// Set this environment's rejection code the provided value, you must also set a rejection 243 | /// message if this is not equal to NoError. 244 | pub fn with_rejection_code(mut self, rejection_code: RejectionCode) -> Self { 245 | self.rejection_code = rejection_code; 246 | self 247 | } 248 | 249 | /// Set the rejection message on this env, only applicable if rejection_code is not zero. 250 | pub fn with_rejection_message>(mut self, rejection_message: S) -> Self { 251 | self.rejection_message = rejection_message.into(); 252 | self 253 | } 254 | } 255 | 256 | impl Env { 257 | /// Return a name we can use to get the method from the symbol table. 258 | pub fn get_entry_point_name(&self) -> String { 259 | match &self.entry_mode { 260 | EntryMode::Init => "canister_init".to_string(), 261 | EntryMode::PreUpgrade => "canister_pre_upgrade".to_string(), 262 | EntryMode::PostUpgrade => "canister_post_upgrade".to_string(), 263 | EntryMode::Heartbeat => "canister_heartbeat".to_string(), 264 | EntryMode::InspectMessage => "canister_inspect_message".to_string(), 265 | EntryMode::Update => { 266 | format!( 267 | "canister_update {}", 268 | self.method_name.as_ref().unwrap_or(&String::new()) 269 | ) 270 | } 271 | EntryMode::Query => format!( 272 | "canister_query {}", 273 | self.method_name.as_ref().unwrap_or(&String::new()) 274 | ), 275 | EntryMode::ReplyCallback => "reply callback".to_string(), 276 | EntryMode::RejectCallback => "reject callback".to_string(), 277 | EntryMode::CleanupCallback => "cleanup callback".to_string(), 278 | EntryMode::CustomTask => "ic-kit: custom".to_string(), 279 | } 280 | } 281 | 282 | /// Returns the second possible name of this entry point. 283 | pub fn get_possible_entry_point_name(&self) -> String { 284 | match &self.entry_mode { 285 | EntryMode::Update => { 286 | format!( 287 | "canister_query {}", 288 | self.method_name.as_ref().unwrap_or(&String::new()) 289 | ) 290 | } 291 | EntryMode::Query => format!( 292 | "canister_update {}", 293 | self.method_name.as_ref().unwrap_or(&String::new()) 294 | ), 295 | _ => self.get_entry_point_name(), 296 | } 297 | } 298 | } 299 | 300 | fn now() -> u64 { 301 | let now = SystemTime::now(); 302 | let unix = now 303 | .duration_since(UNIX_EPOCH) 304 | .expect("ic-kit-runtime: could not retrieve unix time."); 305 | unix.as_nanos() as u64 306 | } 307 | -------------------------------------------------------------------------------- /ic-kit-certified/src/collections/map.rs: -------------------------------------------------------------------------------- 1 | use crate::collections::seq::Seq; 2 | use crate::label::{Label, Prefix}; 3 | use crate::rbtree::entry::Entry; 4 | use crate::rbtree::iterator::RbTreeIterator; 5 | use crate::rbtree::RbTree; 6 | use crate::{AsHashTree, Hash, HashTree}; 7 | use candid::types::{Compound, Field, Label as CLabel, Type}; 8 | use candid::CandidType; 9 | use serde::de::{MapAccess, Visitor}; 10 | use serde::ser::SerializeMap; 11 | use serde::{Deserialize, Deserializer, Serialize, Serializer}; 12 | use std::borrow::Borrow; 13 | use std::fmt::{self, Debug, Formatter}; 14 | use std::iter::FromIterator; 15 | use std::marker::PhantomData; 16 | 17 | #[derive(Default)] 18 | pub struct Map { 19 | pub(crate) inner: RbTree, 20 | } 21 | 22 | impl Map { 23 | #[inline] 24 | pub fn new() -> Self { 25 | Self { 26 | inner: RbTree::new(), 27 | } 28 | } 29 | 30 | /// Returns `true` if the map does not contain any values. 31 | #[inline] 32 | pub fn is_empty(&self) -> bool { 33 | self.inner.is_empty() 34 | } 35 | 36 | /// Returns the number of elements in the map. 37 | #[inline] 38 | pub fn len(&self) -> usize { 39 | self.inner.len() 40 | } 41 | 42 | /// Clear the map. 43 | #[inline] 44 | pub fn clear(&mut self) { 45 | self.inner = RbTree::new(); 46 | } 47 | 48 | /// Insert a key-value pair into the map. Returns [`None`] if the key did not 49 | /// exists in the map, otherwise the previous value associated with the provided 50 | /// key will be returned. 51 | #[inline] 52 | pub fn insert(&mut self, key: K, value: V) -> Option { 53 | self.inner.insert(key, value).0 54 | } 55 | 56 | /// Remove the value associated with the given key from the map, returns the 57 | /// previous value associated with the key. 58 | #[inline] 59 | pub fn remove(&mut self, key: &Q) -> Option 60 | where 61 | K: Borrow, 62 | Q: Ord, 63 | { 64 | self.inner.delete(key).map(|(_, v)| v) 65 | } 66 | 67 | /// Remove an entry from the map and return the key and value. 68 | #[inline] 69 | pub fn remove_entry(&mut self, key: &Q) -> Option<(K, V)> 70 | where 71 | K: Borrow, 72 | Q: Ord, 73 | { 74 | self.inner.delete(key) 75 | } 76 | 77 | #[inline] 78 | pub fn entry(&mut self, key: K) -> Entry { 79 | self.inner.entry(key) 80 | } 81 | 82 | /// Returns a mutable reference to the value corresponding to the key. 83 | #[inline] 84 | pub fn get_mut(&mut self, key: &Q) -> Option<&mut V> 85 | where 86 | K: Borrow, 87 | Q: Ord, 88 | { 89 | self.inner.modify(key, |v| v) 90 | } 91 | 92 | /// Return the value associated with the given key. 93 | #[inline] 94 | pub fn get(&self, key: &Q) -> Option<&V> 95 | where 96 | K: Borrow, 97 | Q: Ord, 98 | { 99 | self.inner.get(key) 100 | } 101 | 102 | /// Return an iterator over the key-values in the map. 103 | #[inline] 104 | pub fn iter(&self) -> RbTreeIterator { 105 | RbTreeIterator::new(&self.inner) 106 | } 107 | 108 | /// Create a HashTree witness for the value associated with given key. 109 | #[inline] 110 | pub fn witness(&self, key: &Q) -> HashTree 111 | where 112 | K: Borrow, 113 | Q: Ord, 114 | { 115 | self.inner.witness(key) 116 | } 117 | 118 | /// Returns a witness enumerating all the keys in this map. The 119 | /// resulting tree doesn't include values, they are replaced with 120 | /// "Pruned" nodes. 121 | pub fn witness_keys(&self) -> HashTree { 122 | self.inner.keys() 123 | } 124 | 125 | /// Returns a witness for the key-value pairs in the specified range. 126 | /// The resulting tree contains both keys and values. 127 | #[inline] 128 | pub fn witness_value_range(&self, first: &K, last: &K) -> HashTree<'_> 129 | where 130 | K: Borrow + Borrow, 131 | Q1: Ord, 132 | Q2: Ord, 133 | { 134 | self.inner.value_range(first, last) 135 | } 136 | 137 | /// Returns a witness for the keys in the specified range. 138 | /// The resulting tree only contains the keys, and the values are replaced with 139 | /// "Pruned" nodes. 140 | #[inline] 141 | pub fn witness_key_range(&self, first: &K, last: &K) -> HashTree<'_> 142 | where 143 | K: Borrow + Borrow, 144 | Q1: Ord, 145 | Q2: Ord, 146 | { 147 | self.inner.key_range(first, last) 148 | } 149 | 150 | /// Returns a witness for the keys with the given prefix, this replaces the values with 151 | /// "Pruned" nodes. 152 | #[inline] 153 | pub fn witness_keys_with_prefix(&self, prefix: &P) -> HashTree<'_> 154 | where 155 | K: Prefix

, 156 | P: Ord, 157 | { 158 | self.inner.keys_with_prefix(prefix) 159 | } 160 | 161 | /// Return the underlying [`RbTree`] for this map. 162 | #[inline] 163 | pub fn as_tree(&self) -> &RbTree { 164 | &self.inner 165 | } 166 | } 167 | 168 | impl Map> { 169 | /// Perform a [`Seq::append`] on the seq associated with the give value, if 170 | /// the seq does not exists, creates an empty one and inserts it to the map. 171 | pub fn append_deep(&mut self, key: K, value: V) { 172 | let mut value = Some(value); 173 | 174 | self.inner.modify(&key, |seq| { 175 | seq.append(value.take().unwrap()); 176 | }); 177 | 178 | if let Some(value) = value.take() { 179 | let mut seq = Seq::new(); 180 | seq.append(value); 181 | self.inner.insert(key, seq); 182 | } 183 | } 184 | 185 | #[inline] 186 | pub fn len_deep(&mut self, key: &Q) -> usize 187 | where 188 | K: Borrow, 189 | Q: Ord, 190 | { 191 | self.inner.get(key).map(|seq| seq.len()).unwrap_or(0) 192 | } 193 | } 194 | 195 | impl AsRef> for Map { 196 | #[inline] 197 | fn as_ref(&self) -> &RbTree { 198 | &self.inner 199 | } 200 | } 201 | 202 | impl Serialize for Map 203 | where 204 | K: Serialize, 205 | V: Serialize, 206 | { 207 | fn serialize(&self, serializer: S) -> Result 208 | where 209 | S: Serializer, 210 | { 211 | let mut s = serializer.serialize_map(Some(self.len()))?; 212 | 213 | for (key, value) in self.iter() { 214 | s.serialize_entry(key, value)?; 215 | } 216 | 217 | s.end() 218 | } 219 | } 220 | 221 | impl<'de, K: 'static + Label, V: AsHashTree + 'static> Deserialize<'de> for Map 222 | where 223 | K: Deserialize<'de>, 224 | V: Deserialize<'de>, 225 | { 226 | fn deserialize(deserializer: D) -> Result 227 | where 228 | D: Deserializer<'de>, 229 | { 230 | deserializer.deserialize_map(MapVisitor(PhantomData::default())) 231 | } 232 | } 233 | 234 | struct MapVisitor(PhantomData<(K, V)>); 235 | 236 | impl<'de, K: 'static + Label, V: AsHashTree + 'static> Visitor<'de> for MapVisitor 237 | where 238 | K: Deserialize<'de>, 239 | V: Deserialize<'de>, 240 | { 241 | type Value = Map; 242 | 243 | fn expecting(&self, formatter: &mut Formatter) -> std::fmt::Result { 244 | write!(formatter, "expected a map") 245 | } 246 | 247 | fn visit_map(self, mut map: A) -> Result 248 | where 249 | A: MapAccess<'de>, 250 | { 251 | let mut result = Map::new(); 252 | 253 | loop { 254 | if let Some((key, value)) = map.next_entry::()? { 255 | result.insert(key, value); 256 | continue; 257 | } 258 | 259 | break; 260 | } 261 | 262 | Ok(result) 263 | } 264 | } 265 | 266 | impl CandidType for Map 267 | where 268 | K: CandidType, 269 | V: CandidType, 270 | { 271 | fn _ty() -> Type { 272 | let tuple = Type::Record(vec![ 273 | Field { 274 | id: CLabel::Id(0), 275 | ty: K::ty(), 276 | }, 277 | Field { 278 | id: CLabel::Id(1), 279 | ty: V::ty(), 280 | }, 281 | ]); 282 | Type::Vec(Box::new(tuple)) 283 | } 284 | 285 | fn idl_serialize(&self, serializer: S) -> Result<(), S::Error> 286 | where 287 | S: candid::types::Serializer, 288 | { 289 | let mut ser = serializer.serialize_vec(self.len())?; 290 | for e in self.iter() { 291 | Compound::serialize_element(&mut ser, &e)?; 292 | } 293 | Ok(()) 294 | } 295 | } 296 | 297 | impl FromIterator<(K, V)> for Map { 298 | fn from_iter>(iter: I) -> Self { 299 | let mut result = Map::new(); 300 | 301 | for (key, value) in iter { 302 | result.insert(key, value); 303 | } 304 | 305 | result 306 | } 307 | } 308 | 309 | impl Debug for Map 310 | where 311 | K: Debug, 312 | V: Debug, 313 | { 314 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { 315 | f.debug_map().entries(self.iter()).finish() 316 | } 317 | } 318 | 319 | impl AsHashTree for Map { 320 | #[inline] 321 | fn root_hash(&self) -> Hash { 322 | self.inner.root_hash() 323 | } 324 | 325 | #[inline] 326 | fn as_hash_tree(&self) -> HashTree<'_> { 327 | self.inner.as_hash_tree() 328 | } 329 | } 330 | 331 | #[cfg(test)] 332 | mod tests { 333 | use super::*; 334 | 335 | #[test] 336 | fn insert() { 337 | let mut map = Map::::new(); 338 | assert_eq!(map.insert("A".into(), 0), None); 339 | assert_eq!(map.insert("A".into(), 1), Some(0)); 340 | assert_eq!(map.insert("B".into(), 2), None); 341 | assert_eq!(map.insert("C".into(), 3), None); 342 | assert_eq!(map.insert("B".into(), 4), Some(2)); 343 | assert_eq!(map.insert("C".into(), 5), Some(3)); 344 | assert_eq!(map.insert("B".into(), 6), Some(4)); 345 | assert_eq!(map.insert("C".into(), 7), Some(5)); 346 | assert_eq!(map.insert("A".into(), 8), Some(1)); 347 | 348 | assert_eq!(map.get("A"), Some(&8)); 349 | assert_eq!(map.get("B"), Some(&6)); 350 | assert_eq!(map.get("C"), Some(&7)); 351 | assert_eq!(map.get("D"), None); 352 | } 353 | 354 | #[test] 355 | fn remove() { 356 | let mut map = Map::::new(); 357 | 358 | for i in 0..200u32 { 359 | map.insert(hex::encode(&i.to_be_bytes()), i); 360 | } 361 | 362 | for i in 0..200u32 { 363 | assert_eq!(map.remove(&hex::encode(&i.to_be_bytes())), Some(i)); 364 | } 365 | 366 | for i in 0..200u32 { 367 | assert_eq!(map.get(&hex::encode(&i.to_be_bytes())), None); 368 | } 369 | } 370 | 371 | #[test] 372 | fn remove_rev() { 373 | let mut map = Map::::new(); 374 | 375 | for i in 0..200u32 { 376 | map.insert(hex::encode(&i.to_be_bytes()), i); 377 | } 378 | 379 | for i in (0..200u32).rev() { 380 | assert_eq!(map.remove(&hex::encode(&i.to_be_bytes())), Some(i)); 381 | } 382 | 383 | for i in 0..200u32 { 384 | assert_eq!(map.get(&hex::encode(&i.to_be_bytes())), None); 385 | } 386 | } 387 | } 388 | --------------------------------------------------------------------------------