kumi 0.0.31 → 0.0.33

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

@@ -23,58 +23,99 @@ module Kumi
 
             def run(_errors)
               ops_by_decl = get_state(:lir_module)
-              # out = {}
+              out = {}
 
-              # ops_by_decl.each do |name, payload|
-              #   out[name] = { operations: optimize_decl(Array(payload[:operations])) }
-              # end
+              ops_by_decl.each do |name, payload|
+                operations = Array(payload[:operations])
+                out[name] = { operations: optimize_decl(operations) }
+              end
 
-              # TODO: - make it work with new gather/length
-              # out.freeze
-              out = ops_by_decl
-              state.with(:lir_06_const_prop, out).with(:lir_module, out.freeze)
+              out.freeze
+              state.with(:lir_module, out).with(:lir_06_const_prop, out)
             end
 
             private
 
             def optimize_decl(ops)
-              known_constants = {}
-              new_ops = []
-
-              ops.each do |ins|
-                # --- START OF THE FIX ---
-
-                new_inputs = []
-                new_immediates = []
-
-                # Process inputs, replacing known constants with a placeholder.
-                Array(ins.inputs).each do |input_reg|
-                  if (literal = known_constants[input_reg])
-                    new_inputs << :__immediate_placeholder__
-                    new_immediates << literal
-                  else
-                    new_inputs << input_reg
+              rewritten = ops.map(&:dup)
+              use_counts = Hash.new(0)
+
+              rewritten.each do |ins|
+                Array(ins.inputs).each do |reg|
+                  use_counts[reg] += 1 if reg
+                end
+              end
+
+              constants = {}
+
+              LIR::Peephole.run(rewritten) do |window|
+                ins = window.current
+                break unless ins
+
+                substitution = substitute_inputs(ins, constants)
+
+                if substitution
+                  new_ins, replaced_regs = substitution
+                  window.replace(1, with: new_ins)
+                  replaced_regs.each do |reg|
+                    use_counts[reg] -= 1
                   end
                 end
 
-                # Append any existing immediates after the new ones.
-                new_immediates.concat(Array(ins.immediates))
+                register = window.result_register
+                constants.delete(register) if register
+
+                if register && window.const?
+                  literal = window.literal
+                  constants[register] = literal if literal
+                end
+
+                window.skip
+              end
 
-                modified_ins = ins.dup
-                modified_ins.inputs = new_inputs
-                modified_ins.immediates = new_immediates
-                new_ops << modified_ins
+              rewritten.reject! do |ins|
+                reg = ins.result_register
+                ins.opcode == :Constant && reg && use_counts[reg].zero?
+              end
 
-                # --- END OF THE FIX ---
+              rewritten
+            end
 
-                # Update the map based on the MODIFIED instruction.
-                known_constants.delete(modified_ins.result_register) if modified_ins.result_register
-                if modified_ins.opcode == :Constant && modified_ins.immediates&.first
-                  known_constants[modified_ins.result_register] = modified_ins.immediates.first
+            def substitute_inputs(ins, constants)
+              original_inputs = Array(ins.inputs)
+              original_imms   = Array(ins.immediates)
+
+              new_inputs = []
+              new_imms   = []
+              changed    = false
+              replaced_regs = []
+
+              original_inputs.each do |reg|
+                if (literal = constants[reg])
+                  new_inputs << :__immediate_placeholder__
+                  new_imms << literal
+                  changed = true
+                  replaced_regs << reg
+                else
+                  new_inputs << reg
                 end
               end
 
-              new_ops
+              new_imms.concat(original_imms)
+              changed ||= new_imms != original_imms
+
+              return unless changed
+
+              instruction = LIR::Instruction.new(
+                opcode: ins.opcode,
+                result_register: ins.result_register,
+                stamp: ins.stamp,
+                inputs: new_inputs.empty? ? [] : new_inputs,
+                immediates: new_imms.empty? ? [] : new_imms,
+                attributes: ins.attributes,
+                location: ins.location
+              )
+              [instruction, replaced_regs]
             end
           end
         end

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

@@ -75,16 +75,17 @@ module Kumi
 
             def lower_expr(node)
               case node
