array: nested map()

vlib/builtin/array_test.v

@@ -384,7 +384,7 @@ mut:
 }
 
 // TODO: default array/struct str methods
-pub fn (ta []Test2) str() string {
+fn (ta []Test2) str() string {
 	mut s := '['
 	for i, t in ta {
 		s += t.str()
@@ -396,11 +396,11 @@ pub fn (ta []Test2) str() string {
 	return s
 }
 
-pub fn (t Test2) str() string {
+fn (t Test2) str() string {
 	return '{$t.one $t.two}'
 }
 
-pub fn (t Test) str() string {
+fn (t Test) str() string {
 	return '{$t.a $t.b}'
 }
 
@@ -524,18 +524,59 @@ fn test_filter() {
 	//assert arr.filter(arr % 2).len == 5
 }
 
+fn map_test_helper_1(i int) int {
+	return i * i
+}
+fn map_test_helper_2(i int, b string) int {
+	return i + b.len
+}
+fn map_test_helper_3(i int, b []string) int {
+	return i + b.map(it.len)[i % b.len]
+}
+
 fn test_map() {
 	nums := [1, 2, 3, 4, 5, 6]
 	strs := ['v', 'is', 'awesome']
 
+	//assert nums.map() == <error>
+	//assert nums.map(it, 'excessive') == <error>
+
+	// identity
+	assert nums.map(it) == [1, 2, 3, 4, 5, 6]
+	assert strs.map(it) == ['v', 'is', 'awesome']
+	assert nums.map(it - it) == [0,0,0,0,0,0]
+	assert nums.map(it - it)[0] == 0
+
+	// type switch
 	assert nums.map(it * 10) == [10, 20, 30, 40, 50, 60]
+	assert nums.map(it * it) == [1, 4, 9, 16, 25, 36]
 	assert nums.map('$it') == ['1', '2', '3', '4', '5', '6']
 	assert nums.map(it % 2 == 0) == [false, true, false, true, false, true]
 
 	assert strs.map(it.to_upper()) == ['V', 'IS', 'AWESOME']
 	assert strs.map(it == 'awesome') == [false, false, true]
+	assert strs.map(it.len in nums) == [true, true, false]
 	assert strs.map(7) == [7, 7, 7]
 
+	// external func
+	assert nums.map(map_test_helper_1(it)) == [1, 4, 9, 16, 25, 36]
+	assert nums.map(map_test_helper_2(it, 'bb')) == [3, 4, 5, 6, 7, 8]
+	assert nums.map(map_test_helper_3(it, strs)) == [3, 9, 4, 6, 12, 7]
+
+	// empty array as input
+	assert []int{len:0}.map(it * 2) == []
+
+	// nested maps (where it is of same type)
+	assert nums.map( strs.map(7) == [7, 7, 7] ) == [true, true, true, true, true, true]
+	assert nums.map( '$it' + strs.map('a')[0] ) == ['1a', '2a', '3a', '4a', '5a', '6a']
+	assert nums.map( it + strs.map(7)[0] ) == [8, 9, 10, 11, 12, 13]
+	assert nums.map( it + strs.map(it.len)[0] ) == [2, 3, 4, 5, 6, 7]
+	assert strs.map( it.len + strs.map(it.len)[0] ) == [2, 3, 8]
+
+	// nested (different it types)
+	assert strs.map( it[ nums.map(it - it)[0] ] ) == [`v`, `i`, `a`]
+	assert nums[0..3].map('$it' + strs.map(it)[it-1]) == ['1v','2is','3awesome']
+
 	assert nums == [1, 2, 3, 4, 5, 6]
 	assert strs == ['v', 'is', 'awesome']
 }

vlib/v/ast/scope.v

@@ -113,6 +113,11 @@ pub fn (s mut Scope) register(name string, obj ScopeObject) {
 	s.objects[name] = obj
 }
 
+pub fn (s mut Scope) register_force(name string, obj ScopeObject) {
+	s.objects[name] = obj
+}
+
+
 pub fn (s &Scope) outermost() &Scope {
 	mut sc := s
 	for !isnil(sc.parent) {

vlib/v/checker/checker.v

@@ -607,13 +607,14 @@ pub fn (mut c Checker) call_method(call_expr mut ast.CallExpr) table.Type {
 			mut scope := c.file.scope.innermost(call_expr.pos.pos)
 			scope.update_var_type('it', array_info.elem_type)
 		}
+		// map/filter are supposed to have 1 arg only
 		mut arg_type := left_type
 		for arg in call_expr.args {
 			arg_type = c.expr(arg.expr)
 		}
 		call_expr.return_type = left_type
 		call_expr.receiver_type = left_type
-		if method_name == 'map' && call_expr.args.len == 1 {
+		if method_name == 'map' {
 			call_expr.return_type = c.table.find_or_register_array(arg_type, 1)
 		} else if method_name == 'clone' {
 			// need to return `array_xxx` instead of `array`

vlib/v/gen/cgen.v

@@ -2544,13 +2544,13 @@ fn (mut g Gen) gen_map(node ast.CallExpr) {
 	g.writeln('.len;')
 	g.writeln('$ret_typ $tmp = __new_array(0, ${tmp}_len, sizeof($ret_elem_type));')
 	g.writeln('for (int i = 0; i < ${tmp}_len; i++) {')
-	g.write('$inp_elem_type it = (($inp_elem_type*) ')
+	g.write('\t$inp_elem_type it = (($inp_elem_type*) ')
 	g.expr(node.left)
 	g.writeln('.data)[i];')
-	g.write('$ret_elem_type ti = ')
+	g.write('\t$ret_elem_type ti = ')
 	g.expr(node.args[0].expr) // the first arg is the filter condition
 	g.writeln(';')
-	g.writeln('array_push(&$tmp, &ti);')
+	g.writeln('\tarray_push(&$tmp, &ti);')
 	g.writeln('}')
 	g.write(s)
 	g.write(tmp)

vlib/v/parser/parser.v

@@ -756,8 +756,9 @@ fn (mut p Parser) index_expr(left ast.Expr) ast.IndexExpr {
 	}
 }
 
-fn (mut p Parser) filter() {
-	p.scope.register('it', ast.Var{
+fn (mut p Parser) scope_register_it() {
+    // force new 'it' even if it exists in parent scope
+	p.scope.register_force('it', ast.Var{
 		name: 'it'
 		pos: p.tok.position()
 		is_used: true
@@ -772,7 +773,7 @@ fn (mut p Parser) dot_expr(left ast.Expr) ast.Expr {
 	if is_filter {
 		p.open_scope()
 		name_pos = p.tok.position()
-		p.filter()
+		p.scope_register_it()
 		// wrong tok position when using defer
 		// defer {
 		// p.close_scope()
@@ -782,6 +783,9 @@ fn (mut p Parser) dot_expr(left ast.Expr) ast.Expr {
 	if p.tok.kind == .lpar {
 		p.next()
 		args := p.call_args()
+		if is_filter && args.len != 1 {
+			p.error('needs exactly 1 argument')
+		}
 		p.check(.rpar)
 		mut or_stmts := []ast.Stmt{}
 		mut is_or_block_used := false