kumi 0.0.22 → 0.0.24

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,148 @@ 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 ---
             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)
-
-              # Hoisting is not allowed at the top level of a declaration
-              raise "Orphaned code was hoisted to top level" unless hoisted_ops.empty?
+              processed, hoisted = process_and_hoist_block(ops, env, reg_map, rename_map)
+              raise "Orphaned code was hoisted to top level" unless hoisted.empty?
 
-              processed_ops
+              processed
             end
 
-            # --- NEW: The core recursive block processor ---
-            # Returns two arrays: [ processed_instructions_for_this_block, instructions_to_hoist_up_one_level ]
+            # ---------------- core ----------------
+
+            # returns [processed_ops, hoisted_ops]
             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 = []
               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, hoisted_from_child =
+                    process_and_hoist_block(loop_body, env, reg_map, child_rename)
                   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)
+                  out.concat(hoisted_from_child)            # emit hoists before loop
+                  rename_map.merge!(child_rename)           # make child renames visible
 
-                  # 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)
+                  out << rewrite(ins, reg_map, rename_map)  # loop shell
+                  out.concat(processed_body)
+                  out << rewrite(block_ops[end_idx], reg_map, rename_map)
 
-                  i = end_idx # Jump iterator past the entire loop block
+                  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, hoist_ops =
+                    handle_load_declaration(ins, env, reg_map, rename_map)
+                  out.concat(inline_ops)
+                  hoisted.concat(hoist_ops)
 
                 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]
             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)
-              callee = ins.immediates.first.value.to_sym
-              decl_axes = ins.attributes.fetch(:axes)
+            # returns [inline_ops, hoisted_ops]
+            def handle_load_declaration(call_ins, env, reg_map, outer_rename_map)
+              raise "LoadDeclaration missing callee" unless call_ins.immediates&.first&.respond_to?(:value)
+
+              callee = call_ins.immediates.first.value.to_sym
+              raise "LoadDeclaration callee #{callee} not found" unless @ops_by_decl.key?(callee)
+
+              decl_axes = fetch_decl_axes(callee, call_ins)
               site_axes = env.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)
+              body, yield_reg, callee_axis_regs = inline_callee_core(callee)
+              axis_map = remap_axes(callee_axis_regs, env)
+
+              local_reg_map = {}
+              _acc, fresh_ops = freshen(body, local_reg_map, pre_map: axis_map)
+
+              nested_rename = {}
+              processed_inner, hoisted_inner =
+                process_and_hoist_block(fresh_ops, env, {}, nested_rename)
+
+              prelim = local_reg_map[yield_reg] || axis_map[yield_reg] || yield_reg
+              mapped_yield = resolve_rename(prelim, nested_rename)
+
+              # def/dominance guard
+              defs_inline  = defs_in(processed_inner)
+              defs_hoisted = defs_in(hoisted_inner)
+              unless defs_inline.include?(mapped_yield) || defs_hoisted.include?(mapped_yield)
+                msg = [
+                  "inliner: mapped yield #{mapped_yield} has no def in emitted ops for #{callee}",
+                  "  original yield: #{yield_reg}",
+                  "  prelim mapping: #{prelim}",
+                  "  nested_rename keys: #{nested_rename.keys.inspect}",
+                  "  inline defs size: #{defs_inline.size}",
+                  "  hoisted defs size: #{defs_hoisted.size}"
+                ].join("\n")
+                raise msg
+              end
 
-              rename_yielded_register(ins, yield_reg, reg_map, remap, rename_map)
+              outer_rename_map[call_ins.result_register] = mapped_yield if call_ins.result_register
 
-              # 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
+                [[], hoisted_inner + processed_inner]         # hoist
               elsif decl_axes == site_axes
-                [fresh_ops, []] # Return ops in the 'inline' bucket
-              # Case 3: Cannot inline
+                [hoisted_inner + processed_inner, []]         # inline in place
               else
-                [[rewrite(ins, reg_map, rename_map)], []]
+                [[rewrite(call_ins, reg_map, outer_rename_map)], []] # cannot inline
               end
             end
 
-            # --- UNCHANGED AND ORIGINAL HELPERS ---
+            # ---------------- helpers ----------------
+
+            def resolve_rename(reg, rename, limit: 64)
+              seen = {}
+              cur = reg
+              limit.times do
+                nxt = rename[cur]
+                break unless nxt
+                raise "inliner: rename cycle at #{cur}" if seen[nxt]
+
+                seen[nxt] = true
+                cur = nxt
+              end
+              cur
+            end
+
+            def defs_in(ops)
+              ops.each_with_object(Set.new) { |ins, s| s << ins.result_register if ins.result_register }
+            end
+
+            def fetch_decl_axes(callee, call_ins)
+              attr_axes  = call_ins.attributes && call_ins.attributes[:axes]
+              gamma_axes = Array(@gamma.fetch(callee)&.axes || [])
+              ax = attr_axes.nil? ? gamma_axes : Array(attr_axes)
+              raise "LoadDeclaration missing :axes" if ax.nil?
+              raise "inliner: non-array axes for #{callee}: #{ax.inspect}" unless ax.is_a?(Array)
+
+              ax
+            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 +187,52 @@ 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 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
+            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
@@ -195,29 +244,30 @@ module Kumi
               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]
+              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 +278,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/core/analyzer/passes/output_schema_pass.rb

@@ -9,15 +9,17 @@ module Kumi
             snast_module = get_state(:snast_module)
             return state unless snast_module
 
-            output_schema = build_output_schema(snast_module)
+            hints = get_state(:hints)
+            output_schema = build_output_schema(snast_module, hints)
             state.with(:output_schema, output_schema.freeze)
           end
 
           private
 
-          def build_output_schema(snast_module)
-            snast_module.decls.transform_values do |decl|
-              build_output_field(decl)
+          def build_output_schema(snast_module, hints)
+            snast_module.decls.each_with_object({}) do |(name, decl), acc|
+              next if hints[name][:inline]
+              acc[name] = build_output_field(decl)
             end
           end
 

data/lib/kumi/core/functions/function_spec.rb

@@ -1,10 +1,16 @@
 # frozen_string_literal: true
 
-# Minimal function specification for NAST dimensional analysis
-FunctionSpec = Struct.new(
-  :id,              # "core.add"
-  :kind,            # :elementwise, :reduce, :constructor
-  :parameter_names, # [:left_operand, :right_operand]
-  :dtype_rule,      # "promote(left_operand,right_operand)" or proc
-  keyword_init: true
-)
+module Kumi
+  module Core
+    module Functions
+      # Minimal function specification for NAST dimensional analysis
+      FunctionSpec = Struct.new(
+        :id,              # "core.add"
+        :kind,            # :elementwise, :reduce, :constructor
+        :parameter_names, # [:left_operand, :right_operand]
+        :dtype_rule,      # "promote(left_operand,right_operand)" or proc
+        keyword_init: true
+      )
+    end
+  end
+end

data/lib/kumi/core/functions/loader.rb

@@ -34,7 +34,7 @@ module Kumi
           parameter_names = (fn_hash["params"] || []).map { |p| p["name"].to_sym }
           dtype_rule_fn = TypeRules.compile_dtype_rule(fn_hash.fetch("dtype"), parameter_names)
 
-          FunctionSpec.new(
+          Functions::FunctionSpec.new(
             id: function_id,
             kind: function_kind,
             parameter_names: parameter_names,

data/lib/kumi/version.rb

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

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: kumi
 version: !ruby/object:Gem::Version
-  version: 0.0.22
+  version: 0.0.24
 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