-              when NAST::Const    then emit_const(node)
-              when NAST::InputRef then emit_input_ref(node)
-              when NAST::Ref      then emit_ref(node)
-              when NAST::Tuple    then emit_tuple(node)
-              when NAST::Select   then emit_select(node)
-              when NAST::Fold     then emit_fold(node)
-              when NAST::Reduce   then emit_reduce(node)
-              when NAST::Call     then call_emit_selection(node)
-              when NAST::Hash     then emit_hash(node)
-              when NAST::IndexRef then emit_index_ref(node)
+              when NAST::Const       then emit_const(node)
+              when NAST::InputRef    then emit_input_ref(node)
+              when NAST::Ref         then emit_ref(node)
+              when NAST::Tuple       then emit_tuple(node)
+              when NAST::Select      then emit_select(node)
+              when NAST::Fold        then emit_fold(node)
+              when NAST::Reduce      then emit_reduce(node)
+              when NAST::Call        then call_emit_selection(node)
+              when NAST::Hash        then emit_hash(node)
+              when NAST::IndexRef    then emit_index_ref(node)
+              when NAST::ImportCall  then emit_import_call(node)
               else raise "unknown node #{node.class}"
               end
             end
@@ -139,6 +140,20 @@ module Kumi
               ins.result_register
             end
 
+            def emit_import_call(n)
+              regs = n.args.map { lower_expr(_1) }
+              ins = Build.import_schema_call(
+                fn_name: n.fn_name,
+                source_module: n.source_module,
+                args: regs,
+                input_mapping_keys: n.input_mapping_keys,
+                out_dtype: dtype_of(n),
+                ids: @ids
+              )
+              @ops << ins
+              ins.result_register
+            end
+
             def emit_select(n)
               ax = axes_of(n)
               ensure_context_for!(ax, anchor_fqn: anchor_fqn_from_node!(n, need_prefix: ax)) unless ax.empty?
@@ -321,7 +336,7 @@ module Kumi
                   walk.call(x.on_false)
                 when NAST::Reduce, NAST::Fold
                   walk.call(x.arg)
-                when NAST::Call, NAST::Tuple
+                when NAST::Call, NAST::Tuple, NAST::ImportCall
                   x.args.each { walk.call(_1) }
                 when NAST::Hash
                   x.pairs.each { walk.call(_1) }

data/lib/kumi/core/analyzer/passes/name_indexer.rb

@@ -8,19 +8,37 @@ module Kumi
         # DEPENDENCIES: None (first pass in pipeline)
         # PRODUCES: :declarations - Hash mapping names to declaration nodes
         #            - annotates hints to declarations (e.g. inlining)
+        #           :imported_declarations - Hash of lazy import references
         # INTERFACE: new(schema, state).run(errors)
         class NameIndexer < PassBase
           def run(errors)
             definitions = {}
+            imported_declarations = {}
             hints = {}
 
+            # Phase 1: Register imports as lazy references
+            (schema.imports || []).each do |import_decl|
+              import_decl.names.each do |name|
+                imported_declarations[name] = {
+                  type: :import,
+                  from_module: import_decl.module_ref,
+                  loc: import_decl.loc
+                }
+              end
+            end
+
+            # Phase 2: Index local declarations
             each_decl do |decl|
-              report_error(errors, "duplicated definition `#{decl.name}`", location: decl.loc) if definitions.key?(decl.name)
+              if definitions.key?(decl.name) || imported_declarations.key?(decl.name)
+                report_error(errors, "duplicated definition `#{decl.name}`", location: decl.loc)
+              end
               definitions[decl.name] = decl
               hints[decl.name] = decl.hints
             end
 
-            state.with(:declarations, definitions.freeze).with(:hints, hints)
+            state.with(:declarations, definitions.freeze)
+                 .with(:imported_declarations, imported_declarations.freeze)
+                 .with(:hints, hints)
           end
         end
       end

data/lib/kumi/core/analyzer/passes/nast_dimensional_analyzer_pass.rb

@@ -54,6 +54,7 @@ module Kumi
           def analyze_expression(expr, errors)
             case expr
             when Kumi::Core::NAST::Call         then analyze_call_expression(expr, errors)
+            when Kumi::Core::NAST::ImportCall   then analyze_import_call(expr, errors)
             when Kumi::Core::NAST::Tuple        then analyze_tuple(expr, errors)
             when Kumi::Core::NAST::InputRef     then analyze_input_ref(expr)
             when Kumi::Core::NAST::IndexRef     then analyze_index_ref(expr, errors)
@@ -163,6 +164,49 @@ module Kumi
             { type: result_type, scope: result_scope }
           end
 
