randomx_rust_wrapper/
dataset.rs

1/*
2 * Copyright 2024 Fluence Labs Limited
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
26#[derive(Debug)]
27pub struct Dataset {
28    inner: Arc<DatasetInner>,
29}
30
31#[derive(Debug)]
32struct DatasetInner {
33    dataset: *mut randomx_dataset,
34}
35
36unsafe impl Send for DatasetInner {}
37unsafe impl Sync for DatasetInner {}
38
39/// Contains a handle of Dataset, can't be created from scratch,
40/// only obtained from already existing Cache.
41#[derive(Clone, Debug)]
42pub struct DatasetHandle {
43    inner: Arc<DatasetInner>,
44}
45
46impl Dataset {
47    /// Allocate and initialize a new database with provided global nonce and flags.
48    /// Only RANDOMX_FLAG_LARGE_PAGES is supported (can be set or unset),
49    /// it forces memory allocation in large pages.
50    pub fn new(global_nonce: &[u8], flags: RandomXFlags) -> RResult<Self> {
51        let cache = Cache::new(global_nonce, flags)?;
52        Self::from_cache(&cache, flags.contains(RandomXFlags::LARGE_PAGES))
53    }
54
55    /// Allocate and initialize a new database with the provided cache,
56    /// large_pages_enabled forces it to allocate memory in large pages.
57    pub fn from_cache(cache: &Cache, large_pages_enabled: bool) -> RResult<Self> {
58        let mut dataset = Self::allocate(large_pages_enabled)?;
59        let items_count = dataset.items_count();
60        dataset.initialize(cache, 0, items_count);
61
62        Ok(dataset)
63    }
64
65    /// Allocate a new dataset, but doesn't initialize it.
66    pub fn allocate(large_pages_enabled: bool) -> RResult<Self> {
67        let flags = if large_pages_enabled {
68            RandomXFlags::LARGE_PAGES
69        } else {
70            RandomXFlags::default()
71        };
72
73        let dataset =
74            try_alloc! { randomx_alloc_dataset(flags.bits()), DatasetAllocationError { flags } };
75        let dataset_inner = DatasetInner { dataset };
76        let dataset = Self {
77            inner: Arc::new(dataset_inner),
78        };
79        Ok(dataset)
80    }
81
82    /// Return a number of elements that a dataset could contain.
83    pub fn items_count(&self) -> u64 {
84        unsafe { randomx_dataset_item_count() }
85    }
86
87    /// Initialize dataset with the provided cache.
88    pub fn initialize(&mut self, cache: &impl CacheRawAPI, start_item: u64, items_count: u64) {
89        unsafe { randomx_init_dataset(self.raw(), cache.raw(), start_item, items_count) };
90    }
91
92    pub fn handle(&self) -> DatasetHandle {
93        DatasetHandle {
94            inner: self.inner.clone(),
95        }
96    }
97
98    pub(crate) fn raw(&self) -> *mut randomx_dataset {
99        self.inner.dataset
100    }
101}
102
103impl DatasetHandle {
104    /// Return a number of elements that a dataset could contain.
105    pub fn items_count(&self) -> u64 {
106        unsafe { randomx_dataset_item_count() }
107    }
108
109    /// Initialize dataset with the provided cache.
110    pub fn initialize(&mut self, cache: &impl CacheRawAPI, start_item: u64, items_count: u64) {
111        unsafe { randomx_init_dataset(self.raw(), cache.raw(), start_item, items_count) };
112    }
113
114    pub(crate) fn raw(&self) -> *mut randomx_dataset {
115        self.inner.dataset
116    }
117}
118
119impl Drop for DatasetInner {
120    fn drop(&mut self) {
121        unsafe { randomx_release_dataset(self.dataset) }
122    }
123}
124
125pub trait DatasetRawAPI {
126    fn raw(&self) -> *mut randomx_dataset;
127}
128
129impl DatasetRawAPI for Dataset {
130    fn raw(&self) -> *mut randomx_dataset {
131        self.raw()
132    }
133}
134
135impl DatasetRawAPI for DatasetHandle {
136    fn raw(&self) -> *mut randomx_dataset {
137        self.raw()
138    }
139}
randomx_rust_wrapper/dataset.rs

randomx_rust_wrapper/
dataset.rs
