xref: /freebsd/contrib/xz/src/liblzma/common/index.c (revision 9398a495eb305d5bf240bc998ee2f6128d75f4ca)
1 // SPDX-License-Identifier: 0BSD
2 
3 ///////////////////////////////////////////////////////////////////////////////
4 //
5 /// \file       index.c
6 /// \brief      Handling of .xz Indexes and some other Stream information
7 //
8 //  Author:     Lasse Collin
9 //
10 ///////////////////////////////////////////////////////////////////////////////
11 
12 #include "common.h"
13 #include "index.h"
14 #include "stream_flags_common.h"
15 
16 
17 /// \brief      How many Records to allocate at once
18 ///
19 /// This should be big enough to avoid making lots of tiny allocations
20 /// but small enough to avoid too much unused memory at once.
21 #define INDEX_GROUP_SIZE 512
22 
23 
24 /// \brief      How many Records can be allocated at once at maximum
25 #define PREALLOC_MAX ((SIZE_MAX - sizeof(index_group)) / sizeof(index_record))
26 
27 
28 /// \brief      Base structure for index_stream and index_group structures
29 typedef struct index_tree_node_s index_tree_node;
30 struct index_tree_node_s {
31 	/// Uncompressed start offset of this Stream (relative to the
32 	/// beginning of the file) or Block (relative to the beginning
33 	/// of the Stream)
34 	lzma_vli uncompressed_base;
35 
36 	/// Compressed start offset of this Stream or Block
37 	lzma_vli compressed_base;
38 
39 	index_tree_node *parent;
40 	index_tree_node *left;
41 	index_tree_node *right;
42 };
43 
44 
45 /// \brief      AVL tree to hold index_stream or index_group structures
46 typedef struct {
47 	/// Root node
48 	index_tree_node *root;
49 
50 	/// Leftmost node. Since the tree will be filled sequentially,
51 	/// this won't change after the first node has been added to
52 	/// the tree.
53 	index_tree_node *leftmost;
54 
55 	/// The rightmost node in the tree. Since the tree is filled
56 	/// sequentially, this is always the node where to add the new data.
57 	index_tree_node *rightmost;
58 
59 	/// Number of nodes in the tree
60 	uint32_t count;
61 
62 } index_tree;
63 
64 
65 typedef struct {
66 	lzma_vli uncompressed_sum;
67 	lzma_vli unpadded_sum;
68 } index_record;
69 
70 
71 typedef struct {
72 	/// Every Record group is part of index_stream.groups tree.
73 	index_tree_node node;
74 
75 	/// Number of Blocks in this Stream before this group.
76 	lzma_vli number_base;
77 
78 	/// Number of Records that can be put in records[].
79 	size_t allocated;
80 
81 	/// Index of the last Record in use.
82 	size_t last;
83 
84 	/// The sizes in this array are stored as cumulative sums relative
85 	/// to the beginning of the Stream. This makes it possible to
86 	/// use binary search in lzma_index_locate().
87 	///
88 	/// Note that the cumulative summing is done specially for
89 	/// unpadded_sum: The previous value is rounded up to the next
90 	/// multiple of four before adding the Unpadded Size of the new
91 	/// Block. The total encoded size of the Blocks in the Stream
92 	/// is records[last].unpadded_sum in the last Record group of
93 	/// the Stream.
94 	///
95 	/// For example, if the Unpadded Sizes are 39, 57, and 81, the
96 	/// stored values are 39, 97 (40 + 57), and 181 (100 + 181).
97 	/// The total encoded size of these Blocks is 184.
98 	///
99 	/// This is a flexible array, because it makes easy to optimize
100 	/// memory usage in case someone concatenates many Streams that
101 	/// have only one or few Blocks.
102 	index_record records[];
103 
104 } index_group;
105 
106 
107 typedef struct {
108 	/// Every index_stream is a node in the tree of Streams.
109 	index_tree_node node;
110 
111 	/// Number of this Stream (first one is 1)
112 	uint32_t number;
113 
114 	/// Total number of Blocks before this Stream
115 	lzma_vli block_number_base;
116 
117 	/// Record groups of this Stream are stored in a tree.
118 	/// It's a T-tree with AVL-tree balancing. There are
119 	/// INDEX_GROUP_SIZE Records per node by default.
120 	/// This keeps the number of memory allocations reasonable
121 	/// and finding a Record is fast.
122 	index_tree groups;
123 
124 	/// Number of Records in this Stream
125 	lzma_vli record_count;
126 
127 	/// Size of the List of Records field in this Stream. This is used
128 	/// together with record_count to calculate the size of the Index
129 	/// field and thus the total size of the Stream.
130 	lzma_vli index_list_size;
131 
132 	/// Stream Flags of this Stream. This is meaningful only if
133 	/// the Stream Flags have been told us with lzma_index_stream_flags().
134 	/// Initially stream_flags.version is set to UINT32_MAX to indicate
135 	/// that the Stream Flags are unknown.
136 	lzma_stream_flags stream_flags;
137 
138 	/// Amount of Stream Padding after this Stream. This defaults to
139 	/// zero and can be set with lzma_index_stream_padding().
140 	lzma_vli stream_padding;
141 
142 } index_stream;
143 
144 
145 struct lzma_index_s {
146 	/// AVL-tree containing the Stream(s). Often there is just one
147 	/// Stream, but using a tree keeps lookups fast even when there
148 	/// are many concatenated Streams.
149 	index_tree streams;
150 
151 	/// Uncompressed size of all the Blocks in the Stream(s)
152 	lzma_vli uncompressed_size;
153 
154 	/// Total size of all the Blocks in the Stream(s)
155 	lzma_vli total_size;
156 
157 	/// Total number of Records in all Streams in this lzma_index
158 	lzma_vli record_count;
159 
160 	/// Size of the List of Records field if all the Streams in this
161 	/// lzma_index were packed into a single Stream (makes it simpler to
162 	/// take many .xz files and combine them into a single Stream).
163 	///
164 	/// This value together with record_count is needed to calculate
165 	/// Backward Size that is stored into Stream Footer.
166 	lzma_vli index_list_size;
167 
168 	/// How many Records to allocate at once in lzma_index_append().
169 	/// This defaults to INDEX_GROUP_SIZE but can be overridden with
170 	/// lzma_index_prealloc().
171 	size_t prealloc;
172 
173 	/// Bitmask indicating what integrity check types have been used
174 	/// as set by lzma_index_stream_flags(). The bit of the last Stream
175 	/// is not included here, since it is possible to change it by
176 	/// calling lzma_index_stream_flags() again.
177 	uint32_t checks;
178 };
179 
180 
181 static void
182 index_tree_init(index_tree *tree)
183 {
184 	tree->root = NULL;
185 	tree->leftmost = NULL;
186 	tree->rightmost = NULL;
187 	tree->count = 0;
188 	return;
189 }
190 
191 
192 /// Helper for index_tree_end()
193 static void
194 index_tree_node_end(index_tree_node *node, const lzma_allocator *allocator,
195 		void (*free_func)(void *node, const lzma_allocator *allocator))
196 {
197 	// The tree won't ever be very huge, so recursion should be fine.
198 	// 20 levels in the tree is likely quite a lot already in practice.
199 	if (node->left != NULL)
200 		index_tree_node_end(node->left, allocator, free_func);
201 
202 	if (node->right != NULL)
203 		index_tree_node_end(node->right, allocator, free_func);
204 
205 	free_func(node, allocator);
206 	return;
207 }
208 
209 
210 /// Free the memory allocated for a tree. Each node is freed using the
211 /// given free_func which is either &lzma_free or &index_stream_end.
212 /// The latter is used to free the Record groups from each index_stream
213 /// before freeing the index_stream itself.
214 static void
215 index_tree_end(index_tree *tree, const lzma_allocator *allocator,
216 		void (*free_func)(void *node, const lzma_allocator *allocator))
217 {
218 	assert(free_func != NULL);
219 
220 	if (tree->root != NULL)
221 		index_tree_node_end(tree->root, allocator, free_func);
222 
223 	return;
224 }
225 
226 
227 /// Add a new node to the tree. node->uncompressed_base and
228 /// node->compressed_base must have been set by the caller already.
229 static void
230 index_tree_append(index_tree *tree, index_tree_node *node)
231 {
232 	node->parent = tree->rightmost;
233 	node->left = NULL;
234 	node->right = NULL;
235 
236 	++tree->count;
237 
238 	// Handle the special case of adding the first node.
239 	if (tree->root == NULL) {
240 		tree->root = node;
241 		tree->leftmost = node;
242 		tree->rightmost = node;
243 		return;
244 	}
245 
246 	// The tree is always filled sequentially.
247 	assert(tree->rightmost->uncompressed_base <= node->uncompressed_base);
248 	assert(tree->rightmost->compressed_base < node->compressed_base);
249 
250 	// Add the new node after the rightmost node. It's the correct
251 	// place due to the reason above.
252 	tree->rightmost->right = node;
253 	tree->rightmost = node;
254 
255 	// Balance the AVL-tree if needed. We don't need to keep the balance
256 	// factors in nodes, because we always fill the tree sequentially,
257 	// and thus know the state of the tree just by looking at the node
258 	// count. From the node count we can calculate how many steps to go
259 	// up in the tree to find the rotation root.
260 	uint32_t up = tree->count ^ (UINT32_C(1) << bsr32(tree->count));
261 	if (up != 0) {
262 		// Locate the root node for the rotation.
263 		up = ctz32(tree->count) + 2;
264 		do {
265 			node = node->parent;
266 		} while (--up > 0);
267 
268 		// Rotate left using node as the rotation root.
269 		index_tree_node *pivot = node->right;
270 
271 		if (node->parent == NULL) {
272 			tree->root = pivot;
273 		} else {
274 			assert(node->parent->right == node);
275 			node->parent->right = pivot;
276 		}
277 
278 		pivot->parent = node->parent;
279 
280 		node->right = pivot->left;
281 		if (node->right != NULL)
282 			node->right->parent = node;
283 
284 		pivot->left = node;
285 		node->parent = pivot;
286 	}
287 
288 	return;
289 }
290 
291 
292 /// Get the next node in the tree. Return NULL if there are no more nodes.
293 static void *
294 index_tree_next(const index_tree_node *node)
295 {
296 	if (node->right != NULL) {
297 		node = node->right;
298 		while (node->left != NULL)
299 			node = node->left;
300 
301 		return (void *)(node);
302 	}
303 
304 	while (node->parent != NULL && node->parent->right == node)
305 		node = node->parent;
306 
307 	return (void *)(node->parent);
308 }
309 
310 
311 /// Locate a node that contains the given uncompressed offset. It is
312 /// caller's job to check that target is not bigger than the uncompressed
313 /// size of the tree (the last node would be returned in that case still).
314 static void *
315 index_tree_locate(const index_tree *tree, lzma_vli target)
316 {
317 	const index_tree_node *result = NULL;
318 	const index_tree_node *node = tree->root;
319 
320 	assert(tree->leftmost == NULL
321 			|| tree->leftmost->uncompressed_base == 0);
322 
323 	// Consecutive nodes may have the same uncompressed_base.
324 	// We must pick the rightmost one.
325 	while (node != NULL) {
326 		if (node->uncompressed_base > target) {
327 			node = node->left;
328 		} else {
329 			result = node;
330 			node = node->right;
331 		}
332 	}
333 
334 	return (void *)(result);
335 }
336 
337 
338 /// Allocate and initialize a new Stream using the given base offsets.
339 static index_stream *
340 index_stream_init(lzma_vli compressed_base, lzma_vli uncompressed_base,
341 		uint32_t stream_number, lzma_vli block_number_base,
342 		const lzma_allocator *allocator)
343 {
344 	index_stream *s = lzma_alloc(sizeof(index_stream), allocator);
345 	if (s == NULL)
346 		return NULL;
347 
348 	s->node.uncompressed_base = uncompressed_base;
349 	s->node.compressed_base = compressed_base;
350 	s->node.parent = NULL;
351 	s->node.left = NULL;
352 	s->node.right = NULL;
353 
354 	s->number = stream_number;
355 	s->block_number_base = block_number_base;
356 
357 	index_tree_init(&s->groups);
358 
359 	s->record_count = 0;
360 	s->index_list_size = 0;
361 	s->stream_flags.version = UINT32_MAX;
362 	s->stream_padding = 0;
363 
364 	return s;
365 }
366 
367 
368 /// Free the memory allocated for a Stream and its Record groups.
369 static void
370 index_stream_end(void *node, const lzma_allocator *allocator)
371 {
372 	index_stream *s = node;
373 	index_tree_end(&s->groups, allocator, &lzma_free);
374 	lzma_free(s, allocator);
375 	return;
376 }
377 
378 
379 static lzma_index *
380 index_init_plain(const lzma_allocator *allocator)
381 {
382 	lzma_index *i = lzma_alloc(sizeof(lzma_index), allocator);
383 	if (i != NULL) {
384 		index_tree_init(&i->streams);
385 		i->uncompressed_size = 0;
386 		i->total_size = 0;
387 		i->record_count = 0;
388 		i->index_list_size = 0;
389 		i->prealloc = INDEX_GROUP_SIZE;
390 		i->checks = 0;
391 	}
392 
393 	return i;
394 }
395 
396 
397 extern LZMA_API(lzma_index *)
398 lzma_index_init(const lzma_allocator *allocator)
399 {
400 	lzma_index *i = index_init_plain(allocator);
401 	if (i == NULL)
402 		return NULL;
403 
404 	index_stream *s = index_stream_init(0, 0, 1, 0, allocator);
405 	if (s == NULL) {
406 		lzma_free(i, allocator);
407 		return NULL;
408 	}
409 
410 	index_tree_append(&i->streams, &s->node);
411 
412 	return i;
413 }
414 
415 
416 extern LZMA_API(void)
417 lzma_index_end(lzma_index *i, const lzma_allocator *allocator)
418 {
419 	// NOTE: If you modify this function, check also the bottom
420 	// of lzma_index_cat().
421 	if (i != NULL) {
422 		index_tree_end(&i->streams, allocator, &index_stream_end);
423 		lzma_free(i, allocator);
424 	}
425 
426 	return;
427 }
428 
429 
430 extern void
431 lzma_index_prealloc(lzma_index *i, lzma_vli records)
432 {
433 	if (records > PREALLOC_MAX)
434 		records = PREALLOC_MAX;
435 
436 	i->prealloc = (size_t)(records);
437 	return;
438 }
439 
440 
441 extern LZMA_API(uint64_t)
442 lzma_index_memusage(lzma_vli streams, lzma_vli blocks)
443 {
444 	// This calculates an upper bound that is only a little bit
445 	// bigger than the exact maximum memory usage with the given
446 	// parameters.
447 
448 	// Typical malloc() overhead is 2 * sizeof(void *) but we take
449 	// a little bit extra just in case. Using LZMA_MEMUSAGE_BASE
450 	// instead would give too inaccurate estimate.
451 	const size_t alloc_overhead = 4 * sizeof(void *);
452 
453 	// Amount of memory needed for each Stream base structures.
454 	// We assume that every Stream has at least one Block and
455 	// thus at least one group.
456 	const size_t stream_base = sizeof(index_stream)
457 			+ sizeof(index_group) + 2 * alloc_overhead;
458 
459 	// Amount of memory needed per group.
460 	const size_t group_base = sizeof(index_group)
461 			+ INDEX_GROUP_SIZE * sizeof(index_record)
462 			+ alloc_overhead;
463 
464 	// Number of groups. There may actually be more, but that overhead
465 	// has been taken into account in stream_base already.
466 	const lzma_vli groups
467 			= (blocks + INDEX_GROUP_SIZE - 1) / INDEX_GROUP_SIZE;
468 
469 	// Memory used by index_stream and index_group structures.
470 	const uint64_t streams_mem = streams * stream_base;
471 	const uint64_t groups_mem = groups * group_base;
472 
473 	// Memory used by the base structure.
474 	const uint64_t index_base = sizeof(lzma_index) + alloc_overhead;
475 
476 	// Validate the arguments and catch integer overflows.
477 	// Maximum number of Streams is "only" UINT32_MAX, because
478 	// that limit is used by the tree containing the Streams.
479 	const uint64_t limit = UINT64_MAX - index_base;
480 	if (streams == 0 || streams > UINT32_MAX || blocks > LZMA_VLI_MAX
481 			|| streams > limit / stream_base
482 			|| groups > limit / group_base
483 			|| limit - streams_mem < groups_mem)
484 		return UINT64_MAX;
485 
486 	return index_base + streams_mem + groups_mem;
487 }
488 
489 
490 extern LZMA_API(uint64_t)
491 lzma_index_memused(const lzma_index *i)
492 {
493 	return lzma_index_memusage(i->streams.count, i->record_count);
494 }
495 
496 
497 extern LZMA_API(lzma_vli)
498 lzma_index_block_count(const lzma_index *i)
499 {
500 	return i->record_count;
501 }
502 
503 
504 extern LZMA_API(lzma_vli)
505 lzma_index_stream_count(const lzma_index *i)
506 {
507 	return i->streams.count;
508 }
509 
510 
511 extern LZMA_API(lzma_vli)
512 lzma_index_size(const lzma_index *i)
513 {
514 	return index_size(i->record_count, i->index_list_size);
515 }
516 
517 
518 extern LZMA_API(lzma_vli)
519 lzma_index_total_size(const lzma_index *i)
520 {
521 	return i->total_size;
522 }
523 
524 
525 extern LZMA_API(lzma_vli)
526 lzma_index_stream_size(const lzma_index *i)
527 {
528 	// Stream Header + Blocks + Index + Stream Footer
529 	return LZMA_STREAM_HEADER_SIZE + i->total_size
530 			+ index_size(i->record_count, i->index_list_size)
531 			+ LZMA_STREAM_HEADER_SIZE;
532 }
533 
534 
535 static lzma_vli
536 index_file_size(lzma_vli compressed_base, lzma_vli unpadded_sum,
537 		lzma_vli record_count, lzma_vli index_list_size,
538 		lzma_vli stream_padding)
539 {
540 	// Earlier Streams and Stream Paddings + Stream Header
541 	// + Blocks + Index + Stream Footer + Stream Padding
542 	//
543 	// This might go over LZMA_VLI_MAX due to too big unpadded_sum
544 	// when this function is used in lzma_index_append().
545 	lzma_vli file_size = compressed_base + 2 * LZMA_STREAM_HEADER_SIZE
546 			+ stream_padding + vli_ceil4(unpadded_sum);
547 	if (file_size > LZMA_VLI_MAX)
548 		return LZMA_VLI_UNKNOWN;
549 
550 	// The same applies here.
551 	file_size += index_size(record_count, index_list_size);
552 	if (file_size > LZMA_VLI_MAX)
553 		return LZMA_VLI_UNKNOWN;
554 
555 	return file_size;
556 }
557 
558 
559 extern LZMA_API(lzma_vli)
560 lzma_index_file_size(const lzma_index *i)
561 {
562 	const index_stream *s = (const index_stream *)(i->streams.rightmost);
563 	const index_group *g = (const index_group *)(s->groups.rightmost);
564 	return index_file_size(s->node.compressed_base,
565 			g == NULL ? 0 : g->records[g->last].unpadded_sum,
566 			s->record_count, s->index_list_size,
567 			s->stream_padding);
568 }
569 
570 
571 extern LZMA_API(lzma_vli)
572 lzma_index_uncompressed_size(const lzma_index *i)
573 {
574 	return i->uncompressed_size;
575 }
576 
577 
578 extern LZMA_API(uint32_t)
579 lzma_index_checks(const lzma_index *i)
580 {
581 	uint32_t checks = i->checks;
582 
583 	// Get the type of the Check of the last Stream too.
584 	const index_stream *s = (const index_stream *)(i->streams.rightmost);
585 	if (s->stream_flags.version != UINT32_MAX)
586 		checks |= UINT32_C(1) << s->stream_flags.check;
587 
588 	return checks;
589 }
590 
591 
592 extern uint32_t
593 lzma_index_padding_size(const lzma_index *i)
594 {
595 	return (LZMA_VLI_C(4) - index_size_unpadded(
596 			i->record_count, i->index_list_size)) & 3;
597 }
598 
599 
600 extern LZMA_API(lzma_ret)
601 lzma_index_stream_flags(lzma_index *i, const lzma_stream_flags *stream_flags)
602 {
603 	if (i == NULL || stream_flags == NULL)
604 		return LZMA_PROG_ERROR;
605 
606 	// Validate the Stream Flags.
607 	return_if_error(lzma_stream_flags_compare(
608 			stream_flags, stream_flags));
609 
610 	index_stream *s = (index_stream *)(i->streams.rightmost);
611 	s->stream_flags = *stream_flags;
612 
613 	return LZMA_OK;
614 }
615 
616 
617 extern LZMA_API(lzma_ret)
618 lzma_index_stream_padding(lzma_index *i, lzma_vli stream_padding)
619 {
620 	if (i == NULL || stream_padding > LZMA_VLI_MAX
621 			|| (stream_padding & 3) != 0)
622 		return LZMA_PROG_ERROR;
623 
624 	index_stream *s = (index_stream *)(i->streams.rightmost);
625 
626 	// Check that the new value won't make the file grow too big.
627 	const lzma_vli old_stream_padding = s->stream_padding;
628 	s->stream_padding = 0;
629 	if (lzma_index_file_size(i) + stream_padding > LZMA_VLI_MAX) {
630 		s->stream_padding = old_stream_padding;
631 		return LZMA_DATA_ERROR;
632 	}
633 
634 	s->stream_padding = stream_padding;
635 	return LZMA_OK;
636 }
637 
638 
639 extern LZMA_API(lzma_ret)
640 lzma_index_append(lzma_index *i, const lzma_allocator *allocator,
641 		lzma_vli unpadded_size, lzma_vli uncompressed_size)
642 {
643 	// Validate.
644 	if (i == NULL || unpadded_size < UNPADDED_SIZE_MIN
645 			|| unpadded_size > UNPADDED_SIZE_MAX
646 			|| uncompressed_size > LZMA_VLI_MAX)
647 		return LZMA_PROG_ERROR;
648 
649 	index_stream *s = (index_stream *)(i->streams.rightmost);
650 	index_group *g = (index_group *)(s->groups.rightmost);
651 
652 	const lzma_vli compressed_base = g == NULL ? 0
653 			: vli_ceil4(g->records[g->last].unpadded_sum);
654 	const lzma_vli uncompressed_base = g == NULL ? 0
655 			: g->records[g->last].uncompressed_sum;
656 	const uint32_t index_list_size_add = lzma_vli_size(unpadded_size)
657 			+ lzma_vli_size(uncompressed_size);
658 
659 	// Check that uncompressed size will not overflow.
660 	if (uncompressed_base + uncompressed_size > LZMA_VLI_MAX)
661 		return LZMA_DATA_ERROR;
662 
663 	// Check that the new unpadded sum will not overflow. This is
664 	// checked again in index_file_size(), but the unpadded sum is
665 	// passed to vli_ceil4() which expects a valid lzma_vli value.
666 	if (compressed_base + unpadded_size > UNPADDED_SIZE_MAX)
667 		return LZMA_DATA_ERROR;
668 
669 	// Check that the file size will stay within limits.
670 	if (index_file_size(s->node.compressed_base,
671 			compressed_base + unpadded_size, s->record_count + 1,
672 			s->index_list_size + index_list_size_add,
673 			s->stream_padding) == LZMA_VLI_UNKNOWN)
674 		return LZMA_DATA_ERROR;
675 
676 	// The size of the Index field must not exceed the maximum value
677 	// that can be stored in the Backward Size field.
678 	if (index_size(i->record_count + 1,
679 			i->index_list_size + index_list_size_add)
680 			> LZMA_BACKWARD_SIZE_MAX)
681 		return LZMA_DATA_ERROR;
682 
683 	if (g != NULL && g->last + 1 < g->allocated) {
684 		// There is space in the last group at least for one Record.
685 		++g->last;
686 	} else {
687 		// We need to allocate a new group.
688 		g = lzma_alloc(sizeof(index_group)
689 				+ i->prealloc * sizeof(index_record),
690 				allocator);
691 		if (g == NULL)
692 			return LZMA_MEM_ERROR;
693 
694 		g->last = 0;
695 		g->allocated = i->prealloc;
696 
697 		// Reset prealloc so that if the application happens to
698 		// add new Records, the allocation size will be sane.
699 		i->prealloc = INDEX_GROUP_SIZE;
700 
701 		// Set the start offsets of this group.
702 		g->node.uncompressed_base = uncompressed_base;
703 		g->node.compressed_base = compressed_base;
704 		g->number_base = s->record_count + 1;
705 
706 		// Add the new group to the Stream.
707 		index_tree_append(&s->groups, &g->node);
708 	}
709 
710 	// Add the new Record to the group.
711 	g->records[g->last].uncompressed_sum
712 			= uncompressed_base + uncompressed_size;
713 	g->records[g->last].unpadded_sum
714 			= compressed_base + unpadded_size;
715 
716 	// Update the totals.
717 	++s->record_count;
718 	s->index_list_size += index_list_size_add;
719 
720 	i->total_size += vli_ceil4(unpadded_size);
721 	i->uncompressed_size += uncompressed_size;
722 	++i->record_count;
723 	i->index_list_size += index_list_size_add;
724 
725 	return LZMA_OK;
726 }
727 
728 
729 /// Structure to pass info to index_cat_helper()
730 typedef struct {
731 	/// Uncompressed size of the destination
732 	lzma_vli uncompressed_size;
733 
734 	/// Compressed file size of the destination
735 	lzma_vli file_size;
736 
737 	/// Same as above but for Block numbers
738 	lzma_vli block_number_add;
739 
740 	/// Number of Streams that were in the destination index before we
741 	/// started appending new Streams from the source index. This is
742 	/// used to fix the Stream numbering.
743 	uint32_t stream_number_add;
744 
745 	/// Destination index' Stream tree
746 	index_tree *streams;
747 
748 } index_cat_info;
749 
750 
751 /// Add the Stream nodes from the source index to dest using recursion.
752 /// Simplest iterative traversal of the source tree wouldn't work, because
753 /// we update the pointers in nodes when moving them to the destination tree.
754 static void
755 index_cat_helper(const index_cat_info *info, index_stream *this)
756 {
757 	index_stream *left = (index_stream *)(this->node.left);
758 	index_stream *right = (index_stream *)(this->node.right);
759 
760 	if (left != NULL)
761 		index_cat_helper(info, left);
762 
763 	this->node.uncompressed_base += info->uncompressed_size;
764 	this->node.compressed_base += info->file_size;
765 	this->number += info->stream_number_add;
766 	this->block_number_base += info->block_number_add;
767 	index_tree_append(info->streams, &this->node);
768 
769 	if (right != NULL)
770 		index_cat_helper(info, right);
771 
772 	return;
773 }
774 
775 
776 extern LZMA_API(lzma_ret)
777 lzma_index_cat(lzma_index *restrict dest, lzma_index *restrict src,
778 		const lzma_allocator *allocator)
779 {
780 	if (dest == NULL || src == NULL)
781 		return LZMA_PROG_ERROR;
782 
783 	const lzma_vli dest_file_size = lzma_index_file_size(dest);
784 
785 	// Check that we don't exceed the file size limits.
786 	if (dest_file_size + lzma_index_file_size(src) > LZMA_VLI_MAX
787 			|| dest->uncompressed_size + src->uncompressed_size
788 				> LZMA_VLI_MAX)
789 		return LZMA_DATA_ERROR;
790 
791 	// Check that the encoded size of the combined lzma_indexes stays
792 	// within limits. In theory, this should be done only if we know
793 	// that the user plans to actually combine the Streams and thus
794 	// construct a single Index (probably rare). However, exceeding
795 	// this limit is quite theoretical, so we do this check always
796 	// to simplify things elsewhere.
797 	{
798 		const lzma_vli dest_size = index_size_unpadded(
799 				dest->record_count, dest->index_list_size);
800 		const lzma_vli src_size = index_size_unpadded(
801 				src->record_count, src->index_list_size);
802 		if (vli_ceil4(dest_size + src_size) > LZMA_BACKWARD_SIZE_MAX)
803 			return LZMA_DATA_ERROR;
804 	}
805 
806 	// Optimize the last group to minimize memory usage. Allocation has
807 	// to be done before modifying dest or src.
808 	{
809 		index_stream *s = (index_stream *)(dest->streams.rightmost);
810 		index_group *g = (index_group *)(s->groups.rightmost);
811 		if (g != NULL && g->last + 1 < g->allocated) {
812 			assert(g->node.left == NULL);
813 			assert(g->node.right == NULL);
814 
815 			index_group *newg = lzma_alloc(sizeof(index_group)
816 					+ (g->last + 1)
817 					* sizeof(index_record),
818 					allocator);
819 			if (newg == NULL)
820 				return LZMA_MEM_ERROR;
821 
822 			newg->node = g->node;
823 			newg->allocated = g->last + 1;
824 			newg->last = g->last;
825 			newg->number_base = g->number_base;
826 
827 			memcpy(newg->records, g->records, newg->allocated
828 					* sizeof(index_record));
829 
830 			if (g->node.parent != NULL) {
831 				assert(g->node.parent->right == &g->node);
832 				g->node.parent->right = &newg->node;
833 			}
834 
835 			if (s->groups.leftmost == &g->node) {
836 				assert(s->groups.root == &g->node);
837 				s->groups.leftmost = &newg->node;
838 				s->groups.root = &newg->node;
839 			}
840 
841 			assert(s->groups.rightmost == &g->node);
842 			s->groups.rightmost = &newg->node;
843 
844 			lzma_free(g, allocator);
845 
846 			// NOTE: newg isn't leaked here because
847 			// newg == (void *)&newg->node.
848 		}
849 	}
850 
851 	// dest->checks includes the check types of all except the last Stream
852 	// in dest. Set the bit for the check type of the last Stream now so
853 	// that it won't get lost when Stream(s) from src are appended to dest.
854 	dest->checks = lzma_index_checks(dest);
855 
856 	// Add all the Streams from src to dest. Update the base offsets
857 	// of each Stream from src.
858 	const index_cat_info info = {
859 		.uncompressed_size = dest->uncompressed_size,
860 		.file_size = dest_file_size,
861 		.stream_number_add = dest->streams.count,
862 		.block_number_add = dest->record_count,
863 		.streams = &dest->streams,
864 	};
865 	index_cat_helper(&info, (index_stream *)(src->streams.root));
866 
867 	// Update info about all the combined Streams.
868 	dest->uncompressed_size += src->uncompressed_size;
869 	dest->total_size += src->total_size;
870 	dest->record_count += src->record_count;
871 	dest->index_list_size += src->index_list_size;
872 	dest->checks |= src->checks;
873 
874 	// There's nothing else left in src than the base structure.
875 	lzma_free(src, allocator);
876 
877 	return LZMA_OK;
878 }
879 
880 
881 /// Duplicate an index_stream.
882 static index_stream *
883 index_dup_stream(const index_stream *src, const lzma_allocator *allocator)
884 {
885 	// Catch a somewhat theoretical integer overflow.
886 	if (src->record_count > PREALLOC_MAX)
887 		return NULL;
888 
889 	// Allocate and initialize a new Stream.
890 	index_stream *dest = index_stream_init(src->node.compressed_base,
891 			src->node.uncompressed_base, src->number,
892 			src->block_number_base, allocator);
893 	if (dest == NULL)
894 		return NULL;
895 
896 	// Copy the overall information.
897 	dest->record_count = src->record_count;
898 	dest->index_list_size = src->index_list_size;
899 	dest->stream_flags = src->stream_flags;
900 	dest->stream_padding = src->stream_padding;
901 
902 	// Return if there are no groups to duplicate.
903 	if (src->groups.leftmost == NULL)
904 		return dest;
905 
906 	// Allocate memory for the Records. We put all the Records into
907 	// a single group. It's simplest and also tends to make
908 	// lzma_index_locate() a little bit faster with very big Indexes.
909 	index_group *destg = lzma_alloc(sizeof(index_group)
910 			+ src->record_count * sizeof(index_record),
911 			allocator);
912 	if (destg == NULL) {
913 		index_stream_end(dest, allocator);
914 		return NULL;
915 	}
916 
917 	// Initialize destg.
918 	destg->node.uncompressed_base = 0;
919 	destg->node.compressed_base = 0;
920 	destg->number_base = 1;
921 	destg->allocated = src->record_count;
922 	destg->last = src->record_count - 1;
923 
924 	// Go through all the groups in src and copy the Records into destg.
925 	const index_group *srcg = (const index_group *)(src->groups.leftmost);
926 	size_t i = 0;
927 	do {
928 		memcpy(destg->records + i, srcg->records,
929 				(srcg->last + 1) * sizeof(index_record));
930 		i += srcg->last + 1;
931 		srcg = index_tree_next(&srcg->node);
932 	} while (srcg != NULL);
933 
934 	assert(i == destg->allocated);
935 
936 	// Add the group to the new Stream.
937 	index_tree_append(&dest->groups, &destg->node);
938 
939 	return dest;
940 }
941 
942 
943 extern LZMA_API(lzma_index *)
944 lzma_index_dup(const lzma_index *src, const lzma_allocator *allocator)
945 {
946 	// Allocate the base structure (no initial Stream).
947 	lzma_index *dest = index_init_plain(allocator);
948 	if (dest == NULL)
949 		return NULL;
950 
951 	// Copy the totals.
952 	dest->uncompressed_size = src->uncompressed_size;
953 	dest->total_size = src->total_size;
954 	dest->record_count = src->record_count;
955 	dest->index_list_size = src->index_list_size;
956 
957 	// Copy the Streams and the groups in them.
958 	const index_stream *srcstream
959 			= (const index_stream *)(src->streams.leftmost);
960 	do {
961 		index_stream *deststream = index_dup_stream(
962 				srcstream, allocator);
963 		if (deststream == NULL) {
964 			lzma_index_end(dest, allocator);
965 			return NULL;
966 		}
967 
968 		index_tree_append(&dest->streams, &deststream->node);
969 
970 		srcstream = index_tree_next(&srcstream->node);
971 	} while (srcstream != NULL);
972 
973 	return dest;
974 }
975 
976 
977 /// Indexing for lzma_index_iter.internal[]
978 enum {
979 	ITER_INDEX,
980 	ITER_STREAM,
981 	ITER_GROUP,
982 	ITER_RECORD,
983 	ITER_METHOD,
984 };
985 
986 
987 /// Values for lzma_index_iter.internal[ITER_METHOD].s
988 enum {
989 	ITER_METHOD_NORMAL,
990 	ITER_METHOD_NEXT,
991 	ITER_METHOD_LEFTMOST,
992 };
993 
994 
995 static void
996 iter_set_info(lzma_index_iter *iter)
997 {
998 	const lzma_index *i = iter->internal[ITER_INDEX].p;
999 	const index_stream *stream = iter->internal[ITER_STREAM].p;
1000 	const index_group *group = iter->internal[ITER_GROUP].p;
1001 	const size_t record = iter->internal[ITER_RECORD].s;
1002 
1003 	// lzma_index_iter.internal must not contain a pointer to the last
1004 	// group in the index, because that may be reallocated by
1005 	// lzma_index_cat().
1006 	if (group == NULL) {
1007 		// There are no groups.
1008 		assert(stream->groups.root == NULL);
1009 		iter->internal[ITER_METHOD].s = ITER_METHOD_LEFTMOST;
1010 
1011 	} else if (i->streams.rightmost != &stream->node
1012 			|| stream->groups.rightmost != &group->node) {
1013 		// The group is not not the last group in the index.
1014 		iter->internal[ITER_METHOD].s = ITER_METHOD_NORMAL;
1015 
1016 	} else if (stream->groups.leftmost != &group->node) {
1017 		// The group isn't the only group in the Stream, thus we
1018 		// know that it must have a parent group i.e. it's not
1019 		// the root node.
1020 		assert(stream->groups.root != &group->node);
1021 		assert(group->node.parent->right == &group->node);
1022 		iter->internal[ITER_METHOD].s = ITER_METHOD_NEXT;
1023 		iter->internal[ITER_GROUP].p = group->node.parent;
1024 
1025 	} else {
1026 		// The Stream has only one group.
1027 		assert(stream->groups.root == &group->node);
1028 		assert(group->node.parent == NULL);
1029 		iter->internal[ITER_METHOD].s = ITER_METHOD_LEFTMOST;
1030 		iter->internal[ITER_GROUP].p = NULL;
1031 	}
1032 
1033 	// NOTE: lzma_index_iter.stream.number is lzma_vli but we use uint32_t
1034 	// internally.
1035 	iter->stream.number = stream->number;
1036 	iter->stream.block_count = stream->record_count;
1037 	iter->stream.compressed_offset = stream->node.compressed_base;
1038 	iter->stream.uncompressed_offset = stream->node.uncompressed_base;
1039 
1040 	// iter->stream.flags will be NULL if the Stream Flags haven't been
1041 	// set with lzma_index_stream_flags().
1042 	iter->stream.flags = stream->stream_flags.version == UINT32_MAX
1043 			? NULL : &stream->stream_flags;
1044 	iter->stream.padding = stream->stream_padding;
1045 
1046 	if (stream->groups.rightmost == NULL) {
1047 		// Stream has no Blocks.
1048 		iter->stream.compressed_size = index_size(0, 0)
1049 				+ 2 * LZMA_STREAM_HEADER_SIZE;
1050 		iter->stream.uncompressed_size = 0;
1051 	} else {
1052 		const index_group *g = (const index_group *)(
1053 				stream->groups.rightmost);
1054 
1055 		// Stream Header + Stream Footer + Index + Blocks
1056 		iter->stream.compressed_size = 2 * LZMA_STREAM_HEADER_SIZE
1057 				+ index_size(stream->record_count,
1058 					stream->index_list_size)
1059 				+ vli_ceil4(g->records[g->last].unpadded_sum);
1060 		iter->stream.uncompressed_size
1061 				= g->records[g->last].uncompressed_sum;
1062 	}
1063 
1064 	if (group != NULL) {
1065 		iter->block.number_in_stream = group->number_base + record;
1066 		iter->block.number_in_file = iter->block.number_in_stream
1067 				+ stream->block_number_base;
1068 
1069 		iter->block.compressed_stream_offset
1070 				= record == 0 ? group->node.compressed_base
1071 				: vli_ceil4(group->records[
1072 					record - 1].unpadded_sum);
1073 		iter->block.uncompressed_stream_offset
1074 				= record == 0 ? group->node.uncompressed_base
1075 				: group->records[record - 1].uncompressed_sum;
1076 
1077 		iter->block.uncompressed_size
1078 				= group->records[record].uncompressed_sum
1079 				- iter->block.uncompressed_stream_offset;
1080 		iter->block.unpadded_size
1081 				= group->records[record].unpadded_sum
1082 				- iter->block.compressed_stream_offset;
1083 		iter->block.total_size = vli_ceil4(iter->block.unpadded_size);
1084 
1085 		iter->block.compressed_stream_offset
1086 				+= LZMA_STREAM_HEADER_SIZE;
1087 
1088 		iter->block.compressed_file_offset
1089 				= iter->block.compressed_stream_offset
1090 				+ iter->stream.compressed_offset;
1091 		iter->block.uncompressed_file_offset
1092 				= iter->block.uncompressed_stream_offset
1093 				+ iter->stream.uncompressed_offset;
1094 	}
1095 
1096 	return;
1097 }
1098 
1099 
1100 extern LZMA_API(void)
1101 lzma_index_iter_init(lzma_index_iter *iter, const lzma_index *i)
1102 {
1103 	iter->internal[ITER_INDEX].p = i;
1104 	lzma_index_iter_rewind(iter);
1105 	return;
1106 }
1107 
1108 
1109 extern LZMA_API(void)
1110 lzma_index_iter_rewind(lzma_index_iter *iter)
1111 {
1112 	iter->internal[ITER_STREAM].p = NULL;
1113 	iter->internal[ITER_GROUP].p = NULL;
1114 	iter->internal[ITER_RECORD].s = 0;
1115 	iter->internal[ITER_METHOD].s = ITER_METHOD_NORMAL;
1116 	return;
1117 }
1118 
1119 
1120 extern LZMA_API(lzma_bool)
1121 lzma_index_iter_next(lzma_index_iter *iter, lzma_index_iter_mode mode)
1122 {
1123 	// Catch unsupported mode values.
1124 	if ((unsigned int)(mode) > LZMA_INDEX_ITER_NONEMPTY_BLOCK)
1125 		return true;
1126 
1127 	const lzma_index *i = iter->internal[ITER_INDEX].p;
1128 	const index_stream *stream = iter->internal[ITER_STREAM].p;
1129 	const index_group *group = NULL;
1130 	size_t record = iter->internal[ITER_RECORD].s;
1131 
1132 	// If we are being asked for the next Stream, leave group to NULL
1133 	// so that the rest of the this function thinks that this Stream
1134 	// has no groups and will thus go to the next Stream.
1135 	if (mode != LZMA_INDEX_ITER_STREAM) {
1136 		// Get the pointer to the current group. See iter_set_inf()
1137 		// for explanation.
1138 		switch (iter->internal[ITER_METHOD].s) {
1139 		case ITER_METHOD_NORMAL:
1140 			group = iter->internal[ITER_GROUP].p;
1141 			break;
1142 
1143 		case ITER_METHOD_NEXT:
1144 			group = index_tree_next(iter->internal[ITER_GROUP].p);
1145 			break;
1146 
1147 		case ITER_METHOD_LEFTMOST:
1148 			group = (const index_group *)(
1149 					stream->groups.leftmost);
1150 			break;
1151 		}
1152 	}
1153 
1154 again:
1155 	if (stream == NULL) {
1156 		// We at the beginning of the lzma_index.
1157 		// Locate the first Stream.
1158 		stream = (const index_stream *)(i->streams.leftmost);
1159 		if (mode >= LZMA_INDEX_ITER_BLOCK) {
1160 			// Since we are being asked to return information
1161 			// about the first a Block, skip Streams that have
1162 			// no Blocks.
1163 			while (stream->groups.leftmost == NULL) {
1164 				stream = index_tree_next(&stream->node);
1165 				if (stream == NULL)
1166 					return true;
1167 			}
1168 		}
1169 
1170 		// Start from the first Record in the Stream.
1171 		group = (const index_group *)(stream->groups.leftmost);
1172 		record = 0;
1173 
1174 	} else if (group != NULL && record < group->last) {
1175 		// The next Record is in the same group.
1176 		++record;
1177 
1178 	} else {
1179 		// This group has no more Records or this Stream has
1180 		// no Blocks at all.
1181 		record = 0;
1182 
1183 		// If group is not NULL, this Stream has at least one Block
1184 		// and thus at least one group. Find the next group.
1185 		if (group != NULL)
1186 			group = index_tree_next(&group->node);
1187 
1188 		if (group == NULL) {
1189 			// This Stream has no more Records. Find the next
1190 			// Stream. If we are being asked to return information
1191 			// about a Block, we skip empty Streams.
1192 			do {
1193 				stream = index_tree_next(&stream->node);
1194 				if (stream == NULL)
1195 					return true;
1196 			} while (mode >= LZMA_INDEX_ITER_BLOCK
1197 					&& stream->groups.leftmost == NULL);
1198 
1199 			group = (const index_group *)(
1200 					stream->groups.leftmost);
1201 		}
1202 	}
1203 
1204 	if (mode == LZMA_INDEX_ITER_NONEMPTY_BLOCK) {
1205 		// We need to look for the next Block again if this Block
1206 		// is empty.
1207 		if (record == 0) {
1208 			if (group->node.uncompressed_base
1209 					== group->records[0].uncompressed_sum)
1210 				goto again;
1211 		} else if (group->records[record - 1].uncompressed_sum
1212 				== group->records[record].uncompressed_sum) {
1213 			goto again;
1214 		}
1215 	}
1216 
1217 	iter->internal[ITER_STREAM].p = stream;
1218 	iter->internal[ITER_GROUP].p = group;
1219 	iter->internal[ITER_RECORD].s = record;
1220 
1221 	iter_set_info(iter);
1222 
1223 	return false;
1224 }
1225 
1226 
1227 extern LZMA_API(lzma_bool)
1228 lzma_index_iter_locate(lzma_index_iter *iter, lzma_vli target)
1229 {
1230 	const lzma_index *i = iter->internal[ITER_INDEX].p;
1231 
1232 	// If the target is past the end of the file, return immediately.
1233 	if (i->uncompressed_size <= target)
1234 		return true;
1235 
1236 	// Locate the Stream containing the target offset.
1237 	const index_stream *stream = index_tree_locate(&i->streams, target);
1238 	assert(stream != NULL);
1239 	target -= stream->node.uncompressed_base;
1240 
1241 	// Locate the group containing the target offset.
1242 	const index_group *group = index_tree_locate(&stream->groups, target);
1243 	assert(group != NULL);
1244 
1245 	// Use binary search to locate the exact Record. It is the first
1246 	// Record whose uncompressed_sum is greater than target.
1247 	// This is because we want the rightmost Record that fulfills the
1248 	// search criterion. It is possible that there are empty Blocks;
1249 	// we don't want to return them.
1250 	size_t left = 0;
1251 	size_t right = group->last;
1252 
1253 	while (left < right) {
1254 		const size_t pos = left + (right - left) / 2;
1255 		if (group->records[pos].uncompressed_sum <= target)
1256 			left = pos + 1;
1257 		else
1258 			right = pos;
1259 	}
1260 
1261 	iter->internal[ITER_STREAM].p = stream;
1262 	iter->internal[ITER_GROUP].p = group;
1263 	iter->internal[ITER_RECORD].s = left;
1264 
1265 	iter_set_info(iter);
1266 
1267 	return false;
1268 }
1269