Project

General

Profile

Feature #16119 ยป 0001-Optimize-Array-flatten-and-flatten-for-already-fl.diff.txt

[PATCH] Optimize Array#flatten and flatten! for already flattened arrays - dylants (Dylan Thacker-Smith), 08/22/2019 11:54 PM

 
1
From a1df4a7b3aff940164ec88e5cbc27c031bafdcb8 Mon Sep 17 00:00:00 2001
2
From: Dylan Thacker-Smith <Dylan.Smith@shopify.com>
3
Date: Thu, 22 Aug 2019 14:08:35 -0400
4
Subject: [PATCH] Optimize Array#flatten and flatten! for already flattened
5
 arrays
6

    
7
---
8
 array.c | 43 +++++++++++++++++++++++++++++++++----------
9
 1 file changed, 33 insertions(+), 10 deletions(-)
10

    
11
diff --git a/array.c b/array.c
12
index bc5c719267..d1947e97fa 100644
13
--- a/array.c
14
+++ b/array.c
15
@@ -5112,18 +5112,40 @@ rb_ary_count(int argc, VALUE *argv, VALUE ary)
16
 }
17
 
18
 static VALUE
19
-flatten(VALUE ary, int level, int *modified)
20
+flatten(VALUE ary, int level)
21
 {
22
-    long i = 0;
23
+    long i;
24
     VALUE stack, result, tmp, elt;
25
     st_table *memo;
26
     st_data_t id;
27
 
28
-    stack = ary_new(0, ARY_DEFAULT_SIZE);
29
+    for (i = 0; i < RARRAY_LEN(ary); i++) {
30
+        elt = RARRAY_AREF(ary, i);
31
+        tmp = rb_check_array_type(elt);
32
+        if (!NIL_P(tmp)) {
33
+            break;
34
+        }
35
+    }
36
+    if (i == RARRAY_LEN(ary)) {
37
+        return ary;
38
+    } else if (tmp == ary) {
39
+        rb_raise(rb_eArgError, "tried to flatten recursive array");
40
+    }
41
+
42
     result = ary_new(0, RARRAY_LEN(ary));
43
+    ary_memcpy(result, 0, i, RARRAY_CONST_PTR_TRANSIENT(ary));
44
+    ARY_SET_LEN(result, i);
45
+
46
+    stack = ary_new(0, ARY_DEFAULT_SIZE);
47
+    rb_ary_push(stack, ary);
48
+    rb_ary_push(stack, LONG2NUM(i + 1));
49
+
50
     memo = st_init_numtable();
51
     st_insert(memo, (st_data_t)ary, (st_data_t)Qtrue);
52
-    *modified = 0;
53
+    st_insert(memo, (st_data_t)tmp, (st_data_t)Qtrue);
54
+
55
+    ary = tmp;
56
+    i = 0;
57
 
58
     while (1) {
59
 	while (i < RARRAY_LEN(ary)) {
60
@@ -5140,7 +5162,6 @@ flatten(VALUE ary, int level, int *modified)
61
 		rb_ary_push(result, elt);
62
 	    }
63
 	    else {
64
-		*modified = 1;
65
 		id = (st_data_t)tmp;
66
 		if (st_lookup(memo, id, 0)) {
67
 		    st_free_table(memo);
68
@@ -5200,9 +5221,8 @@ rb_ary_flatten_bang(int argc, VALUE *argv, VALUE ary)
69
     if (!NIL_P(lv)) level = NUM2INT(lv);
70
     if (level == 0) return Qnil;
71
 
72
-    result = flatten(ary, level, &mod);
73
-    if (mod == 0) {
74
-	ary_discard(result);
75
+    result = flatten(ary, level);
76
+    if (result == ary) {
77
 	return Qnil;
78
     }
79
     if (!(mod = ARY_EMBED_P(result))) rb_obj_freeze(result);
80
@@ -5237,7 +5257,7 @@ rb_ary_flatten_bang(int argc, VALUE *argv, VALUE ary)
81
 static VALUE
82
 rb_ary_flatten(int argc, VALUE *argv, VALUE ary)
83
 {
84
-    int mod = 0, level = -1;
85
+    int level = -1;
86
     VALUE result;
87
 
88
     if (rb_check_arity(argc, 0, 1) && !NIL_P(argv[0])) {
89
@@ -5245,7 +5265,10 @@ rb_ary_flatten(int argc, VALUE *argv, VALUE ary)
90
         if (level == 0) return ary_make_shared_copy(ary);
91
     }
92
 
93
-    result = flatten(ary, level, &mod);
94
+    result = flatten(ary, level);
95
+    if (result == ary) {
96
+        result = ary_make_shared_copy(ary);
97
+    }
98
     OBJ_INFECT(result, ary);
99
 
100
     return result;
101
-- 
102
2.21.0
103