Skip to main content

gear_common/
program_storage.rs

1// Copyright (C) Gear Technologies Inc.
2// SPDX-License-Identifier: GPL-3.0-or-later WITH Classpath-exception-2.0
3
4use gear_core::pages::{WasmPage, numerated::tree::IntervalsTree};
5
6use super::*;
7use crate::storage::{MapStorage, TripleMapStorage};
8use core::fmt::Debug;
9
10/// Trait for ProgramStorage errors.
11///
12/// Contains constructors for all existing errors.
13pub trait Error {
14    /// Program already exists in storage.
15    fn duplicate_item() -> Self;
16
17    /// Program is not found in the storage.
18    fn program_not_found() -> Self;
19
20    /// Program is not an instance of ActiveProgram.
21    fn not_active_program() -> Self;
22
23    /// There is no data for specified `program_id` and `page`.
24    fn cannot_find_page_data() -> Self;
25
26    /// Failed to find the program binary code.
27    fn program_code_not_found() -> Self;
28}
29
30pub type MemoryMap = BTreeMap<GearPage, PageBuf>;
31
32/// Trait to work with program data in a storage.
33pub trait ProgramStorage {
34    type InternalError: Error;
35    type Error: From<Self::InternalError> + Debug;
36    type BlockNumber: Copy + Saturating;
37    type AccountId: Eq + PartialEq;
38
39    type ProgramMap: MapStorage<Key = ActorId, Value = Program<Self::BlockNumber>>;
40    type MemoryPageMap: TripleMapStorage<Key1 = ActorId, Key2 = MemoryInfix, Key3 = GearPage, Value = PageBuf>;
41    type AllocationsMap: MapStorage<Key = ActorId, Value = IntervalsTree<WasmPage>>;
42
43    /// Attempt to remove all items from all the associated maps.
44    fn reset() {
45        Self::ProgramMap::clear();
46        Self::MemoryPageMap::clear();
47        Self::AllocationsMap::clear();
48    }
49
50    /// Store a program to be associated with the given key `program_id` from the map.
51    fn add_program(
52        program_id: ActorId,
53        program: ActiveProgram<Self::BlockNumber>,
54    ) -> Result<(), Self::Error> {
55        Self::ProgramMap::mutate(program_id, |maybe| {
56            if maybe.is_some() {
57                return Err(Self::InternalError::duplicate_item().into());
58            }
59
60            *maybe = Some(Program::Active(program));
61            Ok(())
62        })
63    }
64
65    /// Load the program associated with the given key `program_id` from the map.
66    fn get_program(program_id: ActorId) -> Option<Program<Self::BlockNumber>> {
67        Self::ProgramMap::get(&program_id)
68    }
69
70    /// Does the program (explicitly) exist in storage?
71    fn program_exists(program_id: ActorId) -> bool {
72        Self::ProgramMap::contains_key(&program_id)
73    }
74
75    /// Update the active program under the given key `program_id`.
76    fn update_active_program<F, ReturnType>(
77        program_id: ActorId,
78        update_action: F,
79    ) -> Result<ReturnType, Self::Error>
80    where
81        F: FnOnce(&mut ActiveProgram<Self::BlockNumber>) -> ReturnType,
82    {
83        Self::update_program_if_active(program_id, |program, _bn| match program {
84            Program::Active(active_program) => update_action(active_program),
85            _ => unreachable!("invariant kept by update_program_if_active"),
86        })
87    }
88
89    fn remove_data_for_pages(
90        program_id: ActorId,
91        memory_infix: MemoryInfix,
92        pages: impl Iterator<Item = GearPage>,
93    ) {
94        for page in pages {
95            Self::remove_program_page_data(program_id, memory_infix, page);
96        }
97    }
98
99    fn allocations(program_id: ActorId) -> Option<IntervalsTree<WasmPage>> {
100        Self::AllocationsMap::get(&program_id)
101    }
102
103    fn set_allocations(program_id: ActorId, allocations: IntervalsTree<WasmPage>) {
104        Self::update_active_program(program_id, |program| {
105            program.allocations_tree_len = u32::try_from(allocations.intervals_amount())
106                .unwrap_or_else(|err| {
107                    // This panic is impossible because page numbers are u32.
108                    unreachable!("allocations tree length is too big to fit into u32: {err}")
109                });
110        })
111        .unwrap_or_else(|err| {
112            // set_allocations must be called only for active programs.
113            unreachable!("Failed to update program allocations: {err:?}")
114        });
115        Self::AllocationsMap::insert(program_id, allocations);
116    }
117
118    fn clear_allocations(program_id: ActorId) {
119        Self::AllocationsMap::remove(program_id);
120    }
121
122    fn memory_infix(program_id: ActorId) -> Option<MemoryInfix> {
123        match Self::ProgramMap::get(&program_id) {
124            Some(Program::Active(program)) => Some(program.memory_infix),
125            _ => None,
126        }
127    }
128
129    /// Update the program under the given key `program_id` only if the
130    /// stored program is an active one.
131    fn update_program_if_active<F, ReturnType>(
132        program_id: ActorId,
133        update_action: F,
134    ) -> Result<ReturnType, Self::Error>
135    where
136        F: FnOnce(&mut Program<Self::BlockNumber>, Self::BlockNumber) -> ReturnType,
137    {
138        let mut program =
139            Self::ProgramMap::get(&program_id).ok_or(Self::InternalError::program_not_found())?;
140        let bn = match program {
141            Program::Active(ref p) => p.expiration_block,
142            _ => return Err(Self::InternalError::not_active_program().into()),
143        };
144
145        let result = update_action(&mut program, bn);
146        Self::ProgramMap::insert(program_id, program);
147
148        Ok(result)
149    }
150
151    /// Return data buffer for each memory page, which has data.
152    fn get_program_pages_data(
153        program_id: ActorId,
154        memory_infix: MemoryInfix,
155    ) -> Result<MemoryMap, Self::Error> {
156        Ok(Self::MemoryPageMap::iter_prefix(&program_id, &memory_infix).collect())
157    }
158
159    /// Store a memory page buffer to be associated with the given keys `program_id`, `memory_infix` and `page` from the map.
160    fn set_program_page_data(
161        program_id: ActorId,
162        memory_infix: MemoryInfix,
163        page: GearPage,
164        page_buf: PageBuf,
165    ) {
166        Self::MemoryPageMap::insert(program_id, memory_infix, page, page_buf);
167    }
168
169    /// Remove a memory page buffer under the given keys `program_id`, `memory_infix` and `page`.
170    fn remove_program_page_data(
171        program_id: ActorId,
172        memory_infix: MemoryInfix,
173        page_num: GearPage,
174    ) {
175        Self::MemoryPageMap::remove(program_id, memory_infix, page_num);
176    }
177
178    /// Remove all memory page buffers under the given keys `program_id` and `memory_infix`.
179    fn clear_program_memory(program_id: ActorId, memory_infix: MemoryInfix) {
180        Self::MemoryPageMap::clear_prefix(program_id, memory_infix);
181    }
182
183    /// Final full prefix that prefixes all keys of memory pages.
184    fn pages_final_prefix() -> [u8; 32];
185}