gear_common/
program_storage.rs

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