@@ -38,6 +38,8 @@ pub struct FileEncoder {
38
38
path : PathBuf ,
39
39
#[ cfg( debug_assertions) ]
40
40
finished : bool ,
41
+ panic_at_offset : Option < usize > ,
42
+ flush_every_write : bool ,
41
43
}
42
44
43
45
impl FileEncoder {
@@ -48,6 +50,9 @@ impl FileEncoder {
48
50
let file =
49
51
File :: options ( ) . read ( true ) . write ( true ) . create ( true ) . truncate ( true ) . open ( & path) ?;
50
52
53
+ let panic_at_offset = std:: env:: var_os ( "RUSTC_FILE_ENCODER_PANIC_AT_OFFSET" )
54
+ . map ( |v| v. to_str ( ) . unwrap ( ) . parse :: < usize > ( ) . unwrap ( ) ) ;
55
+
51
56
Ok ( FileEncoder {
52
57
buf : vec ! [ 0u8 ; BUF_SIZE ] . into_boxed_slice ( ) . try_into ( ) . unwrap ( ) ,
53
58
path : path. as_ref ( ) . into ( ) ,
@@ -57,6 +62,8 @@ impl FileEncoder {
57
62
res : Ok ( ( ) ) ,
58
63
#[ cfg( debug_assertions) ]
59
64
finished : false ,
65
+ panic_at_offset,
66
+ flush_every_write : panic_at_offset. is_some_and ( |v| v <= BUF_SIZE ) ,
60
67
} )
61
68
}
62
69
@@ -75,6 +82,19 @@ impl FileEncoder {
75
82
{
76
83
self . finished = false ;
77
84
}
85
+
86
+ if let Some ( panic_offset) = self . panic_at_offset {
87
+ // If we are within the buffer size of the offset we're to panic at, we need to start
88
+ // flushing every write so that we don't panic late.
89
+ if panic_offset - self . flushed <= BUF_SIZE {
90
+ self . flush_every_write = true ;
91
+ }
92
+ // If the offset we want to panic at is in the range we're about to write, panic.
93
+ if ( self . flushed ..self . flushed + self . buffered ) . contains ( & panic_offset) {
94
+ panic ! ( )
95
+ }
96
+ }
97
+
78
98
if self . res . is_ok ( ) {
79
99
self . res = self . file . write_all ( & self . buf [ ..self . buffered ] ) ;
80
100
}
@@ -107,6 +127,12 @@ impl FileEncoder {
107
127
if self . res . is_ok ( ) {
108
128
self . res = self . file . write_all ( buf) ;
109
129
}
130
+ // This write bypasses the buffer, so we need to duplicate the check logic here
131
+ if let Some ( panic_offset) = self . panic_at_offset {
132
+ if ( self . flushed ..self . flushed + buf. len ( ) ) . contains ( & panic_offset) {
133
+ panic ! ( )
134
+ }
135
+ }
110
136
self . flushed += buf. len ( ) ;
111
137
}
112
138
}
@@ -117,7 +143,9 @@ impl FileEncoder {
117
143
{
118
144
self . finished = false ;
119
145
}
120
- if let Some ( dest) = self . buffer_empty ( ) . get_mut ( ..buf. len ( ) ) {
146
+ if !self . flush_every_write
147
+ && let Some ( dest) = self . buffer_empty ( ) . get_mut ( ..buf. len ( ) )
148
+ {
121
149
dest. copy_from_slice ( buf) ;
122
150
self . buffered += buf. len ( ) ;
123
151
} else {
@@ -146,21 +174,39 @@ impl FileEncoder {
146
174
self . finished = false ;
147
175
}
148
176
let flush_threshold = const { BUF_SIZE . checked_sub ( N ) . unwrap ( ) } ;
149
- if std:: intrinsics:: unlikely ( self . buffered > flush_threshold) {
177
+ if std:: intrinsics:: unlikely ( self . flush_every_write || self . buffered > flush_threshold) {
150
178
self . flush ( ) ;
151
179
}
152
180
// SAFETY: We checked above that that N < self.buffer_empty().len(),
153
181
// and if isn't, flush ensures that our empty buffer is now BUF_SIZE.
154
182
// We produce a post-mono error if N > BUF_SIZE.
155
183
let buf = unsafe { self . buffer_empty ( ) . first_chunk_mut :: < N > ( ) . unwrap_unchecked ( ) } ;
156
184
let written = visitor ( buf) ;
157
- // We have to ensure that an errant visitor cannot cause self.buffered to exeed BUF_SIZE.
185
+ // We have to ensure that an errant visitor cannot cause self.buffered to exceed BUF_SIZE.
158
186
if written > N {
159
187
Self :: panic_invalid_write :: < N > ( written) ;
160
188
}
161
189
self . buffered += written;
162
190
}
163
191
192
+ #[ inline]
193
+ pub fn write_with_spare ( & mut self , visitor : impl FnOnce ( & mut [ u8 ] ) -> usize ) {
194
+ #[ cold]
195
+ #[ inline( never) ]
196
+ fn panic_invalid_write_spare ( ) {
197
+ panic ! (
198
+ "FileEncoder::write_with_spare cannot be used to write more bytes than the current buffer size!"
199
+ ) ;
200
+ }
201
+
202
+ let buf = self . buffer_empty ( ) ;
203
+ let written = visitor ( buf) ;
204
+ if written > buf. len ( ) {
205
+ panic_invalid_write_spare ( ) ;
206
+ }
207
+ self . buffered += written;
208
+ }
209
+
164
210
#[ cold]
165
211
#[ inline( never) ]
166
212
fn panic_invalid_write < const N : usize > ( written : usize ) {
0 commit comments