From c9f3a05ef6664a35079260d0199db8f7246db384 Mon Sep 17 00:00:00 2001
From: Tanel Liiv <204019+dkull@users.noreply.github.com>
Date: Wed, 20 May 2020 15:28:33 +0300
Subject: [PATCH] checker: fix map/filter called with a plain fn argument

---
 vlib/builtin/array_test.v |  9 +++++++++
 vlib/v/checker/checker.v  |  9 ++++++++-
 vlib/v/gen/cgen.v         | 22 ++++++++++++++++++++--
 3 files changed, 37 insertions(+), 3 deletions(-)

diff --git a/vlib/builtin/array_test.v b/vlib/builtin/array_test.v
index 4927cce3a6..1f2a481c94 100644
--- a/vlib/builtin/array_test.v
+++ b/vlib/builtin/array_test.v
@@ -500,6 +500,10 @@ fn test_reduce() {
 }
 */
 
+fn filter_test_helper_1(a int) bool {
+	return a > 3
+}
+
 fn test_filter() {
 	a := [1, 2, 3, 4, 5, 6]
 	b := a.filter(it % 2 == 0)
@@ -519,6 +523,8 @@ fn test_filter() {
 	mut mut_arr := [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
 	mut_arr = mut_arr.filter(it < 4)
 	assert mut_arr.len == 3
+	assert a.filter(filter_test_helper_1) == [4,5,6]
+	assert [1,5,10].filter(filter_test_helper_1) == [5,10]
 
 	// TODO
 	//assert arr.filter(arr % 2).len == 5
@@ -577,6 +583,9 @@ fn test_map() {
 	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.map(map_test_helper_1) == [1,4,9,16,25,36]
+	assert [1,5,10].map(map_test_helper_1) == [1,25,100]
+
 	assert nums == [1, 2, 3, 4, 5, 6]
 	assert strs == ['v', 'is', 'awesome']
 }
diff --git a/vlib/v/checker/checker.v b/vlib/v/checker/checker.v
index 549603d0b7..8388376b05 100644
--- a/vlib/v/checker/checker.v
+++ b/vlib/v/checker/checker.v
@@ -711,7 +711,14 @@ pub fn (mut c Checker) call_method(mut call_expr ast.CallExpr) table.Type {
 		call_expr.return_type = left_type
 		call_expr.receiver_type = left_type
 		if method_name == 'map' {
-			call_expr.return_type = c.table.find_or_register_array(arg_type, 1)
+			arg_sym := c.table.get_type_symbol(arg_type)
+			// FIXME: match expr failed for now
+			mut ret_type := 0
+			match arg_sym.info {
+				table.FnType { ret_type = it.func.return_type }
+				else { ret_type = arg_type }
+			}
+			call_expr.return_type = c.table.find_or_register_array(ret_type, 1)
 		} else if method_name == 'clone' {
 			// need to return `array_xxx` instead of `array`
 			// in ['clone', 'str'] {
diff --git a/vlib/v/gen/cgen.v b/vlib/v/gen/cgen.v
index 2d397b8536..54ce09d698 100644
--- a/vlib/v/gen/cgen.v
+++ b/vlib/v/gen/cgen.v
@@ -2855,7 +2855,16 @@ fn (mut g Gen) gen_map(node ast.CallExpr) {
 	g.expr(node.left)
 	g.writeln('.data)[$i];')
 	g.write('\t$ret_elem_type ti = ')
-	g.expr(node.args[0].expr) // the first arg is the filter condition
+	match node.args[0].expr {
+		ast.Ident {
+			if it.kind == .function {
+				g.writeln('${it.name}(it)')
+			} else {
+				g.expr(node.args[0].expr) 
+			}
+		}
+		else { g.expr(node.args[0].expr) }
+	}
 	g.writeln(';')
 	g.writeln('\tarray_push(&$tmp, &ti);')
 	g.writeln('}')
@@ -2884,7 +2893,16 @@ fn (mut g Gen) gen_filter(node ast.CallExpr) {
 	g.expr(node.left)
 	g.writeln('.data)[i];')
 	g.write('if (')
-	g.expr(node.args[0].expr) // the first arg is the filter condition
+	match node.args[0].expr {
+		ast.Ident {
+			if it.kind == .function {
+				g.writeln('${node.args[0]}(it)')
+			} else {
+				g.expr(node.args[0].expr) 
+			}
+		}
+		else { g.expr(node.args[0].expr) }
+	}
 	g.writeln(') array_push(&$tmp, &it); \n }')
 	g.write(s)
 	g.write(' ')