hegel/generators/
strings.rs

1use std::sync::OnceLock;
2
3use super::{Generator, TestCase, labels};
4use crate::control::hegel_internal_assert;
5use crate::ffi;
6use crate::test_case::{full_ranges, invalid_argument};
7
8/// Categories that include surrogate codepoints. Rust strings cannot contain
9/// surrogates, so these are forbidden in `categories()`.
10const SURROGATE_CATEGORIES: &[&str] = &["Cs", "C"];
11
12/// Default upper bound for string/byte sizes when the caller doesn't set one.
13const DEFAULT_MAX_SIZE: usize = 100;
14
15/// Shared character filtering fields used by both [`TextGenerator`] and
16/// [`CharactersGenerator`].
17struct CharacterFields {
18    codec: Option<String>,
19    min_codepoint: Option<u32>,
20    max_codepoint: Option<u32>,
21    categories: Option<Vec<String>>,
22    exclude_categories: Option<Vec<String>>,
23    include_characters: Option<String>,
24    exclude_characters: Option<String>,
25}
26
27impl CharacterFields {
28    fn new() -> Self {
29        CharacterFields {
30            codec: None,
31            min_codepoint: None,
32            max_codepoint: None,
33            categories: None,
34            exclude_categories: None,
35            include_characters: None,
36            exclude_characters: None,
37        }
38    }
39
40    /// Build a libhegel text generator over this alphabet with the given
41    /// size bounds. Surrogates are always excluded: category restrictions
42    /// naming a surrogate category are rejected, and without explicit
43    /// categories the `Cs` category is excluded (Rust strings cannot hold
44    /// surrogates).
45    fn build_text_handle(&self, min_size: u64, max_size: u64) -> ffi::StringGenerator {
46        let (categories, exclude_categories) = if let Some(ref cats) = self.categories {
47            for cat in cats {
48                if SURROGATE_CATEGORIES.contains(&cat.as_str()) {
49                    invalid_argument!(
50                        "Category \"{cat}\" includes surrogate codepoints (Cs), \
51                         which Rust strings cannot represent."
52                    );
53                }
54            }
55            (Some(cats.clone()), None)
56        } else {
57            let mut excl = self.exclude_categories.clone().unwrap_or_default();
58            if !excl.iter().any(|c| c == "Cs") {
59                excl.push("Cs".to_string());
60            }
61            (None, Some(excl))
62        };
63        ffi::StringGenerator::text(
64            min_size,
65            max_size,
66            self.codec.as_deref(),
67            self.min_codepoint.unwrap_or(0),
68            self.max_codepoint,
69            categories.as_deref(),
70            exclude_categories.as_deref(),
71            self.include_characters.as_deref(),
72            self.exclude_characters.as_deref(),
73        )
74        .unwrap_or_else(|msg| invalid_argument!("{msg}"))
75    }
76}
77
78/// Generator for Unicode text strings. Created by [`text()`].
79pub struct TextGenerator {
80    min_size: usize,
81    max_size: Option<usize>,
82    char_fields: CharacterFields,
83    alphabet_called: bool,
84    char_param_called: bool,
85    handle: OnceLock<ffi::StringGenerator>,
86}
87
88impl TextGenerator {
89    /// Set the minimum length in characters.
90    pub fn min_size(mut self, min_size: usize) -> Self {
91        self.handle = OnceLock::new();
92        self.min_size = min_size;
93        self
94    }
95
96    /// Set the maximum length in characters.
97    pub fn max_size(mut self, max_size: usize) -> Self {
98        self.handle = OnceLock::new();
99        self.max_size = Some(max_size);
100        self
101    }
102
103    /// Use a fixed set of characters. Each character in the generated string
104    /// will be a member of the alphabet.
105    ///
106    /// Mutually exclusive with the character filtering methods like `codec`,
107    /// `categories`, `min_codepoint`, etc.
108    pub fn alphabet(mut self, chars: &str) -> Self {
109        self.handle = OnceLock::new();
110        self.char_fields = CharacterFields {
111            codec: None,
112            min_codepoint: None,
113            max_codepoint: None,
114            categories: Some(vec![]),
115            exclude_categories: None,
116            include_characters: Some(chars.to_string()),
117            exclude_characters: None,
118        };
119        self.alphabet_called = true;
120        self
121    }
122
123    /// Restrict to characters encodable in this codec (e.g. `"ascii"`, `"utf-8"`, `"latin-1"`).
124    pub fn codec(mut self, codec: &str) -> Self {
125        self.handle = OnceLock::new();
126        self.char_param_called = true;
127        self.char_fields.codec = Some(codec.to_string());
128        self
129    }
130
131    /// Set the minimum Unicode codepoint.
132    pub fn min_codepoint(mut self, min_codepoint: u32) -> Self {
133        self.handle = OnceLock::new();
134        self.char_param_called = true;
135        self.char_fields.min_codepoint = Some(min_codepoint);
136        self
137    }
138
139    /// Set the maximum Unicode codepoint.
140    pub fn max_codepoint(mut self, max_codepoint: u32) -> Self {
141        self.handle = OnceLock::new();
142        self.char_param_called = true;
143        self.char_fields.max_codepoint = Some(max_codepoint);
144        self
145    }
146
147    /// Include only characters from these Unicode general categories (e.g. `["L", "Nd"]`).
148    ///
149    /// Mutually exclusive with [`exclude_categories`](Self::exclude_categories).
150    pub fn categories(mut self, categories: &[&str]) -> Self {
151        self.handle = OnceLock::new();
152        self.char_param_called = true;
153        self.char_fields.categories = Some(categories.iter().map(|s| s.to_string()).collect());
154        self
155    }
156
157    /// Exclude characters from these Unicode general categories.
158    ///
159    /// Mutually exclusive with [`categories`](Self::categories).
160    pub fn exclude_categories(mut self, exclude_categories: &[&str]) -> Self {
161        self.handle = OnceLock::new();
162        self.char_param_called = true;
163        self.char_fields.exclude_categories =
164            Some(exclude_categories.iter().map(|s| s.to_string()).collect());
165        self
166    }
167
168    /// Always include these specific characters, even if excluded by other filters.
169    pub fn include_characters(mut self, include_characters: &str) -> Self {
170        self.handle = OnceLock::new();
171        self.char_param_called = true;
172        self.char_fields.include_characters = Some(include_characters.to_string());
173        self
174    }
175
176    /// Always exclude these specific characters.
177    pub fn exclude_characters(mut self, exclude_characters: &str) -> Self {
178        self.handle = OnceLock::new();
179        self.char_param_called = true;
180        self.char_fields.exclude_characters = Some(exclude_characters.to_string());
181        self
182    }
183
184    fn handle(&self) -> &ffi::StringGenerator {
185        self.handle.get_or_init(|| {
186            if self.alphabet_called && self.char_param_called {
187                invalid_argument!("Cannot combine .alphabet() with character methods.");
188            }
189            if let Some(max) = self.max_size {
190                if self.min_size > max {
191                    invalid_argument!("Cannot have max_size < min_size");
192                }
193            }
194            let max_size = self
195                .max_size
196                .unwrap_or(if self.min_size > DEFAULT_MAX_SIZE {
197                    self.min_size + DEFAULT_MAX_SIZE
198                } else {
199                    DEFAULT_MAX_SIZE
200                });
201            self.char_fields
202                .build_text_handle(self.min_size as u64, max_size as u64)
203        })
204    }
205}
206
207impl Generator<String> for TextGenerator {
208    fn do_draw(&self, tc: &TestCase) -> String {
209        tc.generate_string(self.handle())
210    }
211}
212
213/// Generate arbitrary Unicode text strings.
214///
215/// See [`TextGenerator`] for builder methods.
216pub fn text() -> TextGenerator {
217    TextGenerator {
218        min_size: 0,
219        max_size: None,
220        char_fields: CharacterFields::new(),
221        alphabet_called: false,
222        char_param_called: false,
223        handle: OnceLock::new(),
224    }
225}
226
227/// Generator for single Unicode characters ([`char`]). Created by [`characters()`].
228pub struct CharactersGenerator {
229    char_fields: CharacterFields,
230    handle: OnceLock<ffi::StringGenerator>,
231}
232
233impl CharactersGenerator {
234    /// Restrict to characters encodable in this codec (e.g. `"ascii"`, `"utf-8"`, `"latin-1"`).
235    pub fn codec(mut self, codec: &str) -> Self {
236        self.handle = OnceLock::new();
237        self.char_fields.codec = Some(codec.to_string());
238        self
239    }
240
241    /// Set the minimum Unicode codepoint.
242    pub fn min_codepoint(mut self, min_codepoint: u32) -> Self {
243        self.handle = OnceLock::new();
244        self.char_fields.min_codepoint = Some(min_codepoint);
245        self
246    }
247
248    /// Set the maximum Unicode codepoint.
249    pub fn max_codepoint(mut self, max_codepoint: u32) -> Self {
250        self.handle = OnceLock::new();
251        self.char_fields.max_codepoint = Some(max_codepoint);
252        self
253    }
254
255    /// Include only characters from these Unicode general categories (e.g. `["L", "Nd"]`).
256    ///
257    /// Mutually exclusive with [`exclude_categories`](Self::exclude_categories).
258    pub fn categories(mut self, categories: &[&str]) -> Self {
259        self.handle = OnceLock::new();
260        self.char_fields.categories = Some(categories.iter().map(|s| s.to_string()).collect());
261        self
262    }
263
264    /// Exclude characters from these Unicode general categories.
265    ///
266    /// Mutually exclusive with [`categories`](Self::categories).
267    pub fn exclude_categories(mut self, exclude_categories: &[&str]) -> Self {
268        self.handle = OnceLock::new();
269        self.char_fields.exclude_categories =
270            Some(exclude_categories.iter().map(|s| s.to_string()).collect());
271        self
272    }
273
274    /// Always include these specific characters, even if excluded by other filters.
275    pub fn include_characters(mut self, include_characters: &str) -> Self {
276        self.handle = OnceLock::new();
277        self.char_fields.include_characters = Some(include_characters.to_string());
278        self
279    }
280
281    /// Always exclude these specific characters.
282    pub fn exclude_characters(mut self, exclude_characters: &str) -> Self {
283        self.handle = OnceLock::new();
284        self.char_fields.exclude_characters = Some(exclude_characters.to_string());
285        self
286    }
287
288    fn handle(&self) -> &ffi::StringGenerator {
289        self.handle
290            .get_or_init(|| self.char_fields.build_text_handle(1, 1))
291    }
292
293    /// Build a standalone text handle carrying this alphabet, for use as a
294    /// regex alphabet constraint. Sized `0..=0` so an empty alphabet — legal
295    /// for regex padding, which simply pads nothing — passes construction.
296    pub(super) fn build_alphabet_handle(&self) -> ffi::StringGenerator {
297        self.char_fields.build_text_handle(0, 0)
298    }
299}
300
301impl Generator<char> for CharactersGenerator {
302    fn do_draw(&self, tc: &TestCase) -> char {
303        let s = tc.generate_string(self.handle());
304        let mut chars = s.chars();
305        let c = chars
306            .next()
307            .expect("expected a single character, got empty string");
308        hegel_internal_assert!(
309            chars.next().is_none(),
310            "expected a single character, got multiple"
311        );
312        c
313    }
314}
315
316/// Generate single Unicode characters ([`char`]).
317///
318/// See [`CharactersGenerator`] for builder methods.
319pub fn characters() -> CharactersGenerator {
320    CharactersGenerator {
321        char_fields: CharacterFields::new(),
322        handle: OnceLock::new(),
323    }
324}
325
326/// Generator for strings matching a regex pattern. Created by [`from_regex()`].
327///
328/// By default the entire string matches the pattern. Use
329/// [`fullmatch(false)`](Self::fullmatch) to generate strings that merely
330/// contain a match.
331pub struct RegexGenerator {
332    pattern: String,
333    fullmatch: bool,
334    alphabet: Option<CharactersGenerator>,
335    handle: OnceLock<ffi::StringGenerator>,
336}
337
338impl RegexGenerator {
339    /// Set whether the entire string must match the pattern (the default), or
340    /// merely contain a match.
341    pub fn fullmatch(mut self, fullmatch: bool) -> Self {
342        self.handle = OnceLock::new();
343        self.fullmatch = fullmatch;
344        self
345    }
346
347    /// Constrain which characters may appear in generated strings.
348    pub fn alphabet(mut self, alphabet: CharactersGenerator) -> Self {
349        self.handle = OnceLock::new();
350        self.alphabet = Some(alphabet);
351        self
352    }
353
354    fn handle(&self) -> &ffi::StringGenerator {
355        self.handle.get_or_init(|| {
356            let alphabet = self.alphabet.as_ref().map(|a| a.build_alphabet_handle());
357            ffi::StringGenerator::regex(&self.pattern, self.fullmatch, alphabet.as_ref())
358                .unwrap_or_else(|msg| invalid_argument!("{msg}"))
359        })
360    }
361}
362
363impl Generator<String> for RegexGenerator {
364    fn do_draw(&self, tc: &TestCase) -> String {
365        tc.generate_string(self.handle())
366    }
367}
368
369/// Generate strings matching a regex pattern.
370///
371/// See [`RegexGenerator`] for builder methods.
372pub fn from_regex(pattern: &str) -> RegexGenerator {
373    RegexGenerator {
374        pattern: pattern.to_string(),
375        fullmatch: true,
376        alphabet: None,
377        handle: OnceLock::new(),
378    }
379}
380
381/// Generator for arbitrary byte sequences. Created by [`binary()`].
382pub struct BinaryGenerator {
383    min_size: usize,
384    max_size: Option<usize>,
385}
386
387impl BinaryGenerator {
388    /// Set the minimum length in bytes.
389    pub fn min_size(mut self, min_size: usize) -> Self {
390        self.min_size = min_size;
391        self
392    }
393
394    /// Set the maximum length in bytes.
395    pub fn max_size(mut self, max_size: usize) -> Self {
396        self.max_size = Some(max_size);
397        self
398    }
399}
400
401impl Generator<Vec<u8>> for BinaryGenerator {
402    fn do_draw(&self, tc: &TestCase) -> Vec<u8> {
403        if let Some(max) = self.max_size {
404            if self.min_size > max {
405                invalid_argument!("Cannot have max_size < min_size");
406            }
407        }
408        let max_size = self
409            .max_size
410            .unwrap_or(if self.min_size > DEFAULT_MAX_SIZE {
411                self.min_size + DEFAULT_MAX_SIZE
412            } else {
413                DEFAULT_MAX_SIZE
414            });
415        tc.generate_bytes(self.min_size, max_size)
416    }
417}
418
419/// Generate arbitrary byte sequences (`Vec<u8>`).
420///
421/// See [`BinaryGenerator`] for builder methods.
422pub fn binary() -> BinaryGenerator {
423    BinaryGenerator {
424        min_size: 0,
425        max_size: None,
426    }
427}
428
429/// Generator for email address strings. Created by [`emails()`].
430pub struct EmailGenerator {
431    handle: OnceLock<ffi::StringGenerator>,
432}
433
434impl Generator<String> for EmailGenerator {
435    fn do_draw(&self, tc: &TestCase) -> String {
436        let handle = self.handle.get_or_init(|| {
437            ffi::StringGenerator::email().unwrap_or_else(|msg| invalid_argument!("{msg}"))
438        });
439        tc.generate_string(handle)
440    }
441}
442
443/// Generate email address strings.
444pub fn emails() -> EmailGenerator {
445    EmailGenerator {
446        handle: OnceLock::new(),
447    }
448}
449
450/// Generator for URL strings. Created by [`urls()`].
451pub struct UrlGenerator {
452    handle: OnceLock<ffi::StringGenerator>,
453}
454
455impl Generator<String> for UrlGenerator {
456    fn do_draw(&self, tc: &TestCase) -> String {
457        let handle = self.handle.get_or_init(|| {
458            ffi::StringGenerator::url().unwrap_or_else(|msg| invalid_argument!("{msg}"))
459        });
460        tc.generate_string(handle)
461    }
462}
463
464/// Generate URL strings.
465pub fn urls() -> UrlGenerator {
466    UrlGenerator {
467        handle: OnceLock::new(),
468    }
469}
470
471/// Generator for domain name strings. Created by [`domains()`].
472pub struct DomainGenerator {
473    max_length: usize,
474    handle: OnceLock<ffi::StringGenerator>,
475}
476
477impl DomainGenerator {
478    /// Set the maximum length (must be between 4 and 255).
479    pub fn max_length(mut self, max_length: usize) -> Self {
480        self.handle = OnceLock::new();
481        self.max_length = max_length;
482        self
483    }
484}
485
486impl Generator<String> for DomainGenerator {
487    fn do_draw(&self, tc: &TestCase) -> String {
488        let handle = self.handle.get_or_init(|| {
489            if !(self.max_length >= 4 && self.max_length <= 255) {
490                invalid_argument!("max_length must be between 4 and 255");
491            }
492            ffi::StringGenerator::domain(self.max_length as u64)
493                .unwrap_or_else(|msg| invalid_argument!("{msg}"))
494        });
495        tc.generate_string(handle)
496    }
497}
498
499/// Generate domain name strings.
500///
501/// See [`DomainGenerator`] for builder methods.
502pub fn domains() -> DomainGenerator {
503    DomainGenerator {
504        max_length: 255,
505        handle: OnceLock::new(),
506    }
507}
508
509/// Generator for IP addresses. Created by [`ip_addresses()`].
510///
511/// Generates both IPv4 and IPv6 addresses.
512pub struct IpAddressGenerator {}
513
514impl IpAddressGenerator {
515    /// Only generate IPv4 addresses.
516    pub fn v4(self) -> Ipv4AddressGenerator {
517        Ipv4AddressGenerator {}
518    }
519
520    /// Only generate IPv6 addresses.
521    pub fn v6(self) -> Ipv6AddressGenerator {
522        Ipv6AddressGenerator {}
523    }
524}
525
526impl Generator<std::net::IpAddr> for IpAddressGenerator {
527    fn do_draw(&self, tc: &TestCase) -> std::net::IpAddr {
528        tc.start_span(labels::ONE_OF);
529        let addr = if tc.generate_integer_i64(0, 1) == 0 {
530            std::net::IpAddr::V4(tc.generate_ipv4())
531        } else {
532            std::net::IpAddr::V6(tc.generate_ipv6())
533        };
534        tc.stop_span(false);
535        addr
536    }
537}
538
539/// Generator for IPv4 addresses. Created by [`IpAddressGenerator::v4`].
540pub struct Ipv4AddressGenerator {}
541
542impl Generator<std::net::Ipv4Addr> for Ipv4AddressGenerator {
543    fn do_draw(&self, tc: &TestCase) -> std::net::Ipv4Addr {
544        tc.generate_ipv4()
545    }
546}
547
548/// Generator for IPv6 addresses. Created by [`IpAddressGenerator::v6`].
549pub struct Ipv6AddressGenerator {}
550
551impl Generator<std::net::Ipv6Addr> for Ipv6AddressGenerator {
552    fn do_draw(&self, tc: &TestCase) -> std::net::Ipv6Addr {
553        tc.generate_ipv6()
554    }
555}
556
557/// Generate IP addresses (IPv4 or IPv6).
558///
559/// See [`IpAddressGenerator`] for builder methods.
560pub fn ip_addresses() -> IpAddressGenerator {
561    IpAddressGenerator {}
562}
563
564/// Format a drawn date as `YYYY-MM-DD`, matching `st.dates().isoformat()`.
565pub(crate) fn format_date(d: hegel_c::hegel_date_t) -> String {
566    format!("{:04}-{:02}-{:02}", d.year, d.month, d.day)
567}
568
569/// Format a drawn time as `HH:MM:SS` or `HH:MM:SS.ffffff`, matching
570/// `st.times().isoformat()`: the fractional part is present iff
571/// `microsecond != 0`.
572pub(crate) fn format_time(t: hegel_c::hegel_time_t) -> String {
573    if t.microsecond == 0 {
574        format!("{:02}:{:02}:{:02}", t.hour, t.minute, t.second)
575    } else {
576        format!(
577            "{:02}:{:02}:{:02}.{:06}",
578            t.hour, t.minute, t.second, t.microsecond
579        )
580    }
581}
582
583/// Generator for date strings in YYYY-MM-DD format. Created by [`dates()`].
584pub struct DateGenerator;
585
586impl Generator<String> for DateGenerator {
587    fn do_draw(&self, tc: &TestCase) -> String {
588        format_date(tc.generate_date(full_ranges::MIN_DATE, full_ranges::MAX_DATE))
589    }
590}
591
592/// Generate date strings in YYYY-MM-DD format.
593pub fn dates() -> DateGenerator {
594    DateGenerator
595}
596
597/// Generator for time strings in HH:MM:SS format. Created by [`times()`].
598pub struct TimeGenerator;
599
600impl Generator<String> for TimeGenerator {
601    fn do_draw(&self, tc: &TestCase) -> String {
602        format_time(tc.generate_time(full_ranges::MIDNIGHT, full_ranges::LAST_MICROSECOND))
603    }
604}
605
606/// Generate time strings in HH:MM:SS format.
607pub fn times() -> TimeGenerator {
608    TimeGenerator
609}
610
611/// Generator for ISO 8601 datetime strings. Created by [`datetimes()`].
612pub struct DateTimeGenerator;
613
614impl Generator<String> for DateTimeGenerator {
615    fn do_draw(&self, tc: &TestCase) -> String {
616        let dt = tc.generate_datetime(full_ranges::MIN_DATETIME, full_ranges::MAX_DATETIME);
617        format!("{}T{}", format_date(dt.date), format_time(dt.time))
618    }
619}
620
621/// Generate ISO 8601 datetime strings.
622pub fn datetimes() -> DateTimeGenerator {
623    DateTimeGenerator
624}
625
626/// Generator for UUID strings in canonical hyphenated form. Created by [`uuids()`].
627///
628/// By default generates UUIDs of any version. Use [`UuidsGenerator::version`]
629/// to restrict to a specific RFC 4122 version (1–5).
630pub struct UuidsGenerator {
631    version: Option<u8>,
632}
633
634impl UuidsGenerator {
635    /// Restrict to UUIDs of a specific version (1–5).
636    pub fn version(mut self, version: u8) -> Self {
637        self.version = Some(version);
638        self
639    }
640}
641
642impl Generator<String> for UuidsGenerator {
643    fn do_draw(&self, tc: &TestCase) -> String {
644        if let Some(v) = self.version {
645            if !(1..=5).contains(&v) {
646                invalid_argument!("UUID version must be between 1 and 5, got {v}");
647            }
648        }
649        let b = tc.generate_uuid(self.version);
650        format!(
651            "{:02x}{:02x}{:02x}{:02x}-{:02x}{:02x}-{:02x}{:02x}-{:02x}{:02x}-{:02x}{:02x}{:02x}{:02x}{:02x}{:02x}",
652            b[0],
653            b[1],
654            b[2],
655            b[3],
656            b[4],
657            b[5],
658            b[6],
659            b[7],
660            b[8],
661            b[9],
662            b[10],
663            b[11],
664            b[12],
665            b[13],
666            b[14],
667            b[15],
668        )
669    }
670}
671
672/// Generate UUID strings in canonical hyphenated form.
673///
674/// See [`UuidsGenerator`] for builder methods.
675pub fn uuids() -> UuidsGenerator {
676    UuidsGenerator { version: None }
677}
hegel/generators/strings.rs

hegel/generators/
strings.rs
