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 }