kumi 0.0.23 → 0.0.25

data/lib/kumi/core/analyzer/passes/lir/inline_declarations_pass.rb

@@ -15,20 +15,18 @@ module Kumi
 
               MAX_PASSES.times do
                 new_ops, changed = run_one_pass(current_ops)
-
                 unless changed
                   new_ops.freeze
-                  return state.with(:lir_module, new_ops).with(:lir_02_inlined_ops_by_decl, new_ops)
+                  return state.with(:lir_module, new_ops)
+                              .with(:lir_02_inlined_ops_by_decl, new_ops)
                 end
                 current_ops = new_ops
               end
-
               raise "LIR inlining did not converge after #{MAX_PASSES} passes."
             end
 
             private
 
-            # --- UNCHANGED: Top-level pass logic ---
             def run_one_pass(ops_by_decl)
               @ops_by_decl = ops_by_decl
               @gamma       = detect_all_gammas(@ops_by_decl)
@@ -36,102 +34,187 @@ module Kumi
               fused = {}
               @ops_by_decl.each do |name, payload|
                 original_ops = Array(payload[:operations])
-
-                # Call the new top-level inliner
-                inlined_ops = inline_top_level_decl(original_ops)
-
-                fused[name] = { operations: inlined_ops }
+                inlined_ops  = inline_top_level_decl(original_ops)
+                fused[name]  = { operations: inlined_ops }
                 changed ||= (inlined_ops != original_ops)
               end
               [fused, changed]
             end
 
-            # --- NEW: Top-level entry point for the recursive processor ---
+            # ---------------- core ----------------
+
+            Hoist = Struct.new(:ops, :target_depth, keyword_init: true)
+
             def inline_top_level_decl(ops)
-              env = Env.new
-              reg_map = {}
+              env        = Env.new
+              reg_map    = {}
               rename_map = {}
-              processed_ops, hoisted_ops = process_and_hoist_block(ops, env, reg_map, rename_map)
+              processed, hoist_pkgs = process_and_hoist_block(ops, env, reg_map, rename_map)
 
-              # Hoisting is not allowed at the top level of a declaration
-              raise "Orphaned code was hoisted to top level" unless hoisted_ops.empty?
+              top_emit, bubble = hoist_pkgs.partition { |p| p.target_depth == 0 }
+              raise "Orphaned code hoist with target depth(s): #{bubble.map(&:target_depth).uniq.inspect}" unless bubble.empty?
 
-              processed_ops
+              top_emit.flat_map(&:ops) + processed
             end
 
-            # --- NEW: The core recursive block processor ---
-            # Returns two arrays: [ processed_instructions_for_this_block, instructions_to_hoist_up_one_level ]
+            # returns [processed_ops, hoist_pkgs]
+            # returns [processed_ops, hoist_pkgs]
             def process_and_hoist_block(block_ops, env, reg_map, rename_map)
-              out_ops = []
-              hoisted_out_ops = [] # Operations to be returned to the parent scope
+              out = []
+              hoisted_pkgs = []
               i = 0
-
               while i < block_ops.length
                 ins = block_ops[i]
                 case ins.opcode
                 when :LoopStart
-                  end_idx = find_matching_loop_end(block_ops, i)
+                  end_idx   = find_matching_loop_end(block_ops, i)
                   loop_body = block_ops[(i + 1)...end_idx]
 
                   env.push(ins)
-                  processed_body, hoisted_from_child = process_and_hoist_block(loop_body, env, reg_map, rename_map)
+                  child_rename = {}
+                  processed_body, child_hoists =
+                    process_and_hoist_block(loop_body, env, reg_map, child_rename)
+
+                  depth_here = env.axes.length
+                  child_el   = ins.attributes[:as_element]
+                  child_ix   = ins.attributes[:as_index]
+
+                  # Partition hoists: those that belong *inside* this loop vs bubble upward
+                  inside_pkgs, bubble_pkgs = child_hoists.partition { |p| p.target_depth == depth_here }
+
+                  # Renames: never let aliases to this loop's el/idx escape upward
+                  safe_pairs = child_rename.reject { |_, v| v == child_el || v == child_ix }
                   env.pop
 
