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 * index.go
12 *
13 *  Created on Apr 23, 2017
14 *      Author Massimiliano Ghilardi
15 */
16
17package fast
18
19:import (
20	"go/ast"
21	r "reflect"
22)
23
24import (
25	"go/ast"
26	r "reflect"
27
28	"github.com/cosmos72/gomacro/base"
29	"github.com/cosmos72/gomacro/base/reflect"
30	xr "github.com/cosmos72/gomacro/xreflect"
31)
32
33func (c *Comp) indexExpr(node *ast.IndexExpr, multivalued bool) *Expr {
34	obj := c.Expr1(node.X, nil)
35	idx := c.Expr1(node.Index, nil)
36	if obj.Untyped() {
37		obj.ConstTo(obj.DefaultType())
38	}
39	t := obj.Type
40	var ret *Expr
41	switch t.Kind() {
42	case r.Array, r.Slice, r.String:
43		ret = c.vectorIndex(node, obj, idx)
44	case r.Map:
45		if multivalued {
46			ret = c.mapIndex(node, obj, idx)
47		} else {
48			ret = c.mapIndex1(node, obj, idx)
49		}
50	case r.Ptr:
51		if t.Elem().Kind() == r.Array {
52			objfun := obj.AsX1()
53			deref := exprFun(t.Elem(), func(env *Env) r.Value {
54				return objfun(env).Elem()
55			})
56			ret = c.vectorIndex(node, deref, idx)
57			break
58		}
59		fallthrough
60	default:
61		c.Errorf("invalid operation: %v (type %v does not support indexing)", node, t)
62		return nil
63	}
64	if obj.Const() && idx.Const() {
65		// constant propagation
66		ret.EvalConst(COptKeepUntyped)
67	}
68	return ret
69}
70
71:func upcasefirstbyte(str string) string {
72	if len(str) > 0 && str[0] >= 'a' && str[0] <= 'z' {
73		bytes := []byte(str)
74		bytes[0] -= 'a' - 'A'
75		return string(bytes)
76	}
77	return str
78}
79
80:func makekind(typ ast.Node) ast.Node {
81	t := EvalType(typ)
82	if t == nil {
83		return nil
84	}
85	// go/ast.SelectorExpr requires the foo in r.foo to be an *ast.Ident, cannot unquote there
86	kind := ~"{r . foo}
87	kind.Sel = &ast.Ident{Name: upcasefirstbyte(t.Name())}
88	return kind
89}
90
91
92:func convertvalue(typ, val ast.Node) (ast.Node, ast.Node) {
93	var t r.Type = EvalType(typ)
94	if t == nil {
95		// keep the result wrapped in a reflect.Value
96		typ = ~'{r.Value}
97	} else {
98		// unwrap the result
99		tname := t.Name()
100		// remove final digits from t.Name()
101		// needed to convert Uint64 -> Uint etc. to calls reflect.Value.{tname}
102		for len(tname) != 0 {
103			ch := tname[len(tname)-1]
104			if ch < '0' || ch > '9' {
105				break
106			}
107			tname = tname[0:len(tname)-1]
108		}
109		if tname == "uintptr" {
110			tname = "uint" // use reflect.Value.Uint()
111		}
112		sel := ~"{~,val . foo} // we modify it destructively
113		sel.Sel = &ast.Ident{Name: upcasefirstbyte(tname)}
114
115		switch t.Kind() {
116		case r.Bool, r.Int64, r.Uint64, r.Float64, r.Complex128, r.String:
117			// result of reflect.Value.{tname} is already the correct type
118			val = ~"{~,sel ()}
119		default:
120			// convert int64, uint64... to the correct type
121			val = ~"{~,typ ( ~,sel () )}
122		}
123	}
124	return typ, val
125}
126
127:macro vec_index_c(typ ast.Node) ast.Node {
128	kind := makekind(typ)
129	typv, val := convertvalue(typ, ~'{objv.Index(i)})
130	fun := ~"{
131		fun = func(env *Env) ~,typv {
132			objv := objfun(env)
133			return ~,val
134		}
135	}
136	if kind == nil {
137		return ~"{ default: ~,fun }
138	} else {
139		return ~"{ case ~,kind: ~,fun }
140	}
141}
142
143:macro vec_index_e(typ ast.Node) ast.Node {
144	kind := makekind(typ)
145	typv, val := convertvalue(typ, ~'{objv.Index(i)})
146	fun := ~"{
147		fun = func(env *Env) ~,typv {
148			objv := objfun(env)
149			i := idxfun(env)
150			return ~,val
151		}
152	}
153	if kind == nil {
154		return ~"{ default: ~,fun }
155	} else {
156		return ~"{ case ~,kind: ~,fun }
157	}
158}
159
160// vectorIndex compiles obj[idx] where obj is an array or slice
161func (c *Comp) vectorIndex(node *ast.IndexExpr, obj *Expr, idx *Expr) *Expr {
162	k := idx.Type.Kind()
163	cat := reflect.Category(k)
164	if cat == r.Int || cat == r.Uint || idx.Untyped() {
165		if !c.TypeOfInt().IdenticalTo(idx.Type) {
166			idx = c.convert(idx, c.TypeOfInt(), node.Index)
167		}
168	} else {
169		c.Errorf("non-integer %s index: %v <%v>", k, node.Index, idx.Type)
170	}
171
172	t := obj.Type
173	if t.Kind() == r.String {
174		return c.stringIndex(node, obj, idx)
175	}
176
177	t = t.Elem()
178	objfun := obj.AsX1()
179	var fun I
180	if idx.Const() {
181		i := idx.Value.(int)
182		switch t.Kind() {
183		{vec_index_c; bool}
184		{vec_index_c; int}
185		{vec_index_c; int8}
186		{vec_index_c; int16}
187		{vec_index_c; int32}
188		{vec_index_c; int64}
189		{vec_index_c; uint}
190		{vec_index_c; uint8}
191		{vec_index_c; uint16}
192		{vec_index_c; uint32}
193		{vec_index_c; uint64}
194		{vec_index_c; uintptr}
195		{vec_index_c; float32}
196		{vec_index_c; float64}
197		{vec_index_c; complex64}
198		{vec_index_c; complex128}
199		{vec_index_c; string}
200		{vec_index_c; nil}
201		}
202	} else {
203		idxfun := idx.WithFun().(func(*Env) int)
204		switch t.Kind() {
205		{vec_index_e; bool}
206		{vec_index_e; int}
207		{vec_index_e; int8}
208		{vec_index_e; int16}
209		{vec_index_e; int32}
210		{vec_index_e; int64}
211		{vec_index_e; uint}
212		{vec_index_e; uint8}
213		{vec_index_e; uint16}
214		{vec_index_e; uint32}
215		{vec_index_e; uint64}
216		{vec_index_e; uintptr}
217		{vec_index_e; float32}
218		{vec_index_e; float64}
219		{vec_index_e; complex64}
220		{vec_index_e; complex128}
221		{vec_index_e; string}
222		{vec_index_e; nil}
223		}
224	}
225	return exprFun(t, fun)
226}
227
228// stringIndex compiles obj[idx] where obj is a string
229func (c *Comp) stringIndex(node *ast.IndexExpr, obj *Expr, idx *Expr) *Expr {
230	idxfun := idx.WithFun().(func(*Env) int)
231	objfun := obj.WithFun().(func(*Env) string)
232	var fun func(env *Env) uint8
233	if obj.Const() {
234		str := obj.Value.(string)
235		fun = func(env *Env) uint8 {
236			i := idxfun(env)
237			return str[i]
238		}
239	} else if idx.Const() {
240		i := idx.Value.(int)
241		fun = func(env *Env) uint8 {
242			str := objfun(env)
243			return str[i]
244		}
245	} else {
246		fun = func(env *Env) uint8 {
247			str := objfun(env)
248			i := idxfun(env)
249			return str[i]
250		}
251	}
252	e := c.exprUint8(fun)
253	if obj.Const() && idx.Const() {
254		panicking := true
255		defer func() {
256			if panicking {
257				recover()
258				c.Errorf("string index out of range: %v", node)
259			}
260		}()
261		e.EvalConst(COptKeepUntyped)
262		panicking = false
263	}
264	return e
265}
266
267// mapIndex compiles obj[idx] where obj is a map
268func (c *Comp) mapIndex(node *ast.IndexExpr, obj *Expr, idx *Expr) *Expr {
269	t := obj.Type
270	tkey := t.Key()
271	tval := t.Elem()
272	idxconst := idx.Const()
273	if idxconst {
274		idx.ConstTo(tkey)
275	} else if idx.Type == nil || !idx.Type.AssignableTo(tkey) {
276		c.Errorf("cannot use %v <%v> as <%v> in map index", node.Index, idx.Type, tkey)
277	}
278	objfun := obj.AsX1()
279	zero := xr.Zero(tval)
280	var fun func(env *Env) (r.Value, []r.Value)
281	if idxconst {
282		key := r.ValueOf(idx.Value)
283		fun = func(env *Env) (r.Value, []r.Value) {
284			obj := objfun(env)
285			val := obj.MapIndex(key)
286			var ok r.Value
287			if val == base.Nil {
288				val = zero // map[key] returns the zero value if key is not present
289				ok = base.False
290			} else {
291				ok = base.True
292			}
293			return val, []r.Value{val, ok}
294		}
295	} else {
296		keyfun := idx.AsX1()
297		fun = func(env *Env) (r.Value, []r.Value) {
298			obj := objfun(env)
299			key := keyfun(env)
300			val := obj.MapIndex(key)
301			var ok r.Value
302			if val == base.Nil {
303				val = zero // map[key] returns the zero value if key is not present
304				ok = base.False
305			} else {
306				ok = base.True
307			}
308			return val, []r.Value{val, ok}
309		}
310	}
311	return exprXV([]xr.Type{tval, c.TypeOfBool()}, fun)
312}
313
314:macro mapindex1_c(typ ast.Node) ast.Node {
315	if EvalType(typ) == nil {
316		return ~'{
317			zero := xr.Zero(tval)
318			fun = func(env *Env) r.Value {
319				obj := objfun(env)
320				result := obj.MapIndex(key)
321				if result == base.Nil {
322					result = zero
323				}
324				return result
325			}
326		}
327	}
328	_, unwrap := convertvalue(typ, ~'v)
329	return ~"{
330		fun = func(env *Env) ~,typ {
331			obj := objfun(env)
332			v := obj.MapIndex(key)
333			var result ~,typ
334			if v != base.Nil {
335				result = ~,unwrap
336			}
337			return result
338		}
339	}
340}
341
342:macro mapindex1_e(typ ast.Node) ast.Node {
343	if EvalType(typ) == nil {
344		return ~'{
345			zero := xr.Zero(tval)
346			fun = func(env *Env) r.Value {
347				obj := objfun(env)
348				key := keyfun(env)
349				result := obj.MapIndex(key)
350				if result == base.Nil {
351					result = zero
352				}
353				return result
354			}
355		}
356	}
357	_, unwrap := convertvalue(typ, ~'v)
358	return ~"{
359		fun = func(env *Env) ~,typ {
360			obj := objfun(env)
361			key := keyfun(env)
362			v := obj.MapIndex(key)
363			var result ~,typ
364			if v != base.Nil {
365				result = ~,unwrap
366			}
367			return result
368		}
369	}
370}
371
372// mapIndex1 compiles obj[idx] where obj is a map, in single-value context
373func (c *Comp) mapIndex1(node *ast.IndexExpr, obj *Expr, idx *Expr) *Expr {
374	t := obj.Type
375	tkey := t.Key()
376	tval := t.Elem()
377	idxconst := idx.Const()
378	if idxconst {
379		idx.ConstTo(tkey)
380	} else if idx.Type == nil || !idx.Type.AssignableTo(tkey) {
381		c.Errorf("cannot use %v <%v> as <%v> in map index", node.Index, idx.Type, tkey)
382	}
383	objfun := obj.AsX1()
384	var fun I
385	if idxconst {
386		key := r.ValueOf(idx.Value)
387		switch tval.Kind() {
388		case r.Bool:       mapindex1_c; bool
389		case r.Int:        mapindex1_c; int
390		case r.Int8:       mapindex1_c; int8
391		case r.Int16:      mapindex1_c; int16
392		case r.Int32:      mapindex1_c; int32
393		case r.Int64:      mapindex1_c; int64
394		case r.Uint:       mapindex1_c; uint
395		case r.Uint8:      mapindex1_c; uint8
396		case r.Uint16:     mapindex1_c; uint16
397		case r.Uint32:     mapindex1_c; uint32
398		case r.Uint64:     mapindex1_c; uint64
399		case r.Uintptr:    mapindex1_c; uintptr
400		case r.Float32:    mapindex1_c; float32
401		case r.Float64:    mapindex1_c; float64
402		case r.Complex64:  mapindex1_c; complex64
403		case r.Complex128: mapindex1_c; complex128
404		case r.String:     mapindex1_c; string
405		default:           mapindex1_c; nil
406		}
407	} else {
408		keyfun := idx.AsX1()
409		switch tval.Kind() {
410		case r.Bool:       mapindex1_e; bool
411		case r.Int:        mapindex1_e; int
412		case r.Int8:       mapindex1_e; int8
413		case r.Int16:      mapindex1_e; int16
414		case r.Int32:      mapindex1_e; int32
415		case r.Int64:      mapindex1_e; int64
416		case r.Uint:       mapindex1_e; uint
417		case r.Uint8:      mapindex1_e; uint8
418		case r.Uint16:     mapindex1_e; uint16
419		case r.Uint32:     mapindex1_e; uint32
420		case r.Uint64:     mapindex1_e; uint64
421		case r.Uintptr:    mapindex1_e; uintptr
422		case r.Float32:    mapindex1_e; float32
423		case r.Float64:    mapindex1_e; float64
424		case r.Complex64:  mapindex1_e; complex64
425		case r.Complex128: mapindex1_e; complex128
426		case r.String:     mapindex1_e; string
427		default:           mapindex1_e; nil
428		}
429	}
430	return exprFun(tval, fun)
431}
432
433// IndexPlace compiles obj[idx] returning a Place, i.e. a settable (and addressable, if possible) reflect.Value
434func (c *Comp) IndexPlace(node *ast.IndexExpr, opt PlaceOption) *Place {
435	obj := c.Expr1(node.X, nil)
436	idx := c.Expr1(node.Index, nil)
437	if obj.Untyped() {
438		obj.ConstTo(obj.DefaultType())
439	}
440	t := obj.Type
441	switch t.Kind() {
442	case r.Array, r.Slice:
443		return c.vectorPlace(node, obj, idx)
444	case r.String:
445		// bytes in a string are not settable nor addressable
446		c.Errorf("%s a byte in a string: %v", opt, node)
447		return nil
448	case r.Map:
449		// elements in a map are settable but not addressable
450		if opt == PlaceAddress {
451			c.Errorf("%s a map element: %v", opt, node)
452			return nil
453		}
454		return c.mapPlace(node, obj, idx)
455	case r.Ptr:
456		if t.Elem().Kind() == r.Array {
457			return c.vectorPtrPlace(node, obj, idx)
458		}
459		fallthrough
460	default:
461		c.Errorf("invalid operation: %v (type %v does not support indexing)", node, t)
462		return nil
463	}
464}
465
466// mapPlace compiles obj[idx] where obj is a map, returning a settable place
467func (c *Comp) mapPlace(node *ast.IndexExpr, obj *Expr, idx *Expr) *Place {
468	tmap := obj.Type
469	tkey := tmap.Key()
470	idxconst := idx.Const()
471	if idxconst {
472		idx.ConstTo(tkey)
473	} else if idx.Type == nil || !idx.Type.AssignableTo(tkey) {
474		c.Errorf("cannot use %v <%v> as type <%v> in map index: %v", node.Index, idx.Type, tkey, node)
475	}
476	return &Place{Var: Var{Type: tmap.Elem()}, Fun: obj.AsX1(), MapKey: idx.AsX1(), MapType: tmap}
477}
478
479// vectorPlace compiles obj[idx] where obj is an array or slice, returning a settable and addressable place
480func (c *Comp) vectorPlace(node *ast.IndexExpr, obj *Expr, idx *Expr) *Place {
481	idxconst := idx.Const()
482	if idxconst {
483		idx.ConstTo(c.TypeOfInt())
484	} else if idx.Type == nil || !idx.Type.AssignableTo(c.TypeOfInt()) {
485		c.Errorf("non-integer %s index: %v <%v>", obj.Type.Kind(), node.Index, idx.Type)
486	}
487	t := obj.Type.Elem()
488	objfun := obj.AsX1()
489	var fun, addr func(env *Env) r.Value
490	if idxconst {
491		i := idx.Value.(int)
492		fun = func(env *Env) r.Value {
493			objv := objfun(env)
494			return objv.Index(i)
495		}
496		addr = func(env *Env) r.Value {
497			objv := objfun(env)
498			return objv.Index(i).Addr()
499		}
500	} else {
501		idxfun := idx.WithFun().(func(*Env) int)
502		fun = func(env *Env) r.Value {
503			objv := objfun(env)
504			i := idxfun(env)
505			return objv.Index(i)
506		}
507		addr = func(env *Env) r.Value {
508			objv := objfun(env)
509			i := idxfun(env)
510			return objv.Index(i).Addr()
511		}
512	}
513	return &Place{Var: Var{Type: t}, Fun: fun, Addr: addr}
514}
515
516// vectorPtrPlace compiles obj[idx] where obj is a pointer to an array, returning a settable and addressable reflect.Value
517func (c *Comp) vectorPtrPlace(node *ast.IndexExpr, obj *Expr, idx *Expr) *Place {
518	idxconst := idx.Const()
519	if idxconst {
520		idx.ConstTo(c.TypeOfInt())
521	} else if idx.Type == nil || !idx.Type.AssignableTo(c.TypeOfInt()) {
522		c.Errorf("non-integer %s index: %v <%v>", obj.Type.Kind(), node.Index, idx.Type)
523	}
524	t := obj.Type.Elem().Elem() // Elem() for the pointer to array, another Elem() for the array element type
525	objfun := obj.AsX1()
526	var fun, addr func(env *Env) r.Value
527	if idxconst {
528		i := idx.Value.(int)
529		fun = func(env *Env) r.Value {
530			objv := objfun(env).Elem()
531			return objv.Index(i)
532		}
533		addr = func(env *Env) r.Value {
534			objv := objfun(env).Elem()
535			return objv.Index(i).Addr()
536		}
537	} else {
538		idxfun := idx.WithFun().(func(*Env) int)
539		fun = func(env *Env) r.Value {
540			objv := objfun(env).Elem()
541			i := idxfun(env)
542			return objv.Index(i)
543		}
544		addr = func(env *Env) r.Value {
545			objv := objfun(env).Elem()
546			i := idxfun(env)
547			return objv.Index(i).Addr()
548		}
549	}
550	return &Place{Var: Var{Type: t}, Fun: fun, Addr: addr}
551}
552