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

2 files changed, 118 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 {
diff --git a/tests/completions_test.odin b/tests/completions_test.odin
index f25fba5..e523592 100644
--- a/tests/completions_test.odin
+++ b/tests/completions_test.odin
@@ -655,6 +655,31 @@ ast_for_in_identifier_completion :: proc(t: ^testing.T) {
 }
 
 @(test)
+ast_for_in_call_expr_completion :: proc(t: ^testing.T) {
+	source := test.Source {
+		main     = `package test	
+		Step :: struct {
+			data: int,
+		}
+
+		main :: proc() {
+			list :: proc() -> []Step {
+				return nil
+			}
+			for zstep in list() {
+				zst{*}
+			}
+		}
+		`,
+		packages = {},
+	}
+
+
+	test.expect_completion_details(t, &source, ".", {"test.zstep: Step"})
+}
+
+
+@(test)
 ast_completion_poly_struct_proc :: proc(t: ^testing.T) {
 	source := test.Source {
 		main     = `package test