From fef476297d762fec1f4c1517895add6b4c342915 Mon Sep 17 00:00:00 2001
From: Eric Biggers <ebiggers3@gmail.com>
Date: Sat, 2 May 2015 13:39:17 -0500
Subject: [PATCH] Win32 capture: sort streams by starting LCN

---
 NEWS                        |  3 ++-
 include/wimlib/blob_table.h |  3 +++
 src/blob_table.c            |  6 +++++-
 src/win32_capture.c         | 36 ++++++++++++++++++++++++++++++++++++
 4 files changed, 46 insertions(+), 2 deletions(-)

diff --git a/NEWS b/NEWS
index ab71b7c8..f5703ca4 100644
--- a/NEWS
+++ b/NEWS
@@ -28,7 +28,8 @@ Version 1.8.1-BETA:
 	In mounted WIM images, the disk usage provided for each file (st_blocks)
 	is now the compressed size rather than the uncompressed size.
 
-	The performance of NTFS-3g capture mode has been slightly improved.
+	The performance of the NTFS-3g and Windows capture modes has been
+	slightly improved.
 
 Version 1.8.0:
 	Improved the LZX compressor.  It is now 15-20% faster than before and
diff --git a/include/wimlib/blob_table.h b/include/wimlib/blob_table.h
index 556187e8..dcf6af7f 100644
--- a/include/wimlib/blob_table.h
+++ b/include/wimlib/blob_table.h
@@ -174,6 +174,9 @@ struct blob_descriptor {
 				struct {
 					tchar *file_on_disk;
 					struct wim_inode *file_inode;
+				#ifdef __WIN32__
+					u64 sort_key;
+				#endif
 				};
 
 				/* BLOB_IN_ATTACHED_BUFFER */
diff --git a/src/blob_table.c b/src/blob_table.c
index 1d87a717..e93fce1e 100644
--- a/src/blob_table.c
+++ b/src/blob_table.c
@@ -395,7 +395,7 @@ cmp_blobs_by_sequential_order(const void *p1, const void *p2)
 
 	v = (int)blob1->blob_location - (int)blob2->blob_location;
 
-	/* Different resource locations?  */
+	/* Different locations?  */
 	if (v)
 		return v;
 
@@ -429,6 +429,10 @@ cmp_blobs_by_sequential_order(const void *p1, const void *p2)
 #ifdef __WIN32__
 	case BLOB_IN_WINNT_FILE_ON_DISK:
 	case BLOB_WIN32_ENCRYPTED:
+		/* Windows: compare by starting LCN (logical cluster number)  */
+		v = cmp_u64(blob1->sort_key, blob2->sort_key);
+		if (v)
+			return v;
 #endif
 		/* Compare files by path: just a heuristic that will place files
 		 * in the same directory next to each other.  */
diff --git a/src/win32_capture.c b/src/win32_capture.c
index 13f63650..3061f6ae 100644
--- a/src/win32_capture.c
+++ b/src/win32_capture.c
@@ -1099,6 +1099,37 @@ out_free_buf:
 	return ret;
 }
 
+static u64
+get_sort_key(HANDLE h)
+{
+	STARTING_VCN_INPUT_BUFFER in = { .StartingVcn.QuadPart = 0 };
+	RETRIEVAL_POINTERS_BUFFER out;
+	DWORD bytesReturned;
+
+	if (!DeviceIoControl(h, FSCTL_GET_RETRIEVAL_POINTERS,
+			     &in, sizeof(in),
+			     &out, sizeof(out),
+			     &bytesReturned, NULL))
+		return 0;
+
+	if (out.ExtentCount < 1)
+		return 0;
+
+	return out.Extents[0].Lcn.QuadPart;
+}
+
+static void
+set_sort_key(struct wim_inode *inode, u64 sort_key)
+{
+	for (unsigned i = 0; i < inode->i_num_streams; i++) {
+		struct wim_inode_stream *strm = &inode->i_streams[i];
+		struct blob_descriptor *blob = stream_blob_resolved(strm);
+		if (blob && (blob->blob_location == BLOB_IN_WINNT_FILE_ON_DISK ||
+			     blob->blob_location == BLOB_WIN32_ENCRYPTED))
+			blob->sort_key = sort_key;
+	}
+}
+
 static int
 winnt_build_dentry_tree_recursive(struct wim_dentry **root_ret,
 				  HANDLE cur_dir,
@@ -1117,6 +1148,7 @@ winnt_build_dentry_tree_recursive(struct wim_dentry **root_ret,
 	NTSTATUS status;
 	FILE_ALL_INFORMATION file_info;
 	ACCESS_MASK requestedPerms;
+	u64 sort_key;
 
 	ret = try_exclude(full_path, full_path_nchars, params);
 	if (ret < 0) /* Excluded? */
@@ -1341,6 +1373,8 @@ retry_open:
 		}
 	}
 
+	sort_key = get_sort_key(h);
+
 	if (unlikely(inode->i_attributes & FILE_ATTRIBUTE_ENCRYPTED)) {
 		/* Load information about the raw encrypted data.  This is
 		 * needed for any directory or non-directory that has
@@ -1378,6 +1412,8 @@ retry_open:
 			goto out;
 	}
 
+	set_sort_key(inode, sort_key);
+
 	if (inode_is_directory(inode)) {
 
 		/* Directory: recurse to children.  */
-- 
2.43.0