tiny.choir.dialects.arith.folds

1const std = @import("std");2const ir = @import("../../core/root.zig");3const predicate_mod = @import("predicate.zig");4const types = @import("types.zig");56const CmpPredicate = predicate_mod.CmpPredicate;78const FoldBinaryOperands = struct {9    lhs: *ir.Value,10    rhs: *ir.Value,11};1213pub fn Folds(comptime Dialect: type) type {14    return struct {15        fn appendFoldValue(16            results: *ir.interfaces.FoldResults,17            value: *ir.Value,18        ) !void {19            try results.append(.{ .value = value });20        }2122        fn appendFoldAttribute(23            results: *ir.interfaces.FoldResults,24            attr: ir.Attribute,25        ) !void {26            try results.append(.{ .attribute = attr });27        }2829        fn foldBinaryOperands(op: *const ir.Operation) ?FoldBinaryOperands {30            if (op.operands.items.len != 2) return null;31            if (op.results.items.len != 1) return null;32            return .{33                .lhs = op.operands.items[0].value,34                .rhs = op.operands.items[1].value,35            };36        }3738        fn appendFoldValueForResult(39            op: *const ir.Operation,40            results: *ir.interfaces.FoldResults,41            value: *ir.Value,42        ) !void {43            if (op.results.items.len != 1) return;44            if (!op.results.items[0].type.eql(value.type)) return;45            try appendFoldValue(results, value);46        }4748        fn appendIntFoldAttributeForResult(49            op: *const ir.Operation,50            results: *ir.interfaces.FoldResults,51            value: i64,52        ) !void {53            if (op.results.items.len != 1) return;54            if (!isIntegerLikeType(op.results.items[0].type)) return;55            const attr = try Dialect.getIntAttr(op.getContext(), value);56            try appendFoldAttribute(results, attr);57        }5859        fn appendBoolFoldAttributeForResult(60            op: *const ir.Operation,61            results: *ir.interfaces.FoldResults,62            value: bool,63        ) !void {64            if (op.results.items.len != 1) return;65            if (!isBoolType(op.results.items[0].type)) return;66            const attr = try Dialect.getBoolAttr(op.getContext(), value);67            try appendFoldAttribute(results, attr);68        }6970        fn constantDefiningOp(value: *ir.Value) ?*ir.Operation {71            const def_any = value.getDefiningOp() orelse return null;72            const def_op: *ir.Operation = @ptrCast(@alignCast(def_any));73            if (!std.mem.eql(u8, def_op.name.name, Dialect.ConstantOp.operation_name)) return null;74            return def_op;75        }7677        pub fn constantIntFromValue(value: *ir.Value) ?i64 {78            if (!isIntegerLikeType(value.type)) return null;79            const def_op = constantDefiningOp(value) orelse return null;80            const int_attr = def_op.getAttrAs(ir.Attribute.IntegerAttr, "value") orelse return null;81            return int_attr.getValue();82        }8384        pub fn constantBoolFromValue(value: *ir.Value) ?bool {85            const def_op = constantDefiningOp(value) orelse return null;86            if (def_op.getAttrAs(ir.Attribute.BoolAttr, "value")) |bool_attr| return bool_attr.getValue();87            if (def_op.getAttrAs(ir.Attribute.IntegerAttr, "value")) |int_attr| {88                const int_value = int_attr.getValue();89                if (!isBoolType(value.type)) return null;90                return int_value != 0;91            }92            return null;93        }9495        pub fn constantIntEquals(value: *ir.Value, expected: i64) bool {96            return (constantIntFromValue(value) orelse return false) == expected;97        }9899        pub fn constantBoolEquals(value: *ir.Value, expected: bool) bool {100            return (constantBoolFromValue(value) orelse return false) == expected;101        }102103        pub fn isBoolType(ty: ir.Type) bool {104            return types.scalarKindFromType(ty) == .bool;105        }106107        pub fn isIntegerLikeType(ty: ir.Type) bool {108            const kind = types.scalarKindFromType(ty) orelse return false;109            return types.scalarKindIsInteger(kind);110        }111112        pub fn foldSameTypeUnary(113            op_ptr: *const anyopaque,114            results: *ir.interfaces.FoldResults,115        ) anyerror!void {116            const op: *const ir.Operation = @ptrCast(@alignCast(op_ptr));117            var declaration = try ir.interfaces.effects.inspect(op.allocator, @constCast(op));118            defer declaration.deinit(op.allocator);119            if (!ir.interfaces.effects.repeatableExpression(declaration.facts)) return;120            if (op.operands.items.len != 1) return;121            if (op.results.items.len != 1) return;122            const input = op.operands.items[0].value;123            if (!op.results.items[0].type.eql(input.type)) return;124            try appendFoldValue(results, input);125        }126127        pub fn foldSelect(128            op_ptr: *const anyopaque,129            results: *ir.interfaces.FoldResults,130        ) anyerror!void {131            const op: *const ir.Operation = @ptrCast(@alignCast(op_ptr));132            var declaration = try ir.interfaces.effects.inspect(op.allocator, @constCast(op));133            defer declaration.deinit(op.allocator);134            if (!ir.interfaces.effects.repeatableExpression(declaration.facts)) return;135            if (op.operands.items.len != 3) return;136            if (op.results.items.len != 1) return;137            const true_value = op.operands.items[1].value;138            const false_value = op.operands.items[2].value;139            if (true_value != false_value) return;140            if (!op.results.items[0].type.eql(true_value.type)) return;141            try appendFoldValue(results, true_value);142        }143144        pub fn foldAdd(145            op_ptr: *const anyopaque,146            results: *ir.interfaces.FoldResults,147        ) anyerror!void {148            const op: *const ir.Operation = @ptrCast(@alignCast(op_ptr));149            var declaration = try ir.interfaces.effects.inspect(op.allocator, @constCast(op));150            defer declaration.deinit(op.allocator);151            if (!ir.interfaces.effects.repeatableExpression(declaration.facts)) return;152            const operands = foldBinaryOperands(op) orelse return;153            if (constantIntEquals(operands.rhs, 0)) return appendFoldValueForResult(op, results, operands.lhs);154            if (constantIntEquals(operands.lhs, 0)) return appendFoldValueForResult(op, results, operands.rhs);155        }156157        pub fn foldSub(158            op_ptr: *const anyopaque,159            results: *ir.interfaces.FoldResults,160        ) anyerror!void {161            const op: *const ir.Operation = @ptrCast(@alignCast(op_ptr));162            var declaration = try ir.interfaces.effects.inspect(op.allocator, @constCast(op));163            defer declaration.deinit(op.allocator);164            if (!ir.interfaces.effects.repeatableExpression(declaration.facts)) return;165            const operands = foldBinaryOperands(op) orelse return;166            if (constantIntEquals(operands.rhs, 0)) return appendFoldValueForResult(op, results, operands.lhs);167            if (operands.lhs == operands.rhs) return appendIntFoldAttributeForResult(op, results, 0);168        }169170        pub fn foldMul(171            op_ptr: *const anyopaque,172            results: *ir.interfaces.FoldResults,173        ) anyerror!void {174            const op: *const ir.Operation = @ptrCast(@alignCast(op_ptr));175            var declaration = try ir.interfaces.effects.inspect(op.allocator, @constCast(op));176            defer declaration.deinit(op.allocator);177            if (!ir.interfaces.effects.repeatableExpression(declaration.facts)) return;178            const operands = foldBinaryOperands(op) orelse return;179            if (constantIntEquals(operands.rhs, 0)) return appendFoldValueForResult(op, results, operands.rhs);180            if (constantIntEquals(operands.lhs, 0)) return appendFoldValueForResult(op, results, operands.lhs);181            if (constantIntEquals(operands.rhs, 1)) return appendFoldValueForResult(op, results, operands.lhs);182            if (constantIntEquals(operands.lhs, 1)) return appendFoldValueForResult(op, results, operands.rhs);183        }184185        pub fn foldDiv(186            op_ptr: *const anyopaque,187            results: *ir.interfaces.FoldResults,188        ) anyerror!void {189            const op: *const ir.Operation = @ptrCast(@alignCast(op_ptr));190            var declaration = try ir.interfaces.effects.inspect(op.allocator, @constCast(op));191            defer declaration.deinit(op.allocator);192            if (!ir.interfaces.effects.repeatableExpression(declaration.facts)) return;193            const operands = foldBinaryOperands(op) orelse return;194            if (constantIntEquals(operands.rhs, 1)) return appendFoldValueForResult(op, results, operands.lhs);195        }196197        pub fn foldAnd(198            op_ptr: *const anyopaque,199            results: *ir.interfaces.FoldResults,200        ) anyerror!void {201            const op: *const ir.Operation = @ptrCast(@alignCast(op_ptr));202            var declaration = try ir.interfaces.effects.inspect(op.allocator, @constCast(op));203            defer declaration.deinit(op.allocator);204            if (!ir.interfaces.effects.repeatableExpression(declaration.facts)) return;205            const operands = foldBinaryOperands(op) orelse return;206            if (operands.lhs == operands.rhs) return appendFoldValueForResult(op, results, operands.lhs);207            if (constantBoolEquals(operands.rhs, false) or constantIntEquals(operands.rhs, 0)) {208                return appendFoldValueForResult(op, results, operands.rhs);209            }210            if (constantBoolEquals(operands.lhs, false) or constantIntEquals(operands.lhs, 0)) {211                return appendFoldValueForResult(op, results, operands.lhs);212            }213            if (constantBoolEquals(operands.rhs, true) or constantIntEquals(operands.rhs, -1)) {214                return appendFoldValueForResult(op, results, operands.lhs);215            }216            if (constantBoolEquals(operands.lhs, true) or constantIntEquals(operands.lhs, -1)) {217                return appendFoldValueForResult(op, results, operands.rhs);218            }219        }220221        pub fn foldOr(222            op_ptr: *const anyopaque,223            results: *ir.interfaces.FoldResults,224        ) anyerror!void {225            const op: *const ir.Operation = @ptrCast(@alignCast(op_ptr));226            var declaration = try ir.interfaces.effects.inspect(op.allocator, @constCast(op));227            defer declaration.deinit(op.allocator);228            if (!ir.interfaces.effects.repeatableExpression(declaration.facts)) return;229            const operands = foldBinaryOperands(op) orelse return;230            if (operands.lhs == operands.rhs) return appendFoldValueForResult(op, results, operands.lhs);231            if (constantBoolEquals(operands.rhs, true) or constantIntEquals(operands.rhs, -1)) {232                return appendFoldValueForResult(op, results, operands.rhs);233            }234            if (constantBoolEquals(operands.lhs, true) or constantIntEquals(operands.lhs, -1)) {235                return appendFoldValueForResult(op, results, operands.lhs);236            }237            if (constantBoolEquals(operands.rhs, false) or constantIntEquals(operands.rhs, 0)) {238                return appendFoldValueForResult(op, results, operands.lhs);239            }240            if (constantBoolEquals(operands.lhs, false) or constantIntEquals(operands.lhs, 0)) {241                return appendFoldValueForResult(op, results, operands.rhs);242            }243        }244245        pub fn foldXor(246            op_ptr: *const anyopaque,247            results: *ir.interfaces.FoldResults,248        ) anyerror!void {249            const op: *const ir.Operation = @ptrCast(@alignCast(op_ptr));250            var declaration = try ir.interfaces.effects.inspect(op.allocator, @constCast(op));251            defer declaration.deinit(op.allocator);252            if (!ir.interfaces.effects.repeatableExpression(declaration.facts)) return;253            const operands = foldBinaryOperands(op) orelse return;254            if (operands.lhs == operands.rhs) {255                if (isBoolType(op.results.items[0].type)) return appendBoolFoldAttributeForResult(op, results, false);256                return appendIntFoldAttributeForResult(op, results, 0);257            }258            if (constantBoolEquals(operands.rhs, false) or constantIntEquals(operands.rhs, 0)) {259                return appendFoldValueForResult(op, results, operands.lhs);260            }261            if (constantBoolEquals(operands.lhs, false) or constantIntEquals(operands.lhs, 0)) {262                return appendFoldValueForResult(op, results, operands.rhs);263            }264        }265266        pub fn foldShift(267            op_ptr: *const anyopaque,268            results: *ir.interfaces.FoldResults,269        ) anyerror!void {270            const op: *const ir.Operation = @ptrCast(@alignCast(op_ptr));271            var declaration = try ir.interfaces.effects.inspect(op.allocator, @constCast(op));272            defer declaration.deinit(op.allocator);273            if (!ir.interfaces.effects.repeatableExpression(declaration.facts)) return;274            const operands = foldBinaryOperands(op) orelse return;275            if (constantIntEquals(operands.rhs, 0)) return appendFoldValueForResult(op, results, operands.lhs);276        }277278        pub fn foldNot(279            op_ptr: *const anyopaque,280            results: *ir.interfaces.FoldResults,281        ) anyerror!void {282            const op: *const ir.Operation = @ptrCast(@alignCast(op_ptr));283            var declaration = try ir.interfaces.effects.inspect(op.allocator, @constCast(op));284            defer declaration.deinit(op.allocator);285            if (!ir.interfaces.effects.repeatableExpression(declaration.facts)) return;286            if (op.operands.items.len != 1) return;287            if (op.results.items.len != 1) return;288            const value = constantBoolFromValue(op.operands.items[0].value) orelse return;289            try appendBoolFoldAttributeForResult(op, results, !value);290        }291292        fn cmpSelfResult(predicate_value: CmpPredicate, ty: ir.Type) ?bool {293            if (isBoolType(ty)) {294                return switch (predicate_value) {295                    .eq => true,296                    .ne => false,297                    else => null,298                };299            }300            if (!isIntegerLikeType(ty)) return null;301            return switch (predicate_value) {302                .eq, .le, .ge, .sle, .sge, .ule, .uge => true,303                .ne, .lt, .gt, .slt, .sgt, .ult, .ugt => false,304            };305        }306307        pub fn foldCmp(308            op_ptr: *const anyopaque,309            results: *ir.interfaces.FoldResults,310        ) anyerror!void {311            const op: *const ir.Operation = @ptrCast(@alignCast(op_ptr));312            var declaration = try ir.interfaces.effects.inspect(op.allocator, @constCast(op));313            defer declaration.deinit(op.allocator);314            if (!ir.interfaces.effects.repeatableExpression(declaration.facts)) return;315            const operands = foldBinaryOperands(op) orelse return;316            if (operands.lhs != operands.rhs) return;317            const cmp = Dialect.CmpOp{ .op = @constCast(op) };318            const predicate = cmp.getPredicate() orelse return;319            const folded = cmpSelfResult(predicate, operands.lhs.type) orelse return;320            try appendBoolFoldAttributeForResult(op, results, folded);321        }322    };323}

1pub const folds = @import("folds.zig");