--- a/src/lib-storage/index/index-mail-headers.c	Mon Apr 04 15:41:19 2011 +0300
+++ b/src/lib-storage/index/index-mail-headers.c	Mon Apr 04 16:02:05 2011 +0300
@@ -152,7 +152,6 @@
 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;
 
@@ -225,13 +224,6 @@
 				      &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.c	Mon Apr 04 15:41:19 2011 +0300
+++ b/src/lib-storage/index/index-search.c	Mon Apr 04 16:02:05 2011 +0300
@@ -1036,41 +1036,55 @@
 static void
 wanted_sort_fields_get(struct mailbox *box,
 		       const enum mail_sort_type *sort_program,
+		       struct mailbox_header_lookup_ctx *wanted_headers,
 		       enum mail_fetch_field *wanted_fields_r,
 		       struct mailbox_header_lookup_ctx **headers_ctx_r)
 {
-	const char *headers[2];
+	ARRAY_TYPE(const_string) headers;
+	const char *header;
+	unsigned int i;
 
 	*wanted_fields_r = 0;
 	*headers_ctx_r = NULL;
 
-	headers[0] = headers[1] = NULL;
+	t_array_init(&headers, 8);
+	header = NULL;
 	switch (sort_program[0] & MAIL_SORT_MASK) {
 	case MAIL_SORT_ARRIVAL:
-		*wanted_fields_r = MAIL_FETCH_RECEIVED_DATE;
+		*wanted_fields_r |= MAIL_FETCH_RECEIVED_DATE;
 		break;
 	case MAIL_SORT_CC:
-		headers[0] = "Cc";
+		header = "Cc";
 		break;
 	case MAIL_SORT_DATE:
-		*wanted_fields_r = MAIL_FETCH_DATE;
+		*wanted_fields_r |= MAIL_FETCH_DATE;
 		break;
 	case MAIL_SORT_FROM:
-		headers[0] = "From";
+		header = "From";
 		break;
 	case MAIL_SORT_SIZE:
-		*wanted_fields_r = MAIL_FETCH_VIRTUAL_SIZE;
+		*wanted_fields_r |= MAIL_FETCH_VIRTUAL_SIZE;
 		break;
 	case MAIL_SORT_SUBJECT:
-		headers[0] = "Subject";
+		header = "Subject";
 		break;
 	case MAIL_SORT_TO:
-		headers[0] = "To";
+		header = "To";
 		break;
 	}
+	if (header != NULL)
+		array_append(&headers, &header, 1);
 
-	if (headers[0] != NULL)
-		*headers_ctx_r = mailbox_header_lookup_init(box, headers);
+	if (wanted_headers != NULL) {
+		for (i = 0; wanted_headers->name[i] != NULL; i++)
+			array_append(&headers, &wanted_headers->name[i], 1);
+	}
+
+	if (array_count(&headers) > 0) {
+		(void)array_append_space(&headers);
+		*headers_ctx_r = mailbox_header_lookup_init(box,
+							array_idx(&headers, 0));
+	}
 }
 
 struct mail_search_context *
@@ -1089,11 +1103,6 @@
 	ctx->view = t->view;
 	ctx->mail_ctx.args = args;
 	ctx->mail_ctx.sort_program = index_sort_program_init(t, sort_program);
-	ctx->mail_ctx.wanted_fields = wanted_fields;
-	if (wanted_headers != NULL) {
-		ctx->mail_ctx.wanted_headers = wanted_headers;
-		mailbox_header_lookup_ref(wanted_headers);
-	}
 
 	ctx->max_mails = t->box->storage->set->mail_prefetch_count + 1;
 	if (ctx->max_mails == 0)
@@ -1119,10 +1128,14 @@
 	}
 
 	if (sort_program != NULL) {
-		wanted_sort_fields_get(ctx->box, sort_program,
-				       &ctx->extra_wanted_fields,
-				       &ctx->extra_wanted_headers);
+		wanted_sort_fields_get(ctx->box, sort_program, wanted_headers,
+				       &ctx->mail_ctx.wanted_fields,
+				       &ctx->mail_ctx.wanted_headers);
+	} else if (wanted_headers != NULL) {
+		ctx->mail_ctx.wanted_headers = wanted_headers;
+		mailbox_header_lookup_ref(wanted_headers);
 	}
