Implement Array.prototype.flat/flatMap (#1132)

boa/src/builtins/array/mod.rs

@@ -16,6 +16,7 @@ mod tests;
 use crate::{
     builtins::array::array_iterator::{ArrayIterationKind, ArrayIterator},
     builtins::BuiltIn,
+    builtins::Number,
     gc::GcObject,
     object::{ConstructorBuilder, FunctionBuilder, ObjectData, PROTOTYPE},
     property::{Attribute, DataDescriptor},
@@ -91,6 +92,8 @@ impl BuiltIn for Array {
         .method(Self::every, "every", 1)
         .method(Self::find, "find", 1)
         .method(Self::find_index, "findIndex", 1)
+        .method(Self::flat, "flat", 0)
+        .method(Self::flat_map, "flatMap", 1)
         .method(Self::slice, "slice", 2)
         .method(Self::some, "some", 2)
         .method(Self::reduce, "reduce", 2)
@@ -885,6 +888,224 @@ impl Array {
         Ok(Value::integer(-1))
     }
 
+    /// `Array.prototype.flat( [depth] )`
+    ///
+    /// This method creates a new array with all sub-array elements concatenated into it
+    /// recursively up to the specified depth.
+    ///
+    /// More information:
+    ///  - [ECMAScript reference][spec]
+    ///  - [MDN documentation][mdn]
+    ///
+    /// [spec]: https://tc39.es/ecma262/#sec-array.prototype.flat
+    /// [mdn]: https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Global_Objects/Array/flat
+    pub(crate) fn flat(this: &Value, args: &[Value], context: &mut Context) -> Result<Value> {
+        // 1. Let O be ToObject(this value)
+        let this: Value = this.to_object(context)?.into();
+
+        // 2. Let sourceLen be LengthOfArrayLike(O)
+        let source_len = this.get_field("length", context)?.to_length(context)? as u32;
+
+        // 3. Let depthNum be 1
+        let depth = args.get(0);
+        let default_depth = Value::Integer(1);
+
+        // 4. If depth is not undefined, then set depthNum to IntegerOrInfinity(depth)
+        // 4.a. Set depthNum to ToIntegerOrInfinity(depth)
+        // 4.b. If depthNum < 0, set depthNum to 0
+        let depth_num = match depth
+            .unwrap_or(&default_depth)
+            .to_integer_or_infinity(context)?
+        {
+            IntegerOrInfinity::Integer(i) if i < 0 => IntegerOrInfinity::Integer(0),
+            num => num,
+        };
+
+        // 5. Let A be ArraySpeciesCreate(O, 0)
+        let new_array = Self::new_array(context);
+
+        // 6. Perform FlattenIntoArray(A, O, sourceLen, 0, depthNum)
+        let len = Self::flatten_into_array(
+            context,
+            &new_array,
+            &this,
+            source_len,
+            0,
+            depth_num,
+            &Value::undefined(),
+            &Value::undefined(),
+        )?;
+        new_array.set_field("length", len.to_length(context)?, context)?;
+
+        Ok(new_array)
+    }
+
+    /// `Array.prototype.flatMap( callback, [ thisArg ] )`
+    ///
+    /// This method returns a new array formed by applying a given callback function to
+    /// each element of the array, and then flattening the result by one level. It is
+    /// identical to a `map()` followed by a `flat()` of depth 1, but slightly more
+    /// efficient than calling those two methods separately.
+    ///
+    /// More information:
+    ///  - [ECMAScript reference][spec]
+    ///  - [MDN documentation][mdn]
+    ///
+    /// [spec]: https://tc39.es/ecma262/#sec-array.prototype.flatMap
+    /// [mdn]: https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Global_Objects/Array/flatMap
+    pub(crate) fn flat_map(this: &Value, args: &[Value], context: &mut Context) -> Result<Value> {
+        // 1. Let O be ToObject(this value)
+        let o: Value = this.to_object(context)?.into();
+
+        // 2. Let sourceLen be LengthOfArrayLike(O)
+        let source_len = this.get_field("length", context)?.to_length(context)? as u32;
+
+        // 3. If IsCallable(mapperFunction) is false, throw a TypeError exception
+        let mapper_function = args.get(0).cloned().unwrap_or_else(Value::undefined);
+        if !mapper_function.is_function() {
+            return context.throw_type_error("flatMap mapper function is not callable");
+        }
+        let this_arg = args.get(1).cloned().unwrap_or(o);
+
+        // 4. Let A be ArraySpeciesCreate(O, 0)
+        let new_array = Self::new_array(context);
+
+        // 5. Perform FlattenIntoArray(A, O, sourceLen, 0, 1, mapperFunction, thisArg)
+        let depth = Value::Integer(1).to_integer_or_infinity(context)?;
+        let len = Self::flatten_into_array(
+            context,
+            &new_array,
+            &this,
+            source_len,
+            0,
+            depth,
+            &mapper_function,
+            &this_arg,
+        )?;
+        new_array.set_field("length", len.to_length(context)?, context)?;
+
+        // 6. Return A
+        Ok(new_array)
+    }
+
+    /// Abstract method `FlattenIntoArray`.
+    ///
+    /// More information:
+    ///  - [ECMAScript reference][spec]
+    ///
+    /// [spec]: https://tc39.es/ecma262/#sec-flattenintoarray
+    #[allow(clippy::too_many_arguments)]
+    fn flatten_into_array(
+        context: &mut Context,
+        target: &Value,
+        source: &Value,
+        source_len: u32,
+        start: u32,
+        depth: IntegerOrInfinity,
+        mapper_function: &Value,
+        this_arg: &Value,
+    ) -> Result<Value> {
+        // 1. Assert target is Object
+        debug_assert!(target.is_object());
+
+        // 2. Assert source is Object
+        debug_assert!(source.is_object());
+
+        // 3. Assert if mapper_function is present, then:
+        // - IsCallable(mapper_function) is true
+        // - thisArg is present
+        // - depth is 1
+
+        // 4. Let targetIndex be start
+        let mut target_index = start;
+
+        // 5. Let sourceIndex be 0
+        let mut source_index = 0;
+
+        // 6. Repeat, while R(sourceIndex) < sourceLen
+        while source_index < source_len {
+            // 6.a. Let P be ToString(sourceIndex)
+            // 6.b. Let exists be HasProperty(source, P)
+            // 6.c. If exists is true, then
+            if source.has_field(source_index) {
+                // 6.c.i. Let element be Get(source, P)
+                let mut element = source.get_field(source_index, context)?;
+
+                // 6.c.ii. If mapperFunction is present, then
+                if !mapper_function.is_undefined() {
+                    // 6.c.ii.1. Set element to Call(mapperFunction, thisArg, <<element, sourceIndex, source>>)
+                    let args = [element, Value::from(source_index), target.clone()];
+                    element = context.call(&mapper_function, &this_arg, &args)?;
+                }
+                let element_as_object = element.as_object();
+
+                // 6.c.iii. Let shouldFlatten be false
+                let mut should_flatten = false;
+
+                // 6.c.iv. If depth > 0, then
+                let depth_is_positive = match depth {
+                    IntegerOrInfinity::PositiveInfinity => true,
+                    IntegerOrInfinity::NegativeInfinity => false,
+                    IntegerOrInfinity::Integer(i) => i > 0,
+                };
+                if depth_is_positive {
+                    // 6.c.iv.1. Set shouldFlatten is IsArray(element)
+                    should_flatten = match element_as_object {
+                        Some(obj) => obj.is_array(),
+                        _ => false,
+                    };
+                }
+                // 6.c.v. If shouldFlatten is true
+                if should_flatten {
+                    // 6.c.v.1. If depth is +Infinity let newDepth be +Infinity
+                    // 6.c.v.2. Else, let newDepth be depth - 1
+                    let new_depth = match depth {
+                        IntegerOrInfinity::PositiveInfinity => IntegerOrInfinity::PositiveInfinity,
+                        IntegerOrInfinity::Integer(d) => IntegerOrInfinity::Integer(d - 1),
+                        IntegerOrInfinity::NegativeInfinity => IntegerOrInfinity::NegativeInfinity,
+                    };
+
+                    // 6.c.v.3. Let elementLen be LengthOfArrayLike(element)
+                    let element_len =
+                        element.get_field("length", context)?.to_length(context)? as u32;
+
+                    // 6.c.v.4. Set targetIndex to FlattenIntoArray(target, element, elementLen, targetIndex, newDepth)
+                    target_index = Self::flatten_into_array(
+                        context,
+                        target,
+                        &element,
+                        element_len,
+                        target_index,
+                        new_depth,
+                        &Value::undefined(),
+                        &Value::undefined(),
+                    )?
+                    .to_u32(context)?;
+
+                // 6.c.vi. Else
+                } else {
+                    // 6.c.vi.1. If targetIndex >= 2^53 - 1, throw a TypeError exception
+                    if target_index.to_f64().ok_or(0)? >= Number::MAX_SAFE_INTEGER {
+                        return context
+                            .throw_type_error("Target index exceeded max safe integer value");
+                    }
+
+                    // 6.c.vi.2. Perform CreateDataPropertyOrThrow(target, targetIndex, element)
+                    target
+                        .set_property(target_index, DataDescriptor::new(element, Attribute::all()));
+
+                    // 6.c.vi.3. Set targetIndex to targetIndex + 1
+                    target_index = target_index.saturating_add(1);
+                }
+            }
+            // 6.d. Set sourceIndex to sourceIndex + 1
+            source_index = source_index.saturating_add(1);
+        }
+
+        // 7. Return targetIndex
+        Ok(Value::Integer(target_index.try_into().unwrap_or(0)))
+    }
+
     /// `Array.prototype.fill( value[, start[, end]] )`
     ///
     /// The method fills (modifies) all the elements of an array from start index (default 0)

boa/src/builtins/array/tests.rs

@@ -231,6 +231,115 @@ fn find_index() {
     assert_eq!(missing, String::from("-1"));
 }
 
+#[test]
+fn flat() {
+    let mut context = Context::new();
+
+    let code = r#"
+        var depth1 = ['a', ['b', 'c']];
+        var flat_depth1 = depth1.flat();
+
+        var depth2 = ['a', ['b', ['c'], 'd']];
+        var flat_depth2 = depth2.flat(2);
+        "#;
+    forward(&mut context, code);
+
+    assert_eq!(forward(&mut context, "flat_depth1[0]"), "\"a\"");
+    assert_eq!(forward(&mut context, "flat_depth1[1]"), "\"b\"");
+    assert_eq!(forward(&mut context, "flat_depth1[2]"), "\"c\"");
+    assert_eq!(forward(&mut context, "flat_depth1.length"), "3");
+
+    assert_eq!(forward(&mut context, "flat_depth2[0]"), "\"a\"");
+    assert_eq!(forward(&mut context, "flat_depth2[1]"), "\"b\"");
+    assert_eq!(forward(&mut context, "flat_depth2[2]"), "\"c\"");
+    assert_eq!(forward(&mut context, "flat_depth2[3]"), "\"d\"");
+    assert_eq!(forward(&mut context, "flat_depth2.length"), "4");
+}
+
+#[test]
+fn flat_empty() {
+    let mut context = Context::new();
+
+    let code = r#"
+        var empty = [[]];
+        var flat_empty = empty.flat();
+        "#;
+    forward(&mut context, code);
+
+    assert_eq!(forward(&mut context, "flat_empty.length"), "0");
+}
+
+#[test]
+fn flat_infinity() {
+    let mut context = Context::new();
+
+    let code = r#"
+        var arr = [[[[[['a']]]]]];
+        var flat_arr = arr.flat(Infinity)
+        "#;
+    forward(&mut context, code);
+
+    assert_eq!(forward(&mut context, "flat_arr[0]"), "\"a\"");
+    assert_eq!(forward(&mut context, "flat_arr.length"), "1");
+}
+
+#[test]
+fn flat_map() {
+    let mut context = Context::new();
+
+    let code = r#"
+        var double = [1, 2, 3];
+        var double_flatmap = double.flatMap(i => [i * 2]);
+
+        var sentence = ["it's Sunny", "in Cali"];
+        var flat_split_sentence = sentence.flatMap(x => x.split(" "));
+    "#;
+    forward(&mut context, code);
+
+    assert_eq!(forward(&mut context, "double_flatmap[0]"), "2");
+    assert_eq!(forward(&mut context, "double_flatmap[1]"), "4");
+    assert_eq!(forward(&mut context, "double_flatmap[2]"), "6");
+    assert_eq!(forward(&mut context, "double_flatmap.length"), "3");
+
+    assert_eq!(forward(&mut context, "flat_split_sentence[0]"), "\"it's\"");
+    assert_eq!(forward(&mut context, "flat_split_sentence[1]"), "\"Sunny\"");
+    assert_eq!(forward(&mut context, "flat_split_sentence[2]"), "\"in\"");
+    assert_eq!(forward(&mut context, "flat_split_sentence[3]"), "\"Cali\"");
+    assert_eq!(forward(&mut context, "flat_split_sentence.length"), "4");
+}
+
+#[test]
+fn flat_map_with_hole() {
+    let mut context = Context::new();
+
+    let code = r#"
+        var arr = [0, 1, 2];
+        delete arr[1];
+        var arr_flattened = arr.flatMap(i => [i * 2]);
+    "#;
+    forward(&mut context, code);
+
+    assert_eq!(forward(&mut context, "arr_flattened[0]"), "0");
+    assert_eq!(forward(&mut context, "arr_flattened[1]"), "4");
+    assert_eq!(forward(&mut context, "arr_flattened.length"), "2");
+}
+
+#[test]
+fn flat_map_not_callable() {
+    let mut context = Context::new();
+
+    let code = r#"
+        try {
+            var array = [1,2,3];
+            array.flatMap("not a function");
+        } catch (err) {
+            err.name === "TypeError"
+        }
+    "#;
+
+    assert_eq!(forward(&mut context, code), "true");
+}
+
 #[test]
 fn push() {
     let mut context = Context::new();