1use crate::util::{as_bytes, as_mut_bytes, read_byte, write_byte, write_exact};
4
5use std::io::{Cursor, Read, Write};
6
7use super::{DecodeError, VertexEncodingVersion};
8
9const VERTEX_HEADER: u8 = 0xa0;
10
11const VERTEX_BLOCK_SIZE_BYTES: usize = 8192;
12const VERTEX_BLOCK_MAX_SIZE: usize = 256;
13const BYTE_GROUP_SIZE: usize = 16;
14const BYTE_GROUP_DECODE_LIMIT: usize = 24;
15const TAIL_MAX_SIZE: usize = 32;
16
17fn get_vertex_block_size(vertex_size: usize) -> usize {
18 let mut result = VERTEX_BLOCK_SIZE_BYTES / vertex_size;
20
21 result &= !(BYTE_GROUP_SIZE - 1);
24
25 result.min(VERTEX_BLOCK_MAX_SIZE)
26}
27
28fn zigzag8(v: u8) -> u8 {
29 ((v as i8) >> 7) as u8 ^ (v << 1)
30}
31
32fn unzigzag8(v: u8) -> u8 {
33 (v & 1).wrapping_neg() ^ (v >> 1)
34}
35
36fn encode_bytes_group_zero(buffer: &[u8]) -> bool {
37 !buffer[0..BYTE_GROUP_SIZE].iter().any(|b| *b > 0)
38}
39
40fn encode_bytes_group_measure(buffer: &[u8], bits: usize) -> usize {
41 assert!(bits >= 1 && bits <= 8);
42
43 match bits {
44 1 => {
45 if encode_bytes_group_zero(buffer) {
46 0
47 } else {
48 usize::MAX
49 }
50 }
51 8 => BYTE_GROUP_SIZE,
52 _ => {
53 let mut result = BYTE_GROUP_SIZE * bits / 8;
54
55 let sentinel = (1 << bits) - 1;
56
57 result += &buffer[0..BYTE_GROUP_SIZE]
58 .iter()
59 .map(|b| (*b >= sentinel) as usize)
60 .sum();
61
62 result
63 }
64 }
65}
66
67fn encode_bytes_group<W: Write>(data: &mut W, buffer: &[u8], bits: usize) -> Option<usize> {
68 assert!(bits >= 1 && bits <= 8);
69
70 match bits {
71 1 => Some(0),
72 8 => write_exact(data, &buffer[0..BYTE_GROUP_SIZE]),
73 _ => {
74 let byte_size = 8 / bits;
75 assert!(BYTE_GROUP_SIZE % byte_size == 0);
76
77 let sentinel = (1 << bits) - 1;
80
81 let mut written = 0;
82
83 for i in (0..BYTE_GROUP_SIZE).step_by(byte_size) {
84 let mut byte = 0;
85
86 for k in 0..byte_size {
87 let enc = if buffer[i + k] >= sentinel {
88 sentinel
89 } else {
90 buffer[i + k]
91 };
92
93 byte <<= bits;
94 byte |= enc;
95 }
96
97 written += write_byte(data, byte);
98 }
99
100 for i in 0..BYTE_GROUP_SIZE {
101 if buffer[i] >= sentinel {
102 written += write_byte(data, buffer[i]);
103 }
104 }
105
106 Some(written)
107 }
108 }
109}
110
111fn encode_bytes(data: &mut [u8], buffer: &[u8]) -> Option<usize> {
112 assert!(buffer.len() % BYTE_GROUP_SIZE == 0);
113
114 let header_size = (buffer.len() / BYTE_GROUP_SIZE + 3) / 4;
116
117 if data.len() < header_size {
118 return None;
119 }
120
121 let (header, mut data) = data.split_at_mut(header_size);
122
123 header.fill(0);
124
125 let mut written = header_size;
126
127 for i in (0..buffer.len()).step_by(BYTE_GROUP_SIZE) {
128 if data.len() < BYTE_GROUP_DECODE_LIMIT {
129 return None;
130 }
131
132 let mut best_bits = 8;
133 let mut best_size = encode_bytes_group_measure(&buffer[i..], 8);
134
135 for bits in [1, 2, 4, 8].iter() {
136 let size = encode_bytes_group_measure(&buffer[i..], *bits);
137
138 if size < best_size {
139 best_bits = *bits;
140 best_size = size;
141 }
142 }
143
144 let bitslog2 = match best_bits {
145 1 => 0,
146 2 => 1,
147 4 => 2,
148 8 => 3,
149 _ => unreachable!(),
150 };
151 assert!((1 << bitslog2) == best_bits);
152
153 let header_offset = i / BYTE_GROUP_SIZE;
154
155 header[header_offset / 4] |= bitslog2 << ((header_offset % 4) * 2);
156
157 let group_written = encode_bytes_group(&mut data, &buffer[i..], best_bits)?;
158
159 assert!(group_written == best_size);
160
161 written += group_written;
162 }
163
164 Some(written)
165}
166
167fn encode_vertex_block(
168 mut data: &mut [u8],
169 vertex_data: &[u8],
170 vertex_count: usize,
171 vertex_size: usize,
172 last_vertex: &mut [u8; 256],
173) -> Option<usize> {
174 assert!(vertex_count > 0 && vertex_count <= VERTEX_BLOCK_MAX_SIZE);
175
176 let mut buffer = [0u8; VERTEX_BLOCK_MAX_SIZE];
178 assert!(VERTEX_BLOCK_MAX_SIZE % BYTE_GROUP_SIZE == 0);
179
180 let mut written_sum = 0;
181
182 for k in 0..vertex_size {
183 let mut vertex_offset = k;
184 let mut p = last_vertex[k];
185
186 for i in 0..vertex_count {
187 buffer[i] = zigzag8(vertex_data[vertex_offset].wrapping_sub(p));
188
189 p = vertex_data[vertex_offset];
190
191 vertex_offset += vertex_size;
192 }
193
194 let written = encode_bytes(
195 data,
196 &buffer[0..(vertex_count + BYTE_GROUP_SIZE - 1) & !(BYTE_GROUP_SIZE - 1)],
197 )?;
198 data = &mut data[written..];
199
200 written_sum += written;
201 }
202
203 let offset = vertex_size * (vertex_count - 1);
204 last_vertex[0..vertex_size].copy_from_slice(&vertex_data[offset..offset + vertex_size]);
205
206 Some(written_sum)
207}
208
209fn decode_bytes_group<W>(data: &mut Cursor<&[u8]>, buffer: &mut W, bitslog2: i32)
210where
211 W: Write,
212{
213 let mut byte = 0;
214
215 let mut data_pos = data.position();
216
217 let read = |data: &mut Cursor<&[u8]>, byte: &mut u8, data_pos: &mut u64| {
218 data.set_position(*data_pos);
219 *byte = read_byte(data);
220 *data_pos += 1;
221 };
222
223 let mut next = |bits, data: &mut Cursor<&[u8]>, byte: &mut u8, data_var_pos: &mut u64| {
224 let enc = *byte >> (8 - bits);
225 *byte <<= bits;
226 data.set_position(*data_var_pos);
227 let encv = read_byte(data);
228 write_byte(buffer, if enc == (1 << bits) as u8 - 1 { encv } else { enc });
229 *data_var_pos += (enc == (1 << bits) as u8 - 1) as u64;
230 };
231
232 let mut buf = [0; BYTE_GROUP_SIZE];
233
234 match bitslog2 {
235 0 => {
236 buffer.write(&[0; BYTE_GROUP_SIZE]).unwrap();
237 }
238 1 => {
239 let mut data_var_pos = data_pos + 4;
240
241 for _ in 0..4 {
243 read(data, &mut byte, &mut data_pos);
244 next(2, data, &mut byte, &mut data_var_pos);
245 next(2, data, &mut byte, &mut data_var_pos);
246 next(2, data, &mut byte, &mut data_var_pos);
247 next(2, data, &mut byte, &mut data_var_pos);
248 }
249
250 data.set_position(data_var_pos);
251 }
252 2 => {
253 let mut data_var_pos = data_pos + 8;
254
255 for _ in 0..8 {
257 read(data, &mut byte, &mut data_pos);
258 next(4, data, &mut byte, &mut data_var_pos);
259 next(4, data, &mut byte, &mut data_var_pos);
260 }
261
262 data.set_position(data_var_pos);
263 }
264 3 => {
265 data.read(&mut buf).unwrap();
266 buffer.write(&buf).unwrap();
267 }
268 _ => unreachable!("Unexpected bit length"), }
270}
271
272fn decode_bytes(data: &mut Cursor<&[u8]>, buffer: &mut [u8]) -> Result<(), DecodeError> {
273 assert!(buffer.len() % BYTE_GROUP_SIZE == 0);
274
275 let header_size = (buffer.len() / BYTE_GROUP_SIZE + 3) / 4;
277
278 let raw_data = &data.get_ref()[data.position() as usize..];
279
280 if raw_data.len() < header_size {
281 return Err(DecodeError::UnexpectedEof);
282 }
283
284 let header = &raw_data[0..header_size];
285
286 data.set_position(data.position() + header_size as u64);
287
288 for i in (0..buffer.len()).step_by(BYTE_GROUP_SIZE) {
289 let raw_data = &data.get_ref()[data.position() as usize..];
290
291 if raw_data.len() < BYTE_GROUP_DECODE_LIMIT {
292 return Err(DecodeError::UnexpectedEof);
293 }
294
295 let header_offset = i / BYTE_GROUP_SIZE;
296
297 let bitslog2 = (header[header_offset / 4] >> ((header_offset % 4) * 2)) & 3;
298
299 let mut b = &mut buffer[i..];
300
301 decode_bytes_group(data, &mut b, bitslog2 as i32);
302 }
303
304 Ok(())
305}
306
307fn decode_vertex_block(
308 data: &mut Cursor<&[u8]>,
309 vertex_data: &mut [u8],
310 vertex_count: usize,
311 vertex_size: usize,
312 last_vertex: &mut [u8; 256],
313) -> Result<(), DecodeError> {
314 assert!(vertex_count > 0 && vertex_count <= VERTEX_BLOCK_MAX_SIZE);
315
316 let mut buffer = [0; VERTEX_BLOCK_MAX_SIZE];
317 let mut transposed = [0; VERTEX_BLOCK_SIZE_BYTES];
318
319 let vertex_count_aligned = (vertex_count + BYTE_GROUP_SIZE - 1) & !(BYTE_GROUP_SIZE - 1);
320
321 for k in 0..vertex_size {
322 decode_bytes(data, &mut buffer[0..vertex_count_aligned])?;
323
324 let mut vertex_offset = k;
325
326 let mut p = last_vertex[k];
327
328 for i in 0..vertex_count {
329 let v = unzigzag8(buffer[i]).wrapping_add(p);
330
331 transposed[vertex_offset] = v;
332 p = v;
333
334 vertex_offset += vertex_size;
335 }
336 }
337
338 vertex_data[0..vertex_count * vertex_size].copy_from_slice(&transposed[0..vertex_count * vertex_size]);
339
340 let offset = vertex_size * (vertex_count - 1);
341 last_vertex[0..vertex_size].copy_from_slice(&transposed[offset..offset + vertex_size]);
342
343 Ok(())
344}
345
346pub fn encode_vertex_buffer<Vertex>(
358 buffer: &mut [u8],
359 vertices: &[Vertex],
360 version: VertexEncodingVersion,
361) -> Option<usize> {
362 let vertex_size = std::mem::size_of::<Vertex>();
363
364 assert!(vertex_size > 0 && vertex_size <= 256);
365 assert!(vertex_size % 4 == 0);
366
367 let vertex_data = as_bytes(vertices);
368
369 let buffer_size = buffer.len();
370 let mut data = buffer;
371
372 if data.len() < 1 + vertex_size {
373 return None;
374 }
375
376 let version: u8 = version.into();
377
378 let mut written_sum = write_byte(&mut data, (VERTEX_HEADER | version) as u8);
379
380 let mut first_vertex = [0; 256];
381 if !vertices.is_empty() {
382 first_vertex[0..vertex_size].copy_from_slice(&vertex_data[0..vertex_size]);
383 }
384
385 let mut last_vertex = [0; 256];
386 last_vertex[0..vertex_size].copy_from_slice(&first_vertex[0..vertex_size]);
387
388 let vertex_block_size = get_vertex_block_size(vertex_size);
389
390 let mut vertex_offset = 0;
391
392 while vertex_offset < vertices.len() {
393 let block_size = if vertex_offset + vertex_block_size < vertices.len() {
394 vertex_block_size
395 } else {
396 vertices.len() - vertex_offset
397 };
398
399 let written = encode_vertex_block(
400 &mut data,
401 &vertex_data[vertex_offset * vertex_size..],
402 block_size,
403 vertex_size,
404 &mut last_vertex,
405 )?;
406 data = &mut data[written..];
407
408 written_sum += written;
409
410 vertex_offset += block_size;
411 }
412
413 let tail_size = vertex_size.max(TAIL_MAX_SIZE);
414
415 if data.len() < tail_size {
416 return None;
417 }
418
419 if vertex_size < TAIL_MAX_SIZE {
421 let written = TAIL_MAX_SIZE - vertex_size;
422 data[0..written].fill(0);
423 data = &mut data[written..];
424 written_sum += written;
425 }
426
427 written_sum += write_exact(&mut data, &first_vertex[0..vertex_size])?;
428
429 assert!(written_sum >= tail_size);
430 assert!(written_sum <= buffer_size);
431
432 Some(written_sum)
433}
434
435pub fn encode_vertex_buffer_bound(vertex_count: usize, vertex_size: usize) -> usize {
437 assert!(vertex_size > 0 && vertex_size <= 256);
438 assert!(vertex_size % 4 == 0);
439
440 let vertex_block_size = get_vertex_block_size(vertex_size);
441 let vertex_block_count = (vertex_count + vertex_block_size - 1) / vertex_block_size;
442
443 let vertex_block_header_size = (vertex_block_size / BYTE_GROUP_SIZE + 3) / 4;
444 let vertex_block_data_size = vertex_block_size;
445
446 let tail_size = vertex_size.max(TAIL_MAX_SIZE);
447
448 1 + vertex_block_count * vertex_size * (vertex_block_header_size + vertex_block_data_size) + tail_size
449}
450
451pub fn decode_vertex_buffer<Vertex>(destination: &mut [Vertex], buffer: &[u8]) -> Result<(), DecodeError> {
461 let vertex_size = std::mem::size_of::<Vertex>();
462 let vertex_count = destination.len();
463
464 assert!(vertex_size > 0 && vertex_size <= 256);
465 assert!(vertex_size % 4 == 0);
466
467 let vertex_data = as_mut_bytes(destination);
468
469 if buffer.len() < 1 + vertex_size {
470 return Err(DecodeError::UnexpectedEof);
471 }
472
473 let mut data = Cursor::new(buffer);
474
475 let data_header = read_byte(&mut data);
476
477 if (data_header & 0xf0) != VERTEX_HEADER {
478 return Err(DecodeError::InvalidHeader);
479 }
480
481 let version = data_header & 0x0f;
482 if version > 0 {
483 return Err(DecodeError::UnsupportedVersion);
484 }
485
486 let mut last_vertex = [0; 256];
487 last_vertex[0..vertex_size].copy_from_slice(&buffer[buffer.len() - vertex_size..]);
488
489 let vertex_block_size = get_vertex_block_size(vertex_size);
490
491 let mut vertex_offset = 0;
492
493 while vertex_offset < vertex_count {
494 let block_size = if vertex_offset + vertex_block_size < vertex_count {
495 vertex_block_size
496 } else {
497 vertex_count - vertex_offset
498 };
499
500 decode_vertex_block(
501 &mut data,
502 &mut vertex_data[vertex_offset * vertex_size..],
503 block_size,
504 vertex_size,
505 &mut last_vertex,
506 )?;
507
508 vertex_offset += block_size;
509 }
510
511 let tail_size = vertex_size.max(TAIL_MAX_SIZE);
512
513 if buffer.len() - data.position() as usize != tail_size {
514 return Err(DecodeError::UnexpectedEof);
515 }
516
517 Ok(())
518}
519
520#[cfg(test)]
521mod test {
522 use super::*;
523
524 #[derive(Clone, Copy, Debug, Default, Eq, Hash, PartialEq)]
525 #[repr(C)]
526 struct PackedVertexOct {
527 p: [u16; 3],
528 n: [i8; 2], t: [u16; 2],
530 }
531
532 const VERTEX_BUFFER: [PackedVertexOct; 4] = [
533 PackedVertexOct {
534 p: [0, 0, 0],
535 n: [0, 0],
536 t: [0, 0],
537 },
538 PackedVertexOct {
539 p: [300, 0, 0],
540 n: [0, 0],
541 t: [500, 0],
542 },
543 PackedVertexOct {
544 p: [0, 300, 0],
545 n: [0, 0],
546 t: [0, 500],
547 },
548 PackedVertexOct {
549 p: [300, 300, 0],
550 n: [0, 0],
551 t: [500, 500],
552 },
553 ];
554
555 const VERTEX_DATA_V0: [u8; 85] = [
556 0xa0, 0x01, 0x3f, 0x00, 0x00, 0x00, 0x58, 0x57, 0x58, 0x01, 0x26, 0x00, 0x00, 0x00, 0x01, 0x0c, 0x00, 0x00,
557 0x00, 0x58, 0x01, 0x08, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x3f, 0x00, 0x00, 0x00, 0x17, 0x18,
558 0x17, 0x01, 0x26, 0x00, 0x00, 0x00, 0x01, 0x0c, 0x00, 0x00, 0x00, 0x17, 0x01, 0x08, 0x00, 0x00, 0x00, 0x00,
559 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
560 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
561 ];
562
563 #[test]
564 fn test_decode_vertex_v0() {
565 let mut decoded = [PackedVertexOct::default(); VERTEX_BUFFER.len()];
566
567 assert!(decode_vertex_buffer(&mut decoded, &VERTEX_DATA_V0).is_ok());
568 assert_eq!(decoded, VERTEX_BUFFER);
569 }
570
571 fn encode_test_vertex() -> Vec<u8> {
572 let mut buffer =
573 vec![0; encode_vertex_buffer_bound(VERTEX_BUFFER.len(), std::mem::size_of::<PackedVertexOct>())];
574
575 let written = encode_vertex_buffer(&mut buffer, &VERTEX_BUFFER, VertexEncodingVersion::default()).unwrap();
576 buffer.resize(written, 0);
577
578 buffer
579 }
580
581 #[test]
582 fn test_encode_vertex_memory_safe() {
583 let mut buffer = encode_test_vertex();
584
585 for i in buffer.len()..=buffer.len() {
587 let result = encode_vertex_buffer(&mut buffer[0..i], &VERTEX_BUFFER, VertexEncodingVersion::default());
588
589 if i == buffer.len() {
590 assert_eq!(result, Some(buffer.len()));
591 } else {
592 assert_eq!(result, None);
593 }
594 }
595 }
596
597 #[test]
598 fn test_decode_vertex_memory_safe() {
599 let buffer = encode_test_vertex();
600
601 let mut decoded = vec![PackedVertexOct::default(); VERTEX_BUFFER.len()];
603
604 for i in buffer.len()..=buffer.len() {
605 let result = decode_vertex_buffer(&mut decoded, &buffer[0..i]);
606
607 if i == buffer.len() {
608 assert!(result.is_ok());
609 } else {
610 assert!(result.is_err());
611 }
612 }
613 }
614
615 #[test]
616 fn test_decode_vertex_reject_extra_bytes() {
617 let mut buffer = encode_test_vertex();
618
619 buffer.push(0);
621
622 let mut decoded = vec![PackedVertexOct::default(); VERTEX_BUFFER.len()];
623 assert!(decode_vertex_buffer(&mut decoded, &buffer).is_err());
624 }
625
626 #[test]
627 fn test_decode_vertex_reject_malformed_headers() {
628 let mut buffer = encode_test_vertex();
629
630 buffer[0] = 0;
632
633 let mut decoded = vec![PackedVertexOct::default(); VERTEX_BUFFER.len()];
634 assert!(decode_vertex_buffer(&mut decoded, &buffer).is_err());
635 }
636
637 #[derive(Clone, Copy, Debug, Default, PartialEq)]
638 #[repr(C)]
639 struct Vertex([u8; 4]);
640
641 #[test]
642 fn test_decode_vertex_bit_groups() {
643 let mut data = [Vertex::default(); 16];
644
645 for (i, v) in data.iter_mut().enumerate() {
647 let i = i as u8;
648 v.0 = [i * 0, i * 1, i * 2, i * 8];
649 }
650
651 let mut buffer = vec![0; encode_vertex_buffer_bound(data.len(), std::mem::size_of::<Vertex>())];
652
653 let written = encode_vertex_buffer(&mut buffer, &data, VertexEncodingVersion::default()).unwrap();
654 buffer.resize(written, 0);
655
656 let mut decoded = [Vertex::default(); 16];
657 assert!(decode_vertex_buffer(&mut decoded, &buffer).is_ok());
658 assert_eq!(decoded, data);
659 }
660
661 #[test]
662 fn test_decode_vertex_bit_group_sentinels() {
663 let mut data = [Vertex::default(); 16];
664
665 for (i, v) in data.iter_mut().enumerate() {
667 let i = i as u8;
668
669 if i == 7 || i == 13 {
670 v.0 = [42; 4];
671 } else {
672 v.0 = [i * 0, i * 1, i * 2, i * 8];
673 }
674 }
675
676 let mut buffer = vec![0; encode_vertex_buffer_bound(data.len(), std::mem::size_of::<Vertex>())];
677
678 let written = encode_vertex_buffer(&mut buffer, &data, VertexEncodingVersion::default()).unwrap();
679 buffer.resize(written, 0);
680
681 let mut decoded = [Vertex::default(); 16];
682 assert!(decode_vertex_buffer(&mut decoded, &buffer).is_ok());
683 assert_eq!(decoded, data);
684 }
685
686 #[test]
687 fn test_decode_vertex_large() {
688 let mut data = [Vertex::default(); 128];
689
690 for (i, v) in data.iter_mut().enumerate() {
692 let i = i as u8;
693 v.0 = [i * 0, i * 1, i.wrapping_mul(2), i.wrapping_mul(8)];
694 }
695
696 let mut buffer = vec![0; encode_vertex_buffer_bound(data.len(), std::mem::size_of::<Vertex>())];
697
698 let written = encode_vertex_buffer(&mut buffer, &data, VertexEncodingVersion::default()).unwrap();
699 buffer.resize(written, 0);
700
701 let mut decoded = [Vertex::default(); 128];
702 assert!(decode_vertex_buffer(&mut decoded, &buffer).is_ok());
703 assert_eq!(decoded, data);
704 }
705
706 #[test]
707 fn test_encode_vertex_empty() {
708 let mut buffer = vec![0; encode_vertex_buffer_bound(0, 16)];
709 let size = encode_vertex_buffer::<PackedVertexOct>(&mut buffer, &[], VertexEncodingVersion::default()).unwrap();
710 buffer.resize(size, 0);
711
712 assert!(decode_vertex_buffer::<PackedVertexOct>(&mut [], &buffer).is_ok());
713 }
714}