+	ctx->mail_ctx.wanted_fields |= wanted_fields;
 
 	search_get_seqset(ctx, status.messages, args->args);
 	(void)mail_search_args_foreach(args->args, search_init_arg, ctx);
@@ -1157,8 +1170,6 @@
 
 	if (ctx->mail_ctx.wanted_headers != NULL)
 		mailbox_header_lookup_unref(&ctx->mail_ctx.wanted_headers);
-	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)
@@ -1456,8 +1467,6 @@
 	imail = (struct index_mail *)mail;
 	imail->search_mail = TRUE;
 	imail->mail.stats_track = TRUE;
-	imail->mail.extra_wanted_fields = ctx->extra_wanted_fields;
-	imail->mail.extra_wanted_headers = ctx->extra_wanted_headers;
 
 	array_append(&ctx->mails, &mail, 1);
 	return mail;

--- a/src/lib-storage/index/index-thread.c	Mon Apr 04 15:41:19 2011 +0300
+++ b/src/lib-storage/index/index-thread.c	Mon Apr 04 16:02:05 2011 +0300
@@ -7,13 +7,10 @@
 #include "bsearch-insert-pos.h"
 #include "hash2.h"
 #include "message-id.h"
-#include "mail-index-private.h"
-#include "mail-index-sync-private.h"
 #include "mail-search.h"
 #include "mail-search-build.h"
 #include "mailbox-search-result-private.h"
 #include "index-storage.h"
-#include "index-search-private.h"
 #include "index-thread-private.h"
 
 #include <stdlib.h>
@@ -298,11 +295,10 @@
 		NULL
 	};
 	struct mail_thread_mailbox *tbox = MAIL_THREAD_CONTEXT(ctx->box);
+	struct mailbox_header_lookup_ctx *headers_ctx;
 	struct mail_search_args *search_args;
 	struct mail_search_context *search_ctx;
-	struct index_search_context *search_ctx_i;
 	struct mail *mail;
-	struct mail_private *pmail;
 	uint32_t last_uid, seq1, seq2;
 	int ret = 0;
 
@@ -318,10 +314,8 @@
 						    &tbox->msgid_hash);
 	}
 
-	ctx->tmp_mail = mail_alloc(ctx->t, 0, NULL);
-	pmail = (struct mail_private *)ctx->tmp_mail;
-	pmail->extra_wanted_headers =
-		mailbox_header_lookup_init(ctx->box, wanted_headers);
+	headers_ctx = mailbox_header_lookup_init(ctx->box, wanted_headers);
+	ctx->tmp_mail = mail_alloc(ctx->t, 0, headers_ctx);
 
 	/* add all missing UIDs */
 	ctx->strmap_sync = mail_index_strmap_view_sync_init(tbox->strmap_view,
@@ -336,11 +330,8 @@
 
 	search_args = mail_search_build_init();
 	mail_search_build_add_seqset(search_args, seq1, seq2);
-	search_ctx = mailbox_search_init(ctx->t, search_args, NULL, 0, NULL);
-
-	search_ctx_i = (struct index_search_context *)search_ctx;
-	search_ctx_i->extra_wanted_headers =
-		mailbox_header_lookup_init(ctx->box, wanted_headers);
+	search_ctx = mailbox_search_init(ctx->t, search_args, NULL,
+					 0, headers_ctx);
 
 	while (mailbox_search_next(search_ctx, &mail)) {
 		if (mail_thread_map_add_mail(ctx, mail) < 0) {
@@ -586,9 +577,6 @@
 
 static void mail_thread_clear(struct mail_thread_context *ctx)
 {
-	struct mail_private *pmail = (struct mail_private *)ctx->tmp_mail;
-
-	mailbox_header_lookup_unref(&pmail->extra_wanted_headers);
 	mail_free(&ctx->tmp_mail);
 	(void)mailbox_transaction_commit(&ctx->t);
 }

--- a/src/lib-storage/mail-storage-private.h	Mon Apr 04 15:41:19 2011 +0300
+++ b/src/lib-storage/mail-storage-private.h	Mon Apr 04 16:02:05 2011 +0300
@@ -336,12 +336,6 @@
 	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;