Add support for call expression "range" loops with custom iterators and non

1 files changed, 93 insertions, 1 deletions

diff --git a/src/server/analysis.odin b/src/server/analysis.odin
index 630f4a1..189bf95 100644
--- a/src/server/analysis.odin
+++ b/src/server/analysis.odin
@@ -3112,8 +3112,100 @@ get_locals_for_range_stmt :: proc(
 		}
 	}
 
-	if symbol, ok := resolve_type_expression(ast_context, stmt.expr); ok {
+	symbol, ok := resolve_type_expression(ast_context, stmt.expr)
+
+	if v, ok := symbol.value.(SymbolProcedureValue); ok {
+		//Not quite sure how the custom iterator is defined, but it seems that it's three arguments. So temporarily just assume three arguments are iterators.
+		if len(v.return_types) != 3 && len(v.return_types) != 0 {
+			if v.return_types[0].type != nil {
+				symbol, ok = resolve_type_expression(ast_context, v.return_types[0].type)
+			} else if v.return_types[0].default_value != nil {
+				symbol, ok = resolve_type_expression(ast_context, v.return_types[0].default_value)
+			}
+		}
+	}
+
+	if ok {
 		#partial switch v in symbol.value {
+		case SymbolProcedureValue:
+			//This can only be custom iterators
+			if len(stmt.vals) >= 1 {
+				if ident, ok := unwrap_ident(stmt.vals[0]); ok {
+					expr: ^ast.Expr
+
+					if v.return_types[0].type != nil {
+						expr = v.return_types[0].type
+					} else if v.return_types[0].default_value != nil {
+						expr = v.return_types[0].default_value
+					}
+
+					store_local(
+						ast_context,
+						ident,
+						expr,
+						ident.pos.offset,
+						ident.name,
+						ast_context.local_id,
+						ast_context.non_mutable_only,
+						false,
+						true,
+						symbol.pkg,
+						false,
+					)
+				}
+			}
+
+			if len(stmt.vals) >= 2 {
+				if ident, ok := unwrap_ident(stmt.vals[1]); ok {
+					expr: ^ast.Expr
+
+					if v.return_types[1].type != nil {
+						expr = v.return_types[1].type
+					} else if v.return_types[1].default_value != nil {
+						expr = v.return_types[1].default_value
+					}
+
+					store_local(
+						ast_context,
+						ident,
+						expr,
+						ident.pos.offset,
+						ident.name,
+						ast_context.local_id,
+						ast_context.non_mutable_only,
+						false,
+						true,
+						symbol.pkg,
+						false,
+					)
+				}
+			}
+
+			if len(stmt.vals) >= 3 {
+				if ident, ok := unwrap_ident(stmt.vals[2]); ok {
+					expr: ^ast.Expr
+
+					if v.return_types[2].type != nil {
+						expr = v.return_types[2].type
+					} else if v.return_types[2].default_value != nil {
+						expr = v.return_types[2].default_value
+					}
+
+					store_local(
+						ast_context,
+						ident,
+						expr,
+						ident.pos.offset,
+						ident.name,
+						ast_context.local_id,
+						ast_context.non_mutable_only,
+						false,
+						true,
+						symbol.pkg,
+						false,
+					)
+				}
+			}
 		case SymbolUntypedValue:
 			if len(stmt.vals) == 1 {
 				if ident, ok := unwrap_ident(stmt.vals[0]); ok {