src.elv.sh@v0.21.0-dev.0.20240515223629-06979efb9a2a/pkg/eval/builtin_fn_stream.go (about)

     1  package eval
     2  
     3  import (
     4  	"errors"
     5  	"fmt"
     6  	"sort"
     7  
     8  	"src.elv.sh/pkg/eval/errs"
     9  	"src.elv.sh/pkg/eval/vals"
    10  )
    11  
    12  // Stream manipulation.
    13  
    14  func init() {
    15  	addBuiltinFns(map[string]any{
    16  		"all": all,
    17  		"one": one,
    18  
    19  		"take":    take,
    20  		"drop":    drop,
    21  		"compact": compact,
    22  
    23  		"count": count,
    24  
    25  		"order": order,
    26  	})
    27  }
    28  
    29  func all(fm *Frame, inputs Inputs) error {
    30  	out := fm.ValueOutput()
    31  	var errOut error
    32  	inputs(func(v any) {
    33  		if errOut != nil {
    34  			return
    35  		}
    36  		errOut = out.Put(v)
    37  	})
    38  	return errOut
    39  }
    40  
    41  func one(fm *Frame, inputs Inputs) error {
    42  	var val any
    43  	n := 0
    44  	inputs(func(v any) {
    45  		if n == 0 {
    46  			val = v
    47  		}
    48  		n++
    49  	})
    50  	if n == 1 {
    51  		return fm.ValueOutput().Put(val)
    52  	}
    53  	return errs.ArityMismatch{What: "values", ValidLow: 1, ValidHigh: 1, Actual: n}
    54  }
    55  
    56  func take(fm *Frame, n int, inputs Inputs) error {
    57  	out := fm.ValueOutput()
    58  	var errOut error
    59  	i := 0
    60  	inputs(func(v any) {
    61  		if errOut != nil {
    62  			return
    63  		}
    64  		if i < n {
    65  			errOut = out.Put(v)
    66  		}
    67  		i++
    68  	})
    69  	return errOut
    70  }
    71  
    72  func drop(fm *Frame, n int, inputs Inputs) error {
    73  	out := fm.ValueOutput()
    74  	var errOut error
    75  	i := 0
    76  	inputs(func(v any) {
    77  		if errOut != nil {
    78  			return
    79  		}
    80  		if i >= n {
    81  			errOut = out.Put(v)
    82  		}
    83  		i++
    84  	})
    85  	return errOut
    86  }
    87  
    88  func compact(fm *Frame, inputs Inputs) error {
    89  	out := fm.ValueOutput()
    90  	first := true
    91  	var errOut error
    92  	var prev any
    93  
    94  	inputs(func(v any) {
    95  		if errOut != nil {
    96  			return
    97  		}
    98  		if first || !vals.Equal(v, prev) {
    99  			errOut = out.Put(v)
   100  			first = false
   101  			prev = v
   102  		}
   103  	})
   104  	return errOut
   105  }
   106  
   107  // The count implementation uses a custom varargs based implementation rather
   108  // than the more common `Inputs` API (see pkg/eval/go_fn.go) because this
   109  // allows the implementation to be O(1) for the common cases rather than O(n).
   110  func count(fm *Frame, args ...any) (int, error) {
   111  	var n int
   112  	switch nargs := len(args); nargs {
   113  	case 0:
   114  		// Count inputs.
   115  		fm.IterateInputs(func(any) {
   116  			n++
   117  		})
   118  	case 1:
   119  		// Get length of argument.
   120  		v := args[0]
   121  		if len := vals.Len(v); len >= 0 {
   122  			n = len
   123  		} else {
   124  			err := vals.Iterate(v, func(any) bool {
   125  				n++
   126  				return true
   127  			})
   128  			if err != nil {
   129  				return 0, fmt.Errorf("cannot get length of a %s", vals.Kind(v))
   130  			}
   131  		}
   132  	default:
   133  		// The error matches what would be returned if the `Inputs` API was
   134  		// used. See GoFn.Call().
   135  		return 0, errs.ArityMismatch{What: "arguments", ValidLow: 0, ValidHigh: 1, Actual: nargs}
   136  	}
   137  	return n, nil
   138  }
   139  
   140  type orderOptions struct {
   141  	Reverse  bool
   142  	Key      Callable
   143  	Total    bool
   144  	LessThan Callable
   145  }
   146  
   147  func (opt *orderOptions) SetDefaultOptions() {}
   148  
   149  // ErrBothTotalAndLessThan is returned by order when both the &total and
   150  // &less-than options are specified.
   151  var ErrBothTotalAndLessThan = errors.New("both &total and &less-than specified")
   152  
   153  func order(fm *Frame, opts orderOptions, inputs Inputs) error {
   154  	if opts.Total && opts.LessThan != nil {
   155  		return ErrBothTotalAndLessThan
   156  	}
   157  	var values, keys []any
   158  	inputs(func(v any) { values = append(values, v) })
   159  	if opts.Key != nil {
   160  		keys = make([]any, len(values))
   161  		for i, value := range values {
   162  			outputs, err := fm.CaptureOutput(func(fm *Frame) error {
   163  				return opts.Key.Call(fm, []any{value}, NoOpts)
   164  			})
   165  			if err != nil {
   166  				return err
   167  			} else if len(outputs) != 1 {
   168  				return errs.ArityMismatch{
   169  					What:     "number of outputs of the &key callback",
   170  					ValidLow: 1, ValidHigh: 1, Actual: len(outputs)}
   171  			}
   172  			keys[i] = outputs[0]
   173  		}
   174  	}
   175  
   176  	s := &slice{fm, opts.Total, opts.LessThan, values, keys, nil}
   177  	if opts.Reverse {
   178  		sort.Stable(sort.Reverse(s))
   179  	} else {
   180  		sort.Stable(s)
   181  	}
   182  	if s.err != nil {
   183  		return s.err
   184  	}
   185  
   186  	out := fm.ValueOutput()
   187  	for _, v := range values {
   188  		err := out.Put(v)
   189  		if err != nil {
   190  			return err
   191  		}
   192  	}
   193  	return nil
   194  }
   195  
   196  type slice struct {
   197  	fm       *Frame
   198  	total    bool
   199  	lessThan Callable
   200  	values   []any
   201  	keys     []any // nil if no keys
   202  	err      error
   203  }
   204  
   205  func (s *slice) Len() int { return len(s.values) }
   206  
   207  func (s *slice) Less(i, j int) bool {
   208  	if s.err != nil {
   209  		return true
   210  	}
   211  
   212  	var a, b any
   213  	if s.keys != nil {
   214  		a, b = s.keys[i], s.keys[j]
   215  	} else {
   216  		a, b = s.values[i], s.values[j]
   217  	}
   218  
   219  	if s.lessThan == nil {
   220  		// Use a builtin comparator depending on s.mixed.
   221  		if s.total {
   222  			return vals.CmpTotal(a, b) == vals.CmpLess
   223  		}
   224  		o := vals.Cmp(a, b)
   225  		if o == vals.CmpUncomparable {
   226  			s.err = ErrUncomparable
   227  			return true
   228  		}
   229  		return o == vals.CmpLess
   230  	}
   231  
   232  	// Use the &less-than callback.
   233  	outputs, err := s.fm.CaptureOutput(func(fm *Frame) error {
   234  		return s.lessThan.Call(fm, []any{a, b}, NoOpts)
   235  	})
   236  	if err != nil {
   237  		s.err = err
   238  		return true
   239  	}
   240  	if len(outputs) != 1 {
   241  		s.err = errs.ArityMismatch{
   242  			What:     "number of outputs of the &less-than callback",
   243  			ValidLow: 1, ValidHigh: 1, Actual: len(outputs)}
   244  		return true
   245  	}
   246  	if b, ok := outputs[0].(bool); ok {
   247  		return b
   248  	}
   249  	s.err = errs.BadValue{
   250  		What:  "output of the &less-than callback",
   251  		Valid: "boolean", Actual: vals.Kind(outputs[0])}
   252  	return true
   253  }
   254  
   255  func (s *slice) Swap(i, j int) {
   256  	s.values[i], s.values[j] = s.values[j], s.values[i]
   257  	if s.keys != nil {
   258  		s.keys[i], s.keys[j] = s.keys[j], s.keys[i]
   259  	}
   260  }