all: implement `ch <- x or {...}` and `ch <- x ?` (#7928)

vlib/sync/channel_push_or_1_test.v

@@ -0,0 +1,67 @@
+const n = 1000
+const c = 100
+
+fn f(ch chan int) {
+	for _ in 0 .. n {
+		_ := <-ch
+	}
+	ch.close()
+}
+
+fn test_push_or_unbuffered() {
+	ch := chan int{}
+	go f(ch)
+	mut j := 0
+	for {
+		ch <- j or {
+			break
+		}
+		j++
+	}
+	assert j == n
+}
+
+fn test_push_or_buffered() {
+	ch := chan int{cap: c}
+	go f(ch)
+	mut j := 0
+	for {
+		ch <- j or {
+			break
+		}
+		j++
+	}
+	// we don't know how many elements are in the buffer when the channel
+	// is closed, so check j against an interval
+	assert j >= n
+	assert j <= n + c
+}
+
+fn g(ch chan int, res chan int) {
+	mut j := 0
+	for {
+		ch <- j or {
+			break
+		}
+		j++
+	}
+	println('done $j')
+	res <- j
+}
+
+fn test_many_senders() {
+	ch := chan int{}
+	res := chan int{}
+	go g(ch, res)
+	go g(ch, res)
+	go g(ch, res)
+	mut k := 0
+	for _ in 0 .. 3 * n {
+		k = <-ch
+	}
+	ch.close()
+	mut sum := <-res
+	sum += <-res
+	sum += <-res
+	assert sum == 3 * n
+}

vlib/sync/channel_push_or_2_test.v

@@ -0,0 +1,27 @@
+const n = 1000
+
+fn f(ch chan int) {
+	mut s := 0
+	for _ in 0 .. n {
+		s += <-ch
+	}
+	assert s == n * (n + 1) / 2
+	ch.close()
+}
+
+fn do_send(ch chan int, val int) ?int {
+	ch <- val ?
+	return val + 1
+}
+
+fn test_push_propargate() {
+	ch := chan int{}
+	go f(ch)
+	mut s := 1
+	for {
+		s = do_send(ch, s) or {
+			break
+		}
+	}
+	assert s == n + 1
+}

vlib/sync/channels.v

@@ -153,6 +153,8 @@ pub fn (mut ch Channel) close() {
 		ch.write_subscriber.sem.post()
 	}
 	C.atomic_store_u16(&ch.write_sub_mtx, u16(0))
+	ch.writesem.post()
+	ch.writesem_im.post()
 }
 
 [inline]
@@ -209,6 +211,10 @@ fn (mut ch Channel) try_push_priv(src voidptr, no_block bool) ChanState {
 			}
 			ch.writesem.wait()
 		}
+		if C.atomic_load_u16(&ch.closed) != 0 {
+			ch.writesem.post()
+			return .closed
+		}
 		if ch.cap == 0 {
 			// try to advertise current object as readable
 			mut read_in_progress := false
@@ -255,6 +261,14 @@ fn (mut ch Channel) try_push_priv(src voidptr, no_block bool) ChanState {
 				} else {
 					ch.writesem_im.wait()
 				}
+				if C.atomic_load_u16(&ch.closed) != 0 {
+					if have_swapped || C.atomic_compare_exchange_strong_ptr(&ch.adr_read, &src2, voidptr(0)) {
+						ch.writesem.post()
+						return .success
+					} else {
+						return .closed
+					}
+				}
 				if have_swapped || C.atomic_compare_exchange_strong_ptr(&ch.adr_read, &src2, voidptr(0)) {
 					ch.writesem.post()
 					break
@@ -323,7 +337,7 @@ fn (mut ch Channel) try_push_priv(src voidptr, no_block bool) ChanState {
 			}
 		}
 	}
-	// this should not happen
+	// we should not get here but the V compiler want's to see a return statement
 	assert false
 	return .success
 }

vlib/v/ast/ast.v

@@ -532,6 +532,7 @@ pub mut:
 	left_type   table.Type
 	right_type  table.Type
 	auto_locked string
+	or_block    OrExpr
 }
 
 // ++, --

vlib/v/checker/checker.v

@@ -902,6 +902,7 @@ pub fn (mut c Checker) infix_expr(mut infix_expr ast.InfixExpr) table.Type {
 					c.error('cannot push non-reference `$right.name` on `$left.name`',
 						right_pos)
 				}
+				c.stmts(infix_expr.or_block.stmts)
 			} else {
 				c.error('cannot push on non-channel `$left.name`', left_pos)
 			}

vlib/v/gen/cgen.v

