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::{WasmPage, numerated::tree::IntervalsTree};
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 = ActorId, Value = Program<Self::BlockNumber>>;
55    type MemoryPageMap: TripleMapStorage<Key1 = ActorId, Key2 = MemoryInfix, Key3 = GearPage, Value = PageBuf>;
56    type AllocationsMap: MapStorage<Key = ActorId, Value = IntervalsTree<WasmPage>>;
57
58    /// Attempt to remove all items from all the associated maps.
59    fn reset() {
60        Self::ProgramMap::clear();
61        Self::MemoryPageMap::clear();
62        Self::AllocationsMap::clear();
63    }
64
65    /// Store a program to be associated with the given key `program_id` from the map.
66    fn add_program(
67        program_id: ActorId,
68        program: ActiveProgram<Self::BlockNumber>,
69    ) -> Result<(), Self::Error> {
70        Self::ProgramMap::mutate(program_id, |maybe| {
71            if maybe.is_some() {
72                return Err(Self::InternalError::duplicate_item().into());
73            }
74
75            *maybe = Some(Program::Active(program));
76            Ok(())
77        })
78    }
79
80    /// Load the program associated with the given key `program_id` from the map.
81    fn get_program(program_id: ActorId) -> Option<Program<Self::BlockNumber>> {
82        Self::ProgramMap::get(&program_id)
83    }
84
85    /// Does the program (explicitly) exist in storage?
86    fn program_exists(program_id: ActorId) -> bool {
87        Self::ProgramMap::contains_key(&program_id)
88    }
89
90    /// Update the active program under the given key `program_id`.
91    fn update_active_program<F, ReturnType>(
92        program_id: ActorId,
93        update_action: F,
94    ) -> Result<ReturnType, Self::Error>
95    where
96        F: FnOnce(&mut ActiveProgram<Self::BlockNumber>) -> ReturnType,
97    {
98        Self::update_program_if_active(program_id, |program, _bn| match program {
99            Program::Active(active_program) => update_action(active_program),
100            _ => unreachable!("invariant kept by update_program_if_active"),
101        })
102    }
103
104    fn remove_data_for_pages(
105        program_id: ActorId,
106        memory_infix: MemoryInfix,
107        pages: impl Iterator<Item = GearPage>,
108    ) {
109        for page in pages {
110            Self::remove_program_page_data(program_id, memory_infix, page);
111        }
112    }
113
114    fn allocations(program_id: ActorId) -> Option<IntervalsTree<WasmPage>> {
115        Self::AllocationsMap::get(&program_id)
116    }
117
118    fn set_allocations(program_id: ActorId, allocations: IntervalsTree<WasmPage>) {
119        Self::update_active_program(program_id, |program| {
120            program.allocations_tree_len = u32::try_from(allocations.intervals_amount())
121                .unwrap_or_else(|err| {
122                    // This panic is impossible because page numbers are u32.
123                    unreachable!("allocations tree length is too big to fit into u32: {err}")
124                });
125        })
126        .unwrap_or_else(|err| {
127            // set_allocations must be called only for active programs.
128            unreachable!("Failed to update program allocations: {err:?}")
129        });
130        Self::AllocationsMap::insert(program_id, allocations);
131    }
132
133    fn clear_allocations(program_id: ActorId) {
134        Self::AllocationsMap::remove(program_id);
135    }
136
137    fn memory_infix(program_id: ActorId) -> Option<MemoryInfix> {
138        match Self::ProgramMap::get(&program_id) {
139            Some(Program::Active(program)) => Some(program.memory_infix),
140            _ => None,
141        }
142    }
143
144    /// Update the program under the given key `program_id` only if the
145    /// stored program is an active one.
146    fn update_program_if_active<F, ReturnType>(
147        program_id: ActorId,
148        update_action: F,
149    ) -> Result<ReturnType, Self::Error>
150    where
151        F: FnOnce(&mut Program<Self::BlockNumber>, Self::BlockNumber) -> ReturnType,
152    {
153        let mut program =
154            Self::ProgramMap::get(&program_id).ok_or(Self::InternalError::program_not_found())?;
155        let bn = match program {
156            Program::Active(ref p) => p.expiration_block,
157            _ => return Err(Self::InternalError::not_active_program().into()),
158        };
159
160        let result = update_action(&mut program, bn);
161        Self::ProgramMap::insert(program_id, program);
162
163        Ok(result)
164    }
165
166    /// Return data buffer for each memory page, which has data.
167    fn get_program_pages_data(
168        program_id: ActorId,
169        memory_infix: MemoryInfix,
170    ) -> Result<MemoryMap, Self::Error> {
171        Ok(Self::MemoryPageMap::iter_prefix(&program_id, &memory_infix).collect())
172    }
173
174    /// Store a memory page buffer to be associated with the given keys `program_id`, `memory_infix` and `page` from the map.
175    fn set_program_page_data(
176        program_id: ActorId,
177        memory_infix: MemoryInfix,
178        page: GearPage,
179        page_buf: PageBuf,
180    ) {
181        Self::MemoryPageMap::insert(program_id, memory_infix, page, page_buf);
182    }
183
184    /// Remove a memory page buffer under the given keys `program_id`, `memory_infix` and `page`.
185    fn remove_program_page_data(
186        program_id: ActorId,
187        memory_infix: MemoryInfix,
188        page_num: GearPage,
189    ) {
190        Self::MemoryPageMap::remove(program_id, memory_infix, page_num);
191    }
192
193    /// Remove all memory page buffers under the given keys `program_id` and `memory_infix`.
194    fn clear_program_memory(program_id: ActorId, memory_infix: MemoryInfix) {
195        Self::MemoryPageMap::clear_prefix(program_id, memory_infix);
196    }
197
198    /// Final full prefix that prefixes all keys of memory pages.
199    fn pages_final_prefix() -> [u8; 32];
200}