1/*
2 * gomacro - A Go interpreter with Lisp-like macros
3 *
4 * Copyright (C) 2017-2019 Massimiliano Ghilardi
5 *
6 *     This Source Code Form is subject to the terms of the Mozilla Public
7 *     License, v. 2.0. If a copy of the MPL was not distributed with this
8 *     file, You can obtain one at http://mozilla.org/MPL/2.0/.
9 *
10 *
11 * range.go
12 *
13 *  Created on Jun 04, 2017
14 *      Author Massimiliano Ghilardi
15 */
16
17package fast
18
19import (
20	"go/ast"
21	"go/token"
22	r "reflect"
23	"sort"
24	"unicode/utf8"
25	"unsafe"
26
27	"github.com/cosmos72/gomacro/base/reflect"
28	xr "github.com/cosmos72/gomacro/xreflect"
29)
30
31type rangeJump struct {
32	Start, Continue, Break int
33}
34
35// Range compiles a "for-range" statement
36func (c *Comp) Range(node *ast.RangeStmt, labels []string) {
37	var nbinds [2]int
38	flag := true // node.Tok == token.DEFINE || (node.Body != nil && containLocalBinds(node.Body.List...))
39
40	c, _ = c.pushEnvIfFlag(&nbinds, flag)
41	erange := c.Expr1(node.X, nil)
42	t := erange.Type
43	if erange.Untyped() {
44		t = erange.DefaultType()
45		erange.ConstTo(t)
46	}
47	var jump rangeJump
48
49	sort.Strings(labels)
50	// we need a fresh Comp here... created above by c.pushEnvIfFlag()
51	c.Loop = &LoopInfo{
52		Continue:   &jump.Continue,
53		Break:      &jump.Break,
54		ThisLabels: labels,
55	}
56
57	switch t.Kind() {
58	case r.Ptr:
59		if t.Elem().Kind() != r.Array {
60			c.Errorf("cannot range over %v <%v>", node.X, t)
61		}
62		// range on pointer to array: dereference it
63		t = t.Elem()
64		efun := erange.AsX1()
65		erange = exprX1(t, func(env *Env) r.Value {
66			return efun(env).Elem()
67		})
68		fallthrough
69	case r.Array, r.Slice:
70		c.rangeSlice(node, erange, &jump)
71	case r.Chan:
72		c.rangeChan(node, erange, &jump)
73	case r.Map:
74		c.rangeMap(node, erange, &jump)
75	case r.String:
76		c.rangeString(node, erange, &jump)
77	default:
78		c.Errorf("cannot range over %v <%v>", node.X, t)
79	}
80
81	jump.Break = c.Code.Len()
82
83	c = c.popEnvIfFlag(&nbinds, flag)
84}
85
86func (c *Comp) rangeChan(node *ast.RangeStmt, erange *Expr, jump *rangeJump) {
87	t := erange.Type
88	telem := t.Elem()
89
90	// unnamed bind, contains channel
91	bindchan := c.DeclVar0("", nil, erange)
92	idxchan := bindchan.Desc.Index()
93
94	placekey, _ := c.rangeVars(node, telem, nil)
95
96	jump.Start = c.Code.Len()
97
98	if placekey == nil {
99		c.append(func(env *Env) (Stmt, *Env) {
100			_, ok := env.Vals[idxchan].Recv()
101			var ip int
102			if ok {
103				ip = env.IP + 1
104			} else {
105				ip = jump.Break
106			}
107			env.IP = ip
108			return env.Code[ip], env
109		})
110	} else {
111		// unnamed bind, contains last received value
112		bindrecv := c.NewBind("", VarBind, c.TypeOfInterface())
113		idxrecv := bindrecv.Desc.Index()
114
115		c.append(func(env *Env) (Stmt, *Env) {
116			v, ok := env.Vals[idxchan].Recv()
117			var ip int
118			if ok {
119				env.Vals[idxrecv] = v
120				ip = env.IP + 1
121			} else {
122				ip = jump.Break
123			}
124			env.IP = ip
125			return env.Code[ip], env
126		})
127		c.SetPlace(placekey, token.ASSIGN, unwrapBind(bindrecv, telem))
128	}
129
130	// compile the body
131	c.Block(node.Body)
132
133	// "continue" is a jump to loop beginning
134	jump.Continue = jump.Start
135
136	// jump back to start
137	c.append(func(env *Env) (Stmt, *Env) {
138		ip := jump.Start
139		env.IP = ip
140		return env.Code[ip], env
141	})
142}
143
144func (c *Comp) rangeMap(node *ast.RangeStmt, erange *Expr, jump *rangeJump) {
145	t := erange.Type
146	tkey, tval := t.Key(), t.Elem()
147	tkeyslice := c.Universe.SliceOf(tkey)
148	rtkeyslice := tkeyslice.ReflectType()
149
150	// unnamed bind, contains map
151	bindmap := c.DeclVar0("", nil, erange)
152	idxmap := bindmap.Desc.Index()
153
154	// unnamed bind, contains map keys
155	bindkeys := c.NewBind("", VarBind, tkeyslice)
156	idxkeys := bindkeys.Desc.Index()
157	c.append(func(env *Env) (Stmt, *Env) {
158		// convert []r.Value slice into a []rtkey slice, to avoid reflect.Value.Interface() while iterating
159		vkeys := env.Vals[idxmap].MapKeys()
160		keys := r.MakeSlice(rtkeyslice, len(vkeys), len(vkeys))
161		for i, vkey := range vkeys {
162			keys.Index(i).Set(vkey)
163		}
164		env.Vals[idxkeys] = keys
165		env.IP++
166		return env.Code[env.IP], env
167	})
168
169	// unnamed bind, contains iteration index
170	bindnext := c.DeclVar0("", c.TypeOfInt(), nil)
171	idxnext := bindnext.Desc.Index()
172
173	placekey, placeval := c.rangeVars(node, tkey, tval)
174
175	var bindkey *Bind
176	var ekey *Expr
177	if placekey != nil || placeval != nil {
178		// unnamed bind, contains iteration map key
179		bindkey = c.DeclVar0("", c.TypeOfInterface(), nil)
180		ekey = unwrapBind(bindkey, tkey)
181	}
182
183	jump.Start = c.Code.Len()
184
185	if bindkey == nil {
186		// check iteration index against # of keys
187		c.append(func(env *Env) (Stmt, *Env) {
188			n := env.Vals[idxkeys].Len()
189			i := *(*int)(unsafe.Pointer(&env.Ints[idxnext]))
190			var ip int
191			if i < n {
192				ip = env.IP + 1
193			} else {
194				ip = jump.Break
195			}
196			env.IP = ip
197			return env.Code[ip], env
198		})
199	} else {
200		// check iteration index against # of keys,
201		// and copy current map key into bindkey
202		idxkey := bindkey.Desc.Index()
203		c.append(func(env *Env) (Stmt, *Env) {
204			vkeys := env.Vals[idxkeys]
205			n := vkeys.Len()
206			i := *(*int)(unsafe.Pointer(&env.Ints[idxnext]))
207			var ip int
208			if i < n {
209				env.Vals[idxkey] = vkeys.Index(i)
210				ip = env.IP + 1
211			} else {
212				ip = jump.Break
213			}
214			env.IP = ip
215			return env.Code[ip], env
216		})
217	}
218
219	if placekey != nil {
220		// copy current map key into placekey
221		c.SetPlace(placekey, token.ASSIGN, ekey)
222	}
223
224	if placeval == nil {
225		// nothing to do
226	} else if placeval.IsVar() && !reflect.IsOptimizedKind(placeval.Type.Kind()) {
227		idxkey := bindkey.Desc.Index()
228		idxval := placeval.Var.Desc.Index()
229		upval := placeval.Var.Upn
230		rtype := tval.ReflectType()
231		zero := r.Zero(rtype)
232		c.append(func(env *Env) (Stmt, *Env) {
233			vmap := env.Vals[idxmap]
234			key := env.Vals[idxkey]
235			o := env
236			for j := 0; j < upval; j++ {
237				o = o.Outer
238			}
239			val := vmap.MapIndex(key)
240			if !val.IsValid() {
241				val = zero
242			} else if val.Type() != rtype {
243				val = convert(val, rtype)
244			}
245			o.Vals[idxval].Set(val)
246			env.IP++
247			return env.Code[env.IP], env
248		})
249	} else {
250		emap := c.Bind(bindmap)
251		c.SetPlace(placeval, token.ASSIGN, c.mapIndex1(nil, emap, ekey))
252	}
253
254	// compile the body
255	c.Block(node.Body)
256
257	// "continue" is a jump to the last statement below
258	jump.Continue = c.Code.Len()
259
260	// increase iteration index and jump back to start
261	c.append(func(env *Env) (Stmt, *Env) {
262		(*(*int)(unsafe.Pointer(&env.Ints[idxnext])))++
263		ip := jump.Start
264		env.IP = ip
265		return env.Code[ip], env
266	})
267}
268
269func (c *Comp) rangeSlice(node *ast.RangeStmt, erange *Expr, jump *rangeJump) {
270	t := erange.Type
271	var constlen int
272	var elen *Expr
273
274	if node.Value != nil || t.Kind() != r.Array {
275		// Go spec: one-variable range on array ONLY evaluates the array length, not the array itself
276		// save range variable in an unnamed bind
277		bind := c.DeclVar0("", nil, erange)
278		erange = c.Bind(bind)
279	}
280
281	if t.Kind() == r.Array {
282		constlen = t.Len()
283	} else {
284		// save range length in an unnamed bind
285		rangefun := erange.AsX1()
286		elen0 := exprFun(c.TypeOfInt(), func(env *Env) int {
287			return rangefun(env).Len()
288		})
289		bindlen := c.DeclVar0("", nil, elen0)
290		elen = c.Bind(bindlen)
291	}
292
293	placekey, placeval := c.rangeVars(node, c.TypeOfInt(), t.Elem())
294
295	if placekey == nil {
296		// we need an interation variable, even if user code ignores it
297		placekey = c.DeclVar0("", c.TypeOfInt(), nil).AsVar(0, PlaceSettable).AsPlace()
298	}
299	if placekey.Desc.Class() != IntBind {
300		c.Errorf("internal error: for-range counter variable allocated with class = %v, expecting class = %v",
301			placekey.Desc.Class(), IntBind)
302	}
303
304	jump.Start = c.Code.Len()
305
306	// compile comparison against range length
307	ekey := c.GetPlace(placekey)
308	funkey := ekey.WithFun().(func(*Env) int)
309
310	if t.Kind() == r.Array {
311		c.append(func(env *Env) (Stmt, *Env) {
312			var ip int
313			if funkey(env) < constlen {
314				ip = env.IP + 1
315			} else {
316				ip = jump.Break
317			}
318			env.IP = ip
319			return env.Code[ip], env
320		})
321	} else {
322		funlen := elen.WithFun().(func(*Env) int)
323		c.append(func(env *Env) (Stmt, *Env) {
324			var ip int
325			if funkey(env) < funlen(env) {
326				ip = env.IP + 1
327			} else {
328				ip = jump.Break
329			}
330			env.IP = ip
331			return env.Code[ip], env
332		})
333	}
334	if placeval != nil {
335		// for error messages
336		indexnode := &ast.IndexExpr{X: node.X, Lbrack: node.X.Pos(), Index: node.Key, Rbrack: node.X.Pos()}
337		eindex := c.vectorIndex(indexnode, erange, ekey)
338		c.SetPlace(placeval, token.ASSIGN, eindex)
339	}
340
341	// compile the body
342	c.Block(node.Body)
343
344	// "continue" is a jump to the increment below
345	jump.Continue = c.Code.Len()
346
347	// increment key
348	c.Pos = node.End() - 1
349	one := c.exprValue(c.TypeOfInt(), 1)
350	c.SetPlace(placekey, token.ADD_ASSIGN, one)
351
352	// jump back to comparison
353	c.append(func(env *Env) (Stmt, *Env) {
354		ip := jump.Start
355		env.IP = ip
356		return env.Code[ip], env
357	})
358}
359
360func (c *Comp) rangeString(node *ast.RangeStmt, erange *Expr, jump *rangeJump) {
361	// save string in an unnamed bind
362	bindrange := c.DeclVar0("", nil, erange)
363	idxrange := bindrange.Desc.Index()
364
365	placekey, placeval := c.rangeVars(node, c.TypeOfInt(), c.TypeOfInt32())
366	bindnext := c.DeclVar0("", c.TypeOfInt(), nil)
367	idxnext := bindnext.Desc.Index()
368
369	var bindrune *Bind
370	if placeval != nil && !placeval.IsVar() {
371		bindrune = c.DeclVar0("", c.TypeOfInt32(), nil)
372	}
373
374	jump.Start = c.Code.Len()
375
376	if placekey != nil {
377		c.SetPlace(placekey, token.ASSIGN, c.Bind(bindnext))
378	}
379	if placeval == nil {
380		c.append(func(env *Env) (Stmt, *Env) {
381			s := env.Vals[idxrange].String()
382			pnext := (*int)(unsafe.Pointer(&env.Ints[idxnext]))
383			next := *pnext
384
385			_, size := utf8.DecodeRuneInString(s[next:])
386			var ip int
387			if size != 0 {
388				next += size
389				*pnext = next
390				ip = env.IP + 1
391			} else {
392				ip = jump.Break
393			}
394			env.IP = ip
395			return env.Code[ip], env
396		})
397	} else if placeval.IsVar() {
398		idxval := placeval.Var.Desc.Index()
399		upval := placeval.Var.Upn
400		c.append(func(env *Env) (Stmt, *Env) {
401			s := env.Vals[idxrange].String()
402			pnext := (*int)(unsafe.Pointer(&env.Ints[idxnext]))
403			next := *pnext
404
405			r, size := utf8.DecodeRuneInString(s[next:])
406			var ip int
407			if size != 0 {
408				next += size
409				*pnext = next
410				o := env
411				for i := 0; i < upval; i++ {
412					o = o.Outer
413				}
414				*(*int32)(unsafe.Pointer(&env.Ints[idxval])) = r
415				ip = env.IP + 1
416			} else {
417				ip = jump.Break
418			}
419			env.IP = ip
420			return env.Code[ip], env
421		})
422	} else {
423		idxrune := bindrune.Desc.Index()
424		c.append(func(env *Env) (Stmt, *Env) {
425			s := env.Vals[idxrange].String()
426			pnext := (*int)(unsafe.Pointer(&env.Ints[idxnext]))
427			next := *pnext
428
429			r, size := utf8.DecodeRuneInString(s[next:])
430			var ip int
431			if size != 0 {
432				next += size
433				*pnext = next
434				*(*int32)(unsafe.Pointer(&env.Ints[idxrune])) = r
435				ip = env.IP + 1
436			} else {
437				ip = jump.Break
438			}
439			env.IP = ip
440			return env.Code[ip], env
441		})
442		c.SetPlace(placeval, token.ASSIGN, c.Bind(bindrune))
443	}
444
445	// compile the body
446	c.Block(node.Body)
447
448	// "continue" is a jump to the iteration above
449	jump.Continue = jump.Start
450
451	// jump back to iteration
452	c.append(func(env *Env) (Stmt, *Env) {
453		ip := jump.Start
454		env.IP = ip
455		return env.Code[ip], env
456	})
457}
458
459// rangeVars compiles the key and value iteration variables in a for-range
460func (c *Comp) rangeVars(node *ast.RangeStmt, tkey xr.Type, tval xr.Type) (*Place, *Place) {
461	place := [2]*Place{nil, nil}
462	t := [2]xr.Type{tkey, tval}
463
464	for i, expr := range [2]ast.Expr{node.Key, node.Value} {
465		if expr == nil {
466			continue
467		} else if t[i] == nil {
468			c.Pos = expr.Pos()
469			c.Errorf("too many variables in range")
470		}
471		c.Pos = expr.Pos()
472		if node.Tok == token.DEFINE {
473			switch expr := expr.(type) {
474			case *ast.Ident:
475				name := expr.Name
476				if name != "_" {
477					place[i] = c.DeclVar0(name, t[i], nil).AsVar(0, PlaceSettable).AsPlace()
478				}
479			default:
480				c.Errorf("non-name %v on left side of :=", expr)
481			}
482		} else {
483			if ident, ok := expr.(*ast.Ident); ok && ident.Name == "_" {
484				// ignore range variable "_"
485				continue
486			}
487			place[i] = c.Place(expr)
488			if !t[i].AssignableTo(place[i].Type) {
489				c.Errorf("cannot assign type <%v> to %v <%v> in range", t[i], expr, place[i].Type)
490			}
491		}
492	}
493	return place[0], place[1]
494}
495