mt-lang 0.3.4 → 0.3.5

data/lib/milk_tea/core/lowering/async/lowering.rb

@@ -150,28 +150,16 @@ module MilkTea
               projection: :result_failure_error,
             )
           end
-          else_body = if statements_contain_await?(statement.else_body, async_info)
-            lower_async_cf_statements(
-              statement.else_body,
-              env: else_env,
-              frame_expr:,
-              raw_frame_expr:,
-              resume_linkage_name:,
-              async_info:,
-              active_defers:,
-              loop_flow:,
-            )
-          else
-            lower_async_non_await_statements(
-              statement.else_body,
-              env: else_env,
-              frame_expr:,
-              raw_frame_expr:,
-              async_info:,
-              active_defers:,
-              loop_flow:,
-            )
-          end
+          else_body = lower_async_body(
+            statement.else_body,
+            env: else_env,
+            frame_expr:,
+            raw_frame_expr:,
+            resume_linkage_name:,
+            async_info:,
+            active_defers:,
+            loop_flow:,
+          )
           lowered << IR::IfStmt.new(
             condition: let_else_failure_condition(target, storage_type),
             then_body: else_body,
@@ -214,6 +202,35 @@ module MilkTea
         end
       end
 
+      def lower_async_body(body, env:, frame_expr:, raw_frame_expr:, async_info:, resume_linkage_name: nil, active_defers: [], loop_flow: nil)
+        if statements_contain_await?(body, async_info)
+          lower_async_cf_statements(body, env:, frame_expr:, raw_frame_expr:, resume_linkage_name:, async_info:, active_defers:, loop_flow:)
+        else
+          lower_async_non_await_statements(body, env:, frame_expr:, raw_frame_expr:, async_info:, active_defers:, loop_flow:)
+        end
+      end
+
+      def build_async_while_with_condition_setup(condition_setup, condition, body, break_label)
+        if condition_setup.empty?
+          stmts = [IR::WhileStmt.new(condition:, body:)]
+          stmts << IR::LabelStmt.new(name: break_label) if contains_label_target?(body, break_label)
+          return stmts
+        end
+
+        loop_body = [
+          *condition_setup,
+          IR::IfStmt.new(
+            condition: IR::Unary.new(operator: "not", operand: condition, type: @ctx.types.fetch("bool")),
+            then_body: [loop_exit_statement(loop_exit_break(break_label), local_defers: [], outer_defers: [])],
+            else_body: nil,
+          ),
+          *body,
+        ]
+        stmts = [IR::WhileStmt.new(condition: IR::BooleanLiteral.new(value: true, type: @ctx.types.fetch("bool")), body: loop_body)]
+        stmts << IR::LabelStmt.new(name: break_label) if contains_label_target?(loop_body, break_label)
+        stmts
+      end
+
       # Lower a list of statements that MAY contain await expressions inside nested control flow.
       # CPS-via-goto: labels placed inside if/while/match bodies, reachable from top-level switch dispatch.
       def lower_async_cf_statements(statements, env:, frame_expr:, raw_frame_expr:, resume_linkage_name:, async_info:, active_defers: [], loop_flow: nil)
@@ -304,20 +321,12 @@ module MilkTea
         branch_entries = statement.branches.map do |branch|
           condition_setup, prepared_cond = prepare_expression_for_inline_lowering(branch.condition, env:)
           condition = lower_contextual_expression(prepared_cond, env:, expected_type: @ctx.types.fetch("bool"))
-          body = if statements_contain_await?(branch.body, async_info)
-            lower_async_cf_statements(branch.body, env:, frame_expr:, raw_frame_expr:, resume_linkage_name:, async_info:, active_defers:, loop_flow:)
-          else
-            lower_async_non_await_statements(branch.body, env:, frame_expr:, raw_frame_expr:, async_info:, active_defers:, loop_flow:)
-          end
+          body = lower_async_body(branch.body, env:, frame_expr:, raw_frame_expr:, resume_linkage_name:, async_info:, active_defers:, loop_flow:)
           { condition_setup:, condition:, body: }
         end
 
         else_body = if statement.else_body
-          if statements_contain_await?(statement.else_body, async_info)
-            lower_async_cf_statements(statement.else_body, env:, frame_expr:, raw_frame_expr:, resume_linkage_name:, async_info:, active_defers:, loop_flow:)
-          else
-            lower_async_non_await_statements(statement.else_body, env:, frame_expr:, raw_frame_expr:, async_info:, active_defers:, loop_flow:)
-          end
+          lower_async_body(statement.else_body, env:, frame_expr:, raw_frame_expr:, resume_linkage_name:, async_info:, active_defers:, loop_flow:)
         end
 
         nested_else = else_body
@@ -336,31 +345,10 @@ module MilkTea
         condition_setup, prepared_cond = prepare_expression_for_inline_lowering(statement.condition, env:)
         condition = lower_contextual_expression(prepared_cond, env:, expected_type: @ctx.types.fetch("bool"))
         inner_loop_flow = loop_flow(break_target: loop_exit_break(break_label), continue_target: loop_exit_continue(continue_label))
-        body = if statements_contain_await?(statement.body, async_info)
-          lower_async_cf_statements(statement.body, env: duplicate_env(env), frame_expr:, raw_frame_expr:, resume_linkage_name:, async_info:, active_defers:, loop_flow: inner_loop_flow)
-        else
-          lower_async_non_await_statements(statement.body, env: duplicate_env(env), frame_expr:, raw_frame_expr:, async_info:, active_defers:, loop_flow: inner_loop_flow)
-        end
+        body = lower_async_body(statement.body, env: duplicate_env(env), frame_expr:, raw_frame_expr:, resume_linkage_name:, async_info:, active_defers:, loop_flow: inner_loop_flow)
         body << IR::LabelStmt.new(name: continue_label) if contains_label_target?(body, continue_label)
 
-        if condition_setup.empty?
-          stmts = [IR::WhileStmt.new(condition:, body:)]
-          stmts << IR::LabelStmt.new(name: break_label) if contains_label_target?(body, break_label)
-          return stmts
-        end
-
-        loop_body = [
-          *condition_setup,
-          IR::IfStmt.new(
-            condition: IR::Unary.new(operator: "not", operand: condition, type: @ctx.types.fetch("bool")),
-            then_body: [loop_exit_statement(loop_exit_break(break_label), local_defers: [], outer_defers: [])],
-            else_body: nil,
-          ),
-          *body,
-        ]
-        stmts = [IR::WhileStmt.new(condition: IR::BooleanLiteral.new(value: true, type: @ctx.types.fetch("bool")), body: loop_body)]
-        stmts << IR::LabelStmt.new(name: break_label) if contains_label_target?(loop_body, break_label)
-        stmts
+        build_async_while_with_condition_setup(condition_setup, condition, body, break_label)
       end
 
       def lower_async_cf_for_stmt(statement, env:, frame_expr:, raw_frame_expr:, resume_linkage_name:, async_info:, active_defers:)
@@ -405,11 +393,7 @@ module MilkTea
         )
         inner_loop_flow = loop_flow(break_target: loop_exit_break(break_label), continue_target: loop_exit_continue(continue_label))
 
