sachy.dev
A pit full of rusty nails
1use std::sync::Arc;
2
3use axum::extract::MatchedPath;
4use bytes::{Bytes, BytesMut};
5use nailconfig::NailConfig;
6use nailkov::NailKov;
7use nailrng::FastRng;
8use rand::{RngExt, Rng, distr::Alphanumeric, seq::IndexedRandom};
9
10/// Provides either the minimum configured size, or a randomised value between
11/// the minimum and maximum configured sizes if a maximum is available.
12#[inline]
13fn get_desired_size(config: &NailConfig, rng: &mut impl Rng) -> usize {
14 match (
15 config.generator.min_paragraph_size,
16 config.generator.max_paragraph_size,
17 ) {
18 (min, None) => min,
19 (min, Some(max)) => rng.random_range(min..=max),
20 }
21}
22
23/// Generates text from the markov chain, using the tokens it outputs to pull
24/// interned text from the interner.
25#[inline]
26pub fn text_generator<'a>(
27 chain: &'a NailKov,
28 size: usize,
29 rng: &'a mut impl Rng,
30) -> impl Iterator<Item = &'a u8> + 'a {
31 chain
32 .generate_tokens(rng)
33 .take(size)
34 // SAFETY: The id comes from the same interner that allocated it
35 .flat_map(|token| token.as_bytes())
36 .skip_while(|&text| !text.is_ascii_alphabetic())
37}
38
39#[inline]
40pub fn static_title<'a>(text: &'a str) -> impl Iterator<Item = &'a u8> + 'a {
41 text.lines()
42 .map(str::trim)
43 .next()
44 .into_iter()
45 .flat_map(str::as_bytes)
46}
47
48#[inline]
49pub fn static_content<'a>(text: &'a str) -> impl Iterator<Item = &'a u8> + 'a {
50 text.lines()
51 .skip(1)
52 .filter_map(|line| {
53 let trimmed = line.trim();
54
55 if line.is_empty() {
56 None
57 } else {
58 Some(trimmed.as_bytes())
59 }
60 })
61 .flat_map(|bytes| b"<p>".iter().chain(bytes).chain(b"</p>\n"))
62}
63
64pub async fn initial_content(
65 buf_mut: BytesMut,
66 chain: Arc<NailKov>,
67 config: Arc<NailConfig>,
68 mut rng: FastRng,
69) -> Bytes {
70 // Randomise how many initial paragraphs we want
71 let max_paras: u32 = rng.random_range(1..=3);
72
73 (0..max_paras)
74 .fold(buf_mut, |mut acc, _| {
75 acc.extend(paragraph(
76 &chain,
77 get_desired_size(&config, &mut rng),
78 &mut rng,
79 ));
80
81 acc
82 })
83 .freeze()
84}
85
86pub async fn main_content(
87 mut buffer: BytesMut,
88 chain: Arc<NailKov>,
89 config: Arc<NailConfig>,
90 mut rng: FastRng,
91) -> Bytes {
92 buffer.reserve(config.generator.chunk_size * 2);
93
94 loop {
95 buffer.extend(header(&chain, config.generator.header_size, &mut rng));
96
97 // Randomise how many paragraphs we want per section
98 let paragraphs = rng.random_range(1..=4);
99
100 (0..paragraphs).for_each(|_| {
101 buffer.extend(paragraph(
102 &chain,
103 get_desired_size(&config, &mut rng),
104 &mut rng,
105 ));
106 });
107
108 // We can generate more before handing it off to be streamed to the client,
109 // A bit more latency, but much more throughput, and friendlier to being compressed.
110 if buffer.len() >= config.generator.chunk_size {
111 return buffer.freeze();
112 }
113
114 // Yield to the runtime to allow other tasks a chance to run before we generate
115 // another chunk of data
116 futures_lite::future::yield_now().await;
117 }
118}
119
120#[inline]
121pub fn extra(buf_mut: &mut BytesMut, config: &NailConfig, rng: &mut FastRng) -> usize {
122 let mut written = 0;
123
124 if let Some(prompt) = match config.generator.prompts.len() {
125 0 => None,
126 1 => config.generator.prompts.first(),
127 _ => config.generator.prompts.choose(rng),
128 } {
129 buf_mut.extend(b"<p>".iter().chain(prompt.as_bytes()).chain(b"</p>"));
130
131 written += prompt.len();
132 }
133
134 written
135}
136
137pub async fn footer(
138 mut buf_mut: BytesMut,
139 chain: Arc<NailKov>,
140 path: MatchedPath,
141 config: Arc<NailConfig>,
142 mut rng: FastRng,
143) -> Bytes {
144 let path = path.as_str();
145
146 let route = path.strip_suffix("/{*generated}").unwrap_or(path);
147
148 let total_links = rng.random_range(1..=config.generator.max_pit_links);
149
150 buf_mut.extend_from_slice(b"<nav style=\"visibility: hidden;\"><ul>");
151
152 for _ in 1..=total_links {
153 buf_mut.extend(b"<li><a href=\"".iter().chain(route.as_bytes()).chain(b"/"));
154 buf_mut.extend((&mut rng).sample_iter(Alphanumeric).take(16));
155 buf_mut.extend(
156 b"\">"
157 .iter()
158 .chain(text_generator(&chain, 8, &mut rng))
159 .chain(b"</a></li>\n"),
160 );
161 }
162
163 buf_mut.extend_from_slice(b"</ul></nav>");
164
165 buf_mut.freeze()
166}
167
168#[inline]
169fn paragraph<'a>(
170 chain: &'a NailKov,
171 size: usize,
172 rng: &'a mut impl Rng,
173) -> impl Iterator<Item = &'a u8> + 'a {
174 b"<p>"
175 .iter()
176 .chain(text_generator(chain, size, rng))
177 .chain(b"</p>\n")
178}
179
180#[inline]
181fn header<'a>(
182 chain: &'a NailKov,
183 size: usize,
184 rng: &'a mut impl Rng,
185) -> impl Iterator<Item = &'a u8> + 'a {
186 b"\n<h2>"
187 .iter()
188 .chain(text_generator(chain, size, rng))
189 .chain(b"</h2>\n")
190}