-                  # This is the crucial step: emit code hoisted from the child loop
-                  # BEFORE emitting the child loop itself.
-                  out_ops.concat(hoisted_from_child)
+                  # Merge only safe renames into outer scope
+                  rename_map.merge!(safe_pairs)
+
+                  # Emit loop shell
+                  out << rewrite(ins, reg_map, rename_map)
+
+                  # Local view inside loop: apply both local and outer renames, with local taking precedence
+                  local_map = child_rename.merge(rename_map)
+
+                  # Emit hoists that belong at this depth *inside* the loop, before the body
+                  inside_ops = inside_pkgs.flat_map(&:ops)
+                  out.concat(rewrite_block(inside_ops, local_map))
+
+                  # Emit rewritten body
+                  out.concat(rewrite_block(processed_body, local_map))
+
+                  # Close loop
+                  out << rewrite(block_ops[end_idx], reg_map, rename_map)
 
-                  # Now, emit the reconstructed loop with its processed body
-                  out_ops << rewrite(ins, reg_map, rename_map)
-                  out_ops.concat(processed_body)
-                  out_ops << rewrite(block_ops[end_idx], reg_map, rename_map)
+                  # Bubble remaining hoists to outer scopes
+                  hoisted_pkgs.concat(bubble_pkgs)
 
-                  i = end_idx # Jump iterator past the entire loop block
+                  # Extra safety: ensure no aliases to this loop's el/idx remain in outer map
+                  rename_map.delete_if { |_, v| v == child_el || v == child_ix }
+
+                  i = end_idx
 
                 when :LoadDeclaration
-                  # This helper now decides if ops go inline or get returned for hoisting
-                  inline_ops, hoist_ops = handle_load_declaration(ins, env, reg_map, rename_map)
-                  out_ops.concat(inline_ops)
-                  hoisted_out_ops.concat(hoist_ops)
+                  inline_ops, new_pkgs = handle_load_declaration(ins, env, reg_map, rename_map)
+                  out.concat(inline_ops)
+                  hoisted_pkgs.concat(new_pkgs)
 
                 else
-                  out_ops << rewrite(ins, reg_map, rename_map)
+                  out << rewrite(ins, reg_map, rename_map)
                 end
                 i += 1
               end
-
-              [out_ops, hoisted_out_ops]
+              [out, hoisted_pkgs]
             end
 
-            # --- NEW: Helper to manage LoadDeclaration logic ---
-            # Returns two arrays: [ instructions_to_add_inline, instructions_to_hoist ]
-            def handle_load_declaration(ins, env, reg_map, rename_map)
+            # returns [inline_ops, hoist_pkgs]
+            def handle_load_declaration(ins, env, _reg_map, rename_map)
               callee = ins.immediates.first.value.to_sym
-              decl_axes = ins.attributes.fetch(:axes)
-              site_axes = env.axes
+
+              # axes presence and agreement with callee gamma
+              decl_axes  = ins.attributes.fetch(:axes) { raise "LoadDeclaration missing :axes for #{callee}" }
+              gamma_axes = @gamma.fetch(callee).axes
+              raise "axes mismatch for #{callee}: decl=#{decl_axes.inspect} gamma=#{gamma_axes.inspect}" unless decl_axes == gamma_axes
 
               body, yield_reg, callee_regs = inline_callee_core(callee)
               remap = remap_axes(callee_regs, env)
-              _acc, fresh_ops = freshen(body, reg_map, pre_map: remap)
 
