From 01caecc284ee9a5246b21603c11112d6863781da Mon Sep 17 00:00:00 2001
From: Swastik Baranwal <swstkbaranwal@gmail.com>
Date: Sun, 2 Apr 2023 19:17:08 +0530
Subject: [PATCH] checker: add check for unwrapped option in array cap (#17851)

---
 vlib/v/checker/containers.v                    |  6 ++++++
 vlib/v/checker/tests/array_init_option_err.out | 14 +++++++++++---
 vlib/v/checker/tests/array_init_option_err.vv  |  5 ++++-
 3 files changed, 21 insertions(+), 4 deletions(-)

diff --git a/vlib/v/checker/containers.v b/vlib/v/checker/containers.v
index c51f7f7e2d..2306aa8266 100644
--- a/vlib/v/checker/containers.v
+++ b/vlib/v/checker/containers.v
@@ -87,6 +87,12 @@ fn (mut c Checker) array_init(mut node ast.ArrayInit) ast.Type {
 			}
 			c.ensure_sumtype_array_has_default_value(node)
 		}
+		if node.has_cap {
+			cap_typ := c.check_expr_opt_call(node.cap_expr, c.expr(node.cap_expr))
+			if cap_typ.has_flag(.option) {
+				c.error('cannot use unwrapped Option as capacity', node.cap_expr.pos())
+			}
+		}
 		c.ensure_type_exists(node.elem_type, node.elem_type_pos) or {}
 		if node.typ.has_flag(.generic) && c.table.cur_fn != unsafe { nil }
 			&& c.table.cur_fn.generic_names.len == 0 {
diff --git a/vlib/v/checker/tests/array_init_option_err.out b/vlib/v/checker/tests/array_init_option_err.out
index fcecf768f3..14412c4b41 100644
--- a/vlib/v/checker/tests/array_init_option_err.out
+++ b/vlib/v/checker/tests/array_init_option_err.out
@@ -6,15 +6,23 @@ vlib/v/checker/tests/array_init_option_err.vv:2:38: error: cannot use unwrapped
     4 |
 vlib/v/checker/tests/array_init_option_err.vv:5:20: error: cannot use unwrapped Option as length
     3 |     dump(arr1)
-    4 | 
+    4 |
     5 |     arr1 = []int{len: get(), init: get()?}
       |                       ~~~~~
     6 |     dump(arr1)
     7 |
 vlib/v/checker/tests/array_init_option_err.vv:8:34: error: cannot use unwrapped Option as initializer
     6 |     dump(arr1)
-    7 | 
+    7 |
     8 |     arr1 = []int{len: get()?, init: get()}
       |                                     ~~~~~
     9 |     dump(arr1)
-   10 | }
+   10 |
+vlib/v/checker/tests/array_init_option_err.vv:11:20: error: cannot use unwrapped Option as capacity
+    9 |     dump(arr1)
+   10 |
+   11 |     arr1 = []int{cap: get(), len: get()?, init: get()?}
+      |                       ~~~~~
+   12 |     dump(arr1)
+   13 | }
+
diff --git a/vlib/v/checker/tests/array_init_option_err.vv b/vlib/v/checker/tests/array_init_option_err.vv
index 87b3ae5cfe..ed315a191d 100644
--- a/vlib/v/checker/tests/array_init_option_err.vv
+++ b/vlib/v/checker/tests/array_init_option_err.vv
@@ -7,8 +7,11 @@ fn main() {
 
 	arr1 = []int{len: get()?, init: get()}
 	dump(arr1)
+
+	arr1 = []int{cap: get(), len: get()?, init: get()?}
+	dump(arr1)
 }
 
 fn get() ?int {
 	return 5
-}
\ No newline at end of file
+}