Support `for` and `inline for` for enumerated arrays

2 files changed, 54 insertions, 0 deletions

diff --git a/src/check_stmt.cpp b/src/check_stmt.cpp
index 7bb6924b4..c365dca84 100644
--- a/src/check_stmt.cpp
+++ b/src/check_stmt.cpp
@@ -652,6 +652,11 @@ void check_inline_range_stmt(CheckerContext *ctx, Ast *node, u32 mod_flags) {
 				val1 = t_int;
 				inline_for_depth = exact_value_i64(t->Array.count);
 				break;
+			case Type_EnumeratedArray:
+				val0 = t->EnumeratedArray.elem;
+				val1 = t->EnumeratedArray.index;
+				inline_for_depth = exact_value_i64(t->EnumeratedArray.count);
+				break;
 			}
 		}
 
@@ -1506,6 +1511,12 @@ void check_stmt_internal(CheckerContext *ctx, Ast *node, u32 flags) {
 						add_package_dependency(ctx, "runtime", "string_decode_rune");
 					}
 					break;
+
+				case Type_EnumeratedArray:
+					val0 = t->EnumeratedArray.elem;
+					val1 = t->EnumeratedArray.index;
+					break;
+
 				case Type_Array:
 					val0 = t->Array.elem;
 					val1 = t_int;
diff --git a/src/ir.cpp b/src/ir.cpp
index 5014bb4bf..31d6c0251 100644
--- a/src/ir.cpp
+++ b/src/ir.cpp
@@ -9137,6 +9137,17 @@ void ir_build_range_indexed(irProcedure *proc, irValue *expr, Type *val_type, ir
 		}
 		break;
 	}
+	case Type_EnumeratedArray: {
+		if (val_type != nullptr) {
+			val = ir_emit_load(proc, ir_emit_array_ep(proc, expr, idx));
+			// NOTE(bill): Override the idx value for the enumeration
+			Type *index_type =expr_type->EnumeratedArray.index;
+			if (compare_exact_values(Token_NotEq, expr_type->EnumeratedArray.min_value, exact_value_u64(0))) {
+				idx = ir_emit_arith(proc, Token_Add, idx, ir_value_constant(index_type, expr_type->EnumeratedArray.min_value), index_type);
+			}
+		}
+		break;
+	}
 	case Type_Slice: {
 		if (val_type != nullptr) {
 			irValue *elem = ir_slice_elem(proc, expr);
@@ -9869,6 +9880,17 @@ void ir_build_stmt_internal(irProcedure *proc, Ast *node) {
 				ir_build_range_indexed(proc, array, val0_type, count_ptr, &val, &key, &loop, &done);
 				break;
 			}
+			case Type_EnumeratedArray: {
+				irValue *count_ptr = nullptr;
+				irValue *array = ir_build_addr_ptr(proc, expr);
+				if (is_type_pointer(type_deref(ir_type(array)))) {
+					array = ir_emit_load(proc, array);
+				}
+				count_ptr = ir_add_local_generated(proc, t_int, false);
+				ir_emit_store(proc, count_ptr, ir_const_int(et->EnumeratedArray.count));
+				ir_build_range_indexed(proc, array, val0_type, count_ptr, &val, &key, &loop, &done);
+				break;
+			}
 			case Type_DynamicArray: {
 				irValue *count_ptr = nullptr;
 				irValue *array = ir_build_addr_ptr(proc, expr);
@@ -10074,6 +10096,27 @@ void ir_build_stmt_internal(irProcedure *proc, Ast *node) {
 
 				}
 				break;
+			case Type_EnumeratedArray:
+				if (t->EnumeratedArray.count > 0) {
+					irValue *val = ir_build_expr(proc, expr);
+					irValue *val_addr = ir_address_from_load_or_generate_local(proc, val);
+
+					for (i64 i = 0; i < t->EnumeratedArray.count; i++) {
+						if (val0_type) {
+							// NOTE(bill): Due to weird legacy issues in LLVM, this needs to be an i32
+							irValue *elem = ir_emit_array_epi(proc, val_addr, cast(i32)i);
+							ir_addr_store(proc, val0_addr, ir_emit_load(proc, elem));
+						}
+						if (val1_type) {
+							ExactValue idx = exact_value_add(exact_value_i64(i), t->EnumeratedArray.min_value);
+							ir_addr_store(proc, val1_addr, ir_value_constant(val1_type, idx));
+						}
+
+						ir_build_stmt(proc, rs->body);
+					}
+
+				}
+				break;
 			default:
 				GB_PANIC("Invalid inline for type");
 				break;