From 69d01653f831052acdc803aad0b78690ad61ca86 Mon Sep 17 00:00:00 2001
From: nagachika <nagachika@ruby-lang.org>
Date: Sat, 20 Feb 2021 12:12:35 +0900
Subject: merge revision(s)
 1b1ea7b3bc9484e6e59d716fce2965a2f39d1e3d,b6e6807993c770c5d2e069d8741f5dadf0b38069:
 [Backport #17092]

	Fix Array#flatten for recursive array when given positive depth [Bug
	 #17092]

	---
	 array.c                 | 44 ++++++++++++++++++++++++++------------------
	 test/ruby/test_array.rb | 14 +++++++++++---
	 2 files changed, 37 insertions(+), 21 deletions(-)

	Initialize memo pointer and use it consistently to silence gcc 7+

	---
	 array.c | 10 +++++-----
	 1 file changed, 5 insertions(+), 5 deletions(-)
---
 array.c                 | 46 +++++++++++++++++++++++++++-------------------
 test/ruby/test_array.rb | 14 +++++++++++---
 version.h               |  4 ++--
 3 files changed, 40 insertions(+), 24 deletions(-)

diff --git a/array.c b/array.c
index 28bd8c866b..90fa2e06d6 100644
--- a/array.c
+++ b/array.c
@@ -5149,7 +5149,7 @@ flatten(VALUE ary, int level)
 {
     long i;
     VALUE stack, result, tmp, elt, vmemo;
-    st_table *memo;
+    st_table *memo = 0;
     st_data_t id;
 
     for (i = 0; i < RARRAY_LEN(ary); i++) {
@@ -5161,8 +5161,6 @@ flatten(VALUE ary, int level)
     }
     if (i == RARRAY_LEN(ary)) {
         return ary;
-    } else if (tmp == ary) {
-        rb_raise(rb_eArgError, "tried to flatten recursive array");
     }
 
     result = ary_new(0, RARRAY_LEN(ary));
@@ -5173,12 +5171,14 @@ flatten(VALUE ary, int level)
     rb_ary_push(stack, ary);
     rb_ary_push(stack, LONG2NUM(i + 1));
 
-    vmemo = rb_hash_new();
-    RBASIC_CLEAR_CLASS(vmemo);
-    memo = st_init_numtable();
-    rb_hash_st_table_set(vmemo, memo);
-    st_insert(memo, (st_data_t)ary, (st_data_t)Qtrue);
-    st_insert(memo, (st_data_t)tmp, (st_data_t)Qtrue);
+    if (level < 0) {
+	vmemo = rb_hash_new();
+	RBASIC_CLEAR_CLASS(vmemo);
+	memo = st_init_numtable();
+	rb_hash_st_table_set(vmemo, memo);
+	st_insert(memo, (st_data_t)ary, (st_data_t)Qtrue);
+	st_insert(memo, (st_data_t)tmp, (st_data_t)Qtrue);
+    }
 
     ary = tmp;
     i = 0;
@@ -5192,20 +5192,24 @@ flatten(VALUE ary, int level)
 	    }
 	    tmp = rb_check_array_type(elt);
 	    if (RBASIC(result)->klass) {
-                RB_GC_GUARD(vmemo);
-                st_clear(memo);
+		if (memo) {
+		    RB_GC_GUARD(vmemo);
+		    st_clear(memo);
+		}
 		rb_raise(rb_eRuntimeError, "flatten reentered");
 	    }
 	    if (NIL_P(tmp)) {
 		rb_ary_push(result, elt);
 	    }
 	    else {
-		id = (st_data_t)tmp;
-		if (st_is_member(memo, id)) {
-                    st_clear(memo);
-		    rb_raise(rb_eArgError, "tried to flatten recursive array");
+		if (memo) {
+		    id = (st_data_t)tmp;
+		    if (st_is_member(memo, id)) {
+			st_clear(memo);
+			rb_raise(rb_eArgError, "tried to flatten recursive array");
+		    }
+		    st_insert(memo, id, (st_data_t)Qtrue);
 		}
-		st_insert(memo, id, (st_data_t)Qtrue);
 		rb_ary_push(stack, ary);
 		rb_ary_push(stack, LONG2NUM(i));
 		ary = tmp;
@@ -5215,14 +5219,18 @@ flatten(VALUE ary, int level)
 	if (RARRAY_LEN(stack) == 0) {
 	    break;
 	}
-	id = (st_data_t)ary;
-	st_delete(memo, &id, 0);
+	if (memo) {
+	    id = (st_data_t)ary;
+	    st_delete(memo, &id, 0);
+	}
 	tmp = rb_ary_pop(stack);
 	i = NUM2LONG(tmp);
 	ary = rb_ary_pop(stack);
     }
 
-    st_clear(memo);
+    if (memo) {
+	st_clear(memo);
+    }
 
     RBASIC_SET_CLASS(result, rb_obj_class(ary));
     return result;
diff --git a/test/ruby/test_array.rb b/test/ruby/test_array.rb
index 476cf795f0..f2956a7fba 100644
--- a/test/ruby/test_array.rb
+++ b/test/ruby/test_array.rb
@@ -886,6 +886,17 @@ class TestArray < Test::Unit::TestCase
     assert_raise(NoMethodError, bug12738) { a.flatten.m }
   end
 
+  def test_flatten_recursive
+    a = []
+    a << a
+    assert_raise(ArgumentError) { a.flatten }
+    b = [1]; c = [2, b]; b << c
+    assert_raise(ArgumentError) { b.flatten }
+
+    assert_equal([1, 2, b], b.flatten(1))
+    assert_equal([1, 2, 1, 2, 1, c], b.flatten(4))
+  end
+
   def test_flatten!
     a1 = @cls[ 1, 2, 3]
     a2 = @cls[ 5, 6 ]
@@ -2624,9 +2635,6 @@ class TestArray < Test::Unit::TestCase
 
   def test_flatten_error
     a = []
-    a << a
-    assert_raise(ArgumentError) { a.flatten }
-
     f = [].freeze
     assert_raise(ArgumentError) { a.flatten!(1, 2) }
     assert_raise(TypeError) { a.flatten!(:foo) }
diff --git a/version.h b/version.h
index e48d1708a0..74122dcab7 100644
--- a/version.h
+++ b/version.h
@@ -2,11 +2,11 @@
 # define RUBY_VERSION_MINOR RUBY_API_VERSION_MINOR
 #define RUBY_VERSION_TEENY 3
 #define RUBY_RELEASE_DATE RUBY_RELEASE_YEAR_STR"-"RUBY_RELEASE_MONTH_STR"-"RUBY_RELEASE_DAY_STR
-#define RUBY_PATCHLEVEL 155
+#define RUBY_PATCHLEVEL 156
 
 #define RUBY_RELEASE_YEAR 2021
 #define RUBY_RELEASE_MONTH 2
-#define RUBY_RELEASE_DAY 11
+#define RUBY_RELEASE_DAY 20
 
 #include "ruby/version.h"
 
-- 
cgit v1.2.3