comptime eval

compiler/checker/checker.odin

@@ -389,6 +389,7 @@ eval_integer_constant_in_context :: proc(
 	pkg: ast.Package_Id,
 	file: ast.File_Id,
 	depth := 0,
+	values: []Comptime_Value = nil,
 ) -> Constant {
 	if depth > 64 || expr_id == ast.INVALID_EXPR || int(expr_id) >= len(checker.ast_module.exprs) {
 		return Constant{kind = .Not_Constant}
@@ -397,8 +398,15 @@ eval_integer_constant_in_context :: proc(
 	#partial switch expr.kind {
 	case .Integer:
 		return Constant{kind = .Value, value = i128(expr.integer)}
+	case .Bool:
+		return Constant{kind = .Value, value = i128(expr.integer)}
 	case .Name:
 		if !symbol.is_valid(expr.qualifier) {
+			if value, ok := find_comptime_value(values, expr.name); ok {
+				if value.kind == .Integer {
+					return Constant{kind = .Value, value = value.value}
+				}
+			}
 			if value, ok := current_comptime_value(checker, expr.name); ok {
 				if value.kind == .Integer {
 					return Constant{kind = .Value, value = value.value}
@@ -417,17 +425,32 @@ eval_integer_constant_in_context :: proc(
 		if g.external || !g.immutable {
 			return Constant{kind = .Not_Constant}
 		}
-		return eval_integer_constant_in_context(checker, g.expr, g.pkg, g.file, depth+1)
+		return eval_integer_constant_in_context(checker, g.expr, g.pkg, g.file, depth+1, values)
+	case .Comptime:
+		if expr.left != ast.INVALID_EXPR {
+			return eval_integer_constant_in_context(checker, expr.left, pkg, file, depth+1, values)
+		}
+		result, produced, ok := eval_comptime_statements(checker, expr.body, pkg, file, depth+1, values, true)
+		if ok && produced {
+			return result
+		}
+		return Constant{kind = .Not_Constant}
 	case .Negate:
-		operand := eval_integer_constant_in_context(checker, expr.left, pkg, file, depth+1)
+		operand := eval_integer_constant_in_context(checker, expr.left, pkg, file, depth+1, values)
 		if operand.kind == .Value {
 			value, overflow := intrinsics.overflow_sub(i128(0), operand.value)
 			return Constant{kind = .Overflow} if overflow else Constant{kind = .Value, value = value}
 		}
 		return operand
+	case .Not:
+		operand := eval_integer_constant_in_context(checker, expr.left, pkg, file, depth+1, values)
+		if operand.kind == .Value {
+			return Constant{kind = .Value, value = 1 if operand.value == 0 else 0}
+		}
+		return operand
 	case .Add, .Sub, .Mul, .Div:
-		left := eval_integer_constant_in_context(checker, expr.left, pkg, file, depth+1)
-		right := eval_integer_constant_in_context(checker, expr.right, pkg, file, depth+1)
+		left := eval_integer_constant_in_context(checker, expr.left, pkg, file, depth+1, values)
+		right := eval_integer_constant_in_context(checker, expr.right, pkg, file, depth+1, values)
 		if left.kind == .Div_By_Zero || right.kind == .Div_By_Zero {
 			return Constant{kind = .Div_By_Zero}
 		}
@@ -453,10 +476,163 @@ eval_integer_constant_in_context :: proc(
 			value, overflow = intrinsics.overflow_add(left.value, right.value)
 		}
 		return Constant{kind = .Overflow} if overflow else Constant{kind = .Value, value = value}
+	case .Eq, .Ne, .Lt, .Le, .Gt, .Ge:
+		left := eval_integer_constant_in_context(checker, expr.left, pkg, file, depth+1, values)
+		right := eval_integer_constant_in_context(checker, expr.right, pkg, file, depth+1, values)
+		if left.kind != .Value || right.kind != .Value {
+			if left.kind == .Div_By_Zero || right.kind == .Div_By_Zero {
+				return Constant{kind = .Div_By_Zero}
+			}
+			if left.kind == .Overflow || right.kind == .Overflow {
+				return Constant{kind = .Overflow}
+			}
+			return Constant{kind = .Not_Constant}
+		}
+		ok := false
+		#partial switch expr.kind {
+		case .Eq: ok = left.value == right.value
+		case .Ne: ok = left.value != right.value
+		case .Lt: ok = left.value < right.value
+		case .Le: ok = left.value <= right.value
+		case .Gt: ok = left.value > right.value
+		case .Ge: ok = left.value >= right.value
+		}
+		return Constant{kind = .Value, value = 1 if ok else 0}
+	case .And, .Or:
+		left := eval_integer_constant_in_context(checker, expr.left, pkg, file, depth+1, values)
+		if left.kind != .Value {
+			return left
+		}
+		if expr.kind == .And && left.value == 0 {
+			return Constant{kind = .Value, value = 0}
+		}
+		if expr.kind == .Or && left.value != 0 {
+			return Constant{kind = .Value, value = 1}
+		}
+		right := eval_integer_constant_in_context(checker, expr.right, pkg, file, depth+1, values)
+		if right.kind == .Value {
+			return Constant{kind = .Value, value = 1 if right.value != 0 else 0}
+		}
+		return right
+	case .Call:
+		return eval_comptime_call(checker, expr, pkg, file, depth+1, values)
 	}
 	return Constant{kind = .Not_Constant}
 }
 
+eval_comptime_statements :: proc(
+	checker: ^Checker,
+	statements: []ast.Stmt_Id,
+	pkg: ast.Package_Id,
+	file: ast.File_Id,
+	depth: int,
+	values: []Comptime_Value,
+	yield_returns: bool,
+) -> (Constant, bool, bool) {
+	env: [dynamic]Comptime_Value
+	env.allocator = checker.allocator
+	append(&env, ..values)
+	defer delete(env)
+	for statement_id in statements {
+		if statement_id == ast.INVALID_STMT || int(statement_id) >= len(checker.ast_module.statements) {
+			return Constant{kind=.Not_Constant}, false, false
+		}
+		statement := checker.ast_module.statements[statement_id]
+		#partial switch statement.kind {
+		case .Declaration:
+			if statement.expr == ast.INVALID_EXPR || !statement.immutable {
+				return Constant{kind=.Not_Constant}, false, false
+			}
+			value := eval_integer_constant_in_context(checker, statement.expr, pkg, file, depth+1, env[:])
+			if value.kind != .Value || statement.name == checker.sink_symbol {
+				return value, false, false
+			}
+			append(&env, Comptime_Value{name=statement.name, type=types.I64, value=value.value})
+		case .Return:
+			if yield_returns || statement.expr == ast.INVALID_EXPR {
+				return Constant{kind=.Not_Constant}, false, false
+			}
+			value := eval_integer_constant_in_context(checker, statement.expr, pkg, file, depth+1, env[:])
+			return value, value.kind == .Value, value.kind == .Value
+		case .Yield:
+			if !yield_returns || statement.expr == ast.INVALID_EXPR {
+				return Constant{kind=.Not_Constant}, false, false
+			}
+			value := eval_integer_constant_in_context(checker, statement.expr, pkg, file, depth+1, env[:])
+			return value, value.kind == .Value, value.kind == .Value
+		case .If:
+			if len(statement.captures) > 0 || statement.guard != ast.INVALID_EXPR {
+				return Constant{kind=.Not_Constant}, false, false
+			}
+			condition := eval_integer_constant_in_context(checker, statement.expr, pkg, file, depth+1, env[:])
+			if condition.kind != .Value {
+				return condition, false, false
+			}
+			body := statement.body if condition.value != 0 else statement.else_body
+			value, produced, ok := eval_comptime_statements(checker, body, pkg, file, depth+1, env[:], yield_returns)
+			if !ok || produced {
+				return value, produced, ok
+			}
+		case:
+			return Constant{kind=.Not_Constant}, false, false
+		}
+	}
+	return Constant{kind=.Not_Constant}, false, true
+}
+
+eval_comptime_call :: proc(
+	checker: ^Checker,
+	expr: ast.Expr,
+	pkg: ast.Package_Id,
+	file: ast.File_Id,
+	depth: int,
+	values: []Comptime_Value,
+) -> Constant {
+	if expr.left != ast.INVALID_EXPR || depth > 64 {
+		return Constant{kind=.Not_Constant}
+	}
+	target_pkg, available := expr_package(checker, expr, pkg, file, false)
+	if !available {
+		return Constant{kind=.Not_Constant}
+	}
+	template := find_template(checker, expr.name, target_pkg)
+	if template == ast.INVALID_FUNCTION || int(template) >= len(checker.ast_module.functions) {
+		return Constant{kind=.Not_Constant}
+	}
+	function := checker.ast_module.functions[template]
+	if function.c_abi || !function.has_body || types.is_valid(function.error) ||
+	   len(function.unsupported_reason) > 0 || !valid_call_arity(function, len(expr.args)) {
+		return Constant{kind=.Not_Constant}
+	}
+	comptime_values, comptime_ok := collect_comptime_values(checker, function, expr.args, pkg, file, false, values)
+	defer delete(comptime_values, checker.allocator)
+	if !comptime_ok {
+		return Constant{kind=.Not_Constant}
+	}
+	env: [dynamic]Comptime_Value
+	env.allocator = checker.allocator
+	defer delete(env)
+	append(&env, ..comptime_values)
+	for param, index in function.params {
+		if param.comptime_value {
+			continue
+		}
+		if index >= len(expr.args) {
+			return Constant{kind=.Not_Constant}
+		}
+		value := eval_integer_constant_in_context(checker, expr.args[index], pkg, file, depth+1, values)
+		if value.kind != .Value {
+			return value
+		}
+		append(&env, Comptime_Value{name=param.name, type=types.I64, value=value.value})
+	}
+	result, produced, ok := eval_comptime_statements(checker, function.body, function.pkg, function.file, depth+1, env[:], false)
+	if !ok || !produced {
+		return Constant{kind=.Not_Constant}
+	}
+	return result
+}
+
 fits_signed_type :: proc(value: i128, value_type: types.Type, selected := target.DEFAULT) -> bool {
 	if !types.is_signed(value_type, selected) {
 		return false
@@ -952,6 +1128,7 @@ collect_comptime_values :: proc(
 	pkg: ast.Package_Id,
 	file: ast.File_Id,
 	diagnose := false,
+	extra_values: []Comptime_Value = nil,
 ) -> ([]Comptime_Value, bool) {
 	if !function_has_comptime_params(function) {
 		return nil, true
@@ -1002,7 +1179,7 @@ collect_comptime_values :: proc(
 		}
 		constant := Constant{kind = .Not_Constant}
 		if index < len(args) {
-			constant = eval_integer_constant_in_context(checker, args[index], pkg, file)
+			constant = eval_integer_constant_in_context(checker, args[index], pkg, file, values=extra_values)
 		}
 		if constant.kind != .Value {
 			if diagnose {
@@ -1145,6 +1322,11 @@ mark_expr_imports_used :: proc(checker: ^Checker, expr_id: ast.Expr_Id, file: as
 			}
 		case .Negate, .Not, .Address, .Deref, .Field, .Unwrap, .Try, .Keyed, .Enum_Literal, .Cast:
 			append(&stack, expr.left)
+		case .Comptime:
+			if expr.left != ast.INVALID_EXPR {
+				append(&stack, expr.left)
+			}
+			mark_block_imports_used(checker, expr.body, file)
 		case .Catch:
 			append(&stack, expr.left)
 			if expr.right != ast.INVALID_EXPR {
@@ -1744,6 +1926,25 @@ infer_compound_expr :: proc(
 ) -> types.Type {
 	store := &checker.module.types
 	#partial switch expr.kind {
+	case .Comptime:
+		constant := Constant{kind=.Not_Constant}
+		if expr.left != ast.INVALID_EXPR {
+			constant = eval_integer_constant_in_context(checker, expr.left, pkg, file)
+		}
+		if expr.left == ast.INVALID_EXPR {
+			value, produced, ok := eval_comptime_statements(checker, expr.body, pkg, file, 0, nil, true)
+			if ok && produced {
+				constant = value
+			}
+		}
+		if constant.kind == .Overflow || constant.kind == .Div_By_Zero ||
+		   (constant.kind == .Value && !fits_i64(constant.value)) {
+			return types.I64
+		}
+		if constant.kind == .Value {
+			return types.smallest_signed_for_literal(i64(constant.value))
+		}
+		return types.INVALID
 	case .Bool:
 		return types.BOOL
 	case .Not:
@@ -1954,7 +2155,7 @@ infer_expr :: proc(
 				_ = pop(&stack)
 			case .String, .Array, .None, .Undefined, .Address, .Deref, .Index, .Slice,
 			     .Field, .Unwrap, .Orelse, .Try, .Catch, .Struct_Literal, .Keyed, .Enum_Literal, .Cast,
-			     .Bool, .Not, .Eq, .Ne, .Lt, .Le, .Gt, .Ge, .And, .Or, .Range:
+			     .Comptime, .Bool, .Not, .Eq, .Ne, .Lt, .Le, .Gt, .Ge, .And, .Or, .Range:
 				last = infer_compound_expr(checker, expr, locals, pkg, file, demanded, local_types)
 				_ = pop(&stack)
 			case .Name:
@@ -2371,7 +2572,7 @@ infer_statements :: proc(
 			open_float := false
 			const_val := i128(0)
 			if !is_runtime_type(checker, declared_local) && !is_undefined_expr(checker, statement.expr) {
-				constant := eval_constant(checker, statement.expr)
+				constant := eval_integer_constant_in_context(checker, statement.expr, pkg, file)
 				if constant.kind == .Value && fits_i64(constant.value) {
 					open = true
 					const_val = constant.value
@@ -2842,7 +3043,7 @@ expr_accepts_numeric_demand :: proc(
 	   expr_id == ast.INVALID_EXPR || int(expr_id) >= len(checker.ast_module.exprs) {
 		return false
 	}
-	if constant := eval_constant(checker, expr_id); constant.kind == .Value {
+	if constant := eval_integer_constant_in_context(checker, expr_id, pkg, file); constant.kind == .Value {
 		return open_integer_accepts_demand(checker, constant.value, demand)
 	}
 	if is_float_constant_expr(checker, expr_id) {
@@ -2945,7 +3146,7 @@ infer_all :: proc(checker: ^Checker) {
 		if global.external {
 			continue
 		}
-		constant := eval_constant(checker, global.expr)
+		constant := eval_integer_constant_in_context(checker, global.expr, global.pkg, global.file)
 		if constant.kind == .Value && fits_i64(constant.value) {
 			checker.global_open_const[index] = true
 			checker.global_const_value[index] = constant.value
@@ -4120,6 +4321,21 @@ build_compound_expr :: proc(
 			target=hir.INVALID_REF, left=hir.INVALID_EXPR, right=hir.INVALID_EXPR,
 			diagnostic=source.INVALID_DIAGNOSTIC,
 		})
+	case .Comptime:
+		constant := Constant{kind=.Not_Constant}
+		if expr.left != ast.INVALID_EXPR {
+			constant = eval_integer_constant_in_context(checker, expr.left, pkg, file)
+		} else {
+			value, produced, ok := eval_comptime_statements(checker, expr.body, pkg, file, 0, nil, true)
+			if ok && produced {
+				constant = value
+			}
+		}
+		if constant.kind == .Value || constant.kind == .Overflow || constant.kind == .Div_By_Zero {
+			return build_constant_expr(checker, expr, constant, expected)
+		}
+		id := source.add(checker.diagnostics, expr.span, "expression cannot be evaluated at comptime")
+		return invalid_hir_expr(checker, expr.span, id)
 	case .Bool:
 		return add_hir_expr(checker, hir.Expr{
 			kind=.Bool, span=expr.span, type=types.BOOL, integer=i64(expr.integer),
@@ -4398,9 +4614,9 @@ build_expr :: proc(
 				continue
 			}
 			switch expr.kind {
-	case .String, .Array, .None, .Undefined, .Address, .Deref, .Index, .Slice,
+			case .String, .Array, .None, .Undefined, .Address, .Deref, .Index, .Slice,
 			     .Field, .Unwrap, .Orelse, .Try, .Catch, .Struct_Literal, .Keyed,
-			     .Bool, .Cast, .Not, .Eq, .Ne, .Lt, .Le, .Gt, .Ge, .And, .Or, .Range,
+			     .Bool, .Cast, .Comptime, .Not, .Eq, .Ne, .Lt, .Le, .Gt, .Ge, .And, .Or, .Range,
 			     .Enum_Literal:
 				last = build_compound_expr(
 					checker, expr, locals, global_reads, calls, frame.expected, pkg, file,
@@ -5092,7 +5308,7 @@ build_block :: proc(
 			// `undefined`, open numeric constants, or arithmetic expressions.
 			open_const_decl := !is_runtime_type(checker, declared) && !is_undefined_expr(checker, statement.expr)
 			if open_const_decl {
-				constant := eval_constant(checker, statement.expr)
+				constant := eval_integer_constant_in_context(checker, statement.expr, ctx.pkg, ctx.file)
 				open_const_decl = constant.kind == .Value && fits_i64(constant.value) ||
 				                  is_float_constant_expr(checker, statement.expr)
 			}
@@ -7650,7 +7866,7 @@ build_globals :: proc(checker: ^Checker) {
 		expected := types.INVALID
 		if is_runtime_type(checker, declared) {
 			expected = declared
-		} else if constant := eval_constant(checker, global.expr);
+		} else if constant := eval_integer_constant_in_context(checker, global.expr, global.pkg, global.file);
 		   constant.kind == .Value && fits_i64(constant.value) &&
 		   is_runtime_type(checker, checker.global_types[global_index]) {
 			// Open integer constant: build against its demanded/defaulted type. Gated