+          def analyze_import_call(call, errors)
+            # Analyze arguments
+            arg_metadata = call.args.map { |arg| analyze_expression(arg, errors) }
+            arg_types = arg_metadata.map { |m| m[:type] }
+            arg_scopes = arg_metadata.map { |m| m[:scope] }
+
+            # Get the imported schemas from state
+            imported_schemas = get_state(:imported_schemas, required: false) || {}
+            import_meta = imported_schemas[call.fn_name]
+
+            unless import_meta
+              report_error(errors, "imported function `#{call.fn_name}` not found", location: call.loc)
+              return { type: Types.scalar(:any), scope: [] }
+            end
+
+            # Get the analyzed state of the source schema
+            analyzed_state = import_meta[:analyzed_state]
+            src_declaration_table = analyzed_state[:declaration_table] || {}
+
+            # Look up the imported declaration in the source schema
+            src_decl_meta = src_declaration_table[call.fn_name]
+            unless src_decl_meta
+              report_error(errors, "declaration `#{call.fn_name}` not found in imported schema", location: call.loc)
+              return { type: Types.scalar(:any), scope: [] }
+            end
+
+            result_type = src_decl_meta[:result_type] || Types.scalar(:any)
+            # ImportCall broadcasts over argument scopes (like elementwise functions)
+            result_scope = lub_by_prefix(arg_scopes)
+
+            @metadata_table[node_id(call)] = {
+              kind: :import_call,
+              result_type: result_type,
+              result_scope: result_scope,
+              imported_fn: call.fn_name,
+              arg_types: arg_types,
+              arg_scopes: arg_scopes
+            }.freeze
+
+            debug "    ImportCall #{call.fn_name}: -> #{result_type} in #{result_scope.inspect}"
+            { type: result_type, scope: result_scope }
+          end
+
           def analyze_tuple(node, errors)
             elems           = node.args.map { |e| analyze_expression(e, errors) }
             element_types   = elems.map { |m| m[:type] }

data/lib/kumi/core/analyzer/passes/normalize_to_nast_pass.rb

@@ -42,6 +42,9 @@ module Kumi
             when Kumi::Syntax::DeclarationReference
               NAST::Ref.new(name: node.name, loc: node.loc)
 
+            when Kumi::Syntax::ImportCall
+              normalize_import_call(node, errors)
+
             when Kumi::Syntax::CallExpression
               # Special handling section - very clear what's happening
               case node.fn_name
@@ -158,6 +161,33 @@ module Kumi
               NAST::Call.new(fn: :"core.and", args: [left, right], loc: loc)
             end
           end
+
+          def normalize_import_call(node, errors)
+            imported_schemas = get_state(:imported_schemas, required: false) || {}
+
+            import_meta = imported_schemas[node.fn_name]
+            unless import_meta
+              add_error(errors, node.loc, "imported function `#{node.fn_name}` not found in imported_schemas")
+              return NAST::Const.new(value: nil, loc: node.loc)
+            end
+
+            # Don't inline the source expression. Instead, create an NAST::ImportCall node
+            # that represents a call to the compiled schema function.
+            # The compiled schema handles its own internal dependencies.
+
+            # Normalize the arguments (the values being passed)
+            args = node.input_mapping.map do |param_name, caller_expr|
+              normalize_expr(caller_expr, errors)
+            end
+
+            NAST::ImportCall.new(
+              fn_name: node.fn_name,
+              args: args,
+              input_mapping_keys: node.input_mapping.keys,
+              source_module: import_meta[:source_module],
+              loc: node.loc
+            )
+          end
         end
       end
     end

data/lib/kumi/core/analyzer/passes/semantic_constraint_validator.rb

@@ -89,7 +89,11 @@ module Kumi
             skip = [:cascade_and] # TODO: - hack
             return if skip.include? fn_name
 
-            # binding.pry
+            # Skip validation for imported functions - they're pure schema methods
+            imported_schemas = state[:imported_schemas] || {}
+            return if imported_schemas.key?(fn_name)
+
+            # Check if it's a built-in function in the registry
             return if @registry.resolve_function(fn_name)
 
             report_error(

data/lib/kumi/core/analyzer/passes/snast_pass.rb

@@ -180,6 +180,21 @@ module Kumi
             stamp!(out, m[:result_scope], m[:result_type])
           end
 
+          def visit_import_call(n)
+            args = n.args.map { _1.accept(self) }
+            m = meta_for(n)
+            out = n.class.new(
+              id: n.id,
+              fn_name: n.fn_name,
+              args: args,
+              input_mapping_keys: n.input_mapping_keys,
+              source_module: n.source_module,
+              loc: n.loc,
+              meta: n.meta.dup
+            )
+            stamp!(out, m[:result_scope], m[:result_type])
+          end
+
           # ---------- Helpers ----------
 
           def stamp!(node, axes, dtype)

data/lib/kumi/core/lir/build.rb

@@ -146,6 +146,33 @@ module Kumi
           )
         end
 
