github.com/MontFerret/ferret@v0.18.0/pkg/stdlib/arrays/flatten.go (about)

     1  package arrays
     2  
     3  import (
     4  	"context"
     5  
     6  	"github.com/MontFerret/ferret/pkg/runtime/core"
     7  	"github.com/MontFerret/ferret/pkg/runtime/values"
     8  	"github.com/MontFerret/ferret/pkg/runtime/values/types"
     9  )
    10  
    11  // FLATTEN turns an array of arrays into a flat array.
    12  // All array elements in array will be expanded in the result array.
    13  // Non-array elements are added as they are.
    14  // The function will recurse into sub-arrays up to the specified depth.
    15  // Duplicates will not be removed.
    16  // @param {Any[]} arr - Target array.
    17  // @param {Int} [depth] - Depth level.
    18  // @return {Any[]} - Flat array.
    19  func Flatten(_ context.Context, args ...core.Value) (core.Value, error) {
    20  	err := core.ValidateArgs(args, 1, 2)
    21  
    22  	if err != nil {
    23  		return values.None, err
    24  	}
    25  
    26  	err = core.ValidateType(args[0], types.Array)
    27  
    28  	if err != nil {
    29  		return values.None, err
    30  	}
    31  
    32  	arr := args[0].(*values.Array)
    33  	level := 1
    34  
    35  	if len(args) > 1 {
    36  		err = core.ValidateType(args[1], types.Int)
    37  
    38  		if err != nil {
    39  			return values.None, err
    40  		}
    41  
    42  		level = int(args[1].(values.Int))
    43  	}
    44  
    45  	currentLevel := 0
    46  	result := values.NewArray(int(arr.Length()) * 2)
    47  	var unwrap func(input *values.Array)
    48  
    49  	unwrap = func(input *values.Array) {
    50  		currentLevel++
    51  
    52  		input.ForEach(func(value core.Value, idx int) bool {
    53  			if value.Type() != types.Array || currentLevel > level {
    54  				result.Push(value)
    55  			} else {
    56  				unwrap(value.(*values.Array))
    57  				currentLevel--
    58  			}
    59  
    60  			return true
    61  		})
    62  	}
    63  
    64  	unwrap(arr)
    65  
    66  	return result, nil
    67  }