Implement flatten/0 in terms of flatten/1

2 files changed, 12 insertions, 4 deletions

diff --git a/builtin.c b/builtin.c
index 7201cf6d..52f469d9 100644
--- a/builtin.c
+++ b/builtin.c
@@ -1420,8 +1420,9 @@ static const char* const jq_builtins[] = {
   "def scalars_or_empty: select(. == null or . == true or . == false or type == \"number\" or type == \"string\" or ((type==\"array\" or type==\"object\") and length==0));",
   "def leaf_paths: paths(scalars);",
   "def join($x): reduce .[] as $i (null; (.//\"\") + (if . == null then $i else $x + $i end))//\"\";",
-  "def flatten: reduce .[] as $i ([]; if $i | type == \"array\" then . + ($i | flatten) else . + [$i] end);",
-  "def flatten($x): reduce .[] as $i ([]; if $i | type == \"array\" and $x > 0 then . + ($i | flatten($x-1)) else . + [$i] end);",
+  "def _flatten($x): reduce .[] as $i ([]; if $i | type == \"array\" and $x != 0 then . + ($i | _flatten($x-1)) else . + [$i] end);",
+  "def flatten($x): if $x < 0 then error(\"flatten depth must not be negative\") else _flatten($x) end;",
+  "def flatten: _flatten(-1);",
   "def range($x): range(0;$x);",
   "def fromdateiso8601: strptime(\"%Y-%m-%dT%H:%M:%SZ\")|mktime;",
   "def todateiso8601: strftime(\"%Y-%m-%dT%H:%M:%SZ\");",
diff --git a/tests/jq.test b/tests/jq.test
index 941d825a..e6f5671c 100644
--- a/tests/jq.test
+++ b/tests/jq.test
@@ -328,9 +328,8 @@ join(",","/")
 ["","","a","aa"]
 
 # Same check for flatten/1
-flatten(-1,3,2,1)
+flatten(3,2,1)
 [0, [1], [[2]], [[[3]]]]
-[0,[1],[[2]],[[[3]]]]
 [0,1,2,3]
 [0,1,2,[3]]
 [0,1,[2],[[3]]]
@@ -1117,6 +1116,10 @@ flatten
 [0, [1], [[2]], [[[3]]]]
 [0, 1, 2, 3]
 
+flatten(0)
+[0, [1], [[2]], [[[3]]]]
+[0, [1], [[2]], [[[3]]]]
+
 flatten(2)
 [0, [1], [[2]], [[[3]]]]
 [0, 1, 2, [3]]
@@ -1125,6 +1128,10 @@ flatten(2)
 [0, [1, [2]], [1, [[3], 2]]]
 [0, 1, 2, 1, [3], 2]
 
+try flatten(-1) catch .
+[0, [1], [[2]], [[[3]]]]
+"flatten depth must not be negative"
+
 transpose
 [[1], [2,3]]
 [[1,2],[null,3]]