]> wimlib.net Git - wimlib/blobdiff - src/write.c
Comment lines that may produce compiler warnings
[wimlib] / src / write.c
index 36b82b453fc9d1ca5c85079de479ffa728ec6779..86730ff7fec7a1158b6ed8bed305ec784b9375aa 100644 (file)
@@ -6,7 +6,6 @@
  */
 
 /*
- * Copyright (C) 2010 Carl Thijssen
  * Copyright (C) 2012 Eric Biggers
  *
  * This file is part of wimlib, a library for working with WIM files.
@@ -35,7 +34,7 @@
 
 #include "list.h"
 #include "wimlib_internal.h"
-#include "io.h"
+#include "buffer_io.h"
 #include "dentry.h"
 #include "lookup_table.h"
 #include "xml.h"
@@ -99,7 +98,7 @@ struct chunk_table {
  * output file.
  */
 static int
-begin_wim_resource_chunk_tab(const struct lookup_table_entry *lte,
+begin_wim_resource_chunk_tab(const struct wim_lookup_table_entry *lte,
                             FILE *out_fp,
                             off_t file_offset,
                             struct chunk_table **chunk_tab_ret)
@@ -251,7 +250,7 @@ finish_wim_resource_chunk_tab(struct chunk_table *chunk_tab,
 
 /* Prepare for multiple reads to a resource by caching a FILE * or NTFS
  * attribute pointer in the lookup table entry. */
-static int prepare_resource_for_read(struct lookup_table_entry *lte
+static int prepare_resource_for_read(struct wim_lookup_table_entry *lte
 
                                        #ifdef WITH_NTFS_3G
                                        , ntfs_inode **ni_ret
@@ -300,7 +299,7 @@ static int prepare_resource_for_read(struct lookup_table_entry *lte
 
 /* Undo prepare_resource_for_read() by closing the cached FILE * or NTFS
  * attribute. */
-static void end_wim_resource_read(struct lookup_table_entry *lte
+static void end_wim_resource_read(struct wim_lookup_table_entry *lte
                                #ifdef WITH_NTFS_3G
                                        , ntfs_inode *ni
                                #endif
@@ -324,7 +323,7 @@ static void end_wim_resource_read(struct lookup_table_entry *lte
 }
 
 static int
-write_uncompressed_resource_and_truncate(struct lookup_table_entry *lte,
+write_uncompressed_resource_and_truncate(struct wim_lookup_table_entry *lte,
                                         FILE *out_fp,
                                         off_t file_offset,
                                         struct resource_entry *out_res_entry)
@@ -365,7 +364,7 @@ write_uncompressed_resource_and_truncate(struct lookup_table_entry *lte,
  *
  * Returns 0 on success; nonzero on failure.
  */
-int write_wim_resource(struct lookup_table_entry *lte,
+int write_wim_resource(struct wim_lookup_table_entry *lte,
                       FILE *out_fp, int out_ctype,
                       struct resource_entry *out_res_entry,
                       int flags)
@@ -549,6 +548,8 @@ out:
 }
 
 #ifdef ENABLE_MULTITHREADED_COMPRESSION
+
+/* Blocking shared queue (solves the producer-consumer problem) */
 struct shared_queue {
        unsigned size;
        unsigned front;
@@ -562,6 +563,7 @@ struct shared_queue {
 
 static int shared_queue_init(struct shared_queue *q, unsigned size)
 {
+       wimlib_assert(size != 0);
        q->array = CALLOC(sizeof(q->array[0]), size);
        if (!q->array)
                return WIMLIB_ERR_NOMEM;
@@ -624,7 +626,7 @@ struct compressor_thread_params {
 #define MAX_CHUNKS_PER_MSG 2
 
 struct message {
-       struct lookup_table_entry *lte;
+       struct wim_lookup_table_entry *lte;
        u8 *uncompressed_chunks[MAX_CHUNKS_PER_MSG];
        u8 *out_compressed_chunks[MAX_CHUNKS_PER_MSG];
        u8 *compressed_chunks[MAX_CHUNKS_PER_MSG];
@@ -653,6 +655,11 @@ static void compress_chunks(struct message *msg, compress_func_t compress)
        }
 }
 
+/* Compressor thread routine.  This is a lot simpler than the main thread
+ * routine: just repeatedly get a group of chunks from the
+ * res_to_compress_queue, compress them, and put them in the
+ * compressed_res_queue.  A NULL pointer indicates that the thread should stop.
+ * */
 static void *compressor_thread_proc(void *arg)
 {
        struct compressor_thread_params *params = arg;
@@ -679,7 +686,7 @@ static int do_write_stream_list(struct list_head *my_resources,
                                int write_resource_flags)
 {
        int ret;
-       struct lookup_table_entry *lte, *tmp;
+       struct wim_lookup_table_entry *lte, *tmp;
 
        list_for_each_entry_safe(lte, tmp, my_resources, staging_list) {
                ret = write_wim_resource(lte,
@@ -762,19 +769,25 @@ static int main_writer_thread_proc(struct list_head *stream_list,
                                   int out_ctype,
                                   struct shared_queue *res_to_compress_queue,
                                   struct shared_queue *compressed_res_queue,
-                                  size_t queue_size,
+                                  size_t num_messages,
                                   int write_flags,
                                   wimlib_progress_func_t progress_func,
                                   union wimlib_progress_info *progress)
 {
        int ret;
-
-       struct message msgs[queue_size];
-       ZERO_ARRAY(msgs);
+       struct chunk_table *cur_chunk_tab = NULL;
+       struct message *msgs = CALLOC(num_messages, sizeof(struct message));
+       struct wim_lookup_table_entry *next_lte = NULL;
 
        // Initially, all the messages are available to use.
        LIST_HEAD(available_msgs);
-       for (size_t i = 0; i < ARRAY_LEN(msgs); i++)
+
+       if (!msgs) {
+               ret = WIMLIB_ERR_NOMEM;
+               goto out;
+       }
+
+       for (size_t i = 0; i < num_messages; i++)
                list_add(&msgs[i].list, &available_msgs);
 
        // outstanding_resources is the list of resources that currently have
@@ -794,7 +807,6 @@ static int main_writer_thread_proc(struct list_head *stream_list,
        //
        LIST_HEAD(outstanding_resources);
        struct list_head *next_resource = stream_list->next;
-       struct lookup_table_entry *next_lte = NULL;
        u64 next_chunk = 0;
        u64 next_num_chunks = 0;
 
@@ -806,8 +818,7 @@ static int main_writer_thread_proc(struct list_head *stream_list,
        // list and written directly by the main thread.
        LIST_HEAD(my_resources);
 
-       struct lookup_table_entry *cur_lte = NULL;
-       struct chunk_table *cur_chunk_tab = NULL;
+       struct wim_lookup_table_entry *cur_lte = NULL;
        struct message *msg;
 
 #ifdef WITH_NTFS_3G
@@ -816,11 +827,11 @@ static int main_writer_thread_proc(struct list_head *stream_list,
 
        DEBUG("Initializing buffers for uncompressed "
              "and compressed data (%zu bytes needed)",
-             queue_size * MAX_CHUNKS_PER_MSG * WIM_CHUNK_SIZE * 2);
+             num_messages * MAX_CHUNKS_PER_MSG * WIM_CHUNK_SIZE * 2);
 
        // Pre-allocate all the buffers that will be needed to do the chunk
        // compression.
-       for (size_t i = 0; i < ARRAY_LEN(msgs); i++) {
+       for (size_t i = 0; i < num_messages; i++) {
                for (size_t j = 0; j < MAX_CHUNKS_PER_MSG; j++) {
                        msgs[i].compressed_chunks[j] = MALLOC(WIM_CHUNK_SIZE);
 
@@ -832,8 +843,6 @@ static int main_writer_thread_proc(struct list_head *stream_list,
                        if (msgs[i].compressed_chunks[j] == NULL ||
                            msgs[i].uncompressed_chunks[j] == NULL)
                        {
-                               ERROR("Could not allocate enough memory for "
-                                     "multi-threaded compression");
                                ret = WIMLIB_ERR_NOMEM;
                                goto out;
                        }
@@ -893,11 +902,13 @@ static int main_writer_thread_proc(struct list_head *stream_list,
                                // list, or we get to a resource that needs compression.
                                while (1) {
                                        if (next_resource == stream_list) {
+                                               // No more resources to send for
+                                               // compression
                                                next_lte = NULL;
                                                break;
                                        }
                                        next_lte = container_of(next_resource,
-                                                               struct lookup_table_entry,
+                                                               struct wim_lookup_table_entry,
                                                                staging_list);
                                        next_resource = next_resource->next;
                                        if ((!(write_flags & WIMLIB_WRITE_FLAG_RECOMPRESS)
@@ -921,15 +932,20 @@ static int main_writer_thread_proc(struct list_head *stream_list,
 
                                                if (ret != 0)
                                                        goto out;
-                                               if (cur_lte == NULL)
+                                               if (cur_lte == NULL) {
+                                                       // Set cur_lte for the
+                                                       // first time
                                                        cur_lte = next_lte;
+                                               }
                                                break;
                                        }
                                }
                        }
 
-                       if (next_lte == NULL)
+                       if (next_lte == NULL) {
+                               // No more resources to send for compression
                                break;
+                       }
 
                        // Get a message from the available messages
                        // list
@@ -956,13 +972,11 @@ static int main_writer_thread_proc(struct list_head *stream_list,
                                // message so that a compressor thread can
                                // compress it.
 
-                               if (next_chunk == next_num_chunks - 1 &&
-                                    wim_resource_size(next_lte) % WIM_CHUNK_SIZE != 0)
-                               {
-                                       size = wim_resource_size(next_lte) % WIM_CHUNK_SIZE;
+                               if (next_chunk == next_num_chunks - 1) {
+                                       size = MODULO_NONZERO(wim_resource_size(next_lte),
+                                                             WIM_CHUNK_SIZE);
                                }
 
-
                                DEBUG2("Read resource (size=%u, offset=%zu)",
                                      size, next_chunk * WIM_CHUNK_SIZE);
 
@@ -1005,10 +1019,11 @@ static int main_writer_thread_proc(struct list_head *stream_list,
                // message around until all earlier chunks are received.
                //
                // Otherwise, write all the chunks we can.
-               while (cur_lte != NULL && !list_empty(&cur_lte->msg_list)
-                       && (msg = container_of(cur_lte->msg_list.next,
-                                              struct message,
-                                              list))->complete)
+               while (cur_lte != NULL &&
+                      !list_empty(&cur_lte->msg_list) &&
+                      (msg = container_of(cur_lte->msg_list.next,
+                                          struct message,
+                                          list))->complete)
                {
                        DEBUG2("Complete msg (begin_chunk=%"PRIu64")", msg->begin_chunk);
                        if (msg->begin_chunk == 0) {
@@ -1040,8 +1055,8 @@ static int main_writer_thread_proc(struct list_head *stream_list,
                        // now.
                        list_add(&msg->list, &available_msgs);
 
-                       // Was this the last chunk of the stream?  If so,
-                       // finish it.
+                       // Was this the last chunk of the stream?  If so, finish
+                       // it.
                        if (list_empty(&cur_lte->msg_list) &&
                            msg->begin_chunk + msg->num_chunks == cur_chunk_tab->num_chunks)
                        {
@@ -1094,21 +1109,12 @@ static int main_writer_thread_proc(struct list_head *stream_list,
                                struct list_head *next = cur_lte->staging_list.next;
                                list_del(&cur_lte->staging_list);
 
-                               if (next == &outstanding_resources) {
-                                       if (next_lte == NULL) {
-                                               DEBUG("No more outstanding resources");
-                                               ret = 0;
-                                               goto out;
-                                       } else {
-                                               DEBUG("No more outstanding resources---"
-                                                     "but still more to compress!");
-                                               cur_lte = NULL;
-                                       }
-                               } else {
+                               if (next == &outstanding_resources)
+                                       cur_lte = NULL;
+                               else
                                        cur_lte = container_of(cur_lte->staging_list.next,
-                                                              struct lookup_table_entry,
+                                                              struct wim_lookup_table_entry,
                                                               staging_list);
-                               }
 
                                // Since we just finished writing a stream,
                                // write any streams that have been added to the
@@ -1130,6 +1136,11 @@ static int main_writer_thread_proc(struct list_head *stream_list,
        }
 
 out:
+       if (ret == WIMLIB_ERR_NOMEM) {
+               ERROR("Could not allocate enough memory for "
+                     "multi-threaded compression");
+       }
+
        if (next_lte) {
 #ifdef WITH_NTFS_3G
                end_wim_resource_read(next_lte, ni);
@@ -1143,28 +1154,32 @@ out:
                                           out_ctype, progress_func,
                                           progress, 0);
        } else {
-               size_t num_available_msgs = 0;
-               struct list_head *cur;
+               if (msgs) {
+                       size_t num_available_msgs = 0;
+                       struct list_head *cur;
 
-               list_for_each(cur, &available_msgs) {
-                       num_available_msgs++;
-               }
+                       list_for_each(cur, &available_msgs) {
+                               num_available_msgs++;
+                       }
 
-               while (num_available_msgs < ARRAY_LEN(msgs)) {
-                       shared_queue_get(compressed_res_queue);
-                       num_available_msgs++;
+                       while (num_available_msgs < num_messages) {
+                               shared_queue_get(compressed_res_queue);
+                               num_available_msgs++;
+                       }
                }
        }
 
-       for (size_t i = 0; i < ARRAY_LEN(msgs); i++) {
-               for (size_t j = 0; j < MAX_CHUNKS_PER_MSG; j++) {
-                       FREE(msgs[i].compressed_chunks[j]);
-                       FREE(msgs[i].uncompressed_chunks[j]);
+       if (msgs) {
+               for (size_t i = 0; i < num_messages; i++) {
+                       for (size_t j = 0; j < MAX_CHUNKS_PER_MSG; j++) {
+                               FREE(msgs[i].compressed_chunks[j]);
+                               FREE(msgs[i].uncompressed_chunks[j]);
+                       }
                }
+               FREE(msgs);
        }
 
-       if (cur_chunk_tab != NULL)
-               FREE(cur_chunk_tab);
+       FREE(cur_chunk_tab);
        return ret;
 }
 
@@ -1214,6 +1229,10 @@ static int write_stream_list_parallel(struct list_head *stream_list,
        params.compress = get_compress_func(out_ctype);
 
        compressor_threads = MALLOC(num_threads * sizeof(pthread_t));
+       if (!compressor_threads) {
+               ret = WIMLIB_ERR_NOMEM;
+               goto out_destroy_compressed_res_queue;
+       }
 
        for (unsigned i = 0; i < num_threads; i++) {
                DEBUG("pthread_create thread %u", i);
@@ -1251,6 +1270,7 @@ out_join:
                }
        }
        FREE(compressor_threads);
+out_destroy_compressed_res_queue:
        shared_queue_destroy(&compressed_res_queue);
 out_destroy_res_to_compress_queue:
        shared_queue_destroy(&res_to_compress_queue);
@@ -1277,12 +1297,11 @@ static int write_stream_list(struct list_head *stream_list, FILE *out_fp,
                             unsigned num_threads,
                             wimlib_progress_func_t progress_func)
 {
-       struct lookup_table_entry *lte;
+       struct wim_lookup_table_entry *lte;
        size_t num_streams = 0;
        u64 total_bytes = 0;
        u64 total_compression_bytes = 0;
        union wimlib_progress_info progress;
-       int ret;
 
        list_for_each_entry(lte, stream_list, staging_list) {
                num_streams++;
@@ -1301,54 +1320,96 @@ static int write_stream_list(struct list_head *stream_list, FILE *out_fp,
        progress.write_streams.num_threads       = num_threads;
        progress.write_streams.compression_type  = out_ctype;
 
-       if (num_streams == 0) {
-               ret = 0;
-               goto out;
-       }
-
 #ifdef ENABLE_MULTITHREADED_COMPRESSION
-       if (total_compression_bytes >= 1000000 && num_threads != 1) {
-               ret = write_stream_list_parallel(stream_list,
-                                                out_fp,
-                                                out_ctype,
-                                                write_flags,
-                                                num_threads,
-                                                progress_func,
-                                                &progress);
-       }
+       if (total_compression_bytes >= 1000000 && num_threads != 1)
+               return write_stream_list_parallel(stream_list,
+                                                 out_fp,
+                                                 out_ctype,
+                                                 write_flags,
+                                                 num_threads,
+                                                 progress_func,
+                                                 &progress);
        else
 #endif
+               return write_stream_list_serial(stream_list,
+                                               out_fp,
+                                               out_ctype,
+                                               write_flags,
+                                               progress_func,
+                                               &progress);
+}
+
+struct lte_overwrite_prepare_args {
+       WIMStruct *wim;
+       off_t end_offset;
+       struct list_head *stream_list;
+};
+
+static int lte_overwrite_prepare(struct wim_lookup_table_entry *lte, void *arg)
+{
+       struct lte_overwrite_prepare_args *args = arg;
+
+       if (lte->resource_location == RESOURCE_IN_WIM &&
+           lte->wim == args->wim &&
+           lte->resource_entry.offset + lte->resource_entry.size > args->end_offset)
        {
-               ret = write_stream_list_serial(stream_list,
-                                              out_fp,
-                                              out_ctype,
-                                              write_flags,
-                                              progress_func,
-                                              &progress);
+               ERROR("The following resource is after the XML data:");
+               print_lookup_table_entry(lte);
+               return WIMLIB_ERR_RESOURCE_ORDER;
        }
-out:
-       return ret;
+
+       lte->out_refcnt = lte->refcnt;
+       memcpy(&lte->output_resource_entry, &lte->resource_entry,
+              sizeof(struct resource_entry));
+       if (!(lte->resource_entry.flags & WIM_RESHDR_FLAG_METADATA)) {
+               wimlib_assert(lte->resource_location != RESOURCE_NONEXISTENT);
+               if (lte->resource_location != RESOURCE_IN_WIM || lte->wim != args->wim)
+                       list_add(&lte->staging_list, args->stream_list);
+       }
+       return 0;
 }
 
+static int wim_find_new_streams(WIMStruct *wim, off_t end_offset,
+                               struct list_head *stream_list)
+{
+       struct lte_overwrite_prepare_args args = {
+               .wim         = wim,
+               .end_offset  = end_offset,
+               .stream_list = stream_list,
+       };
+
+       return for_lookup_table_entry(wim->lookup_table,
+                                     lte_overwrite_prepare, &args);
+}
 
-static int dentry_find_streams_to_write(struct dentry *dentry,
-                                       void *wim)
+static int inode_find_streams_to_write(struct wim_inode *inode,
+                                      struct wim_lookup_table *table,
+                                      struct list_head *stream_list)
 {
-       WIMStruct *w = wim;
-       struct list_head *stream_list = w->private;
-       struct lookup_table_entry *lte;
-       for (unsigned i = 0; i <= dentry->d_inode->num_ads; i++) {
-               lte = inode_stream_lte(dentry->d_inode, i, w->lookup_table);
-               if (lte && ++lte->out_refcnt == 1)
-                       list_add_tail(&lte->staging_list, stream_list);
+       struct wim_lookup_table_entry *lte;
+       for (unsigned i = 0; i <= inode->i_num_ads; i++) {
+               lte = inode_stream_lte(inode, i, table);
+               if (lte) {
+                       if (lte->out_refcnt == 0)
+                               list_add_tail(&lte->staging_list, stream_list);
+                       lte->out_refcnt += inode->i_nlink;
+               }
        }
        return 0;
 }
 
-static int find_streams_to_write(WIMStruct *w)
+static int image_find_streams_to_write(WIMStruct *w)
 {
-       return for_dentry_in_tree(wim_root_dentry(w),
-                                 dentry_find_streams_to_write, w);
+       struct wim_inode *inode;
+       struct hlist_node *cur;
+       struct hlist_head *inode_list;
+
+       inode_list = &wim_get_current_image_metadata(w)->inode_list;
+       hlist_for_each_entry(inode, cur, inode_list, i_hlist) {
+               inode_find_streams_to_write(inode, w->lookup_table,
+                                           (struct list_head*)w->private);
+       }
+       return 0;
 }
 
 static int write_wim_streams(WIMStruct *w, int image, int write_flags,
@@ -1359,7 +1420,7 @@ static int write_wim_streams(WIMStruct *w, int image, int write_flags,
        for_lookup_table_entry(w->lookup_table, lte_zero_out_refcnt, NULL);
        LIST_HEAD(stream_list);
        w->private = &stream_list;
-       for_image(w, image, find_streams_to_write);
+       for_image(w, image, image_find_streams_to_write);
        return write_stream_list(&stream_list, w->out_fp,
                                 wimlib_get_compression_type(w), write_flags,
                                 num_threads, progress_func);
@@ -1426,6 +1487,8 @@ int finish_write(WIMStruct *w, int image, int write_flags,
                        memcpy(&checkpoint_hdr, &hdr, sizeof(struct wim_header));
                        memset(&checkpoint_hdr.integrity, 0, sizeof(struct resource_entry));
                        if (fseeko(out, 0, SEEK_SET) != 0) {
+                               ERROR_WITH_ERRNO("Failed to seek to beginning "
+                                                "of WIM being written");
                                ret = WIMLIB_ERR_WRITE;
                                goto out;
                        }
@@ -1440,6 +1503,8 @@ int finish_write(WIMStruct *w, int image, int write_flags,
                        }
 
                        if (fseeko(out, 0, SEEK_END) != 0) {
+                               ERROR_WITH_ERRNO("Failed to seek to end "
+                                                "of WIM being written");
                                ret = WIMLIB_ERR_WRITE;
                                goto out;
                        }
@@ -1474,16 +1539,6 @@ int finish_write(WIMStruct *w, int image, int write_flags,
         * it should be a copy of the resource entry for the image that is
         * marked as bootable.  This is not well documented...
         */
-       if (hdr.boot_idx == 0 || !w->image_metadata
-                       || (image != WIMLIB_ALL_IMAGES && image != hdr.boot_idx)) {
-               memset(&hdr.boot_metadata_res_entry, 0,
-                      sizeof(struct resource_entry));
-       } else {
-               memcpy(&hdr.boot_metadata_res_entry,
-                      &w->image_metadata[
-                         hdr.boot_idx - 1].metadata_lte->output_resource_entry,
-                      sizeof(struct resource_entry));
-       }
 
        /* Set image count and boot index correctly for single image writes */
        if (image != WIMLIB_ALL_IMAGES) {
@@ -1494,7 +1549,19 @@ int finish_write(WIMStruct *w, int image, int write_flags,
                        hdr.boot_idx = 0;
        }
 
+       if (hdr.boot_idx == 0) {
+               memset(&hdr.boot_metadata_res_entry, 0,
+                      sizeof(struct resource_entry));
+       } else {
+               memcpy(&hdr.boot_metadata_res_entry,
+                      &w->image_metadata[
+                         hdr.boot_idx - 1].metadata_lte->output_resource_entry,
+                      sizeof(struct resource_entry));
+       }
+
        if (fseeko(out, 0, SEEK_SET) != 0) {
+               ERROR_WITH_ERRNO("Failed to seek to beginning of WIM "
+                                "being written");
                ret = WIMLIB_ERR_WRITE;
                goto out;
        }
@@ -1522,22 +1589,24 @@ out:
 }
 
 #if defined(HAVE_SYS_FILE_H) && defined(HAVE_FLOCK)
-int lock_wim(FILE *fp, const char *path)
+int lock_wim(WIMStruct *w, FILE *fp)
 {
        int ret = 0;
-       if (fp) {
+       if (fp && !w->wim_locked) {
                ret = flock(fileno(fp), LOCK_EX | LOCK_NB);
                if (ret != 0) {
                        if (errno == EWOULDBLOCK) {
                                ERROR("`%s' is already being modified or has been "
                                      "mounted read-write\n"
-                                     "        by another process!", path);
+                                     "        by another process!", w->filename);
                                ret = WIMLIB_ERR_ALREADY_LOCKED;
                        } else {
                                WARNING("Failed to lock `%s': %s",
-                                       path, strerror(errno));
+                                       w->filename, strerror(errno));
                                ret = 0;
                        }
+               } else {
+                       w->wim_locked = 1;
                }
        }
        return ret;
@@ -1597,7 +1666,7 @@ WIMLIBAPI int wimlib_write(WIMStruct *w, const char *path,
 {
        int ret;
 
-       if (!w || !path)
+       if (!path)
                return WIMLIB_ERR_INVALID_PARAM;
 
        write_flags &= WIMLIB_WRITE_MASK_PUBLIC;
@@ -1636,37 +1705,12 @@ out:
        return ret;
 }
 
-static int lte_overwrite_prepare(struct lookup_table_entry *lte,
-                                void *ignore)
-{
-       memcpy(&lte->output_resource_entry, &lte->resource_entry,
-              sizeof(struct resource_entry));
-       lte->out_refcnt = 0;
-       return 0;
-}
-
-static int check_resource_offset(struct lookup_table_entry *lte, void *arg)
-{
-       off_t end_offset = *(u64*)arg;
-
-       wimlib_assert(lte->out_refcnt <= lte->refcnt);
-       if (lte->out_refcnt < lte->refcnt) {
-               if (lte->resource_entry.offset + lte->resource_entry.size > end_offset) {
-                       ERROR("The following resource is after the XML data:");
-                       print_lookup_table_entry(lte);
-                       return WIMLIB_ERR_RESOURCE_ORDER;
-               }
-       }
-       return 0;
-}
-
-static int find_new_streams(struct lookup_table_entry *lte, void *arg)
+static bool any_images_modified(WIMStruct *w)
 {
-       if (lte->out_refcnt == lte->refcnt)
-               list_add(&lte->staging_list, (struct list_head*)arg);
-       else
-               lte->out_refcnt = lte->refcnt;
-       return 0;
+       for (int i = 0; i < w->hdr.image_count; i++)
+               if (w->image_metadata[i].modified)
+                       return true;
+       return false;
 }
 
 /*
@@ -1724,16 +1768,16 @@ static int find_new_streams(struct lookup_table_entry *lte, void *arg)
  * is is crash-safe except in the case of write re-ordering, but the
  * disadvantage is that a small hole is left in the WIM where the old lookup
  * table, xml data, and integrity table were.  (These usually only take up a
- * small amount of space compared to the streams, however.
+ * small amount of space compared to the streams, however.)
  */
 static int overwrite_wim_inplace(WIMStruct *w, int write_flags,
                                 unsigned num_threads,
-                                wimlib_progress_func_t progress_func,
-                                int modified_image_idx)
+                                wimlib_progress_func_t progress_func)
 {
        int ret;
        struct list_head stream_list;
        off_t old_wim_end;
+       bool found_modified_image;
 
        DEBUG("Overwriting `%s' in-place", w->filename);
 
@@ -1752,27 +1796,13 @@ static int overwrite_wim_inplace(WIMStruct *w, int write_flags,
                return WIMLIB_ERR_RESOURCE_ORDER;
        }
 
-       DEBUG("Identifying newly added streams");
-       for_lookup_table_entry(w->lookup_table, lte_overwrite_prepare, NULL);
-       INIT_LIST_HEAD(&stream_list);
-       for (int i = modified_image_idx; i < w->hdr.image_count; i++) {
-               DEBUG("Identifiying streams in image %d", i + 1);
-               w->private = &stream_list;
-               for_dentry_in_tree(w->image_metadata[i].root_dentry,
-                                  dentry_find_streams_to_write, w);
-       }
 
        if (w->hdr.integrity.offset)
                old_wim_end = w->hdr.integrity.offset + w->hdr.integrity.size;
        else
                old_wim_end = w->hdr.xml_res_entry.offset + w->hdr.xml_res_entry.size;
 
-       ret = for_lookup_table_entry(w->lookup_table, check_resource_offset,
-                                    &old_wim_end);
-       if (ret != 0)
-               return ret;
-
-       if (modified_image_idx == w->hdr.image_count && !w->deletion_occurred) {
+       if (!w->deletion_occurred && !any_images_modified(w)) {
                /* If no images have been modified and no images have been
                 * deleted, a new lookup table does not need to be written. */
                old_wim_end = w->hdr.lookup_table_res_entry.offset +
@@ -1780,17 +1810,17 @@ static int overwrite_wim_inplace(WIMStruct *w, int write_flags,
                write_flags |= WIMLIB_WRITE_FLAG_NO_LOOKUP_TABLE |
                               WIMLIB_WRITE_FLAG_CHECKPOINT_AFTER_XML;
        }
-
        INIT_LIST_HEAD(&stream_list);
-       for_lookup_table_entry(w->lookup_table, find_new_streams,
-                              &stream_list);
+       ret = wim_find_new_streams(w, old_wim_end, &stream_list);
+       if (ret != 0)
+               return ret;
 
        ret = open_wim_writable(w, w->filename, false,
                                (write_flags & WIMLIB_WRITE_FLAG_CHECK_INTEGRITY) != 0);
        if (ret != 0)
                return ret;
 
-       ret = lock_wim(w->out_fp, w->filename);
+       ret = lock_wim(w, w->out_fp);
        if (ret != 0) {
                fclose(w->out_fp);
                w->out_fp = NULL;
@@ -1799,8 +1829,10 @@ static int overwrite_wim_inplace(WIMStruct *w, int write_flags,
 
        if (fseeko(w->out_fp, old_wim_end, SEEK_SET) != 0) {
                ERROR_WITH_ERRNO("Can't seek to end of WIM");
-               ret = WIMLIB_ERR_WRITE;
-               goto out_ftruncate;
+               fclose(w->out_fp);
+               w->out_fp = NULL;
+               w->wim_locked = 0;
+               return WIMLIB_ERR_WRITE;
        }
 
        if (!list_empty(&stream_list)) {
@@ -1816,22 +1848,30 @@ static int overwrite_wim_inplace(WIMStruct *w, int write_flags,
                DEBUG("No new streams were added");
        }
 
-       for (int i = modified_image_idx; i < w->hdr.image_count; i++) {
-               select_wim_image(w, i + 1);
-               ret = write_metadata_resource(w);
-               if (ret != 0)
-                       goto out_ftruncate;
+       found_modified_image = false;
+       for (int i = 0; i < w->hdr.image_count; i++) {
+               if (!found_modified_image)
+                       found_modified_image = w->image_metadata[i].modified;
+               if (found_modified_image) {
+                       select_wim_image(w, i + 1);
+                       ret = write_metadata_resource(w);
+                       if (ret != 0)
+                               goto out_ftruncate;
+               }
        }
        write_flags |= WIMLIB_WRITE_FLAG_REUSE_INTEGRITY_TABLE;
        ret = finish_write(w, WIMLIB_ALL_IMAGES, write_flags,
                           progress_func);
 out_ftruncate:
        close_wim_writable(w);
-       if (ret != 0) {
+       if (ret != 0 && !(write_flags & WIMLIB_WRITE_FLAG_NO_LOOKUP_TABLE)) {
                WARNING("Truncating `%s' to its original size (%"PRIu64" bytes)",
                        w->filename, old_wim_end);
-               truncate(w->filename, old_wim_end);
+               /* Return value of truncate() is ignored because this is already
+                * an error path. */
+               (void)truncate(w->filename, old_wim_end);
        }
+       w->wim_locked = 0;
        return ret;
 }
 
@@ -1907,9 +1947,6 @@ WIMLIBAPI int wimlib_overwrite(WIMStruct *w, int write_flags,
                               unsigned num_threads,
                               wimlib_progress_func_t progress_func)
 {
-       if (!w)
-               return WIMLIB_ERR_INVALID_PARAM;
-
        write_flags &= WIMLIB_WRITE_MASK_PUBLIC;
 
        if (!w->filename)
@@ -1923,18 +1960,13 @@ WIMLIBAPI int wimlib_overwrite(WIMStruct *w, int write_flags,
        if ((!w->deletion_occurred || (write_flags & WIMLIB_WRITE_FLAG_SOFT_DELETE))
            && !(write_flags & WIMLIB_WRITE_FLAG_REBUILD))
        {
-               int i, modified_image_idx;
-               for (i = 0; i < w->hdr.image_count && !w->image_metadata[i].modified; i++)
-                       ;
-               modified_image_idx = i;
-               for (; i < w->hdr.image_count && w->image_metadata[i].modified &&
-                       !w->image_metadata[i].has_been_mounted_rw; i++)
-                       ;
-               if (i == w->hdr.image_count) {
-                       return overwrite_wim_inplace(w, write_flags, num_threads,
-                                                    progress_func,
-                                                    modified_image_idx);
-               }
+               int ret;
+               ret = overwrite_wim_inplace(w, write_flags, num_threads,
+                                           progress_func);
+               if (ret == WIMLIB_ERR_RESOURCE_ORDER)
+                       WARNING("Falling back to re-building entire WIM");
+               else
+                       return ret;
        }
        return overwrite_wim_via_tmpfile(w, write_flags, num_threads,
                                         progress_func);