Implement -limit N

Closes: https://github.com/tavianator/bfs/issues/133

15 files changed, 103 insertions, 3 deletions

diff --git a/completions/bfs.bash b/completions/bfs.bash
index 2f52e8d..db582da 100644
--- a/completions/bfs.bash
+++ b/completions/bfs.bash
@@ -37,6 +37,7 @@ _bfs() {
         -ipath
         -iregex
         -iwholename
+        -limit
         -links
         -lname
         -maxdepth
@@ -94,8 +95,6 @@ _bfs() {
         -depth
         -follow
         -ignore_readdir_race
-        -maxdepth
-        -mindepth
         -mount
         -nocolor
         -noignore_readdir_race
diff --git a/completions/bfs.fish b/completions/bfs.fish
index 3f399e7..1303639 100644
--- a/completions/bfs.fish
+++ b/completions/bfs.fish
@@ -133,6 +133,7 @@ complete -c bfs -o fls -d "Like -ls, but write to specified file" -F
 complete -c bfs -o fprint -d "Like -print, but write to specified file" -F
 complete -c bfs -o fprint0 -d "Like -print0, but write to specified file" -F
 complete -c bfs -o fprintf -d "Like -printf, but write to specified file" -F
+complete -c bfs -o limit -d "Limit the number of results" -x
 complete -c bfs -o ls -d "List files like ls -dils"
 complete -c bfs -o print -d "Print the path to the found file"
 complete -c bfs -o print0 -d "Like -print, but use the null character as a separator rather than newlines"
diff --git a/completions/bfs.zsh b/completions/bfs.zsh
index 3d7dc3a..51b5029 100644
--- a/completions/bfs.zsh
+++ b/completions/bfs.zsh
@@ -133,6 +133,7 @@ args=(
     '*-fprint0[print the path to the found file using null character as separator, but write to FILE instead of standard output]:output file:_files'
     '*-fprintf[print according to format string, but write to FILE instead of standard output]:output file:_files:output format'
 
+    '*-limit[quit after N results]:maximum result count'
     '*-ls[list files like ls -dils]'
     '*-print[print the path to the found file]'
     '*-print0[print the path to the found file using null character as separator]'
diff --git a/docs/USAGE.md b/docs/USAGE.md
index 3efdee0..071c95b 100644
--- a/docs/USAGE.md
+++ b/docs/USAGE.md
@@ -130,6 +130,40 @@ Unlike `-prune`, `-exclude` even works in combination with `-depth`/`-delete`.
 
 ---
 
+### `-limit`
+
+The `-limit N` action makes `bfs` quit once it gets evaluated `N` times.
+Placing it after an action like `-print` limits the number of results that get printed, for example:
+
+```console
+$ bfs -s -type f -name '*.txt'
+./1.txt
+./2.txt
+./3.txt
+./4.txt
+$ bfs -s -type f -name '*.txt' -print -limit 2
+./1.txt
+./2.txt
+```
+
+This is similar to
+
+```console
+$ bfs -s -type f -name '*.txt' | head -n2
+```
+
+but more powerful because you can apply separate limits to different expressions:
+
+```console
+$ bfs \( -name '*.txt' -print -limit 3 -o -name '*.log' -print -limit 4 \) -limit 5
+[At most 3 .txt files, at most 4 .log files, and at most 5 in total]
+```
+
+and more efficient because it will quit immediately.
+When piping to `head`, `bfs` will only quit *after* it tries to output too many results.
+
+---
+
 ### `-hidden`/`-nohidden`
 
 `-hidden` matches "hidden" files (dotfiles).
diff --git a/docs/bfs.1 b/docs/bfs.1
index 2ecb891..3a4f15a 100644
--- a/docs/bfs.1
+++ b/docs/bfs.1
@@ -725,6 +725,11 @@ but write to
 instead of standard output.
 .RE
 .TP
+\fB\-limit \fIN\fR
+Quit once this action is evaluated
+.I N
+times.
+.TP
 .B \-ls
 List files like
 .B ls
diff --git a/src/eval.c b/src/eval.c
index 9e55964..2f06858 100644
--- a/src/eval.c
+++ b/src/eval.c
@@ -841,6 +841,19 @@ error:
 }
 
 /**
+ * -limit action.
+ */
+bool eval_limit(const struct bfs_expr *expr, struct bfs_eval *state) {
+	long long evals = expr->evaluations + 1;
+	if (evals >= expr->num) {
+		state->action = BFTW_STOP;
+		state->quit = true;
+	}
+
+	return true;
+}
+
+/**
  * -prune action.
  */
 bool eval_prune(const struct bfs_expr *expr, struct bfs_eval *state) {
diff --git a/src/eval.h b/src/eval.h
index 98bbc08..f7f6c77 100644
--- a/src/eval.h
+++ b/src/eval.h
@@ -88,6 +88,7 @@ bool eval_fprint(const struct bfs_expr *expr, struct bfs_eval *state);
 bool eval_fprint0(const struct bfs_expr *expr, struct bfs_eval *state);
 bool eval_fprintf(const struct bfs_expr *expr, struct bfs_eval *state);
 bool eval_fprintx(const struct bfs_expr *expr, struct bfs_eval *state);
+bool eval_limit(const struct bfs_expr *expr, struct bfs_eval *state);
 bool eval_prune(const struct bfs_expr *expr, struct bfs_eval *state);
 bool eval_quit(const struct bfs_expr *expr, struct bfs_eval *state);
 
diff --git a/src/opt.c b/src/opt.c
index 76965de..a470d25 100644
--- a/src/opt.c
+++ b/src/opt.c
@@ -1181,6 +1181,7 @@ static struct bfs_expr *annotate_visit(struct bfs_opt *opt, struct bfs_expr *exp
 		eval_fprint0,
 		eval_fprintf,
 		eval_fprintx,
+		eval_limit,
 		eval_prune,
 		eval_true,
 		// Non-returning
diff --git a/src/parse.c b/src/parse.c
index 3b7386d..2dfcab2 100644
--- a/src/parse.c
+++ b/src/parse.c
@@ -97,6 +97,8 @@ struct bfs_parser {
 	char **last_arg;
 	/** A "-depth"-type argument, if any. */
 	char **depth_arg;
+	/** A "-limit" argument, if any. */
+	char **limit_arg;
 	/** A "-prune" argument, if any. */
 	char **prune_arg;
 	/** A "-mount" argument, if any. */
@@ -733,7 +735,7 @@ static struct bfs_expr *parse_action(struct bfs_parser *parser, bfs_eval_fn *eva
 		return NULL;
 	}
 
-	if (eval_fn != eval_prune && eval_fn != eval_quit) {
+	if (eval_fn != eval_limit && eval_fn != eval_prune && eval_fn != eval_quit) {
 		parser->implicit_print = false;
 	}
 
@@ -1570,6 +1572,29 @@ static struct bfs_expr *parse_jobs(struct bfs_parser *parser, int arg1, int arg2
 }
 
 /**
+ * Parse -limit N.
+ */
+static struct bfs_expr *parse_limit(struct bfs_parser *parser, int arg1, int arg2) {
+	struct bfs_expr *expr = parse_unary_action(parser, eval_limit);
+	if (!expr) {
+		return NULL;
+	}
+
+	char **arg = &expr->argv[1];
+	if (!parse_int(parser, arg, *arg, &expr->num, IF_LONG_LONG)) {
+		return NULL;
+	}
+
+	if (expr->num <= 0) {
+		parse_expr_error(parser, expr, "The ${blu}%s${rs} must be at least ${bld}1${rs}.\n", expr->argv[0]);
+		return NULL;
+	}
+
+	parser->limit_arg = expr->argv;
+	return expr;
+}
+
+/**
  * Parse -links N.
  */
 static struct bfs_expr *parse_links(struct bfs_parser *parser, int arg1, int arg2) {
@@ -2845,6 +2870,8 @@ static struct bfs_expr *parse_help(struct bfs_parser *parser, int arg1, int arg2
 	cfprintf(cout, "  ${blu}-fprintf${rs} ${bld}FILE${rs} ${bld}FORMAT${rs}\n");
 	cfprintf(cout, "      Like ${blu}-ls${rs}/${blu}-print${rs}/${blu}-print0${rs}/${blu}-printf${rs}, but write to ${bld}FILE${rs} instead of standard\n"
 	               "      output\n");
+	cfprintf(cout, "  ${blu}-limit${rs} ${bld}N${rs}\n");
+	cfprintf(cout, "      Quit after this action is evaluated ${bld}N${rs} times\n");
 	cfprintf(cout, "  ${blu}-ls${rs}\n");
 	cfprintf(cout, "      List files like ${ex}ls${rs} ${bld}-dils${rs}\n");
 	cfprintf(cout, "  ${blu}-print${rs}\n");
@@ -2968,6 +2995,7 @@ static const struct table_entry parse_table[] = {
 	{"-iregex", T_TEST, parse_regex, BFS_REGEX_ICASE},
 	{"-iwholename", T_TEST, parse_path, true},
 	{"-j", T_FLAG, parse_jobs, 0, 0, true},
+	{"-limit", T_ACTION, parse_limit},
 	{"-links", T_TEST, parse_links},
 	{"-lname", T_TEST, parse_lname, false},
 	{"-ls", T_ACTION, parse_ls},
@@ -3330,6 +3358,14 @@ static struct bfs_expr *parse_whole_expr(struct bfs_parser *parser) {
 	}
 
 	if (parser->implicit_print) {
+		char **limit = parser->limit_arg;
+		if (limit) {
+			parse_argv_error(parser, parser->limit_arg, 2,
+				"With ${blu}%s${rs}, you must specify an action explicitly; for example, ${blu}-print${rs} ${blu}%s${rs} ${bld}%s${rs}.\n",
+				limit[0], limit[0], limit[1]);
+			return NULL;
+		}
+
 		struct bfs_expr *print = parse_new_expr(parser, eval_fprint, 1, &fake_print_arg);
 		if (!print) {
 			return NULL;
diff --git a/tests/bfs/limit.out b/tests/bfs/limit.out
new file mode 100644
index 0000000..ea94276
--- /dev/null
+++ b/tests/bfs/limit.out
@@ -0,0 +1,4 @@
+basic/a
+basic/b
+basic/c/d
+basic/e/f
diff --git a/tests/bfs/limit.sh b/tests/bfs/limit.sh
new file mode 100644
index 0000000..84b605f
--- /dev/null
+++ b/tests/bfs/limit.sh
@@ -0,0 +1 @@
+bfs_diff -s basic -type f -print -limit 4
diff --git a/tests/bfs/limit_0.sh b/tests/bfs/limit_0.sh
new file mode 100644
index 0000000..3ce26de
--- /dev/null
+++ b/tests/bfs/limit_0.sh
@@ -0,0 +1 @@
+! invoke_bfs basic -print -limit 0
diff --git a/tests/bfs/limit_implicit_print.sh b/tests/bfs/limit_implicit_print.sh
new file mode 100644
index 0000000..cdb059d
--- /dev/null
+++ b/tests/bfs/limit_implicit_print.sh
@@ -0,0 +1 @@
+! invoke_bfs basic -type f -limit 1
diff --git a/tests/bfs/limit_incomplete.sh b/tests/bfs/limit_incomplete.sh
new file mode 100644
index 0000000..2d1e842
--- /dev/null
+++ b/tests/bfs/limit_incomplete.sh
@@ -0,0 +1 @@
+! invoke_bfs basic -print -limit
diff --git a/tests/bfs/limit_one.sh b/tests/bfs/limit_one.sh
new file mode 100644
index 0000000..3f8181c
--- /dev/null
+++ b/tests/bfs/limit_one.sh
@@ -0,0 +1 @@
+! invoke_bfs basic -print -limit one