1// Copyright 2019 Gregory Petrosyan <gregory.petrosyan@gmail.com>
2//
3// This Source Code Form is subject to the terms of the Mozilla Public
4// License, v. 2.0. If a copy of the MPL was not distributed with this
5// file, You can obtain one at https://mozilla.org/MPL/2.0/.
6
7package rapid
8
9import (
10	"fmt"
11	"reflect"
12)
13
14func SliceOf(elem *Generator) *Generator {
15	return SliceOfN(elem, -1, -1)
16}
17
18func SliceOfN(elem *Generator, minLen int, maxLen int) *Generator {
19	assertValidRange(minLen, maxLen)
20
21	return newGenerator(&sliceGen{
22		typ:    reflect.SliceOf(elem.type_()),
23		minLen: minLen,
24		maxLen: maxLen,
25		elem:   elem,
26	})
27}
28
29func SliceOfDistinct(elem *Generator, keyFn interface{}) *Generator {
30	return SliceOfNDistinct(elem, -1, -1, keyFn)
31}
32
33func SliceOfNDistinct(elem *Generator, minLen int, maxLen int, keyFn interface{}) *Generator {
34	assertValidRange(minLen, maxLen)
35
36	keyTyp := elem.type_()
37	if keyFn != nil {
38		t := reflect.TypeOf(keyFn)
39		assertCallable(t, elem.type_(), "keyFn")
40		keyTyp = t.Out(0)
41	}
42	assertf(keyTyp.Comparable(), "key type should be comparable (got %v)", keyTyp)
43
44	return newGenerator(&sliceGen{
45		typ:    reflect.SliceOf(elem.type_()),
46		minLen: minLen,
47		maxLen: maxLen,
48		elem:   elem,
49		keyTyp: keyTyp,
50		keyFn:  reflect.ValueOf(keyFn),
51	})
52}
53
54type sliceGen struct {
55	typ    reflect.Type
56	minLen int
57	maxLen int
58	elem   *Generator
59	keyTyp reflect.Type
60	keyFn  reflect.Value
61}
62
63func (g *sliceGen) String() string {
64	if g.keyTyp == nil {
65		if g.minLen < 0 && g.maxLen < 0 {
66			return fmt.Sprintf("SliceOf(%v)", g.elem)
67		} else {
68			return fmt.Sprintf("SliceOfN(%v, minLen=%v, maxLen=%v)", g.elem, g.minLen, g.maxLen)
69		}
70	} else {
71		key := ""
72		if g.keyFn.IsValid() {
73			key = fmt.Sprintf(", key=func(%v) %v", g.elem.type_(), g.keyTyp)
74		}
75
76		if g.minLen < 0 && g.maxLen < 0 {
77			return fmt.Sprintf("SliceOfDistinct(%v%v)", g.elem, key)
78		} else {
79			return fmt.Sprintf("SliceOfNDistinct(%v, minLen=%v, maxLen=%v%v)", g.elem, g.minLen, g.maxLen, key)
80		}
81	}
82}
83
84func (g *sliceGen) type_() reflect.Type {
85	return g.typ
86}
87
88func (g *sliceGen) value(t *T) value {
89	repeat := newRepeat(g.minLen, g.maxLen, -1)
90
91	var seen reflect.Value
92	if g.keyTyp != nil {
93		seen = reflect.MakeMapWithSize(reflect.MapOf(g.keyTyp, emptyStructType), repeat.avg())
94	}
95
96	sl := reflect.MakeSlice(g.typ, 0, repeat.avg())
97	for repeat.more(t.s, g.elem.String()) {
98		e := reflect.ValueOf(g.elem.value(t))
99		if g.keyTyp == nil {
100			sl = reflect.Append(sl, e)
101		} else {
102			k := e
103			if g.keyFn.IsValid() {
104				k = g.keyFn.Call([]reflect.Value{k})[0]
105			}
106
107			if seen.MapIndex(k).IsValid() {
108				repeat.reject()
109			} else {
110				seen.SetMapIndex(k, emptyStructValue)
111				sl = reflect.Append(sl, e)
112			}
113		}
114	}
115
116	return sl.Interface()
117}
118
119func MapOf(key *Generator, val *Generator) *Generator {
120	return MapOfN(key, val, -1, -1)
121}
122
123func MapOfN(key *Generator, val *Generator, minLen int, maxLen int) *Generator {
124	assertValidRange(minLen, maxLen)
125	assertf(key.type_().Comparable(), "key type should be comparable (got %v)", key.type_())
126
127	return newGenerator(&mapGen{
128		typ:    reflect.MapOf(key.type_(), val.type_()),
129		minLen: minLen,
130		maxLen: maxLen,
131		key:    key,
132		val:    val,
133	})
134}
135
136func MapOfValues(val *Generator, keyFn interface{}) *Generator {
137	return MapOfNValues(val, -1, -1, keyFn)
138}
139
140func MapOfNValues(val *Generator, minLen int, maxLen int, keyFn interface{}) *Generator {
141	assertValidRange(minLen, maxLen)
142
143	keyTyp := val.type_()
144	if keyFn != nil {
145		t := reflect.TypeOf(keyFn)
146		assertCallable(t, val.type_(), "keyFn")
147		keyTyp = t.Out(0)
148	}
149	assertf(keyTyp.Comparable(), "key type should be comparable (got %v)", keyTyp)
150
151	return newGenerator(&mapGen{
152		typ:    reflect.MapOf(keyTyp, val.type_()),
153		minLen: minLen,
154		maxLen: maxLen,
155		val:    val,
156		keyTyp: keyTyp,
157		keyFn:  reflect.ValueOf(keyFn),
158	})
159}
160
161type mapGen struct {
162	typ    reflect.Type
163	minLen int
164	maxLen int
165	key    *Generator
166	val    *Generator
167	keyTyp reflect.Type
168	keyFn  reflect.Value
169}
170
171func (g *mapGen) String() string {
172	if g.keyTyp == nil {
173		if g.minLen < 0 && g.maxLen < 0 {
174			return fmt.Sprintf("MapOf(%v, %v)", g.key, g.val)
175		} else {
176			return fmt.Sprintf("MapOfN(%v, %v, minLen=%v, maxLen=%v)", g.key, g.val, g.minLen, g.maxLen)
177		}
178	} else {
179		key := ""
180		if g.keyFn.IsValid() {
181			key = fmt.Sprintf(", key=func(%v) %v", g.val.type_(), g.keyTyp)
182		}
183
184		if g.minLen < 0 && g.maxLen < 0 {
185			return fmt.Sprintf("MapOfValues(%v%v)", g.val, key)
186		} else {
187			return fmt.Sprintf("MapOfNValues(%v, minLen=%v, maxLen=%v%v)", g.val, g.minLen, g.maxLen, key)
188		}
189	}
190}
191
192func (g *mapGen) type_() reflect.Type {
193	return g.typ
194}
195
196func (g *mapGen) value(t *T) value {
197	label := g.val.String()
198	if g.key != nil {
199		label = g.key.String() + "," + label
200	}
201
202	repeat := newRepeat(g.minLen, g.maxLen, -1)
203
204	m := reflect.MakeMapWithSize(g.typ, repeat.avg())
205	for repeat.more(t.s, label) {
206		var k, v reflect.Value
207		if g.keyTyp == nil {
208			k = reflect.ValueOf(g.key.value(t))
209			v = reflect.ValueOf(g.val.value(t))
210		} else {
211			v = reflect.ValueOf(g.val.value(t))
212			k = v
213			if g.keyFn.IsValid() {
214				k = g.keyFn.Call([]reflect.Value{v})[0]
215			}
216		}
217
218		if m.MapIndex(k).IsValid() {
219			repeat.reject()
220		} else {
221			m.SetMapIndex(k, v)
222		}
223	}
224
225	return m.Interface()
226}
227
228func ArrayOf(count int, elem *Generator) *Generator {
229	assertf(count >= 0 && count < 1024, "array element count should be in [0, 1024] (got %v)", count)
230
231	return newGenerator(&arrayGen{
232		typ:   reflect.ArrayOf(count, elem.type_()),
233		count: count,
234		elem:  elem,
235	})
236}
237
238type arrayGen struct {
239	typ   reflect.Type
240	count int
241	elem  *Generator
242}
243
244func (g *arrayGen) String() string {
245	return fmt.Sprintf("ArrayOf(%v, %v)", g.count, g.elem)
246}
247
248func (g *arrayGen) type_() reflect.Type {
249	return g.typ
250}
251
252func (g *arrayGen) value(t *T) value {
253	a := reflect.Indirect(reflect.New(g.typ))
254
255	if g.count == 0 {
256		t.s.drawBits(0)
257	} else {
258		for i := 0; i < g.count; i++ {
259			e := reflect.ValueOf(g.elem.value(t))
260			a.Index(i).Set(e)
261		}
262	}
263
264	return a.Interface()
265}
266