-              rename_yielded_register(ins, yield_reg, reg_map, remap, rename_map)
+              # per-callsite freshening
+              local_reg_map = {}
+              _acc, fresh_ops = freshen(body, local_reg_map, pre_map: remap)
+
+              # recursively process nested calls
+              processed_inline, nested_pkgs = process_and_hoist_block(fresh_ops, env, {}, rename_map)
+
+              # compute yielded register mapping, then resolve through any renames created by nested inlines
+              mapped_yield =
+                local_reg_map[yield_reg] || remap[yield_reg] ||
+                (raise "inliner: yielded reg #{yield_reg} not produced in inlined body for #{callee}")
+              resolved_yield = resolve_rename(mapped_yield, rename_map)
+
+              # sanity: resolved_yield must be definable at site
+              emitted_defs = processed_inline.map(&:result_register).compact +
+                             nested_pkgs.flat_map { |p| p.ops }.map(&:result_register).compact
+              unless emitted_defs.include?(resolved_yield) || env.ambient_regs.include?(resolved_yield)
+                raise "inliner: mapped yield #{resolved_yield} has no def in emitted ops for #{callee}\n" \
+                      "original yield: #{yield_reg}\n" \
+                      "inline defs size: #{processed_inline.count { |x| x.result_register }}\n" \
+                      "nested hoist defs size: #{nested_pkgs.flat_map { |p| p.ops }.count { |x| x.result_register }}"
+              end
+
+              # final rename for call site result uses the resolved register
+              rename_map[ins.result_register] = resolved_yield
+
+              # decide placement by depth
+              site_depth   = env.axes.length
+              callee_depth = decl_axes.length
+
+              if callee_depth < site_depth
+                forb = forbidden_ambient_after(callee_depth, env)
+                used = uses_of(processed_inline)
+                bad  = used & forb
+                unless bad.empty?
+                  raise "scope error: would hoist ops using deeper-axis regs #{bad.inspect} " \
+                        "(callee_depth=#{callee_depth}, site_depth=#{site_depth})"
+                end
+                pkgs = nested_pkgs + [Hoist.new(ops: processed_inline, target_depth: callee_depth)]
+                [[], pkgs]
+
+              elsif callee_depth == site_depth
+                emit, bubble = nested_pkgs.partition { |p| p.target_depth == site_depth }
+                [(emit.flat_map(&:ops) + processed_inline), bubble]
 
-              # Case 1: Hoisting
-              if prefix?(decl_axes, site_axes) && decl_axes.length < site_axes.length
-                [[], fresh_ops] # Return ops in the 'hoist' bucket
-              # Case 2: In-place
-              elsif decl_axes == site_axes
-                [fresh_ops, []] # Return ops in the 'inline' bucket
-              # Case 3: Cannot inline
               else
-                [[rewrite(ins, reg_map, rename_map)], []]
+                [[rewrite(ins, {}, rename_map)], []]
               end
             end
 
-            # --- UNCHANGED AND ORIGINAL HELPERS ---
+            # ---------------- helpers ----------------
+            def rewrite_block(ops, rename)
+              # Ensure late-added renames apply to a block we built earlier.
+              ops.map { |ins| rewrite(ins, {}, rename) }
+            end
+
+            def resolve_rename(reg, rename)
+              seen = {}
+              cur = reg
+              while (n = rename[cur]) && !seen[n]
+                seen[cur] = true
+                cur = n
+              end
+              cur
+            end
+
+            def uses_of(ops)
+              ops.flat_map { |x| Array(x.inputs) }.compact
+            end
+
+            def forbidden_ambient_after(depth, env)
+              env.frames_after(depth).flat_map { |f| [f[:el], f[:idx]] }
+            end
 
             def find_matching_loop_end(ops, start_index)
-              depth = 1; (start_index + 1...ops.length).each do |i|
+              depth = 1
+              (start_index + 1...ops.length).each do |i|
                 op = ops[i].opcode
                 depth += 1 if op == :LoopStart
                 depth -= 1 if op == :LoopEnd
@@ -143,47 +226,57 @@ module Kumi
             def remap_axes(callee_axis_regs, env)
               callee_axis_regs.each_with_object({}) do |r, h|
                 caller = env.reg_for_axis(r[:axis])