+        # ImportSchemaCall
+        # Params:
+        #   fn_name:           imported function name
+        #   source_module:     the module containing the imported schema
+        #   args:              argument registers (mapped parameter values)
+        #   input_mapping_keys:keys of input mapping in order of args
+        #   out_dtype:         dtype of the result
+        #   as:                result register
+        #   location:          optional Location
+        # Result: produces
+        def import_schema_call(fn_name:, source_module:, args:, input_mapping_keys:, out_dtype:, as: nil, ids: nil, location: nil)
+          as ||= ids.generate_temp
+          Instruction.new(
+            opcode: :ImportSchemaCall,
+            result_register: as,
+            stamp: Stamp.new(dtype: out_dtype),
+            inputs: args,
+            immediates: [],
+            attributes: {
+              fn_name: fn_name.to_s,
+              source_module: source_module.to_s,
+              input_mapping_keys: Array(input_mapping_keys).map(&:to_s)
+            },
+            location:
+          )
+        end
+
         # Fold
         # Params:
         #   function:   String function id (e.g., "core.mul")

data/lib/kumi/core/lir/peephole.rb

@@ -0,0 +1,164 @@
+# frozen_string_literal: true
+
+module Kumi
+  module Core
+    module LIR
+      # Peephole provides a small helper for implementing local LIR rewrites.
+      # It iterates over an instruction array and yields a Window object that
+      # exposes the current instruction alongside a handful of mutation helpers.
+      #
+      # Example:
+      #   Peephole.run(ops) do |window|
+      #     next unless window.match?(:Constant, :Constant, :KernelCall)
+      #
+      #     last = window.instruction(2)
+      #     const = Build.constant(
+      #       value: 3,
+      #       dtype: last.stamp.dtype,
+      #       as: last.result_register,
+      #       ids: ids
+      #     )
+      #     window.replace(3, with: const)
+      #   end
+      class Peephole
+        attr_reader :ops
+
+        def self.run(ops, &)
+          new(ops).run(&)
+        end
+
+        def initialize(ops)
+          @ops = ops
+        end
+
+        def run
+          index = 0
+          while index < @ops.length
+            window = Window.new(@ops, index)
+            yield window
+            index = window.next_index
+          end
+          @ops
+        end
+
+        # A mutable view over the instruction stream anchored at a given index.
+        class Window
+          attr_reader :index
+
+          def initialize(ops, index)
+            @ops = ops
+            @index = index
+            @next_index = index + 1
+          end
+
+          def current = instruction(0)
+
+          def instruction(offset = 0) = @ops[@index + offset]
+
+          def opcode(offset = 0)
+            instruction(offset)&.opcode
+          end
+
+          def instructions(count)
+            @ops[@index, count].compact
+          end
+
+          def match?(*opcodes)
+            opcodes.each_with_index.all? do |opcode, off|
+              ins = instruction(off)
+              ins&.respond_to?(:opcode) && ins.opcode == opcode
+            end
+          end
+
+          def result_register(offset = 0)
+            instruction(offset)&.result_register
+          end
+
+          def const?(offset = 0, value: nil)
+            ins = instruction(offset)
+            return false unless ins&.opcode == :Constant
+
+            return true if value.nil?
+
+            literal_value(offset) == value
+          end
+
+          def zero?(offset = 0)
+            const?(offset, value: 0)
+          end
+
+          def literal(offset = 0)
+            ins = instruction(offset)
+            Array(ins&.immediates).first
+          end
+
+          def literal_value(offset = 0)
+            literal(offset)&.value
+          end
+
+          def replace(count, with:)
+            replacements = normalize(with)
+            @ops[@index, count] = replacements
+            @next_index = @index
+            replacements
+          end
+
+          def delete(count = 1)
+            replace(count, with: [])
+          end
+
+          def insert_before(*new_ops)
+            new_ops = normalize(new_ops)
+            return @next_index = @index if new_ops.empty?
+
+            @ops.insert(@index, *new_ops)
+            @index += new_ops.length
+            @next_index = @index
+            new_ops
+          end
+
+          def insert_after(*new_ops)
+            new_ops = normalize(new_ops)
+            return @next_index = @index + 1 if new_ops.empty?
+
+            @ops.insert(@index + 1, *new_ops)
+            @next_index = @index + 1
+            new_ops
+          end
+
+          def skip(count = 1)
+            count = 1 if count.nil? || count < 1
+            @next_index = @index + count
+            nil
+          end
+
+          def rewind(count = 1)
+            count = 1 if count.nil? || count < 1
+            @next_index = [@index - count, 0].max
+            nil
+          end
+
+          def next_index
+            @next_index = 0 if @next_index.negative?
+            @next_index = @ops.length if @next_index > @ops.length
+            @next_index
+          end
+
+          def size
+            @ops.length
+          end
+
+          private
+
+          def normalize(value)
+            if value.is_a?(Array)
+              value.compact
+            else
+              value ? [value] : []
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/lib/kumi/core/nast.rb

