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
29typedef struct index_tree_node_s index_tree_node;
30struct 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
46typedef 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
65typedef struct {
66	lzma_vli uncompressed_sum;
67	lzma_vli unpadded_sum;
68} index_record;
69
70
71typedef 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
107typedef 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
145struct 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
181static void
182index_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()
193static void
194index_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.
214static void
215index_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.
229static void
230index_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.
293static void *
294index_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).
314static void *
315index_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.
339static index_stream *
340index_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.
369static void
370index_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
379static lzma_index *
380index_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
397extern LZMA_API(lzma_index *)
398lzma_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
416extern LZMA_API(void)
417lzma_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
430extern void
431lzma_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
441extern LZMA_API(uint64_t)
442lzma_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
490extern LZMA_API(uint64_t)
491lzma_index_memused(const lzma_index *i)
492{
493	return lzma_index_memusage(i->streams.count, i->record_count);
494}
495
496
497extern LZMA_API(lzma_vli)
498lzma_index_block_count(const lzma_index *i)
499{
500	return i->record_count;
501}
502
503
504extern LZMA_API(lzma_vli)
505lzma_index_stream_count(const lzma_index *i)
506{
507	return i->streams.count;
508}
509
510
511extern LZMA_API(lzma_vli)
512lzma_index_size(const lzma_index *i)
513{
514	return index_size(i->record_count, i->index_list_size);
515}
516
517
518extern LZMA_API(lzma_vli)
519lzma_index_total_size(const lzma_index *i)
520{
521	return i->total_size;
522}
523
524
525extern LZMA_API(lzma_vli)
526lzma_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
535static lzma_vli
536index_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
559extern LZMA_API(lzma_vli)
560lzma_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
571extern LZMA_API(lzma_vli)
572lzma_index_uncompressed_size(const lzma_index *i)
573{
574	return i->uncompressed_size;
575}
576
577
578extern LZMA_API(uint32_t)
579lzma_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
592extern uint32_t
593lzma_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
600extern LZMA_API(lzma_ret)
601lzma_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
617extern LZMA_API(lzma_ret)
618lzma_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
639extern LZMA_API(lzma_ret)
640lzma_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()
730typedef 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.
754static void
755index_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
776extern LZMA_API(lzma_ret)
777lzma_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.
882static index_stream *
883index_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
943extern LZMA_API(lzma_index *)
944lzma_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[]
978enum {
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
988enum {
989	ITER_METHOD_NORMAL,
990	ITER_METHOD_NEXT,
991	ITER_METHOD_LEFTMOST,
992};
993
994
995static void
996iter_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
1100extern LZMA_API(void)
1101lzma_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
1109extern LZMA_API(void)
1110lzma_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
1120extern LZMA_API(lzma_bool)
1121lzma_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
1154again:
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
1227extern LZMA_API(lzma_bool)
1228lzma_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