--- a/src/imap/imap-search.c	Mon Jan 31 01:36:29 2011 +0200
+++ b/src/imap/imap-search.c	Mon Jan 31 02:13:06 2011 +0200
@@ -474,55 +474,11 @@
 	return 1;
 }
 
-static void wanted_fields_get(struct mailbox *box,
-			      const enum mail_sort_type *sort_program,
-			      enum mail_fetch_field *wanted_fields_r,
-			      struct mailbox_header_lookup_ctx **headers_ctx_r)
-{
-	const char *headers[2];
-
-	*wanted_fields_r = 0;
-	*headers_ctx_r = NULL;
-
-	if (sort_program == NULL)
-		return;
-
-	headers[0] = headers[1] = NULL;
-	switch (sort_program[0] & MAIL_SORT_MASK) {
-	case MAIL_SORT_ARRIVAL:
-		*wanted_fields_r = MAIL_FETCH_RECEIVED_DATE;
-		break;
-	case MAIL_SORT_CC:
-		headers[0] = "Cc";
-		break;
-	case MAIL_SORT_DATE:
-		*wanted_fields_r = MAIL_FETCH_DATE;
-		break;
-	case MAIL_SORT_FROM:
-		headers[0] = "From";
-		break;
-	case MAIL_SORT_SIZE:
-		*wanted_fields_r = MAIL_FETCH_VIRTUAL_SIZE;
-		break;
-	case MAIL_SORT_SUBJECT:
-		headers[0] = "Subject";
-		break;
-	case MAIL_SORT_TO:
-		headers[0] = "To";
-		break;
-	}
-
-	if (headers[0] != NULL)
-		*headers_ctx_r = mailbox_header_lookup_init(box, headers);
-}
-
 bool imap_search_start(struct imap_search_context *ctx,
 		       struct mail_search_args *sargs,
 		       const enum mail_sort_type *sort_program)
 {
 	struct client_command_context *cmd = ctx->cmd;
-	enum mail_fetch_field wanted_fields;
-	struct mailbox_header_lookup_ctx *wanted_headers;
 
 	imap_search_args_check(ctx, sargs->args);
 
@@ -532,13 +488,10 @@
 	}
 
 	ctx->box = cmd->client->mailbox;
-	wanted_fields_get(ctx->box, sort_program,
-			  &wanted_fields, &wanted_headers);
-
 	ctx->trans = mailbox_transaction_begin(ctx->box, 0);
 	ctx->sargs = sargs;
 	ctx->search_ctx = mailbox_search_init(ctx->trans, sargs, sort_program);
-	ctx->mail = mail_alloc(ctx->trans, wanted_fields, wanted_headers);
+	ctx->mail = mail_alloc(ctx->trans, 0, NULL);
 	ctx->sorting = sort_program != NULL;
 	(void)gettimeofday(&ctx->start_time, NULL);
 	i_array_init(&ctx->result, 128);

--- a/src/lib-storage/index/index-mail-headers.c	Mon Jan 31 01:36:29 2011 +0200
+++ b/src/lib-storage/index/index-mail-headers.c	Mon Jan 31 02:13:06 2011 +0200
@@ -152,6 +152,7 @@
 bool index_mail_want_parse_headers(struct index_mail *mail)
 {
 	if (mail->wanted_headers != NULL ||
+	    mail->mail.extra_wanted_headers != NULL ||
 	    mail->data.save_bodystructure_header)
 		return TRUE;
 
@@ -224,6 +225,13 @@
 				      &mail->header_match_value);
 		}
 	}
+	if (mail->mail.extra_wanted_headers != NULL) {
+		headers = mail->mail.extra_wanted_headers;
+		for (i = 0; i < headers->count; i++) {
+			array_idx_set(&mail->header_match, headers->idx[i],
+				      &mail->header_match_value);
+		}
+	}
 
 	/* register also all the other headers that exist in cache file */
 	T_BEGIN {

--- a/src/lib-storage/index/index-search-private.h	Mon Jan 31 01:36:29 2011 +0200
+++ b/src/lib-storage/index/index-search-private.h	Mon Jan 31 02:13:06 2011 +0200
@@ -8,6 +8,9 @@
 	struct mail_index_view *view;
 	struct mailbox *box;
 
+	enum mail_fetch_field extra_wanted_fields;
+	struct mailbox_header_lookup_ctx *extra_wanted_headers;
+
 	uint32_t seq1, seq2;
 	struct mail *mail;
 	struct index_mail *imail;

--- a/src/lib-storage/index/index-search.c	Mon Jan 31 01:36:29 2011 +0200
+++ b/src/lib-storage/index/index-search.c	Mon Jan 31 02:13:06 2011 +0200
@@ -1040,6 +1040,46 @@
 	return ret;
 }
 