@@ -67,6 +67,7 @@ mut:
 	vlines_path                      string   // set to the proper path for generating #line directives
 	optionals                        []string // to avoid duplicates TODO perf, use map
 	chan_pop_optionals               []string // types for `x := <-ch or {...}`
+	chan_push_optionals              []string // types for `ch <- x or {...}`
 	shareds                          []int    // types with hidden mutex for which decl has been emitted
 	inside_ternary                   int      // ?: comma separated statements on a single line
 	inside_map_postfix               bool     // inside map++/-- postfix expr
@@ -555,6 +556,21 @@ static inline $opt_el_type __Option_${styp}_popval($styp ch) {
 	}
 }
 
+fn (mut g Gen) register_chan_push_optional_call(el_type string, styp string) {
+	if styp !in g.chan_push_optionals {
+		g.chan_push_optionals << styp
+		g.channel_definitions.writeln('
+static inline Option_void __Option_${styp}_pushval($styp ch, $el_type e) {
+	if (sync__Channel_try_push_priv(ch, &e, false)) {
+		Option _tmp2 = v_error(_SLIT("channel closed"));
+		return *(Option_void*)&_tmp2;
+	}
+	Option_void _tmp = {.ok = true, .is_none = false, .v_error = (string){.str=(byteptr)""}, .ecode = 0};
+	return _tmp;
+}')
+	}
+}
+
 // TODO: merge cc_type and cc_type2
 // cc_type but without the `struct` prefix
 fn (g &Gen) cc_type2(t table.Type) string {
@@ -3167,12 +3183,25 @@ fn (mut g Gen) infix_expr(node ast.InfixExpr) {
 		}
 	} else if node.op == .arrow {
 		// chan <- val
+		gen_or := node.or_block.kind != .absent
 		styp := left_sym.cname
-		g.write('__${styp}_pushval(')
+		mut left_inf := left_sym.info as table.Chan
+		elem_type := left_inf.elem_type
+		tmp_opt := if gen_or { g.new_tmp_var() } else { '' }
+		if gen_or {
+			elem_styp := g.typ(elem_type)
+			g.register_chan_push_optional_call(elem_styp, styp)
+			g.write('Option_void $tmp_opt = __Option_${styp}_pushval(')
+		} else {
+			g.write('__${styp}_pushval(')
+		}
 		g.expr(node.left)
 		g.write(', ')
 		g.expr(node.right)
 		g.write(')')
+		if gen_or {
+			g.or_block(tmp_opt, node.or_block, table.void_type)
+		}
 	} else if unaliased_left.idx() in [table.u32_type_idx, table.u64_type_idx] && unaliased_right.is_signed() &&
 		node.op in [.eq, .ne, .gt, .lt, .ge, .le] {
 		bitsize := if unaliased_left.idx() == table.u32_type_idx &&

vlib/v/parser/pratt.v

@@ -360,6 +360,7 @@ pub fn (mut p Parser) expr_with_left(left ast.Expr, precedence int, is_stmt_iden
 fn (mut p Parser) infix_expr(left ast.Expr) ast.Expr {
 	op := p.tok.kind
 	if op == .arrow {
+		p.or_is_handled = true
 		p.register_auto_import('sync')
 	}
 	// mut typ := p.
@@ -378,11 +379,47 @@ fn (mut p Parser) infix_expr(left ast.Expr) ast.Expr {
 		1 {
 		p.vet_error('Use `var == value` instead of `var in [value]`', pos.line_nr)
 	}
+	mut or_stmts := []ast.Stmt{}
+	mut or_kind := ast.OrKind.absent
+	mut or_pos := p.tok.position()
+	// allow `x := <-ch or {...}` to handle closed channel
+	if op == .arrow {
+		if p.tok.kind == .key_orelse {
+			p.next()
+			p.open_scope()
+			p.scope.register(ast.Var{
+				name: 'errcode'
+				typ: table.int_type
+				pos: p.tok.position()
+				is_used: true
+			})
+			p.scope.register(ast.Var{
+				name: 'err'
+				typ: table.string_type
+				pos: p.tok.position()
+				is_used: true
+			})
+			or_kind = .block
+			or_stmts = p.parse_block_no_scope(false)
+			or_pos = or_pos.extend(p.prev_tok.position())
+			p.close_scope()
+		}
+		if p.tok.kind == .question {
+			p.next()
+			or_kind = .propagate
+		}
+		p.or_is_handled = false
+	}
 	return ast.InfixExpr{
 		left: left
 		right: right
 		op: op
 		pos: pos
+		or_block: ast.OrExpr{
+			stmts: or_stmts
+			kind: or_kind
+			pos: or_pos
+		}
 	}
 }