lib/chant/src/lower/statement/loop.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const ast = @import("../../ast/root.zig");
2 const lower_mod = @import("../root.zig");
3 const expression = @import("../expression/root.zig");
4 const canonical_mod = @import("canonical.zig");
5 const types = @import("types.zig");
6
7 const dialects = @import("choir").dialects;
8 const ArithDialect = dialects.ArithDialect;
9 const ScfDialect = dialects.ScfDialect;
10 const Error = lower_mod.Error;
11 const Lowerer = lower_mod.Lowerer;
12
13 pub fn lower(lowerer: *Lowerer, for_stmt: ast.Stmt.For, comptime recurse: types.LowerStatement) Error!void {
14 const canonical = canonical_mod.recognize(lowerer, for_stmt) orelse return error.UnsupportedConstruct;
15
16 try lower_mod.scope.push(lowerer);
17 defer lower_mod.scope.pop(lowerer);
18
19 const initial = try expression.lowerExpression(lowerer, canonical.initial);
20 const lower_bound = try lower_mod.convert.toIndex(
21 lowerer,
22 try lower_mod.convert.convert(lowerer, initial.value, initial.c_type, canonical.iv_type),
23 canonical.iv_type,
24 );
25
26 const limit = try expression.lowerExpression(lowerer, canonical.limit);
27 var upper_bound = try lower_mod.convert.toIndex(
28 lowerer,
29 try lower_mod.convert.convert(lowerer, limit.value, limit.c_type, canonical.iv_type),
30 canonical.iv_type,
31 );
32 if (canonical.inclusive) {
33 const one = try lower_mod.emit.indexConstant(lowerer, 1);
34 const bumped = ArithDialect.AddOp.create(lowerer.ctx, lowerer.loc, upper_bound, one) catch return error.OutOfMemory;
35 try lower_mod.emit.append(lowerer, bumped.op);
36 var bumped_mutable = bumped;
37 upper_bound = bumped_mutable.getResult();
38 }
39
40 const step = try lower_mod.emit.indexConstant(lowerer, @intCast(canonical.step));
41
42 const for_op = ScfDialect.ForOp.create(lowerer.ctx, lowerer.loc, lower_bound, upper_bound, step, &.{}, &.{}) catch return error.OutOfMemory;
43 try lower_mod.emit.append(lowerer, for_op.op);
44
45 const body = for_op.getBodyBlock();
46 const iv = body.arguments.items[0];
47
48 const saved = lowerer.block;
49 lowerer.block = body;
50 defer lowerer.block = saved;
51
52 const iv_cast = ArithDialect.CastOp.create(lowerer.ctx, lowerer.loc, iv, try lower_mod.convert.scalarType(lowerer, canonical.iv_type)) catch return error.OutOfMemory;
53 try lower_mod.emit.append(lowerer, iv_cast.op);
54 var iv_mutable = iv_cast;
55 try lower_mod.scope.bindLoopVariable(lowerer, canonical.name, iv_mutable.getResult(), canonical.iv_type);
56
57 try lower_mod.scope.push(lowerer);
58 defer lower_mod.scope.pop(lowerer);
59 try recurse(lowerer, for_stmt.body);
60 try lower_mod.emit.emitYield(lowerer, &.{});
61 }