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 }