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