improve flatten performance by reducing copy

itchyny · Oct 18, 2021 · a1a2550 · a1a2550
1 parent 1ba8b03
commit a1a2550
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 7 deletions.
diff --git a/builtin.go b/builtin.go
diff --git a/builtin.jq b/builtin.jq
@@ -25,12 +25,7 @@ def range($start; $end): _range($start; $end; 1);
 def range($start; $end; $step): _range($start; $end; $step);
 
 def _flatten($x):
-  reduce .[] as $i
-    ( [];
-      if $i | type == "array" and $x != 0
-      then . + ($i | _flatten($x-1))
-      else . + [$i]
-      end);
+  map(if type == "array" and $x != 0 then _flatten($x - 1) else [.] end) | add;
 def flatten($x):
   if $x < 0
   then error("flatten depth must not be negative")

diff --git a/cli/test.yaml b/cli/test.yaml
@@ -1170,9 +1170,11 @@
   input: |
     [0, [[[[[1]]]]]]
     [0, [1, [2]], [1, [[3], 2]]]
+    {"x": [1], "y": [[2]], "z": [[[3]]]}
   expected: |
     [0,1]
     [0,1,2,1,3,2]
+    [1,2,3]
 
 - name: flatten/1 function
   args:
@@ -1181,9 +1183,11 @@
   input: |
     [0, [[[[[1]]]]]]
     [0, [1, [2]], [1, [[3], 2]]]
+    {"x": [1], "y": [[2]], "z": [[[3]]]}
   expected: |
     [0,[[[1]]]]
     [0,1,2,1,[3],2]
+    [1,2,[3]]
 
 - name: min, min_by, max, max_by functions
   args: