瀏覽代碼

Add `rand.choice_enum`

gingerBill 1 年之前
父節點
當前提交
e201a2fabb
共有 1 個文件被更改,包括 20 次插入0 次删除
  1. 20 0
      core/math/rand/rand.odin

+ 20 - 0
core/math/rand/rand.odin

@@ -834,3 +834,23 @@ choice :: proc(array: $T/[]$E, r: ^Rand = nil) -> (res: E) {
 	}
 	return array[int63_max(n, r)]
 }
+
+
+@(require_results)
+choice_enum :: proc($T: typeid, r: ^Rand = nil) -> T
+	where
+		intrinsics.type_is_enum(T),
+		size_of(T) <= 8,
+		len(T) == cap(T) /* Only allow contiguous enum types */
+{
+	when intrinsics.type_is_unsigned(intrinsics.type_core_type(T)) &&
+	     u64(max(T)) > u64(max(i64)) {
+		i := uint64(r) % u64(len(T))
+		i += u64(min(T))
+		return T(i)
+	} else {
+		i := int63_max(i64(len(T)), r)
+		i += i64(min(T))
+		return T(i)
+	}
+}