[PATCH 77/82] imap: wire up Xapian, MSN SEARCH and multi sequence-sets

[Eric Wong <e@yhbt.net>]

Simple queries work, more complex queries involving parentheses,
"OR", "NOT" don't work, yet.

Tested with "=b", "=B", and "=H" search and limits in mutt
on both v1 and v2 with multiple Xapian shards.
---
 lib/PublicInbox/IMAP.pm   | 142 ++++++++++++++++++++++++++++++--------
 lib/PublicInbox/Search.pm |  19 ++++-
 t/imapd.t                 |  46 +++++++++++-
 3 files changed, 175 insertions(+), 32 deletions(-)

diff --git a/lib/PublicInbox/IMAP.pm b/lib/PublicInbox/IMAP.pm
index f6106a1e806..67bc32ab59b 100644
--- a/lib/PublicInbox/IMAP.pm
+++ b/lib/PublicInbox/IMAP.pm
@@ -24,6 +24,8 @@ use Errno qw(EAGAIN);
 use Time::Local qw(timegm);
 use POSIX qw(strftime);
 use Hash::Util qw(unlock_hash); # dependency of fields for perl 5.10+, anyways
+use PublicInbox::Search;
+*mdocid = \&PublicInbox::Search::mdocid;
 
 my $Address;
 for my $mod (qw(Email::Address::XS Mail::Address)) {
@@ -592,22 +594,26 @@ sub range_step ($$) {
 		$range = $$range_csv;
 		$$range_csv = undef;
 	}
+	my $uid_base = $self->{uid_base};
+	my $uid_end = $uid_base + UID_BLOCK;
 	if ($range =~ /\A([0-9]+):([0-9]+)\z/) {
 		($beg, $end) = ($1 + 0, $2 + 0);
+		uid_clamp($self, \$beg, \$end);
 	} elsif ($range =~ /\A([0-9]+):\*\z/) {
 		$beg = $1 + 0;
 		$end = $self->{ibx}->over->max;
-		my $uid_end = $self->{uid_base} + UID_BLOCK;
 		$end = $uid_end if $end > $uid_end;
 		$beg = $end if $beg > $end;
+		uid_clamp($self, \$beg, \$end);
 	} elsif ($range =~ /\A[0-9]+\z/) {
 		$beg = $end = $range + 0;
-		undef $range;
+		# just let the caller do an out-of-range query if a single
+		# UID is out-of-range
+		++$beg if ($beg <= $uid_base || $end > $uid_end);
 	} else {
 		return 'BAD fetch range';
 	}
-	uid_clamp($self, \$beg, \$end) if defined($range);
-	my $msn = $beg - $self->{uid_base};
+	my $msn = $beg - $uid_base;
 	[ $beg, $end, $$range_csv, \$msn ];
 }
 
@@ -971,15 +977,22 @@ sub parse_date ($) { # 02-Oct-1993
 	timegm(0, 0, 0, $dd, $mm, $yyyy);
 }
 