-        body = if statements_contain_await?(statement.body, async_info)
-          lower_async_cf_statements(statement.body, env: inner_env, frame_expr:, raw_frame_expr:, resume_linkage_name:, async_info:, active_defers:, loop_flow: inner_loop_flow)
-        else
-          lower_async_non_await_statements(statement.body, env: inner_env, frame_expr:, raw_frame_expr:, async_info:, active_defers:, loop_flow: inner_loop_flow)
-        end
+        body = lower_async_body(statement.body, env: inner_env, frame_expr:, raw_frame_expr:, resume_linkage_name:, async_info:, active_defers:, loop_flow: inner_loop_flow)
         body << IR::LabelStmt.new(name: continue_label) if contains_label_target?(body, continue_label)
 
         cmp_op = inclusive ? "<=" : "<"
@@ -461,11 +445,7 @@ module MilkTea
         inner_loop_flow = loop_flow(break_target: loop_exit_break(break_label), continue_target: loop_exit_continue(continue_label))
 
         assign_item = IR::Assignment.new(target: loop_var_expr, operator: "=", value: item_value)
-        body_stmts = if statements_contain_await?(statement.body, async_info)
-          lower_async_cf_statements(statement.body, env: inner_env, frame_expr:, raw_frame_expr:, resume_linkage_name:, async_info:, active_defers:, loop_flow: inner_loop_flow)
-        else
-          lower_async_non_await_statements(statement.body, env: inner_env, frame_expr:, raw_frame_expr:, async_info:, active_defers:, loop_flow: inner_loop_flow)
-        end
+        body_stmts = lower_async_body(statement.body, env: inner_env, frame_expr:, raw_frame_expr:, resume_linkage_name:, async_info:, active_defers:, loop_flow: inner_loop_flow)
         body_stmts << IR::LabelStmt.new(name: continue_label) if contains_label_target?(body_stmts, continue_label)
 
         stmts = [
@@ -533,11 +513,7 @@ module MilkTea
           IR::Assignment.new(target: binding_target, operator: "=", value: loop_item_value)
         end
         inner_loop_flow = loop_flow(break_target: loop_exit_break(break_label), continue_target: loop_exit_continue(continue_label))
-        body_stmts = if statements_contain_await?(statement.body, async_info)
-          lower_async_cf_statements(statement.body, env: inner_env, frame_expr:, raw_frame_expr:, resume_linkage_name:, async_info:, active_defers:, loop_flow: inner_loop_flow)
-        else
-          lower_async_non_await_statements(statement.body, env: inner_env, frame_expr:, raw_frame_expr:, async_info:, active_defers:, loop_flow: inner_loop_flow)
-        end
+        body_stmts = lower_async_body(statement.body, env: inner_env, frame_expr:, raw_frame_expr:, resume_linkage_name:, async_info:, active_defers:, loop_flow: inner_loop_flow)
         body_stmts << IR::LabelStmt.new(name: continue_label) if contains_label_target?(body_stmts, continue_label)
 
         length_checks = iterable_entries.drop(1).each_with_index.map do |entry, offset|
@@ -593,11 +569,7 @@ module MilkTea
           kind_expr = IR::Member.new(receiver: match_expr, member: "kind", type: kind_type)
           cases = statement.arms.map do |arm|
             arm_env, binding_decl = async_variant_match_arm_binding(arm, match_expr, match_type, env:, frame_expr:, local_fields: async_info[:local_fields])
-            arm_body = if statements_contain_await?(arm.body, async_info)
-                         lower_async_cf_statements(arm.body, env: arm_env, frame_expr:, raw_frame_expr:, resume_linkage_name:, async_info:, active_defers:, loop_flow: arm_loop_flow)
-                       else
-                         lower_async_non_await_statements(arm.body, env: arm_env, frame_expr:, raw_frame_expr:, async_info:, active_defers:, loop_flow: arm_loop_flow)
-                       end
+            arm_body = lower_async_body(arm.body, env: arm_env, frame_expr:, raw_frame_expr:, resume_linkage_name:, async_info:, active_defers:, loop_flow: arm_loop_flow)
             body = [binding_decl, *arm_body].compact + [IR::BreakStmt.new]
             if wildcard_arm_pattern?(arm.pattern)
               IR::SwitchDefaultCase.new(body: body)
@@ -611,11 +583,7 @@ module MilkTea
         end
 
         cases = statement.arms.map do |arm|
-          arm_body = if statements_contain_await?(arm.body, async_info)
-            lower_async_cf_statements(arm.body, env:, frame_expr:, raw_frame_expr:, resume_linkage_name:, async_info:, active_defers:, loop_flow: arm_loop_flow)
-          else
-            lower_async_non_await_statements(arm.body, env:, frame_expr:, raw_frame_expr:, async_info:, active_defers:, loop_flow: arm_loop_flow)
-          end
+          arm_body = lower_async_body(arm.body, env:, frame_expr:, raw_frame_expr:, resume_linkage_name:, async_info:, active_defers:, loop_flow: arm_loop_flow)
           if wildcard_arm_pattern?(arm.pattern)
             IR::SwitchDefaultCase.new(body: arm_body + [IR::BreakStmt.new])
           else
@@ -804,23 +772,7 @@ module MilkTea
         body << IR::LabelStmt.new(name: continue_label) if contains_label_target?(body, continue_label)
         cond = lower_expression(prepared_cond, env:, expected_type: @ctx.types.fetch("bool"))
 
-        if condition_setup.empty?
-          stmts = [IR::WhileStmt.new(condition: cond, body:)]
-          stmts << IR::LabelStmt.new(name: break_label) if contains_label_target?(body, break_label)
-          return IR::BlockStmt.new(body: stmts)
-        end
-
-        loop_body = [
-          *condition_setup,
-          IR::IfStmt.new(
-            condition: IR::Unary.new(operator: "not", operand: cond, type: @ctx.types.fetch("bool")),
-            then_body: [loop_exit_statement(loop_exit_break(break_label), local_defers: [], outer_defers: [])],
-            else_body: nil,
-          ),
-          *body,
-        ]
-        stmts = [IR::WhileStmt.new(condition: IR::BooleanLiteral.new(value: true, type: @ctx.types.fetch("bool")), body: loop_body)]
-        stmts << IR::LabelStmt.new(name: break_label) if contains_label_target?(loop_body, break_label)
+        stmts = build_async_while_with_condition_setup(condition_setup, cond, body, break_label)
         IR::BlockStmt.new(body: stmts)
       end
 
@@ -1242,28 +1194,16 @@ module MilkTea
                 projection: :result_failure_error,
               )
             end