-                h[r[:el]] = caller[:el]
+                h[r[:el]]  = caller[:el]
                 h[r[:idx]] = caller[:idx]
               end
             end
 
-            def rename_yielded_register(ins, yielded_reg, reg_map, axis_remap, rename)
-              return unless ins.result_register && yielded_reg
-
-              mapped = reg_map.fetch(yielded_reg, axis_remap.fetch(yielded_reg, yielded_reg))
-              rename[ins.result_register] = mapped
-            end
-
-            # (Your original `detect_all_gammas`, `detect_gamma`, `inline_callee_core`,
-            # `Env`, `freshen`, `rewrite`, and `prefix?` methods go here, unchanged)
             class Env
               def initialize = @frames = []
               def axes = @frames.map { _1[:axis] }
+              def ambient_regs = @frames.flat_map { |f| [f[:el], f[:idx]] }
 
               def push(loop_ins)
-                @frames << { axis: loop_ins.attributes[:axis], el: loop_ins.attributes[:as_element], idx: loop_ins.attributes[:as_index] }
+                @frames << {
+                  axis: loop_ins.attributes[:axis],
+                  el: loop_ins.attributes[:as_element],
+                  idx: loop_ins.attributes[:as_index]
+                }
               end
 
               def pop = @frames.pop
-              def reg_for_axis(axis) = @frames.reverse.find { _1[:axis] == axis } || raise("no element for #{axis}")
+
+              def reg_for_axis(axis)
+                @frames.reverse.find { _1[:axis] == axis } ||
+                  raise("no element for axis=#{axis.inspect}")
+              end
+
+              def frames_after(depth)
+                @frames[depth..] || []
+              end
+            end
+
+            def detect_all_gammas(ops_by_decl)
+              ops_by_decl.transform_values { |p| detect_gamma(Array(p[:operations])) }
             end
 
-            def detect_all_gammas(ops_by_decl) = ops_by_decl.transform_values { |p| detect_gamma(Array(p[:operations])) }
             GammaInfo = Struct.new(:start_idx, :axes, :axis_regs, keyword_init: true)
+
             def detect_gamma(ops)
               frames = []
               ops.each do |ins|
                 case ins.opcode
                 when :LoopStart
-                  frames << { axis: ins.attributes[:axis], el: ins.attributes[:as_element], idx: ins.attributes[:as_index] }
+                  frames << {
+                    axis: ins.attributes[:axis],
+                    el: ins.attributes[:as_element],
+                    idx: ins.attributes[:as_index]
+                  }
                 when :LoopEnd
                   frames.pop
                 when :Yield
                   axes = frames.map { _1[:axis] }
-                  axis_regs = frames.map do |f|
-                    { axis: f[:axis], el: f[:el], idx: f[:idx] }
-                  end
+                  axis_regs = frames.map { |f| { axis: f[:axis], el: f[:el], idx: f[:idx] } }
                   return GammaInfo.new(start_idx: nil, axes: axes, axis_regs: axis_regs)
                 end
               end
@@ -191,33 +284,34 @@ module Kumi
             end
 
             def inline_callee_core(callee_name)
-              ops = Array(@ops_by_decl.fetch(callee_name)[:operations])
+              ops  = Array(@ops_by_decl.fetch(callee_name)[:operations])
               info = @gamma.fetch(callee_name)
               axes = info.axes
-              k = axes.length
-              yield_index = ops.rindex { |ins| ins.opcode == :Yield } or raise "callee #{callee_name} has no Yield"
-              yielded_reg = Array(ops[yield_index].inputs).first
-              first_loop_index = ops.index { |ins| ins.opcode == :LoopStart }
-              return [ops[0...yield_index], yielded_reg, info.axis_regs] unless first_loop_index
-
-              prologue = ops[0...first_loop_index]
-              main_part = ops[first_loop_index...yield_index]
+              k    = axes.length
+
+              yi = ops.rindex { |x| x.opcode == :Yield } or raise "callee #{callee_name} has no Yield"
+              yielded_reg = Array(ops[yi].inputs).first
+
+              first_loop = ops.index { |x| x.opcode == :LoopStart }
+              return [ops[0...yi], yielded_reg, info.axis_regs] unless first_loop
+
+              prologue   = ops[0...first_loop]
+              main       = ops[first_loop...yi]
               inner_body = []
               open_gamma = 0
