1f6384cfcf6d9f9d945d88b2d6406b7d394a068a
[wimlib] / src / lookup_table.c
1 /*
2  * lookup_table.c
3  *
4  * Lookup table, implemented as a hash table, that maps dentries to file
5  * resources.
6  */
7
8 /*
9  * Copyright (C) 2012 Eric Biggers
10  *
11  * This file is part of wimlib, a library for working with WIM files.
12  *
13  * wimlib is free software; you can redistribute it and/or modify it under the
14  * terms of the GNU Lesser General Public License as published by the Free
15  * Software Foundation; either version 2.1 of the License, or (at your option)
16  * any later version.
17  *
18  * wimlib is distributed in the hope that it will be useful, but WITHOUT ANY
19  * WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR
20  * A PARTICULAR PURPOSE. See the GNU Lesser General Public License for more
21  * details.
22  *
23  * You should have received a copy of the GNU Lesser General Public License
24  * along with wimlib; if not, see http://www.gnu.org/licenses/.
25  */
26
27 #include "wimlib_internal.h"
28 #include "lookup_table.h"
29 #include "io.h"
30
31 struct lookup_table *new_lookup_table(size_t capacity)
32 {
33         struct lookup_table *table;
34         struct lookup_table_entry **array;
35
36         table = CALLOC(1, sizeof(struct lookup_table));
37         if (!table)
38                 goto err;
39         array = CALLOC(capacity, sizeof(array[0]));
40         if (!array) {
41                 FREE(table);
42                 goto err;
43         }
44         table->num_entries = 0;
45         table->capacity = capacity;
46         table->array = array;
47         return table;
48 err:
49         ERROR("Failed to allocate memory for lookup table with capacity "
50                         "%zu\n", capacity);
51         return NULL;
52 }
53
54 struct lookup_table_entry *new_lookup_table_entry()
55 {
56         struct lookup_table_entry *lte;
57         
58         lte = MALLOC(sizeof(struct lookup_table_entry));
59         if (!lte) {
60                 ERROR("Out of memory (tried to allocate %zu bytes for "
61                                 "lookup table entry)\n", 
62                                 sizeof(struct lookup_table_entry));
63                 return NULL;
64         }
65
66         lte->file_on_disk = NULL;
67         lte->other_wim_fp = NULL;
68         lte->next         = NULL;
69         lte->part_number  = 1;
70         lte->refcnt       = 1;
71         lte->staging_num_times_opened = 0;
72         return lte;
73 }
74
75
76
77 /*
78  * Inserts an entry into the lookup table.
79  *
80  * @lookup_table:       A pointer to the lookup table.
81  * @entry:              A pointer to the entry to insert.
82  */
83 void lookup_table_insert(struct lookup_table *table, 
84                          struct lookup_table_entry *lte)
85 {
86         size_t pos;
87         pos = lte->hash_short % table->capacity;
88         lte->next = table->array[pos];
89         table->array[pos] = lte;
90         /* XXX Make the table grow when too many entries have been inserted. */
91         table->num_entries++;
92 }
93
94
95 /* Unlinks a lookup table entry from the table; does not free it. */
96 void lookup_table_unlink(struct lookup_table *table, 
97                          struct lookup_table_entry *lte)
98 {
99         size_t pos;
100         struct lookup_table_entry *prev, *cur_entry, *next;
101
102         pos = lte->hash_short % table->capacity;
103         prev = NULL;
104         cur_entry = table->array[pos];
105
106         while (cur_entry) {
107                 next = cur_entry->next;
108                 if (cur_entry == lte) {
109                         if (prev)
110                                 prev->next = next;
111                         else
112                                 table->array[pos] = next;
113                         table->num_entries--;
114                         return;
115                 }
116                 prev = cur_entry;
117                 cur_entry = next;
118         }
119 }
120
121
122 /* Decrement the reference count for the dentry having hash value @hash in the
123  * lookup table.  The lookup table entry is unlinked and freed if there are no
124  * references to in remaining.  */
125 void lookup_table_decrement_refcnt(struct lookup_table* table, const u8 hash[])
126 {
127         size_t pos = *(size_t*)hash % table->capacity;
128         struct lookup_table_entry *prev = NULL;
129         struct lookup_table_entry *entry = table->array[pos];
130         struct lookup_table_entry *next;
131         while (entry) {
132                 next = entry->next;
133                 if (memcmp(hash, entry->hash, WIM_HASH_SIZE) == 0) {
134                         wimlib_assert(entry->refcnt != 0);
135                         if (--entry->refcnt == 0) {
136                                 free_lookup_table_entry(entry);
137                                 if (prev)
138                                         prev->next = next;
139                                 else
140                                         table->array[pos] = next;
141                         }
142                 }
143                 prev = entry;
144                 entry = next;
145         }
146 }
147
148 /* 
149  * Looks up an entry in the lookup table.
150  */
151 struct lookup_table_entry *lookup_resource(const struct lookup_table *lookup_table, 
152                                      const u8 hash[])
153 {
154         size_t pos;
155         struct lookup_table_entry *lte;
156
157         pos = *(size_t*)hash % lookup_table->capacity;
158         lte = lookup_table->array[pos];
159         while (lte) {
160                 if (memcmp(hash, lte->hash, WIM_HASH_SIZE) == 0)
161                         return lte;
162                 lte = lte->next;
163         }
164         return NULL;
165 }
166
167 /* 
168  * Calls a function on all the entries in the lookup table.  Stop early and
169  * return nonzero if any call to the function returns nonzero.
170  */
171 int for_lookup_table_entry(struct lookup_table *table, 
172                 int (*visitor)(struct lookup_table_entry *, void *), void *arg)
173 {
174         struct lookup_table_entry *entry, *next;
175         size_t i;
176         int ret;
177
178         for (i = 0; i < table->capacity; i++) {
179                 entry = table->array[i];
180                 while (entry) {
181                         next = entry->next;
182                         ret = visitor(entry, arg);
183                         if (ret != 0)
184                                 return ret;
185                         entry = next;
186                 }
187         }
188         return 0;
189 }
190
191
192 /*
193  * Reads the lookup table from a WIM file.
194  *
195  * @fp:                 The FILE* for the WIM file.
196  * @offset:             The offset of the lookup table resource.
197  * @size:               The size of the lookup table resource.
198  * @lookup_table_ret:   A pointer to a struct lookup_table structure into which the
199  *                              lookup table will be returned.
200  * @return:             True on success, false on failure.
201  */
202 int read_lookup_table(FILE *fp, u64 offset, u64 size, 
203                       struct lookup_table **table_ret)
204 {
205         size_t num_entries;
206         u8     buf[WIM_LOOKUP_TABLE_ENTRY_DISK_SIZE];
207         int    ret;
208         struct lookup_table *table;
209         const u8 *p;
210         struct lookup_table_entry *cur_entry;
211
212         DEBUG("Reading lookup table: offset %"PRIu64", size %"PRIu64"\n",
213                         offset, size);
214
215         if (fseeko(fp, offset, SEEK_SET) != 0) {
216                 ERROR("Failed to seek to byte %"PRIu64" to read lookup table: "
217                                 "%m\n", offset);
218                 return WIMLIB_ERR_READ;
219         }
220
221         num_entries = size / WIM_LOOKUP_TABLE_ENTRY_DISK_SIZE;
222         table = new_lookup_table(num_entries * 2 + 1);
223         if (!table)
224                 return WIMLIB_ERR_NOMEM;
225
226         while (num_entries--) {
227                 if (fread(buf, 1, sizeof(buf), fp) != sizeof(buf)) {
228                         if (feof(fp)) {
229                                 ERROR("Unexpected EOF in lookup table!\n");
230                         } else {
231                                 ERROR("Stream read error while reading "
232                                                 "lookup table!\n");
233                         }
234                         ret = WIMLIB_ERR_READ;
235                         goto err;
236                 }
237                 cur_entry = new_lookup_table_entry();
238                 if (!cur_entry) {
239                         ret = WIMLIB_ERR_NOMEM;
240                         goto err;
241                 }
242                          
243                 p = get_resource_entry(buf, &cur_entry->resource_entry);
244                 p = get_u16(p, &cur_entry->part_number);
245                 p = get_u32(p, &cur_entry->refcnt);
246                 p = get_bytes(p, sizeof(cur_entry->hash), cur_entry->hash);
247                 lookup_table_insert(table, cur_entry);
248         }
249         DEBUG("Done reading lookup table.\n");
250         *table_ret = table;
251         return 0;
252 err:
253         free_lookup_table(table);
254         return ret;
255 }
256
257
258 /* 
259  * Writes a lookup table entry to the output file.
260  */
261 int write_lookup_table_entry(struct lookup_table_entry *lte, void *__out)
262 {
263         FILE *out;
264         u8 buf[WIM_LOOKUP_TABLE_ENTRY_DISK_SIZE];
265         u8 *p;
266
267         out = __out;
268
269         /* do not write lookup table entries for empty files */
270         if (lte->output_resource_entry.original_size == 0)
271                 return 0;
272
273         /* Don't write entries that have not had file resources or metadata
274          * resources written for them. */
275         if (lte->out_refcnt == 0)
276                 return 0;
277
278         if (lte->output_resource_entry.flags & WIM_RESHDR_FLAG_METADATA)
279                 DEBUG("Writing metadata entry at %lu\n", ftello(out));
280
281         p = put_resource_entry(buf, &lte->output_resource_entry);
282         p = put_u16(p, lte->part_number);
283         p = put_u32(p, lte->out_refcnt);
284         p = put_bytes(p, WIM_HASH_SIZE, lte->hash);
285         if (fwrite(buf, 1, sizeof(buf), out) != sizeof(buf)) {
286                 ERROR("Failed to write lookup table entry: %m\n");
287                 return WIMLIB_ERR_WRITE;
288         }
289         return 0;
290 }
291
292 static int do_free_lookup_table_entry(struct lookup_table_entry *entry, void *ignore)
293 {
294         free_lookup_table_entry(entry);
295         return 0;
296 }
297
298 void free_lookup_table(struct lookup_table *table)
299 {
300         if (table) {
301                 if (table->array) {
302                         for_lookup_table_entry(table, 
303                                                do_free_lookup_table_entry, 
304                                                NULL);
305                         FREE(table->array);
306                 }
307                 FREE(table);
308         }
309 }
310
311 int zero_out_refcnts(struct lookup_table_entry *entry, void *ignore)
312 {
313         entry->out_refcnt = 0;
314         return 0;
315 }
316
317 int print_lookup_table_entry(struct lookup_table_entry *entry, void *ignore)
318 {
319         printf("Offset            = %"PRIu64" bytes\n", 
320                         entry->resource_entry.offset);
321         printf("Size              = %"PRIu64" bytes\n", 
322                         (u64)entry->resource_entry.size);
323         printf("Original size     = %"PRIu64" bytes\n", 
324                         entry->resource_entry.original_size);
325         printf("Part Number       = %hu\n", entry->part_number);
326         printf("Reference Count   = %u\n", entry->refcnt);
327         printf("Hash              = ");
328         print_hash(entry->hash);
329         putchar('\n');
330         printf("Flags             = ");
331         u8 flags = entry->resource_entry.flags;
332         if (flags & WIM_RESHDR_FLAG_COMPRESSED)
333                 fputs("WIM_RESHDR_FLAG_COMPRESSED, ", stdout);
334         if (flags & WIM_RESHDR_FLAG_FREE)
335                 fputs("WIM_RESHDR_FLAG_FREE, ", stdout);
336         if (flags & WIM_RESHDR_FLAG_METADATA)
337                 fputs("WIM_RESHDR_FLAG_METADATA, ", stdout);
338         if (flags & WIM_RESHDR_FLAG_SPANNED)
339                 fputs("WIM_RESHDR_FLAG_SPANNED, ", stdout);
340         putchar('\n');
341         if (entry->file_on_disk)
342                 printf("File on Disk      = \"%s\"\n", entry->file_on_disk);
343         putchar('\n');
344         return 0;
345 }
346
347 /*
348  * Prints the lookup table of a WIM file. 
349  */
350 WIMLIBAPI void wimlib_print_lookup_table(WIMStruct *w)
351 {
352         for_lookup_table_entry(w->lookup_table, 
353                                print_lookup_table_entry, NULL);
354 }