-            else_body = if statements_contain_await?(statement.else_body, async_info)
-              lower_async_cf_statements(
-                statement.else_body,
-                env: else_env,
-                frame_expr:,
-                raw_frame_expr:,
-                resume_linkage_name:,
-                async_info:,
-                active_defers:,
-                loop_flow:,
-              )
-            else
-              lower_async_non_await_statements(
-                statement.else_body,
-                env: else_env,
-                frame_expr:,
-                raw_frame_expr:,
-                async_info:,
-                active_defers:,
-                loop_flow:,
-              )
-            end
+            else_body = lower_async_body(
+              statement.else_body,
+              env: else_env,
+              frame_expr:,
+              raw_frame_expr:,
+              resume_linkage_name:,
+              async_info:,
+              active_defers:,
+              loop_flow:,
+            )
             lowered << IR::IfStmt.new(
               condition: let_else_failure_condition(target, storage_type),
               then_body: else_body,
@@ -1315,28 +1255,16 @@ module MilkTea
           next [] if cleanup_entry[:body].empty?
 
           cleanup_env = duplicate_env(cleanup_entry[:env])
-          if statements_contain_await?(cleanup_entry[:body], async_info)
-            lower_async_cf_statements(
-              cleanup_entry[:body],
-              env: cleanup_env,
-              frame_expr:,
-              raw_frame_expr:,
-              resume_linkage_name: async_info.fetch(:resume_linkage_name),
-              async_info:,
-              active_defers: [],
-              loop_flow: nil,
-            )
-          else
-            lower_async_non_await_statements(
-              cleanup_entry[:body],
-              env: cleanup_env,
-              frame_expr:,
-              raw_frame_expr:,
-              async_info:,
-              active_defers: [],
-              loop_flow: nil,
-            )
-          end
+          lower_async_body(
+            cleanup_entry[:body],
+            env: cleanup_env,
+            frame_expr:,
+            raw_frame_expr:,
+            resume_linkage_name: async_info.fetch(:resume_linkage_name),
+            async_info:,
+            active_defers: [],
+            loop_flow: nil,
+          )
         end
       end
 
@@ -1403,5 +1331,83 @@ module MilkTea
         lowered << IR::ReturnStmt.new(value: nil)
         lowered
       end
+
+      def async_variant_match_arm_binding(arm, scrutinee_expr, scrutinee_type, env:, frame_expr: nil, local_fields: nil)
+        arm_env = duplicate_env(env)
+        binding_decl = nil
+
+        if arm.binding_name && !wildcard_arm_pattern?(arm.pattern)
+          arm_name = variant_match_arm_name_from_pattern(arm.pattern)
+          if arm_name && scrutinee_type.has_payload?(arm_name)
+            fields = scrutinee_type.arm(arm_name)
+            payload_type = Types::VariantArmPayload.new(scrutinee_type, arm_name, fields)
+
+            field_key = async_match_binding_field_key(arm)
+            field_info = local_fields&.fetch(field_key, nil)
+            if field_info && frame_expr
+              target = async_frame_field_expression(frame_expr, field_info[:field_name], field_info[:storage_type])
+              binding_c = async_frame_field_c_name(field_info[:field_name])
+              arm_env[:scopes].last[arm.binding_name] = local_binding(type: payload_type, linkage_name: binding_c, mutable: false, pointer: false)
+              data_expr = IR::Member.new(receiver: scrutinee_expr, member: "data", type: nil)
+              arm_expr = IR::Member.new(receiver: data_expr, member: arm_name, type: payload_type)
+              binding_decl = IR::Assignment.new(target:, operator: "=", value: arm_expr)
+            else
+              data_expr = IR::Member.new(receiver: scrutinee_expr, member: "data", type: nil)
+              arm_expr = IR::Member.new(receiver: data_expr, member: arm_name, type: payload_type)
+              binding_c = c_local_name(arm.binding_name)
+              arm_env[:scopes].last[arm.binding_name] = local_binding(type: payload_type, linkage_name: binding_c, mutable: false, pointer: false)
+              binding_decl = IR::LocalDecl.new(name: arm.binding_name, linkage_name: binding_c, type: payload_type, value: arm_expr)
+            end
+          end
+        end
+
+        [arm_env, binding_decl]
+      end
+
+      def bind_async_variant_match_arm_env!(arm_env, scrutinee_type, arm)
+        return unless scrutinee_type.is_a?(Types::Variant)
+        return unless arm.binding_name && !wildcard_arm_pattern?(arm.pattern)
+
+        arm_name = variant_match_arm_name_from_pattern(arm.pattern)
+        return unless arm_name && scrutinee_type.has_payload?(arm_name)
+
+        fields = scrutinee_type.arm(arm_name)
+        payload_type = Types::VariantArmPayload.new(scrutinee_type, arm_name, fields)
+        arm_env[:scopes].last[arm.binding_name] = local_binding(type: payload_type, linkage_name: c_local_name(arm.binding_name), mutable: true, pointer: false)
+      end
+
+      def lower_async_loop_exit(target, local_defers, outer_defers, frame_expr:, raw_frame_expr:, async_info:)
+        cleanup = lower_async_cleanup_entries(local_defers, outer_defers, frame_expr:, raw_frame_expr:, async_info:)
+        if cleanup.empty?
+          [loop_exit_statement(target, local_defers:, outer_defers:)]
+        else
+          label = target[:label]
+          raise LoweringError, "structured loop exits with cleanup are unsupported" unless label
+
+          cleanup + [IR::GotoStmt.new(label:)]
+        end
+      end
+
+      def async_local_decl_field_key(statement)
+        return "__discard_#{statement.line}" if statement.name == "_"
+
+        statement.name
+      end
+
+      def async_local_decl_field_name(statement)
+        return "local_discard_#{statement.line}" if statement.name == "_"
+
+        "local_#{statement.name}"
+      end
+
+      def async_match_binding_field_key(arm)
+        @async_binding_counter ||= 0
+        @async_binding_counter += 1
+        "match_binding_#{@async_binding_counter}"
+      end
+
+      def async_match_binding_field_name(arm)
+        "local_match_binding_#{@async_binding_counter}"
+      end
   end
 end