+static void
+wanted_sort_fields_get(struct mailbox *box,
+		       const enum mail_sort_type *sort_program,
+		       enum mail_fetch_field *wanted_fields_r,
+		       struct mailbox_header_lookup_ctx **headers_ctx_r)
+{
+	const char *headers[2];
+
+	*wanted_fields_r = 0;
+	*headers_ctx_r = NULL;
+
+	headers[0] = headers[1] = NULL;
+	switch (sort_program[0] & MAIL_SORT_MASK) {
+	case MAIL_SORT_ARRIVAL:
+		*wanted_fields_r = MAIL_FETCH_RECEIVED_DATE;
+		break;
+	case MAIL_SORT_CC:
+		headers[0] = "Cc";
+		break;
+	case MAIL_SORT_DATE:
+		*wanted_fields_r = MAIL_FETCH_DATE;
+		break;
+	case MAIL_SORT_FROM:
+		headers[0] = "From";
+		break;
+	case MAIL_SORT_SIZE:
+		*wanted_fields_r = MAIL_FETCH_VIRTUAL_SIZE;
+		break;
+	case MAIL_SORT_SUBJECT:
+		headers[0] = "Subject";
+		break;
+	case MAIL_SORT_TO:
+		headers[0] = "To";
+		break;
+	}
+
+	if (headers[0] != NULL)
+		*headers_ctx_r = mailbox_header_lookup_init(box, headers);
+}
+
 void index_storage_search_init_context(struct index_search_context *ctx,
 				       struct mailbox_transaction_context *t,
 				       struct mail_search_args *args,
@@ -1071,6 +1111,12 @@
 			ctx->failed = TRUE;
 	}
 
+	if (sort_program != NULL) {
+		wanted_sort_fields_get(ctx->box, sort_program,
+				       &ctx->extra_wanted_fields,
+				       &ctx->extra_wanted_headers);
+	}
+
 	search_get_seqset(ctx, status.messages, args->args);
 	(void)mail_search_args_foreach(args->args, search_init_arg, ctx);
 
@@ -1117,6 +1163,8 @@
 	(void)mail_search_args_foreach(ctx->mail_ctx.args->args,
 				       search_arg_deinit, NULL);
 
+	if (ctx->extra_wanted_headers != NULL)
+		mailbox_header_lookup_unref(&ctx->extra_wanted_headers);
 	if (ctx->mail_ctx.sort_program != NULL)
 		index_sort_program_deinit(&ctx->mail_ctx.sort_program);
 	if (ctx->thread_ctx != NULL)
@@ -1347,6 +1395,8 @@
 	}
 
 	ctx->mail = mail;
+	mail_private->extra_wanted_fields = ctx->extra_wanted_fields;
+	mail_private->extra_wanted_headers = ctx->extra_wanted_headers;
 
 	if (ioloop_time - ctx->last_notify.tv_sec >=
 	    SEARCH_NOTIFY_INTERVAL_SECS)
@@ -1396,6 +1446,8 @@
 	ctx->mail = NULL;
 	ctx->imail = NULL;
 	mail_private->stats_track = old_stats_track;
+	mail_private->extra_wanted_fields = 0;
+	mail_private->extra_wanted_headers = NULL;
 
 	if (!match && _ctx->sort_program != NULL &&
 	    !*tryagain_r && !ctx->failed) {

--- a/src/lib-storage/index/index-sort.c	Mon Jan 31 01:36:29 2011 +0200
+++ b/src/lib-storage/index/index-sort.c	Mon Jan 31 02:13:06 2011 +0200
@@ -108,7 +108,9 @@
 {
 	i_assert(mail->transaction == program->t);
 
-	program->sort_list_add(program, mail);
+	T_BEGIN {
+		program->sort_list_add(program, mail);
+	} T_END;
 }
 
 static int sort_node_date_cmp(const struct mail_sort_node_date *n1,

--- a/src/lib-storage/mail-storage-private.h	Mon Jan 31 01:36:29 2011 +0200
+++ b/src/lib-storage/mail-storage-private.h	Mon Jan 31 02:13:06 2011 +0200
@@ -319,6 +319,12 @@
 	pool_t pool;
 	ARRAY_DEFINE(module_contexts, union mail_module_context *);
 
+	/* temporary extra wanted fields/headers that should be looked up
+	   automatically if possible. for example sort_program fields during
+	   search. */
+	enum mail_fetch_field extra_wanted_fields;
+	struct mailbox_header_lookup_ctx *extra_wanted_headers;
+
 	/* these statistics are never reset by mail-storage API: */
 
 	unsigned long stats_open_lookup_count;