Add underlying type for `bit_set`

8 files changed, 135 insertions, 72 deletions

diff --git a/src/check_expr.cpp b/src/check_expr.cpp
index 3b7571f74..b74658b6e 100644
--- a/src/check_expr.cpp
+++ b/src/check_expr.cpp
@@ -836,6 +836,9 @@ bool is_polymorphic_type_assignable(CheckerContext *c, Type *poly, Type *source,
 		return false;
 
 	case Type_BitSet:
+		if (source->kind == Type_BitSet) {
+			return is_polymorphic_type_assignable(c, poly->BitSet.elem, source->BitSet.elem, true, modify_type);
+		}
 		return false;
 
 	case Type_Union:
@@ -2040,7 +2043,7 @@ void check_binary_expr(CheckerContext *c, Operand *x, Ast *node, bool use_lhs_as
 			add_package_dependency(c, "runtime", "__dynamic_map_get");
 		} else if (is_type_bit_set(y->type)) {
 			Type *yt = base_type(y->type);
-			check_assignment(c, x, yt->BitSet.base, str_lit("bit_set 'in'"));
+			check_assignment(c, x, yt->BitSet.elem, str_lit("bit_set 'in'"));
 
 			if (x->mode == Addressing_Constant && y->mode == Addressing_Constant) {
 				ExactValue k = exact_value_to_integer(x->value);
@@ -5552,7 +5555,7 @@ ExprKind check_expr_base_internal(CheckerContext *c, Operand *o, Ast *node, Type
 			if (cl->elems.count == 0) {
 				break; // NOTE(bill): No need to init
 			}
-			Type *et = base_type(t->BitSet.base);
+			Type *et = base_type(t->BitSet.elem);
 			isize field_count = 0;
 			if (et->kind == Type_Enum) {
 				field_count = et->Enum.fields.count;
@@ -5576,7 +5579,7 @@ ExprKind check_expr_base_internal(CheckerContext *c, Operand *o, Ast *node, Type
 						is_constant = o->mode == Addressing_Constant;
 					}
 
-					check_assignment(c, o, t->BitSet.base, str_lit("bit_set literal"));
+					check_assignment(c, o, t->BitSet.elem, str_lit("bit_set literal"));
 				}
 			}
 			break;
@@ -6318,7 +6321,7 @@ gbString write_expr_to_string(gbString str, Ast *node) {
 
 	case_ast_node(bs, BitSetType, node);
 		str = gb_string_appendc(str, "bit_set[");
-		str = write_expr_to_string(str, bs->base);
+		str = write_expr_to_string(str, bs->elem);
 		str = gb_string_appendc(str, "]");
 	case_end;
 
diff --git a/src/check_type.cpp b/src/check_type.cpp
index dccaafae0..083a9e7c7 100644
--- a/src/check_type.cpp
+++ b/src/check_type.cpp
@@ -696,7 +696,7 @@ void check_bit_set_type(CheckerContext *c, Type *type, Ast *node) {
 
 	i64 const MAX_BITS = 64;
 
-	Ast *base = unparen_expr(bs->base);
+	Ast *base = unparen_expr(bs->elem);
 	if (is_ast_range(base)) {
 		ast_node(be, BinaryExpr, base);
 		Operand lhs = {};
@@ -719,8 +719,8 @@ void check_bit_set_type(CheckerContext *c, Type *type, Ast *node) {
 			    rhs.type != t_invalid) {
 				gbString xt = type_to_string(lhs.type);
 				gbString yt = type_to_string(rhs.type);
-				gbString expr_str = expr_to_string(bs->base);
-				error(bs->base, "Mismatched types in range '%s' : '%s' vs '%s'", expr_str, xt, yt);
+				gbString expr_str = expr_to_string(bs->elem);
+				error(bs->elem, "Mismatched types in range '%s' : '%s' vs '%s'", expr_str, xt, yt);
 				gb_string_free(expr_str);
 				gb_string_free(yt);
 				gb_string_free(xt);
@@ -730,13 +730,13 @@ void check_bit_set_type(CheckerContext *c, Type *type, Ast *node) {
 
 		if (!is_type_valid_bit_set_range(lhs.type)) {
 			gbString str = type_to_string(lhs.type);
-			error(bs->base, "'%s' is invalid for an interval expression, expected an integer or rune", str);
+			error(bs->elem, "'%s' is invalid for an interval expression, expected an integer or rune", str);
 			gb_string_free(str);
 			return;
 		}
 
 		if (lhs.mode != Addressing_Constant || rhs.mode != Addressing_Constant) {
-			error(bs->base, "Intervals must be constant values");
+			error(bs->elem, "Intervals must be constant values");
 			return;
 		}
 
@@ -751,19 +751,29 @@ void check_bit_set_type(CheckerContext *c, Type *type, Ast *node) {
 			gbAllocator a = heap_allocator();
 			String si = big_int_to_string(a, &i);
 			String sj = big_int_to_string(a, &j);
-			error(bs->base, "Lower interval bound larger than upper bound, %.*s .. %.*s", LIT(si), LIT(sj));
+			error(bs->elem, "Lower interval bound larger than upper bound, %.*s .. %.*s", LIT(si), LIT(sj));
 			gb_free(a, si.text);
 			gb_free(a, sj.text);
 			return;
 		}
 
 		Type *t = default_type(lhs.type);
+		if (bs->underlying != nullptr) {
+			Type *u = check_type(c, bs->underlying);
+			if (!is_type_integer(u)) {
+				gbString ts = type_to_string(u);
+				error(bs->underlying, "Expected an underlying integer for the bit set, got %s", ts);
+				gb_string_free(ts);
+				return;
+			}
+			type->BitSet.underlying = u;
+		}
 
 		if (!check_representable_as_constant(c, iv, t, nullptr)) {
 			gbAllocator a = heap_allocator();
 			String s = big_int_to_string(a, &i);
 			gbString ts = type_to_string(t);
-			error(bs->base, "%.*s is not representable by %s", LIT(s), ts);
+			error(bs->elem, "%.*s is not representable by %s", LIT(s), ts);
 			gb_string_free(ts);
 			gb_free(a, s.text);
 			return;
@@ -772,7 +782,7 @@ void check_bit_set_type(CheckerContext *c, Type *type, Ast *node) {
 			gbAllocator a = heap_allocator();
 			String s = big_int_to_string(a, &j);
 			gbString ts = type_to_string(t);
-			error(bs->base, "%.*s is not representable by %s", LIT(s), ts);
+			error(bs->elem, "%.*s is not representable by %s", LIT(s), ts);
 			gb_string_free(ts);
 			gb_free(a, s.text);
 			return;
@@ -780,50 +790,68 @@ void check_bit_set_type(CheckerContext *c, Type *type, Ast *node) {
 		i64 lower = big_int_to_i64(&i);
 		i64 upper = big_int_to_i64(&j);
 
-		if (upper - lower > MAX_BITS) {
-			error(bs->base, "bit_set range is greater than %lld bits, %lld bits are required", MAX_BITS, (upper-lower+1));
+		i64 bits = MAX_BITS;
+		if (type->BitSet.underlying != nullptr) {
+			bits = 8*type_size_of(type->BitSet.underlying);
+		}
+
+		if (upper - lower >= bits) {
+			error(bs->elem, "bit_set range is greater than %lld bits, %lld bits are required", bits, (upper-lower+1));
 		}
-		type->BitSet.base  = t;
+		type->BitSet.elem  = t;
 		type->BitSet.lower = lower;
 		type->BitSet.upper = upper;
-
 	} else {
-		Type *bt = check_type_expr(c, bs->base, nullptr);
+		Type *elem = check_type_expr(c, bs->elem, nullptr);
 
-		type->BitSet.base = bt;
-		if (!is_type_enum(bt)) {
-			error(bs->base, "Expected an enum type for a bit_set");
+		type->BitSet.elem = elem;
+		if (!is_type_valid_bit_set_elem(elem)) {
+			error(bs->elem, "Expected an enum type for a bit_set");
 		} else {
-			Type *et = base_type(bt);
-			GB_ASSERT(et->kind == Type_Enum);
-			if (!is_type_integer(et->Enum.base_type)) {
-				error(bs->base, "Enum type for bit_set must be an integer");
-				return;
-			}
-			i64 lower = 0;
-			i64 upper = 0;
+			Type *et = base_type(elem);
+			if (et->kind == Type_Enum) {
+				if (!is_type_integer(et->Enum.base_type)) {
+					error(bs->elem, "Enum type for bit_set must be an integer");
+					return;
+				}
+				i64 lower = 0;
+				i64 upper = 0;
 
-			for_array(i, et->Enum.fields) {
-				Entity *e = et->Enum.fields[i];
-				if (e->kind != Entity_Constant) {
-					continue;
+				for_array(i, et->Enum.fields) {
+					Entity *e = et->Enum.fields[i];
+					if (e->kind != Entity_Constant) {
+						continue;
+					}
+					ExactValue value = exact_value_to_integer(e->Constant.value);
+					GB_ASSERT(value.kind == ExactValue_Integer);
+					// NOTE(bill): enum types should be able to store i64 values
+					i64 x = big_int_to_i64(&value.value_integer);
+					lower = gb_min(lower, x);
+					upper = gb_max(upper, x);
 				}
-				ExactValue value = exact_value_to_integer(e->Constant.value);
-				GB_ASSERT(value.kind == ExactValue_Integer);
-				// NOTE(bill): enum types should be able to store i64 values
-				i64 x = big_int_to_i64(&value.value_integer);
-				lower = gb_min(lower, x);
-				upper = gb_max(upper, x);
-			}
 
-			GB_ASSERT(lower <= upper);
+				GB_ASSERT(lower <= upper);
 
-			if (upper - lower > MAX_BITS) {
-				error(bs->base, "bit_set range is greater than %lld bits, %lld bits are required", MAX_BITS, (upper-lower+1));
-			}
+				i64 bits = MAX_BITS;
+				if (bs->underlying != nullptr) {
+					Type *u = check_type(c, bs->underlying);
+					if (!is_type_integer(u)) {
+						gbString ts = type_to_string(u);
+						error(bs->underlying, "Expected an underlying integer for the bit set, got %s", ts);
+						gb_string_free(ts);
+						return;
+					}
+					type->BitSet.underlying = u;
+					bits = 8*type_size_of(u);
+				}
+
+				if (upper - lower >= MAX_BITS) {
+					error(bs->elem, "bit_set range is greater than %lld bits, %lld bits are required", MAX_BITS, (upper-lower+1));
+				}
 
-			type->BitSet.lower = lower;
-			type->BitSet.upper = upper;
+				type->BitSet.lower = lower;
+				type->BitSet.upper = upper;
+			}
 		}
 	}
 }
diff --git a/src/checker.cpp b/src/checker.cpp
index 20c37fbb4..c42413718 100644
--- a/src/checker.cpp
+++ b/src/checker.cpp
@@ -1044,7 +1044,7 @@ void add_type_info_type(CheckerContext *c, Type *t) {
 		break;
 
 	case Type_BitSet:
-		add_type_info_type(c, bt->BitSet.base);
+		add_type_info_type(c, bt->BitSet.elem);
 		break;
 
 	case Type_Union:
diff --git a/src/ir.cpp b/src/ir.cpp
index 67fee8556..c824ec1ef 100644
--- a/src/ir.cpp
+++ b/src/ir.cpp
@@ -4909,7 +4909,7 @@ irValue *ir_build_expr_internal(irProcedure *proc, Ast *expr) {
 				{
 					ir_emit_comment(proc, str_lit("bit_set in"));
 
-					Type *key_type = rt->BitSet.base;
+					Type *key_type = rt->BitSet.elem;
 					GB_ASSERT(are_types_identical(ir_type(left), key_type));
 
 					Type *it = bit_set_to_int(rt);
@@ -5663,7 +5663,7 @@ irAddr ir_build_addr(irProcedure *proc, Ast *expr) {
 		switch (bt->kind) {
 		case Type_Array:  et = bt->Array.elem;  break;
 		case Type_Slice:  et = bt->Slice.elem;  break;
-		case Type_BitSet: et = bt->BitSet.base; break;
+		case Type_BitSet: et = bt->BitSet.elem; break;
 		}
 
 		String proc_name = {};
@@ -8118,10 +8118,13 @@ void ir_setup_type_info_data(irProcedure *proc) { // NOTE(bill): Setup type_info
 			ir_emit_comment(proc, str_lit("Type_Info_Bit_Set"));
 			tag = ir_emit_conv(proc, variant_ptr, t_type_info_bit_set_ptr);
 
-			GB_ASSERT(is_type_typed(t->BitSet.base));
-			ir_emit_store(proc, ir_emit_struct_ep(proc, tag, 0), ir_get_type_info_ptr(proc, t->BitSet.base));
-			ir_emit_store(proc, ir_emit_struct_ep(proc, tag, 1), ir_const_i64(t->BitSet.lower));
-			ir_emit_store(proc, ir_emit_struct_ep(proc, tag, 2), ir_const_i64(t->BitSet.upper));
+			GB_ASSERT(is_type_typed(t->BitSet.elem));
+			ir_emit_store(proc, ir_emit_struct_ep(proc, tag, 0), ir_get_type_info_ptr(proc, t->BitSet.elem));
+			if (t->BitSet.underlying != nullptr) {
+				ir_emit_store(proc, ir_emit_struct_ep(proc, tag, 1), ir_get_type_info_ptr(proc, t->BitSet.underlying));
+			}
+			ir_emit_store(proc, ir_emit_struct_ep(proc, tag, 2), ir_const_i64(t->BitSet.lower));
+			ir_emit_store(proc, ir_emit_struct_ep(proc, tag, 3), ir_const_i64(t->BitSet.upper));
 			break;
 
 
diff --git a/src/ir_print.cpp b/src/ir_print.cpp
index 85efd12ba..547235dd0 100644
--- a/src/ir_print.cpp
+++ b/src/ir_print.cpp
@@ -508,16 +508,13 @@ void ir_print_type(irFileBuffer *f, irModule *m, Type *t, bool in_struct) {
 	}
 
 	case Type_BitSet: {
-		i64 align = type_align_of(t);
-		i64 size  = type_size_of(t);
-		switch (size) {
-		case 0: ir_write_str_lit(f, "{}");  return;
-		case 1: ir_write_str_lit(f, "i8");  return;
-		case 2: ir_write_str_lit(f, "i16"); return;
-		case 4: ir_write_str_lit(f, "i32"); return;
-		case 8: ir_write_str_lit(f, "i64"); return;
-		default: GB_PANIC("Unknown bit_set size"); break;
+		i64 size = type_size_of(t);
+		if (size == 0) {
+			ir_write_str_lit(f, "{}");
+			return;
 		}
+		ir_print_type(f, m, bit_set_to_int(t));
+		return;
 	}
 	}
 }
diff --git a/src/parser.cpp b/src/parser.cpp
index 94d87ecee..cce1729c1 100644
--- a/src/parser.cpp
+++ b/src/parser.cpp
@@ -347,7 +347,8 @@ Ast *clone_ast(Ast *node) {
 		n->BitFieldType.align = clone_ast(n->BitFieldType.align);
 		break;
 	case Ast_BitSetType:
-		n->BitSetType.base = clone_ast(n->BitSetType.base);
+		n->BitSetType.elem       = clone_ast(n->BitSetType.elem);
+		n->BitSetType.underlying = clone_ast(n->BitSetType.underlying);
 		break;
 	case Ast_MapType:
 		n->MapType.count = clone_ast(n->MapType.count);
@@ -927,10 +928,11 @@ Ast *ast_bit_field_type(AstFile *f, Token token, Array<Ast *> fields, Ast *align
 	return result;
 }
 
-Ast *ast_bit_set_type(AstFile *f, Token token, Ast *base) {
+Ast *ast_bit_set_type(AstFile *f, Token token, Ast *elem, Ast *underlying) {
 	Ast *result = alloc_ast_node(f, Ast_BitSetType);
 	result->BitSetType.token = token;
-	result->BitSetType.base = base;
+	result->BitSetType.elem = elem;
+	result->BitSetType.underlying = underlying;
 	return result;
 }
 
@@ -1963,14 +1965,20 @@ Ast *parse_operand(AstFile *f, bool lhs) {
 		Token token = expect_token(f, Token_bit_set);
 		Token open  = expect_token(f, Token_OpenBracket);
 
+		Ast *elem = nullptr;
+		Ast *underlying = nullptr;
+
 		bool prev_allow_range = f->allow_range;
 		f->allow_range = true;
-		Ast *base = parse_expr(f, false);
+		elem = parse_expr(f, false);
 		f->allow_range = prev_allow_range;
+		if (allow_token(f, Token_Semicolon)) {
+			underlying = parse_type(f);
+		}
 
 		Token close = expect_token(f, Token_CloseBracket);
 
-		return ast_bit_set_type(f, token, base);
+		return ast_bit_set_type(f, token, elem, underlying);
 	}
 
 	default: {
diff --git a/src/parser.hpp b/src/parser.hpp
index b25740315..0cc77d2cc 100644
--- a/src/parser.hpp
+++ b/src/parser.hpp
@@ -483,7 +483,8 @@ AST_KIND(_TypeBegin, "", bool) \
 	}) \
 	AST_KIND(BitSetType, "bit set type", struct { \
 		Token token; \
-		Ast * base;  \
+		Ast * elem;  \
+		Ast * underlying; \
 	}) \
 	AST_KIND(MapType, "map type", struct { \
 		Token token; \
diff --git a/src/types.cpp b/src/types.cpp
index 406b55b82..695187f05 100644
--- a/src/types.cpp
+++ b/src/types.cpp
@@ -179,7 +179,8 @@ struct TypeStruct {
 		i64             custom_align;                     \
 	})                                                    \
 	TYPE_KIND(BitSet, struct {                            \
-		Type *base;                                       \
+		Type *elem;                                       \
+		Type *underlying;                                 \
 		i64   lower;                                      \
 		i64   upper;                                      \
 	})                                                    \
@@ -981,10 +982,28 @@ bool is_type_valid_for_keys(Type *t) {
 	return false;
 }
 
+bool is_type_valid_bit_set_elem(Type *t) {
+	if (is_type_enum(t)) {
+		return true;
+	}
+	t = core_type(t);
+	if (t->kind == Type_Generic) {
+		return true;
+	}
+	return false;
+}
+
 Type *bit_set_to_int(Type *t) {
 	GB_ASSERT(is_type_bit_set(t));
+	Type *bt = base_type(t);
+	Type *underlying = bt->BitSet.underlying;
+	if (underlying != nullptr && is_type_integer(underlying)) {
+		return underlying;
+	}
+
 	i64 sz = type_size_of(t);
 	switch (sz) {
+	case 0: return t_u8;
 	case 1: return t_u8;
 	case 2: return t_u16;
 	case 4: return t_u32;
@@ -1277,7 +1296,7 @@ bool are_types_identical(Type *x, Type *y) {
 
 	case Type_BitSet:
 		if (y->kind == Type_BitSet) {
-			return are_types_identical(x->BitSet.base, y->BitSet.base);
+			return are_types_identical(x->BitSet.elem, y->BitSet.elem);
 		}
 		break;
 
@@ -2069,8 +2088,10 @@ i64 type_align_of_internal(Type *t, TypePath *path) {
 	} break;
 
 	case Type_BitSet: {
+		if (t->BitSet.underlying != nullptr) {
+			return type_align_of(t->BitSet.underlying);
+		}
 		i64 bits = t->BitSet.upper - t->BitSet.lower + 1;
-		if (bits == 0)  return 0;
 		if (bits <= 8)  return 1;
 		if (bits <= 16) return 2;
 		if (bits <= 32) return 4;
@@ -2296,8 +2317,10 @@ i64 type_size_of_internal(Type *t, TypePath *path) {
 	} break;
 
 	case Type_BitSet: {
+		if (t->BitSet.underlying != nullptr) {
+			return type_size_of(t->BitSet.underlying);
+		}
 		i64 bits = t->BitSet.upper - t->BitSet.lower + 1;
-		if (bits == 0)  return 0;
 		if (bits <= 8)  return 1;
 		if (bits <= 16) return 2;
 		if (bits <= 32) return 4;
@@ -2617,7 +2640,7 @@ gbString write_type_to_string(gbString str, Type *type) {
 
 	case Type_BitSet:
 		str = gb_string_appendc(str, "bit_set[");
-		str = write_type_to_string(str, type->BitSet.base);
+		str = write_type_to_string(str, type->BitSet.elem);
 		str = gb_string_appendc(str, "]");
 		break;
 	}