-              kind_stack = []
-              main_part.each do |ins|
+              stack = []
+              main.each do |ins|
                 case ins.opcode
                 when :LoopStart
                   if open_gamma < k && ins.attributes[:axis] == axes[open_gamma]
-                    kind_stack << :gamma
+                    stack << :gamma
                     open_gamma += 1
                   else
-                    kind_stack << :inner
+                    stack << :inner
                     inner_body << ins
                   end
                 when :LoopEnd
-                  kind = kind_stack.pop or raise "unbalanced loops in #{callee_name}"
-                  inner_body << ins if kind == :inner
+                  inner_body << ins if stack.pop == :inner
                 else
                   inner_body << ins
                 end
@@ -228,44 +322,70 @@ module Kumi
             def freshen(block_ops, reg_map, pre_map: {})
               acc_map = {}
               new_ops = block_ops.map do |ins|
+                attrs = (ins.attributes || {}).dup
+
+                if ins.opcode == :LoopStart
+                  attrs[:id] = @ids.generate_loop_id
+                  new_el  = @ids.generate_temp
+                  new_idx = @ids.generate_temp
+                  reg_map[attrs[:as_element]] = new_el
+                  reg_map[attrs[:as_index]]   = new_idx
+                  attrs[:as_element] = new_el
+                  attrs[:as_index]   = new_idx
+                end
+
                 res = ins.result_register
                 reg_map[res] ||= @ids.generate_temp if res
+
                 new_inputs = Array(ins.inputs).map do |r|
                   r1 = pre_map.fetch(r, r)
                   reg_map.fetch(r1, r1)
                 end
-                attrs = (ins.attributes || {}).dup
+
                 case ins.opcode
                 when :DeclareAccumulator
-                  original_acc = ins.result_register
-                  acc_map[original_acc] ||= @ids.generate_acc
-                  res = acc_map[original_acc]
+                  orig = ins.result_register
+                  acc_map[orig] ||= @ids.generate_acc
+                  res = acc_map[orig]
                 when :Accumulate
-                  original_acc = ins.result_register
-                  acc_map[original_acc] ||= @ids.generate_acc
-                  res = acc_map[original_acc]
+                  orig = ins.result_register
+                  acc_map[orig] ||= @ids.generate_acc
+                  res = acc_map[orig]
                 when :LoadAccumulator
-                  original_acc = ins.inputs.first
-                  acc_map[original_acc] ||= @ids.generate_acc
-                  new_inputs[0] = acc_map[original_acc]
-                when :LoopStart
-                  attrs[:id] = @ids.generate_loop_id
+                  orig = ins.inputs.first
+                  acc_map[orig] ||= @ids.generate_acc
+                  new_inputs[0] = acc_map[orig]
                 end
-                LIR::Instruction.new(opcode: ins.opcode, result_register: res ? reg_map.fetch(res, res) : nil, stamp: ins.stamp, inputs: new_inputs,
-                                     immediates: ins.immediates, attributes: attrs, location: ins.location)
+
+                LIR::Instruction.new(
+                  opcode: ins.opcode,
+                  result_register: res ? reg_map.fetch(res, res) : nil,
+                  stamp: ins.stamp,
+                  inputs: new_inputs,
+                  immediates: ins.immediates,
+                  attributes: attrs,
+                  location: ins.location
+                )
               end
               [acc_map, new_ops]
             end
 
             def rewrite(ins, _reg_map, rename)
