ccp_randomx/
dataset.rs

1/*
2 * Copyright 2024 Fluence DAO
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 *     http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17use std::sync::Arc;
18
19use crate::bindings::dataset::*;
20use crate::cache::{Cache, CacheRawAPI};
21use crate::errors::RandomXError::DatasetAllocationError;
22use crate::flags::RandomXFlags;
23use crate::try_alloc;
24use crate::RResult;
25
26pub const DATASET_ITEM_SIZE: u64 = 64;
27
28#[derive(Debug)]
29pub struct Dataset {
30    inner: Arc<DatasetInner>,
31}
32
33#[derive(Debug)]
34struct DatasetInner {
35    dataset: *mut randomx_dataset,
36}
37
38unsafe impl Send for DatasetInner {}
39unsafe impl Sync for DatasetInner {}
40
41/// Contains a handle of Dataset, can't be created from scratch,
42/// only obtained from already existing Cache.
43#[derive(Clone, Debug)]
44pub struct DatasetHandle {
45    inner: Arc<DatasetInner>,
46}
47
48impl Dataset {
49    /// Allocate and initialize a new database with provided global nonce and flags.
50    /// Only RANDOMX_FLAG_LARGE_PAGES is supported (can be set or unset),
51    /// it forces memory allocation in large pages.
52    pub fn new(global_nonce: &[u8], flags: RandomXFlags) -> RResult<Self> {
53        let cache = Cache::new(global_nonce, flags)?;
54        Self::from_cache(&cache, flags.contains(RandomXFlags::LARGE_PAGES))
55    }
56
57    /// Allocate and initialize a new database with the provided cache,
58    /// large_pages_enabled forces it to allocate memory in large pages.
59    pub fn from_cache(cache: &Cache, large_pages_enabled: bool) -> RResult<Self> {
60        let mut dataset = Self::allocate(large_pages_enabled)?;
61        let items_count = dataset.items_count();
62        dataset.initialize(cache, 0, items_count);
63
64        Ok(dataset)
65    }
66
67    /// Allocate a new dataset, but doesn't initialize it.
68    pub fn allocate(large_pages_enabled: bool) -> RResult<Self> {
69        let flags = if large_pages_enabled {
70            RandomXFlags::LARGE_PAGES
71        } else {
72            RandomXFlags::default()
73        };
74
75        let dataset =
76            try_alloc! { randomx_alloc_dataset(flags.bits()), DatasetAllocationError { flags } };
77        let dataset_inner = DatasetInner { dataset };
78        let dataset = Self {
79            inner: Arc::new(dataset_inner),
80        };
81        Ok(dataset)
82    }
83
84    /// Return a number of elements that a dataset could contain.
85    pub fn items_count(&self) -> u64 {
86        unsafe { randomx_dataset_item_count() }
87    }
88
89    /// Initialize dataset with the provided cache.
90    pub fn initialize(&mut self, cache: &impl CacheRawAPI, start_item: u64, items_count: u64) {
91        unsafe { randomx_init_dataset(self.raw(), cache.raw(), start_item, items_count) };
92    }
93
94    pub fn handle(&self) -> DatasetHandle {
95        DatasetHandle {
96            inner: self.inner.clone(),
97        }
98    }
99
100    pub(crate) fn raw(&self) -> *mut randomx_dataset {
101        self.inner.dataset
102    }
103}
104
105impl DatasetHandle {
106    /// Return a number of elements that a dataset could contain.
107    pub fn items_count(&self) -> u64 {
108        unsafe { randomx_dataset_item_count() }
109    }
110
111    /// Initialize dataset with the provided cache.
112    pub fn initialize(&mut self, cache: &impl CacheRawAPI, start_item: u64, items_count: u64) {
113        unsafe { randomx_init_dataset(self.raw(), cache.raw(), start_item, items_count) };
114    }
115
116    pub(crate) fn raw(&self) -> *mut randomx_dataset {
117        self.inner.dataset
118    }
119
120    pub fn memory(&self) -> *mut ::std::os::raw::c_void {
121        unsafe { randomx_get_dataset_memory(self.raw()) }
122    }
123}
124
125impl Drop for DatasetInner {
126    fn drop(&mut self) {
127        unsafe { randomx_release_dataset(self.dataset) }
128    }
129}
130
131pub trait DatasetRawAPI {
132    fn raw(&self) -> *mut randomx_dataset;
133}
134
135impl DatasetRawAPI for Dataset {
136    fn raw(&self) -> *mut randomx_dataset {
137        self.raw()
138    }
139}
140
141impl DatasetRawAPI for DatasetHandle {
142    fn raw(&self) -> *mut randomx_dataset {
143        self.raw()
144    }
145}