all: implement iterators in for loops (#7867)
							parent
							
								
									82162b8ff8
								
							
						
					
					
						commit
						efb80bdffd
					
				| 
						 | 
				
			
			@ -2824,32 +2824,53 @@ fn (mut c Checker) stmt(node ast.Stmt) {
 | 
			
		|||
				}
 | 
			
		||||
			} else {
 | 
			
		||||
				sym := c.table.get_type_symbol(typ)
 | 
			
		||||
				if sym.kind == .map && !(node.key_var.len > 0 && node.val_var.len > 0) {
 | 
			
		||||
					c.error('declare a key and a value variable when ranging a map: `for key, val in map {`\n' +
 | 
			
		||||
						'use `_` if you do not need the variable', node.pos)
 | 
			
		||||
				}
 | 
			
		||||
				if node.key_var.len > 0 {
 | 
			
		||||
					key_type := match sym.kind {
 | 
			
		||||
						.map { sym.map_info().key_type }
 | 
			
		||||
						else { table.int_type }
 | 
			
		||||
					}
 | 
			
		||||
					node.key_type = key_type
 | 
			
		||||
					node.scope.update_var_type(node.key_var, key_type)
 | 
			
		||||
				}
 | 
			
		||||
				mut value_type := c.table.value_type(typ)
 | 
			
		||||
				if value_type == table.void_type || typ.has_flag(.optional) {
 | 
			
		||||
					if typ != table.void_type {
 | 
			
		||||
						c.error('for in: cannot index `${c.table.type_to_str(typ)}`',
 | 
			
		||||
				if sym.kind == .struct_ {
 | 
			
		||||
					// iterators
 | 
			
		||||
					next_fn := sym.find_method('next') or {
 | 
			
		||||
						c.error('a struct must have a `next()` method to be an iterator',
 | 
			
		||||
							node.cond.position())
 | 
			
		||||
						return
 | 
			
		||||
					}
 | 
			
		||||
					if !next_fn.return_type.has_flag(.optional) {
 | 
			
		||||
						c.error('iterator method `next()` must return an optional', node.cond.position())
 | 
			
		||||
					}
 | 
			
		||||
					// the receiver
 | 
			
		||||
					if next_fn.params.len != 1 {
 | 
			
		||||
						c.error('iterator method `next()` must have 0 parameters', node.cond.position())
 | 
			
		||||
					}
 | 
			
		||||
					val_type := next_fn.return_type.clear_flag(.optional)
 | 
			
		||||
					node.cond_type = typ
 | 
			
		||||
					node.kind = sym.kind
 | 
			
		||||
					node.val_type = val_type
 | 
			
		||||
					node.scope.update_var_type(node.val_var, val_type)
 | 
			
		||||
				} else {
 | 
			
		||||
					if sym.kind == .map && !(node.key_var.len > 0 && node.val_var.len > 0) {
 | 
			
		||||
						c.error('declare a key and a value variable when ranging a map: `for key, val in map {`\n' +
 | 
			
		||||
							'use `_` if you do not need the variable', node.pos)
 | 
			
		||||
					}
 | 
			
		||||
					if node.key_var.len > 0 {
 | 
			
		||||
						key_type := match sym.kind {
 | 
			
		||||
							.map { sym.map_info().key_type }
 | 
			
		||||
							else { table.int_type }
 | 
			
		||||
						}
 | 
			
		||||
						node.key_type = key_type
 | 
			
		||||
						node.scope.update_var_type(node.key_var, key_type)
 | 
			
		||||
					}
 | 
			
		||||
					mut value_type := c.table.value_type(typ)
 | 
			
		||||
					if value_type == table.void_type || typ.has_flag(.optional) {
 | 
			
		||||
						if typ != table.void_type {
 | 
			
		||||
							c.error('for in: cannot index `${c.table.type_to_str(typ)}`',
 | 
			
		||||
								node.cond.position())
 | 
			
		||||
						}
 | 
			
		||||
					}
 | 
			
		||||
					if node.val_is_mut {
 | 
			
		||||
						value_type = value_type.to_ptr()
 | 
			
		||||
					}
 | 
			
		||||
					node.cond_type = typ
 | 
			
		||||
					node.kind = sym.kind
 | 
			
		||||
					node.val_type = value_type
 | 
			
		||||
					node.scope.update_var_type(node.val_var, value_type)
 | 
			
		||||
				}
 | 
			
		||||
				if node.val_is_mut {
 | 
			
		||||
					value_type = value_type.to_ptr()
 | 
			
		||||
				}
 | 
			
		||||
				node.cond_type = typ
 | 
			
		||||
				node.kind = sym.kind
 | 
			
		||||
				node.val_type = value_type
 | 
			
		||||
				node.scope.update_var_type(node.val_var, value_type)
 | 
			
		||||
			}
 | 
			
		||||
			c.check_loop_label(node.label, node.pos)
 | 
			
		||||
			c.stmts(node.stmts)
 | 
			
		||||
| 
						 | 
				
			
			
 | 
			
		|||
| 
						 | 
				
			
			@ -1305,6 +1305,27 @@ fn (mut g Gen) for_in(it ast.ForInStmt) {
 | 
			
		|||
			g.expr(it.cond)
 | 
			
		||||
			g.writeln('.str[$i];')
 | 
			
		||||
		}
 | 
			
		||||
	} else if it.kind == .struct_ {
 | 
			
		||||
		cond_type_sym := g.table.get_type_symbol(it.cond_type)
 | 
			
		||||
		next_fn := cond_type_sym.find_method('next') or {
 | 
			
		||||
			verror('`next` method not found')
 | 
			
		||||
			return
 | 
			
		||||
		}
 | 
			
		||||
		ret_typ := next_fn.return_type
 | 
			
		||||
		g.writeln('while (1) {')
 | 
			
		||||
		t := g.new_tmp_var()
 | 
			
		||||
		receiver_styp := g.typ(next_fn.params[0].typ)
 | 
			
		||||
		fn_name := receiver_styp.replace_each(['*', '', '.', '__']) + '_next'
 | 
			
		||||
		g.write('\t${g.typ(ret_typ)} $t = ${fn_name}(')
 | 
			
		||||
		if !it.cond_type.is_ptr() {
 | 
			
		||||
			g.write('&')
 | 
			
		||||
		}
 | 
			
		||||
		g.expr(it.cond)
 | 
			
		||||
		g.writeln(');')
 | 
			
		||||
		g.writeln('\tif (!${t}.ok) { break; }')
 | 
			
		||||
		val := if it.val_var in ['', '_'] { g.new_tmp_var() } else { it.val_var }
 | 
			
		||||
		val_styp := g.typ(it.val_type)
 | 
			
		||||
		g.writeln('\t$val_styp $val = *($val_styp*)${t}.data;')
 | 
			
		||||
	} else {
 | 
			
		||||
		s := g.table.type_to_str(it.cond_type)
 | 
			
		||||
		g.error('for in: unhandled symbol `$it.cond` of type `$s`', it.pos)
 | 
			
		||||
| 
						 | 
				
			
			
 | 
			
		|||
| 
						 | 
				
			
			@ -0,0 +1,36 @@
 | 
			
		|||
struct Doubler {
 | 
			
		||||
mut:
 | 
			
		||||
	val int
 | 
			
		||||
	until int
 | 
			
		||||
}
 | 
			
		||||
 | 
			
		||||
fn (mut it Doubler) next() ?int {
 | 
			
		||||
	v := it.val
 | 
			
		||||
	if v > it.until {
 | 
			
		||||
		return none
 | 
			
		||||
	}
 | 
			
		||||
	it.val *= 2
 | 
			
		||||
	return v
 | 
			
		||||
}
 | 
			
		||||
 | 
			
		||||
fn doubler(start int, until int) Doubler {
 | 
			
		||||
	return Doubler{start, until}
 | 
			
		||||
}
 | 
			
		||||
 | 
			
		||||
fn test_for_in_iterator() {
 | 
			
		||||
	mut d := doubler(5, 30)
 | 
			
		||||
	mut vals := []int{}
 | 
			
		||||
	for val in d {
 | 
			
		||||
		vals << val
 | 
			
		||||
	}
 | 
			
		||||
	assert vals == [5, 10, 20]
 | 
			
		||||
}
 | 
			
		||||
 | 
			
		||||
fn test_for_in_empty_iterator() {
 | 
			
		||||
	mut d := doubler(5, 2)
 | 
			
		||||
	mut vals := []int{}
 | 
			
		||||
	for val in d {
 | 
			
		||||
		vals << val
 | 
			
		||||
	}
 | 
			
		||||
	assert vals == []
 | 
			
		||||
}
 | 
			
		||||
		Loading…
	
		Reference in New Issue