From 9b6e07e2a65b9477cbe55ce165b0144bc28cc3e7 Mon Sep 17 00:00:00 2001
From: spaceface <spaceface7777@gmail.com>
Date: Sun, 3 Oct 2021 14:44:13 +0200
Subject: [PATCH] cgen: support `index` and `in` for sumtype/interface arrays
 (#12051)

---
 vlib/v/gen/c/array.v                        | 18 +++++++++--
 vlib/v/tests/interface_array_methods_test.v | 36 +++++++++++++++++++++
 vlib/v/tests/sumtype_array_methods_test.v   | 34 +++++++++++++++++++
 3 files changed, 85 insertions(+), 3 deletions(-)
 create mode 100644 vlib/v/tests/interface_array_methods_test.v
 create mode 100644 vlib/v/tests/sumtype_array_methods_test.v

diff --git a/vlib/v/gen/c/array.v b/vlib/v/gen/c/array.v
index 1e99e13f49..c148e03d05 100644
--- a/vlib/v/gen/c/array.v
+++ b/vlib/v/gen/c/array.v
@@ -530,6 +530,12 @@ fn (mut g Gen) gen_array_contains_methods() {
 		} else if elem_sym.kind == .struct_ && left_info.elem_type.nr_muls() == 0 {
 			ptr_typ := g.equality_fn(left_info.elem_type)
 			fn_builder.writeln('\t\tif (${ptr_typ}_struct_eq((($elem_type_str*)a.data)[i], v)) {')
+		} else if elem_sym.kind == .interface_ && left_info.elem_type.nr_muls() == 0 {
+			ptr_typ := g.equality_fn(left_info.elem_type)
+			fn_builder.writeln('\t\tif (${ptr_typ}_interface_eq((($elem_type_str*)a.data)[i], v)) {')
+		} else if elem_sym.kind == .sum_type && left_info.elem_type.nr_muls() == 0 {
+			ptr_typ := g.equality_fn(left_info.elem_type)
+			fn_builder.writeln('\t\tif (${ptr_typ}_sumtype_eq((($elem_type_str*)a.data)[i], v)) {')
 		} else {
 			fn_builder.writeln('\t\tif ((($elem_type_str*)a.data)[i] == v) {')
 		}
@@ -591,15 +597,21 @@ fn (mut g Gen) gen_array_index_methods() {
 			fn_builder.writeln('\t\tif (fast_string_eq(*pelem, v)) {')
 		} else if elem_sym.kind == .array && !info.elem_type.is_ptr() {
 			ptr_typ := g.equality_fn(info.elem_type)
-			fn_builder.writeln('\t\tif (${ptr_typ}_arr_eq( *pelem, v)) {')
+			fn_builder.writeln('\t\tif (${ptr_typ}_arr_eq(*pelem, v)) {')
 		} else if elem_sym.kind == .function && !info.elem_type.is_ptr() {
 			fn_builder.writeln('\t\tif ( pelem == v) {')
 		} else if elem_sym.kind == .map && !info.elem_type.is_ptr() {
 			ptr_typ := g.equality_fn(info.elem_type)
-			fn_builder.writeln('\t\tif (${ptr_typ}_map_eq(( *pelem, v))) {')
+			fn_builder.writeln('\t\tif (${ptr_typ}_map_eq((*pelem, v))) {')
 		} else if elem_sym.kind == .struct_ && !info.elem_type.is_ptr() {
 			ptr_typ := g.equality_fn(info.elem_type)
-			fn_builder.writeln('\t\tif (${ptr_typ}_struct_eq( *pelem, v)) {')
+			fn_builder.writeln('\t\tif (${ptr_typ}_struct_eq(*pelem, v)) {')
+		} else if elem_sym.kind == .interface_ {
+			ptr_typ := g.equality_fn(info.elem_type)
+			fn_builder.writeln('\t\tif (${ptr_typ}_interface_eq(*pelem, v)) {')
+		} else if elem_sym.kind == .sum_type {
+			ptr_typ := g.equality_fn(info.elem_type)
+			fn_builder.writeln('\t\tif (${ptr_typ}_sumtype_eq(*pelem, v)) {')
 		} else {
 			fn_builder.writeln('\t\tif (*pelem == v) {')
 		}
diff --git a/vlib/v/tests/interface_array_methods_test.v b/vlib/v/tests/interface_array_methods_test.v
new file mode 100644
index 0000000000..f21319a819
--- /dev/null
+++ b/vlib/v/tests/interface_array_methods_test.v
@@ -0,0 +1,36 @@
+interface IA {
+	a int
+}
+
+struct AA {
+	a int
+}
+
+struct AB {
+	a int
+}
+
+fn test_interface_array_index() {
+	mut ia_ary := []IA{}
+	aa, ab := AA{12}, AB{13}
+	ia_ary << aa
+	ia_ary << ab
+
+	abi := IA(AB{13})
+	aci := IA(AB{14})
+	assert ia_ary.index(abi) == 1
+	assert ia_ary.index(aci) == -1
+}
+
+fn test_interface_array_contains() {
+	mut ia_ary := []IA{}
+	aa, ab := AA{12}, AB{13}
+	ia_ary << aa
+	ia_ary << ab
+
+	abi := IA(AB{13})
+	aci := IA(AB{14})
+
+	assert abi in ia_ary
+	assert aci !in ia_ary
+}
diff --git a/vlib/v/tests/sumtype_array_methods_test.v b/vlib/v/tests/sumtype_array_methods_test.v
new file mode 100644
index 0000000000..71da49500d
--- /dev/null
+++ b/vlib/v/tests/sumtype_array_methods_test.v
@@ -0,0 +1,34 @@
+type IA = AA | AB
+
+struct AA {
+	a int
+}
+
+struct AB {
+	a int
+}
+
+fn test_interface_array_index() {
+	mut ia_ary := []IA{}
+	aa, ab := AA{12}, AB{13}
+	ia_ary << aa
+	ia_ary << ab
+
+	abi := IA(AB{13})
+	aci := IA(AB{14})
+	assert ia_ary.index(abi) == 1
+	assert ia_ary.index(aci) == -1
+}
+
+fn test_interface_array_contains() {
+	mut ia_ary := []IA{}
+	aa, ab := AA{12}, AB{13}
+	ia_ary << aa
+	ia_ary << ab
+
+	abi := IA(AB{13})
+	aci := IA(AB{14})
+
+	assert abi in ia_ary
+	assert aci !in ia_ary
+}