-              new_inputs = Array(ins.inputs).map do |r|
-                rename.fetch(r, r)
-              end
-              LIR::Instruction.new(opcode: ins.opcode, result_register: ins.result_register, stamp: ins.stamp, inputs: new_inputs,
-                                   immediates: ins.immediates, attributes: ins.attributes, location: ins.location)
+              new_inputs = Array(ins.inputs).map { |r| rename.fetch(r, r) }
+              LIR::Instruction.new(
+                opcode: ins.opcode,
+                result_register: ins.result_register,
+                stamp: ins.stamp,
+                inputs: new_inputs,
+                immediates: ins.immediates,
+                attributes: ins.attributes,
+                location: ins.location
+              )
             end
 
-            def prefix?(pre, full) = pre.each_with_index.all? { |tok, i| full[i] == tok }
+            def prefix?(pre, full)
+              pre.each_with_index.all? { |tok, i| full[i] == tok }
+            end
           end
         end
       end

data/lib/kumi/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module Kumi
-  VERSION = "0.0.23"
+  VERSION = "0.0.25"
 end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: kumi
 version: !ruby/object:Gem::Version
-  version: 0.0.23
+  version: 0.0.25
 platform: ruby
 authors:
 - André Muta
@@ -265,6 +265,22 @@ files:
 - golden/hierarchical_complex/expected/snast.txt
 - golden/hierarchical_complex/input.json
 - golden/hierarchical_complex/schema.kumi
+- golden/inline_rename_scope_leak/expected.json
+- golden/inline_rename_scope_leak/expected/ast.txt
+- golden/inline_rename_scope_leak/expected/input_plan.txt
+- golden/inline_rename_scope_leak/expected/lir_00_unoptimized.txt
+- golden/inline_rename_scope_leak/expected/lir_01_hoist_scalar_references.txt
+- golden/inline_rename_scope_leak/expected/lir_02_inlined.txt
+- golden/inline_rename_scope_leak/expected/lir_03_cse.txt
+- golden/inline_rename_scope_leak/expected/lir_04_1_loop_fusion.txt
+- golden/inline_rename_scope_leak/expected/lir_04_loop_invcm.txt
+- golden/inline_rename_scope_leak/expected/lir_06_const_prop.txt
+- golden/inline_rename_scope_leak/expected/nast.txt
+- golden/inline_rename_scope_leak/expected/schema_javascript.mjs
+- golden/inline_rename_scope_leak/expected/schema_ruby.rb
+- golden/inline_rename_scope_leak/expected/snast.txt
+- golden/inline_rename_scope_leak/input.json
+- golden/inline_rename_scope_leak/schema.kumi
 - golden/input_reference/expected.json
 - golden/input_reference/expected/ast.txt
 - golden/input_reference/expected/input_plan.txt
@@ -521,6 +537,22 @@ files:
 - golden/tuples_and_arrays/expected/snast.txt
 - golden/tuples_and_arrays/input.json
 - golden/tuples_and_arrays/schema.kumi
+- golden/us_tax_2024/expected.json
+- golden/us_tax_2024/expected/ast.txt
+- golden/us_tax_2024/expected/input_plan.txt
+- golden/us_tax_2024/expected/lir_00_unoptimized.txt
+- golden/us_tax_2024/expected/lir_01_hoist_scalar_references.txt
+- golden/us_tax_2024/expected/lir_02_inlined.txt
+- golden/us_tax_2024/expected/lir_03_cse.txt
+- golden/us_tax_2024/expected/lir_04_1_loop_fusion.txt
+- golden/us_tax_2024/expected/lir_04_loop_invcm.txt
+- golden/us_tax_2024/expected/lir_06_const_prop.txt
+- golden/us_tax_2024/expected/nast.txt
+- golden/us_tax_2024/expected/schema_javascript.mjs
+- golden/us_tax_2024/expected/schema_ruby.rb
+- golden/us_tax_2024/expected/snast.txt
+- golden/us_tax_2024/input.json
+- golden/us_tax_2024/schema.kumi
 - golden/with_constants/expected/ast.txt
 - golden/with_constants/expected/input_plan.txt
 - golden/with_constants/expected/lir_00_unoptimized.txt