1use std::future::IntoFuture;
2
3use n0_future::{stream, StreamExt};
4use rand::{RngCore, SeedableRng};
5
6use crate::{
7 api::{blobs::AddBytesOptions, tags::TagInfo, RequestResult, Store},
8 hashseq::HashSeq,
9 BlobFormat,
10};
11
12pub async fn create_random_blobs<R: rand::Rng>(
13 store: &Store,
14 num_blobs: usize,
15 blob_size: impl Fn(usize, &mut R) -> usize,
16 mut rand: R,
17) -> anyhow::Result<Vec<TagInfo>> {
18 let sizes = (0..num_blobs)
20 .map(|n| (blob_size(n, &mut rand), rand.r#gen::<u64>()))
21 .collect::<Vec<_>>();
22 let infos = stream::iter(sizes)
24 .then(|(size, seed)| {
25 let mut rand = rand::rngs::StdRng::seed_from_u64(seed);
26 let mut data = vec![0u8; size];
27 rand.fill_bytes(&mut data);
28 store.add_bytes(data).into_future()
29 })
30 .collect::<Vec<_>>()
31 .await;
32 let infos = infos.into_iter().collect::<RequestResult<Vec<_>>>()?;
33 Ok(infos)
34}
35
36pub async fn add_hash_sequences<R: rand::Rng>(
37 store: &Store,
38 tags: &[TagInfo],
39 num_seqs: usize,
40 seq_size: impl Fn(usize, &mut R) -> usize,
41 mut rand: R,
42) -> anyhow::Result<Vec<TagInfo>> {
43 let infos = stream::iter(0..num_seqs)
44 .then(|n| {
45 let size = seq_size(n, &mut rand);
46 let hs = (0..size)
47 .map(|_| {
48 let j = rand.gen_range(0..tags.len());
49 tags[j].hash
50 })
51 .collect::<HashSeq>();
52 store
53 .add_bytes_with_opts(AddBytesOptions {
54 data: hs.into(),
55 format: BlobFormat::HashSeq,
56 })
57 .into_future()
58 })
59 .collect::<Vec<_>>()
60 .await;
61 let infos = infos.into_iter().collect::<RequestResult<Vec<_>>>()?;
62 Ok(infos)
63}