qubit_budget/string/budgeted_string_writer.rs
1// =============================================================================
2// Copyright (c) 2025 - 2026 Haixing Hu.
3//
4// SPDX-License-Identifier: Apache-2.0
5//
6// Licensed under the Apache License, Version 2.0.
7// =============================================================================
8//! Single-pass string rendering with transactional budget accounting.
9
10use std::fmt;
11use std::io;
12
13use super::BudgetedStringError;
14use super::internal::FmtWriter;
15use super::internal::IoWriter;
16use super::internal::WriterFailure;
17use crate::resource::ResourceBudget;
18use crate::resource::ResourceQuantity;
19
20/// Collects a UTF-8 string while checking a finite byte budget incrementally.
21///
22/// The writer is constructed and committed by
23/// [`ResourceBudget::try_write_string`]. A failed render drops the buffered
24/// prefix and leaves the budget unchanged.
25///
26/// Use [`Self::as_fmt`] for formatting APIs or [`Self::as_io`] for
27/// byte-oriented I/O. The returned adapter borrows this writer, so the render
28/// callback must finish using it before returning.
29///
30/// # Type Parameters
31///
32/// * `R` - Caller-defined resource identity retained by limits and errors.
33/// * `Q` - Exact unsigned quantity used for measurements and accounting.
34///
35/// # Examples
36///
37/// ```
38/// use std::fmt::Write as _;
39/// use qubit_budget::ResourceBudget;
40///
41/// let mut budget = ResourceBudget::new("response bytes", 16_u64);
42/// let output = budget
43/// .try_write_string(|writer| {
44/// let mut formatted = writer.as_fmt();
45/// write!(&mut formatted, "status={}", 200)
46/// })
47/// .expect("the rendered response should fit");
48///
49/// assert_eq!(output, "status=200");
50/// assert_eq!(budget.used(), 10);
51/// ```
52pub struct BudgetedStringWriter<'a, R, Q = u64>
53where
54 Q: ResourceQuantity,
55{
56 /// Immutable budget snapshot used to validate the final output length.
57 budget: &'a ResourceBudget<R, Q>,
58 /// Bytes staged until the entire rendering transaction succeeds.
59 output: Vec<u8>,
60 /// First writer-side failure retained for deterministic error precedence.
61 failure: Option<WriterFailure<R, Q>>,
62}
63
64impl<'a, R, Q> BudgetedStringWriter<'a, R, Q>
65where
66 R: Clone,
67 Q: ResourceQuantity,
68{
69 /// Creates an empty writer backed by an immutable budget snapshot.
70 ///
71 /// # Parameters
72 ///
73 /// * `budget` - Immutable capacity snapshot used to validate output.
74 ///
75 /// # Returns
76 ///
77 /// Creates an empty writer backed by an immutable budget snapshot.
78 fn new(budget: &'a ResourceBudget<R, Q>) -> Self {
79 Self {
80 budget,
81 output: Vec::new(),
82 failure: None,
83 }
84 }
85
86 /// Separates rendered bytes from the first captured failure.
87 ///
88 /// # Returns
89 ///
90 /// Separates rendered bytes from the first captured failure.
91 ///
92 /// A `None` failure component indicates that no writer-side failure was
93 /// captured.
94 fn into_parts(self) -> (Vec<u8>, Option<WriterFailure<R, Q>>) {
95 (self.output, self.failure)
96 }
97
98 /// Appends bytes after checking the next cumulative length and budget.
99 ///
100 /// Returns `true` when the bytes were appended, or `false` after storing
101 /// the first failure for the enclosing transaction.
102 ///
103 /// # Parameters
104 ///
105 /// * `bytes` - Bytes to append to the transactional output.
106 ///
107 /// # Returns
108 ///
109 /// Appends bytes after checking the next cumulative length and budget.
110 pub(crate) fn append(&mut self, bytes: &[u8]) -> bool {
111 if self.failure.is_some() {
112 return false;
113 }
114 let Some(next_len) = checked_output_len(self.output.len(), bytes.len()) else {
115 self.failure = Some(WriterFailure::LengthOverflow);
116 return false;
117 };
118 let next_length = match Q::try_from_usize(next_len) {
119 Ok(value) => value,
120 Err(source) => {
121 self.failure = Some(WriterFailure::Quantity {
122 resource: self.budget.resource().clone(),
123 source,
124 });
125 return false;
126 }
127 };
128 if let Err(error) = self.budget.check_available(next_length) {
129 self.failure = Some(WriterFailure::Budget(error));
130 return false;
131 }
132 if next_len > self.output.capacity() {
133 let target = self.output.capacity().saturating_mul(2).max(next_len);
134 if let Err(source) = self.output.try_reserve_exact(target.saturating_sub(self.output.len())) {
135 self.failure = Some(WriterFailure::Allocation(source));
136 return false;
137 }
138 }
139 self.output.extend_from_slice(bytes);
140 true
141 }
142
143 /// Returns a formatting writer view over the current transaction.
144 ///
145 /// # Returns
146 ///
147 /// Returns a formatting writer view over the current transaction.
148 #[must_use]
149 #[inline]
150 pub fn as_fmt(&mut self) -> impl fmt::Write + '_ {
151 FmtWriter { writer: self }
152 }
153
154 /// Returns an I/O writer view over the current transaction.
155 ///
156 /// # Returns
157 ///
158 /// Returns an I/O writer view over the current transaction.
159 #[must_use]
160 #[inline]
161 pub fn as_io(&mut self) -> impl io::Write + '_ {
162 IoWriter { writer: self }
163 }
164}
165
166/// Adds two output lengths while detecting `usize` overflow.
167///
168/// # Parameters
169///
170/// * `current` - Bytes already staged in the output buffer.
171/// * `additional` - Bytes requested by the next append.
172///
173/// # Returns
174///
175/// Adds two output lengths while detecting `usize` overflow.
176///
177/// `None` indicates that the arithmetic sum would overflow `usize`.
178const fn checked_output_len(current: usize, additional: usize) -> Option<usize> {
179 current.checked_add(additional)
180}
181
182/// Renders a UTF-8 string and commits its byte length to a resource budget.
183///
184/// This crate-internal implementation keeps string buffering and writer error
185/// precedence in the string domain. [`ResourceBudget::try_write_string`]
186/// exposes the public type-owned forwarding method.
187///
188/// # Type Parameters
189///
190/// * `R` - Caller-defined resource identity retained by limits and errors.
191/// * `Q` - Exact unsigned quantity used for byte accounting.
192/// * `E` - Error type returned by the caller-provided renderer.
193/// * `F` - Closure that renders into the transactional writer.
194///
195/// # Parameters
196///
197/// * `budget` - Finite byte budget charged only after rendering succeeds.
198/// * `render` - Caller-provided renderer writing into the transactional
199/// adapter.
200///
201/// # Returns
202///
203/// `Ok(rendered)` after the complete UTF-8 output is charged and committed.
204///
205/// # Errors
206///
207/// Returns [`BudgetedStringError`] when rendering, allocation, UTF-8
208/// validation, measurement, or budget accounting fails.
209pub(crate) fn render_budgeted_string<R, Q, E, F>(
210 budget: &mut ResourceBudget<R, Q>,
211 render: F,
212) -> Result<String, BudgetedStringError<R, E, Q>>
213where
214 R: Clone + fmt::Debug,
215 Q: ResourceQuantity,
216 E: fmt::Debug + fmt::Display,
217 F: FnOnce(&mut BudgetedStringWriter<'_, R, Q>) -> Result<(), E>,
218{
219 let mut writer = BudgetedStringWriter::new(budget);
220 let rendered = render(&mut writer);
221 let (bytes, failure) = writer.into_parts();
222 match failure {
223 Some(WriterFailure::Budget(error)) => {
224 return Err(BudgetedStringError::Budget(error));
225 }
226 Some(WriterFailure::Quantity { resource, source }) => {
227 return Err(BudgetedStringError::Quantity { resource, source });
228 }
229 Some(WriterFailure::LengthOverflow) => {
230 return Err(BudgetedStringError::LengthOverflow);
231 }
232 Some(WriterFailure::Allocation(source)) => {
233 return Err(BudgetedStringError::Allocation(source));
234 }
235 None => {}
236 }
237 if let Err(error) = rendered {
238 return Err(BudgetedStringError::Render(error));
239 }
240 let output = String::from_utf8(bytes).map_err(BudgetedStringError::InvalidUtf8)?;
241 let output_length = Q::try_from_usize(output.len()).map_err(|source| BudgetedStringError::Quantity {
242 resource: budget.resource().clone(),
243 source,
244 })?;
245 budget.try_consume(output_length).map_err(BudgetedStringError::Budget)?;
246 Ok(output)
247}