-sub uid_search_uid_range { # long_response
-	my ($self, $tag, $uids, $sql, $range_info) = @_;
+sub msn_convert ($$) {
+	my ($self, $uids) = @_;
+	my $adj = $self->{uid_base};
+	$_ -= $adj for @$uids;
+}
+
+sub search_uid_range { # long_response
+	my ($self, $tag, $sql, $range_info, $want_msn) = @_;
+	my $uids = [];
 	if (defined(my $err = refill_uids($self, $uids, $range_info, $sql))) {
 		$err ||= 'OK Search done';
 		$self->write("\r\n$tag $err\r\n");
 		return;
 	}
+	msn_convert($self, $uids) if $want_msn;
 	$self->msg_more(join(' ', '', @$uids));
-	@$uids = ();
 	1; # more
 }
 
@@ -1029,6 +1042,21 @@ my %I2X = (
 	# KEYWORD # TODO ? dfpre,dfpost,...
 );
 
+# IMAP allows searching arbitrary headers via "HEADER $HDR_NAME $HDR_VAL"
+# which gets silly expensive.  We only allow the headers we already index.
+my %H2X = (%I2X, 'MESSAGE-ID' => 'm:', 'LIST-ID' => 'l:');
+
+sub xap_append ($$$$) {
+	my ($q, $rest, $k, $xk) = @_;
+	delete $q->{sql}; # can't use over.sqlite3
+	defined(my $arg = shift @$rest) or return "BAD $k no arg";
+
+	# AFAIK Xapian can't handle [*"] in probabilistic terms
+	$arg =~ tr/*"//d;
+	${$q->{xap}} .= qq[ $xk"$arg"];
+	undef;
+}
+
 sub parse_query {
 	my ($self, $rest) = @_;
 	if (uc($rest->[0]) eq 'CHARSET') {
@@ -1038,7 +1066,8 @@ sub parse_query {
 	}
 
 	my $sql = ''; # date conditions, {sql} deleted if Xapian is needed
-	my $q = { xap => '', sql => \$sql };
+	my $xap = '';
+	my $q = { sql => \$sql, xap => \$xap };
 	while (@$rest) {
 		my $k = uc(shift @$rest);
 		# default criteria
@@ -1059,17 +1088,18 @@ sub parse_query {
 			delete $q->{sql}; # can't use over.sqlite3
 			my $bytes = shift(@$rest) // '';
 			$bytes =~ /\A[0-9]+\z/ or return "BAD $k not a number";
-			$q->{xap} .= ' bytes:' . ($k eq 'SMALLER' ?
+			$xap .= ' bytes:' . ($k eq 'SMALLER' ?
 							'..'.(--$bytes) :
 							(++$bytes).'..');
+		} elsif ($k eq 'HEADER') {
+			$k = uc(shift(@$rest) // '');
+			my $xk = $H2X{$k} or
+				return "BAD HEADER $k not supported";
+			my $err = xap_append($q, $rest, $k, $xk);
+			return $err if $err;
 		} elsif (defined(my $xk = $I2X{$k})) {
-			delete $q->{sql}; # can't use over.sqlite3
-			my $arg = shift @$rest;
-			defined($arg) or return "BAD $k no arg";
-
-			# Xapian can't handle [*"] in probabilistic terms
-			$arg =~ tr/*"//d;
-			$q->{xap} .= qq[ $xk:"$arg"];
+			my $err = xap_append($q, $rest, $k, $xk);
+			return $err if $err;
 		} else {
 			# TODO: parentheses, OR, NOT ...
 			return "BAD $k not supported (yet?)";
@@ -1083,31 +1113,87 @@ sub parse_query {
 	} elsif (!$self->{ibx}->search) {
 		return 'BAD Xapian not configured for mailbox';
 	}
-
+	my $max = $self->{ibx}->over->max;
 	if (my $uid = delete $q->{uid}) {
-		$q->{uid} = join(',', @$uid);
+		my $range_csv = join(',', @$uid);
+		do {
+			my $nxt = range_step($self, \$range_csv);
+			my ($beg, $end) = @$nxt;
+			if ($xap) {
+				$xap .= " uid:$beg..$end";
+			} elsif ($beg == $end) {
+				$sql .= " AND num = $beg";
+			} else {
+				$sql .= " AND num >= $beg AND num <= $end";
+			}
+		} while ($range_csv);
 	}
+	my $beg = 1;
+	uid_clamp($self, \$beg, \$max);
+	$q->{range_info} = [ $beg, $max ];
 	$q;
 }
 
-sub cmd_uid_search ($$$;) {
-	my ($self, $tag) = splice(@_, 0, 2);
+sub refill_xap ($$$$) {
+	my ($self, $uids, $range_info, $q) = @_;
+	my ($beg, $end) = @$range_info;
+	my $srch = $self->{ibx}->search;
+	my $opt = { mset => 2, limit => 1000 };
+	my $nshard = $srch->{nshard} // 1;
+	while (1) {
+		my $mset = $srch->query("$$q uid:$beg..$end", $opt);
+		@$uids = map { mdocid($nshard, $_) } $mset->items;
+		if (@$uids) {
+			$range_info->[0] = $uids->[-1] + 1; # update $beg
+			return;
+		} else { # all done
+			return 0;
+		}
+	}
+}
+
+sub search_xap_range { # long_response
+	my ($self, $tag, $q, $range_info, $want_msn) = @_;
+	my $uids = [];
+	if (defined(my $err = refill_xap($self, $uids, $range_info, $q))) {
+		$err ||= 'OK Search done';
+		$self->write("\r\n$tag $err\r\n");
+		return;
+	}
+	msn_convert($self, $uids) if $want_msn;
+	$self->msg_more(join(' ', '', @$uids));
+	1; # more
+}
+
+sub search_common {
+	my ($self, $tag, $rest, $want_msn) = @_;
 	my $ibx = $self->{ibx} or return "$tag BAD No mailbox selected\r\n";
-	my $q = parse_query($self, \@_);
+	my $q = parse_query($self, $rest);
 	return "$tag $q\r\n" if !ref($q);
-	my $sql = delete $q->{sql};
-	my $range_csv = delete $q->{uid} // '1:*';
-	my $range_info = range_step($self, \$range_csv);
-	return "$tag $range_info\r\n" if !ref($range_info);
-	if (!scalar(keys %$q)) {
+	my ($sql, $range_info) = delete @$q{qw(sql range_info)};
+	if (!scalar(keys %$q)) { # overview.sqlite3
+		$self->msg_more('* SEARCH');
+		long_response($self, \&search_uid_range,
+				$tag, $sql, $range_info, $want_msn);
+	} elsif ($q = $q->{xap}) {
 		$self->msg_more('* SEARCH');
-		long_response($self, \&uid_search_uid_range,
-				$tag, [], $sql, $range_info);
+		long_response($self, \&search_xap_range,
+				$tag, $q, $range_info, $want_msn);
 	} else {
 		"$tag BAD Error\r\n";
 	}
 }
 
+sub cmd_uid_search ($$$;) {
+	my ($self, $tag) = splice(@_, 0, 2);
+	search_common($self, $tag, \@_);
+}
+
+sub cmd_search ($$$;) {
+	my ($self, $tag) = splice(@_, 0, 2);
+	search_common($self, $tag, \@_, 1);
+}
+
 sub args_ok ($$) { # duplicated from PublicInbox::NNTP
 	my ($cb, $argc) = @_;
 	my $tot = prototype $cb;
diff --git a/lib/PublicInbox/Search.pm b/lib/PublicInbox/Search.pm
index c54cf7b9911..55eee41ca4a 100644
--- a/lib/PublicInbox/Search.pm
+++ b/lib/PublicInbox/Search.pm
@@ -182,6 +182,7 @@ sub _xdb ($) {
 	my ($xdb, $slow_phrase);
 	my $qpf = \($self->{qp_flags} ||= $QP_FLAGS);
 	if ($self->{ibx_ver} >= 2) {
+		my $n = 0;
 		foreach my $shard (<$dir/*>) {
 			-d $shard && $shard =~ m!/[0-9]+\z! or next;
 			my $sub = $X{Database}->new($shard);
@@ -191,7 +192,9 @@ sub _xdb ($) {
 				$xdb = $sub;
 			}
 			$slow_phrase ||= -f "$shard/iamchert";
+			++$n;
 		}
+		$self->{nshard} = $n;
 	} else {
 		$slow_phrase = -f "$dir/iamchert";
 		$xdb = $X{Database}->new($dir);
@@ -200,6 +203,15 @@ sub _xdb ($) {
 	$xdb;
 }
 
+# v2 Xapian docids don't conflict, so they're identical to
+# NNTP article numbers and IMAP UIDs.
+# https://trac.xapian.org/wiki/FAQ/MultiDatabaseDocumentID
+sub mdocid {
+	my ($nshard, $mitem) = @_;
+	my $docid = $mitem->get_docid;
+	int(($docid - 1) / $nshard) + 1;
+}
+
 sub xdb ($) {
 	my ($self) = @_;
 	$self->{xdb} ||= do {
@@ -283,7 +295,7 @@ sub _enquire_once { # retry_reopen callback
 	$enquire->set_query($query);
 	$opts ||= {};
         my $desc = !$opts->{asc};
-	if (($opts->{mset} || 0) == 2) {
+	if (($opts->{mset} || 0) == 2) { # mset == 2: ORDER BY docid/UID
 		$enquire->set_docid_order($ENQ_ASCENDING);
 		$enquire->set_weighting_scheme($X{BoolWeight}->new);
 	} elsif ($opts->{relevance}) {
@@ -322,6 +334,11 @@ sub qp {
 	$qp->add_valuerangeprocessor($nvrp->new(YYYYMMDD, 'd:'));
 	$qp->add_valuerangeprocessor($nvrp->new(DT, 'dt:'));
 
+	# for IMAP, undocumented for WWW and may be split off go away
+	$qp->add_valuerangeprocessor($nvrp->new(BYTES, 'bytes:'));
+	$qp->add_valuerangeprocessor($nvrp->new(TS, 'ts:'));
+	$qp->add_valuerangeprocessor($nvrp->new(UID, 'uid:'));
+
 	while (my ($name, $prefix) = each %bool_pfx_external) {
 		$qp->add_boolean_prefix($name, $_) foreach split(/ /, $prefix);
 	}
diff --git a/t/imapd.t b/t/imapd.t
index fdab074a249..e3cce2d30f0 100644
--- a/t/imapd.t
+++ b/t/imapd.t
@@ -19,10 +19,10 @@ if ($can_compress) { # hope this gets fixed upstream, soon
 require_ok 'PublicInbox::IMAP';
 my $first_range = '0';
 
-my $level = '-Lbasic';
+my $level = 'basic';
 SKIP: {
 	require_mods('Search::Xapian', 1);
-	$level = '-Lmedium';
+	$level = 'medium';
 };
 
 my @V = (1);
@@ -38,7 +38,7 @@ for my $V (@V) {
 	my $url = "http://example.com/i$V";
 	my $inboxdir = "$tmpdir/$name";
 	my $folder = "inbox.i$V";
-	my $cmd = ['-init', "-V$V", $level, $name, $inboxdir, $url, $addr];
+	my $cmd = ['-init', "-V$V", "-L$level", $name, $inboxdir, $url, $addr];
 	run_script($cmd) or BAIL_OUT("init $name");
 	xsys(qw(git config), "--file=$ENV{HOME}/.public-inbox/config",
 			"publicinbox.$name.newsgroup", $folder) == 0 or
@@ -120,6 +120,24 @@ is_deeply($ret, [ 1 ], 'search UID 1:1 works');
 $ret = $mic->search('uid 1:*') or BAIL_OUT "SEARCH FAIL $@";
 is_deeply($ret, [ 1 ], 'search UID 1:* works');
 
+SKIP: {
+	skip 'Xapian missing', 6 if $level eq 'basic';
+	my $x = $mic->search(qw(smaller 99999));
+	is_deeply($x, [1], 'SMALLER works with Xapian (hit)');
+	$x = $mic->search(qw(smaller 9));
+	is_deeply($x, [], 'SMALLER works with Xapian (miss)');
+
+	$x = $mic->search(qw(larger 99999));
+	is_deeply($x, [], 'LARGER works with Xapian (miss)');
+	$x = $mic->search(qw(larger 9));
+	is_deeply($x, [1], 'LARGER works with Xapian (hit)');
+
+	$x = $mic->search(qw(HEADER Message-ID testmessage@example.com));
+	is_deeply($x, [1], 'HEADER Message-ID works');
+	$x = $mic->search(qw(HEADER Message-ID miss));
+	is_deeply($x, [], 'HEADER Message-ID can miss');
+}
+
 is_deeply(scalar $mic->flags('1'), [], '->flags works');
 {
 	# RFC 3501 section 6.4.8 states:
@@ -341,12 +359,34 @@ $ret = $mic->fetch_hash('1,2:3', 'RFC822') or BAIL_OUT "FETCH $@";
 is(scalar keys %$ret, 3, 'got all 3 messages with comma-separated sequence');
 $ret = $mic->fetch_hash('1:*', 'RFC822') or BAIL_OUT "FETCH $@";
 is(scalar keys %$ret, 3, 'got all 3 messages');
+
+SKIP: {
+	# do any clients use non-UID IMAP SEARCH?
+	skip 'Xapian missing', 2 if $level eq 'basic';
+	my $x = $mic->search('all');
+	is_deeply($x, [1, 2, 3], 'MSN SEARCH works before rm');
+	$x = $mic->search(qw(header subject embedded));
+	is_deeply($x, [2], 'MSN SEARCH on Subject works before rm');
+}
+
 {
 	my $rdr = { 0 => \($ret->{1}->{RFC822}) };
 	my $env = { HOME => $ENV{HOME} };
 	my @cmd = qw(-learn rm --all);
 	run_script(\@cmd, $env, $rdr) or BAIL_OUT('-learn rm');
 }
+
+SKIP: {
+	# do any clients use non-UID IMAP SEARCH?  We only ensure
+	# MSN "SEARCH" can return a result which can be retrieved
+	# via MSN "FETCH"
+	skip 'Xapian missing', 3 if $level eq 'basic';
+	my $x = $mic->search(qw(header subject embedded));
+	is(scalar(@$x), 1, 'MSN SEARCH on Subject works after rm');
+	$x = $mic->message_string($x->[0]);
+	is($x, $ret->{2}->{RFC822}, 'message 2 unchanged');
+}
+
 my $r2 = $mic->fetch_hash('1:*', 'BODY.PEEK[]') or BAIL_OUT "FETCH $@";
 is(scalar keys %$r2, 2, 'did not get all 3 messages');
 is($r2->{1}->{'BODY[]'}, $ret->{2}->{RFC822}, 'message 2 unchanged');