@@ -101,6 +101,22 @@ module Kumi
         end
       end
 
+      class ImportCall < Node
+        attr_reader :fn_name, :args, :input_mapping_keys, :source_module
+
+        def initialize(fn_name:, args:, input_mapping_keys:, source_module:, **k)
+          super(**k)
+          @fn_name = fn_name.to_sym
+          @args = args
+          @input_mapping_keys = input_mapping_keys
+          @source_module = source_module
+        end
+
+        def accept(visitor)
+          visitor.respond_to?(:visit_import_call) ? visitor.visit_import_call(self) : super
+        end
+      end
+
       class Tuple < Node
         attr_reader :args
 

data/lib/kumi/core/ruby_parser/build_context.rb

@@ -4,13 +4,15 @@ module Kumi
   module Core
     module RubyParser
       class BuildContext
-        attr_reader :inputs, :values, :traits
+        attr_reader :inputs, :values, :traits, :imports, :imported_names
         attr_accessor :current_location
 
         def initialize
           @inputs = []
           @values = []
           @traits = []
+          @imports = []
+          @imported_names = Set.new
           @input_block_defined = false
         end
 

data/lib/kumi/core/ruby_parser/parser.rb

@@ -45,7 +45,7 @@ module Kumi
         end
 
         def build_syntax_tree
-          Root.new(@context.inputs, @context.values, @context.traits)
+          Root.new(@context.inputs, @context.values, @context.traits, @context.imports)
         end
 
         def handle_parse_error(error)

data/lib/kumi/core/ruby_parser/schema_builder.rb

@@ -8,7 +8,7 @@ module Kumi
         include Syntax
         include ErrorReporting
 
-        DSL_METHODS = %i[value trait input ref literal fn select shift roll].freeze
+        DSL_METHODS = %i[value trait input ref literal fn select shift roll import].freeze
 
         def initialize(context)
           @context = context
@@ -57,10 +57,31 @@ module Kumi
           Kumi::Syntax::Literal.new(value, loc: @context.current_location)
         end
 
+        def import(*names, from:)
+          update_location
+
+          unless from.is_a?(Module) || from.is_a?(Class)
+            raise_syntax_error(
+              "import 'from:' must reference a module or class",
+              location: @context.current_location
+            )
+          end
+
+          import_decl = Kumi::Syntax::ImportDeclaration.new(names, from, loc: @context.current_location)
+          @context.imports << import_decl
+          @context.imported_names.merge(names)
+        end
+
         def fn(fn_name, *args, **kwargs)
           update_location
-          expr_args = args.map { ensure_syntax(_1) }
-          Kumi::Syntax::CallExpression.new(fn_name, expr_args, kwargs, loc: @context.current_location)
+
+          if args.empty? && !kwargs.empty? && @context.imported_names.include?(fn_name)
+            mapping = build_import_mapping(kwargs)
+            Kumi::Syntax::ImportCall.new(fn_name, mapping, loc: @context.current_location)
+          else
+            expr_args = args.map { ensure_syntax(_1) }
+            Kumi::Syntax::CallExpression.new(fn_name, expr_args, kwargs, loc: @context.current_location)
+          end
         end
 
         def select(condition, value_when_true, value_when_false)
@@ -97,6 +118,15 @@ module Kumi
 
         private
 
+        def build_import_mapping(kwargs)
+          converter = ExpressionConverter.new(@context)
+          mapping = {}
+          kwargs.each do |field_name, expr|
+            mapping[field_name] = converter.ensure_syntax(expr)
+          end
+          mapping
+        end
+
         def update_location
           # Use caller_locations(2, 1) to skip the DSL method and get the actual user code location
           # Stack: [0] update_location, [1] DSL method (value/trait/etc), [2] user's DSL code