[PATCH 2/4] each: support passing user args to callback

[Eric Wong <e@80x24.org>]

This allows users to avoid creating expensive and cycle-prone
anonymous subs.  This API promote safer usage and reduces the
likelyhood of users introducing leaks.
---
 Mwrap.xs  |  8 ++++++--
 t/mwrap.t | 30 ++++++++++++++++++++++--------
 2 files changed, 28 insertions(+), 10 deletions(-)

diff --git a/Mwrap.xs b/Mwrap.xs
index c93223e..54daf45 100644
--- a/Mwrap.xs
+++ b/Mwrap.xs
@@ -810,9 +810,10 @@ OUTPUT:
 	RETVAL
 
 void
-mwrap_each(min, cb)
+mwrap_each(min, cb, arg = &PL_sv_undef)
 	size_t min;
 	SV *cb;
+	SV *arg;
 PREINIT:
 	struct cds_lfht *t;
 	struct cds_lfht_iter iter;
@@ -836,6 +837,7 @@ CODE:
 				PUSHMARK(SP);
 				loc = sv_newmortal();
 				sv_setref_pv(loc, "Devel::Mwrap::SrcLoc", l);
+				XPUSHs(arg);
 				XPUSHs(loc);
 				PUTBACK;
 
@@ -996,10 +998,11 @@ CLEANUP:
 	--locating;
 
 void
-src_loc_each(self, min, cb)
+src_loc_each(self, min, cb, arg = &PL_sv_undef)
 	Devel::Mwrap::SrcLoc self;
 	size_t min;
 	SV *cb;
+	SV *arg;
 PREINIT:
 	struct alloc_hdr *h;
 	bool err = false;
@@ -1021,6 +1024,7 @@ CODE:
 			 * note: we MUST NOT expose alloc_hdr to Perl code
 			 * since that opens us up to use-after-free
 			 */
+			XPUSHs(arg);
 			XPUSHs(sv_2mortal(newSVuv(size)));
 			XPUSHs(sv_2mortal(newSVuv(gen)));
 			PUTBACK;
diff --git a/t/mwrap.t b/t/mwrap.t
index fdae523..91bae8a 100644
--- a/t/mwrap.t
+++ b/t/mwrap.t
@@ -61,26 +61,40 @@ mwrap_run('Devel::Mwrap::each', {}, '-e', <<'EOF');
 open my $zero, '<', '/dev/zero' or die "open /dev/zero: $!";
 my $nbytes = 1024 * 512;
 sysread($zero, my $before, $nbytes);
-my @keep;
-eval { Devel::Mwrap::each(0, sub { die "died in each"; push @keep, @_; }) };
-"$@" =~ /died in each at -e line 5/ or die "did not propagate die in each";
+my (@keep, @uargs);
+my $uarg = 'user-arg';
+eval {
+	Devel::Mwrap::each(0, sub {
+		my $ua = shift;
+		push @uargs, $ua;
+		die "died in each";
+		push @keep, @_;
+	}, $uarg) };
+"$@" =~ /died in each at -e line \d+/ or die "did not propagate die in each";
 scalar(@keep) == 0 or die "nothing if died in sub";
+$uarg eq join('-', @uargs) or die "uarg passed to sub";
 my %h;
 my $iter;
+my $arg = "FOO";
 Devel::Mwrap::each(1, sub {
-	my ($src_loc, @more) = @_;
+	my ($argh, $src_loc, @more) = @_;
 	my $name = $src_loc->name;
 	$h{$name} = $src_loc;
 	sysread($zero, my $iter, $nbytes);
 	die 'extra args from ::each' if @more;
-});
+	die "arg mismatch $arg / $argh" if $argh ne $arg;
+}, $arg);
 my $src_loc = $h{"-e:3"} or die "missing src_loc for -e:3";
 $src_loc->total >= $nbytes or die "allocated <= $nbytes";
 $src_loc = Devel::Mwrap::get("-e:8");
 if ($src_loc && $src_loc->total >= $nbytes) {
 	die "large allocation inside :each should not be tracked";
 }
-Devel::Mwrap::each($nbytes, sub { push @keep, @_ });
+my $nr = 0;
+Devel::Mwrap::each($nbytes, sub {
+	$foo += shift;
+	push @keep, @_;
+}, 1);
 scalar(@keep) == 1 or die "got multiple matches for > $nbytes";
 $keep[0]->name eq '-e:3' or die "unexpected location for > $nbytes";
 @keep = ();
@@ -106,7 +120,7 @@ $loc && $loc->total >= ($nbytes * $nr) or die "wrong line or bad stats";
 my $ml = $loc->mean_lifespan;
 $ml >= 0.0 or die "mean_lifespan broken";
 my @sl_each;
-$loc->each($nbytes, sub { push @sl_each, \@_ });
+$loc->each($nbytes, sub { shift; push @sl_each, \@_ });
 my $n = scalar(@sl_each);
 $n == 1 or die "SrcLoc::each returned unexpected: $n";
 $sl_each[0]->[0] >= $nbytes or die "$sl_each[0]->[0] < $nbytes";
@@ -115,7 +129,7 @@ my $age_before = $sl_each[0]->[1];
 $nbytes = 1024 * 1024 * 8;
 do_read() until Devel::Mwrap::current_age() > ($age_before + $nbytes);
 @sl_each = ();
-$loc->each($nbytes, sub { push @sl_each, \@_ });
+$loc->each($nbytes, sub { shift; push @sl_each, \@_ });
 $n = scalar(@sl_each);
 $n == 1 or die "SrcLoc::each returned unexpected: $n";
 $sl_each[0]->[0] >= $nbytes or die "$sl_each[0]->[0] < $nbytes";