kumi 0.0.18 → 0.0.20

data/lib/kumi/frontends.rb

@@ -3,27 +3,25 @@
 module Kumi
   module Frontends
     def self.load(path:, inputs: {})
-      mode = (ENV["KUMI_PARSER"] || "auto") # auto|text|ruby
+      mode = ENV["KUMI_PARSER"] || "auto" # auto|text|ruby
       ext  = File.extname(path)
 
       # Explicit mode selection
       return Text.load(path:, inputs:) if mode == "text"
       return Ruby.load(path:, inputs:) if mode == "ruby"
-      
+
       # Auto mode: prefer .kumi if present
       if mode == "auto" && ext == ".rb"
         kumi_path = path.sub(/\.rb\z/, ".kumi")
-        if File.exist?(kumi_path)
-          return Text.load(path: kumi_path, inputs: inputs)
-        end
+        return Text.load(path: kumi_path, inputs: inputs) if File.exist?(kumi_path)
       end
-      
+
       # File extension based selection
       return Text.load(path:, inputs:) if ext == ".kumi"
       return Ruby.load(path:, inputs:) if ext == ".rb"
-      
+
       # Default fallback
       Ruby.load(path:, inputs:)
     end
   end
-end
+end

data/lib/kumi/kernel_registry.rb

@@ -0,0 +1,59 @@
+require "yaml"
+require "json"
+require "digest"
+
+module Kumi
+  class KernelRegistry
+    KERNEL_DIRS = {
+      ruby: "data/kernels/ruby"
+    }
+    Entry = Struct.new(:id, :fn, :impl, :identity, keyword_init: true)
+
+    def self.load_ruby
+      load_dir(KERNEL_DIRS[:ruby])
+    end
+
+    def self.load_dir(dir)
+      files = Dir.glob(File.join(dir, "**", "*.y{a,}ml")).sort
+      entries = files.flat_map { |p| (YAML.load_file(p) || {}).fetch("kernels", []) }
+                     .map { |h| Entry.new(id: h["id"], fn: h["fn"], impl: h["impl"], identity: h["identity"]) }
+      new(entries)
+    end
+
+    def initialize(entries)
+      @by_fn = Hash.new { |h, k| h[k] = [] }
+      @by_id = {}
+      entries.each { |e| (@by_fn[e.fn] << e) && (@by_id[e.id] = e) }
+      @by_fn.each { |fn, arr| raise "multiple kernels for #{fn}" if arr.size > 1 }
+    end
+
+    def pick(fn)
+      (@by_fn[fn]&.first || raise("no kernel for #{fn}")).id
+    end
+
+    def registry_ref
+      stable = { kernels: @by_id.values.map { |e| { "id" => e.id, "fn" => e.fn, "impl" => e.impl } }.sort_by { _1["id"] } }
+      "sha256:#{Digest::SHA256.hexdigest(JSON.generate(stable))}"
+    end
+
+    def identity(kernel_id, dtype)
+      e = @by_id[kernel_id] or raise "unknown kernel id #{kernel_id}"
+      identity_map = e.identity or raise "no identity map for #{kernel_id}"
+
+      (identity_map[dtype.to_s] || identity_map["any"]) or raise "no identity with dtype `#{dtype}` for #{kernel_id}"
+    end
+
+    # Late-bind the Ruby implementation
+    # Returns implementation structure for the kernel
+    def impl_for(kernel_id)
+      e = @by_id[kernel_id] or raise "unknown kernel id #{kernel_id}"
+      e.impl
+    end
+
+    private
+
+    def constantize(path)
+      path.split("::").inject(Object) { |ctx, name| ctx.const_get(name) }
+    end
+  end
+end

data/lib/kumi/pack/builder.rb

@@ -0,0 +1,229 @@
+# frozen_string_literal: true
+
+require "json"
+require "digest"
+require "fileutils"
+
+module Kumi
+  module Pack
+    module Builder
+      VERSION = "0.1"
+
+      module_function
+
+      def build(schema:, out_dir:, targets: %w[ruby], include_ir: false)
+        ir, planning, bindings, inputs, module_id = generate_artifacts(schema)
+        FileUtils.mkdir_p(out_dir)
+        write_json("#{out_dir}/irv2.json", ir)
+        write_json("#{out_dir}/planning.json", planning)
+        write_json("#{out_dir}/bindings.json", bindings)
+        write_json("#{out_dir}/inputs.json", inputs)
+
+        pack = assemble_pack(module_id, ir, planning, bindings, inputs, targets, include_ir)
+        write_json("#{out_dir}/pack.json", pack)
+        canonical_json(pack)
+      end
+
+      def print(schema:, targets: %w[ruby], include_ir: false)
+        ir, planning, bindings, inputs, module_id = generate_artifacts(schema)
+        pack = assemble_pack(module_id, ir, planning, bindings, inputs, targets, include_ir)
+        canonical_json(pack)
+      end
+
+      def build_for_golden(schema_path, golden_dir, targets: %w[ruby])
+        ir, planning, bindings, inputs, module_id = generate_artifacts(schema_path)
+
+        targets.each do |target|
+          pack = assemble_pack(module_id, ir, planning, bindings, inputs, [target], false)
+
+          filename = targets.size == 1 ? "pack.json" : "pack_#{target}.json"
+          pack_file = File.join(golden_dir, filename)
+          File.write(pack_file, canonical_json(pack))
+        end
+      end
+
+      def generate_artifacts(schema_path)
+        schema, = Kumi::Frontends.load(path: schema_path)
+        res = Kumi::Analyzer.analyze!(schema, side_tables: true)
+
+        irv2 = stringify_keys(res.state[:irv2]) or raise "No IRV2"
+        module_id = irv2["module"] || irv2["module_id"] || "kumi_module"
+
+        require_relative "../codegen/planning"
+        plan_bundle = Kumi::Codegen::Planning.from_ir(res.state[:irv2])
+        planning = Kumi::Codegen::Planning.to_json(plan_bundle)
+
+        bindings = stringify_keys(res.state[:binding_manifest] || {})
+        inputs = stringify_keys(res.state[:input_table] || {})
+
+        [irv2, planning, bindings, inputs, module_id]
+      end
+
+      def assemble_pack(module_id, ir, planning, bindings, inputs, _targets, include_ir)
+        plan_obj = planning["plan"] || planning
+        plan_declarations = plan_obj["declarations"] || {}
+
+        pack = {
+          "pack_version" => VERSION,
+          "module_id" => module_id,
+          "declarations" => extract_ops_by_decl_with_planning(ir, plan_declarations),
+          "inputs" => extract_inputs(inputs),
+          "bindings" => format_bindings_for_pack(bindings),
+          "capabilities" => { "layout" => "nested_array" }
+        }
+        pack["ir_debug"] = ir if include_ir
+        pack["hashes"] = compute_hashes(pack)
+        pack
+      end
+
+      def extract_ops_by_decl_with_planning(ir, plan_declarations)
+        declarations = ir["declarations"] || {}
+        declarations.keys.map do |name|
+          d = declarations[name]
+          ops = (d["operations"] || []).map do |op|
+            {
+              "id" => op["id"],
+              "op" => op["op"] || op["kind"],
+              "args" => op["args"] || [],
+              "attrs" => op["attrs"] || {}
+            }
+          end
+
+          # Start with IRV2 operations
+          result = { "name" => name, "operations" => ops }
+          result["result_op_id"] = d["result"] if d.key?("result")
+          result["axes"] = d["axes"] if d.key?("axes")
+
+          # Merge in planning metadata in sorted key order
+          plan_data = plan_declarations[name] || {}
+          plan_data.keys.sort.each do |key|
+            # Don't override IRV2 fields, but add planning fields
+            result[key] = plan_data[key] unless result.key?(key)
+          end
+
+          result
+        end
+      end
+
+      def extract_ops_by_decl(ir)
+        declarations = ir["declarations"] || {}
+        declarations.keys.sort.map do |name|
+          d = declarations[name]
+          ops = (d["operations"] || []).map do |op|
+            {
+              "id" => op["id"],
+              "op" => op["op"] || op["kind"],
+              "args" => op["args"] || [],
+              "attrs" => op["attrs"] || {}
+            }
+          end
+          result = { "name" => name, "operations" => ops }
+          result["result_op_id"] = d["result"] if d.key?("result")
+          result["axes"] = d["axes"] if d.key?("axes")
+          result
+        end
+      end
+
+      def extract_inputs(input_table)
+        # input_table is already in the new format with navigation_steps
+        input_table.map do |entry|
+          path_fqn = entry.path_fqn
+          navigation_steps = stringify_keys(entry.navigation_steps)
+
+          {
+            "name" => path_fqn,
+            "path_fqn" => path_fqn,
+            "axes" => entry.axes.map(&:to_s),
+            "dtype" => entry.dtype.to_s,
+            "accessor_name" => accessor_name_for(path_fqn),
+            "navigation_steps" => navigation_steps
+          }
+        end
+      end
+
+      def format_bindings_for_pack(bindings)
+        return {} unless bindings.is_a?(Hash)
+
+        formatted = {}
+        target = bindings["target"]
+        kernels_array = bindings["kernels"].map do |kernel_data|
+          kernel_id = kernel_data["kernel_id"]
+          impl = kernel_data["impl"]
+          attrs = kernel_data["attrs"] || {}
+
+          # Extract the function name from kernel_id (e.g., "core.add:ruby:v1" -> "core.add")
+          fn_name = kernel_id.split(":").first
+          # Format the implementation as proper Ruby lambda
+          ruby_impl = format_ruby_lambda(impl)
+
+          {
+            "kernel_id" => fn_name,
+            "impl" => ruby_impl,
+            "attrs" => attrs
+          }
+        end
+        formatted[target] = { "kernels" => kernels_array }
+
+        formatted
+      end
+
+      def format_ruby_lambda(impl)
+        # Convert "(a, b)\n  a + b" to "->(a, b) { a + b }"
+        lines = impl.strip.split("\n")
+        if lines.length >= 2
+          params = lines[0].strip.gsub(/^\(|\)$/, "") # Remove outer parentheses
+          body = lines[1..-1].map(&:strip).join("; ")
+          "->(#{params}) { #{body} }"
+        else
+          # Fallback for single line
+          "->(#{impl.strip})"
+        end
+      end
+
+      def stringify_keys(obj)
+        case obj
+        when Hash
+          obj.transform_keys(&:to_s).transform_values { |v| stringify_keys(v) }
+        when Array
+          obj.map { |v| stringify_keys(v) }
+        else
+          obj
+        end
+      end
+
+      def accessor_name_for(name)
+        name.gsub(/[^a-zA-Z0-9_]/, "_").downcase
+      end
+
+      def write_json(file_path, data)
+        File.write(file_path, JSON.pretty_generate(data))
+      end
+
+      # Compute canonical section hashes (sorted-key JSON)
+      def compute_hashes(pack)
+        keys = %w[declarations inputs bindings]
+        keys << "ir_debug" if pack.key?("ir_debug")
+        keys.to_h { |k| [k, sha256(pack[k])] }
+      end
+
+      # Stable JSON for hashing
+      def canonical_json(obj)
+        case obj
+        when Hash
+          "{#{obj.keys.map(&:to_s).sort.map { |k| "\"#{k}\":#{canonical_json(obj[k])}" }.join(',')}}"
+        when Array
+          "[#{obj.map { |v| canonical_json(v) }.join(',')}]"
+        when String   then JSON.generate(obj)
+        when Numeric  then obj.to_s
+        when true, false then obj.to_s
+        when NilClass then "null"
+        else JSON.generate(obj)
+        end
+      end
+
+      def sha256(obj)
+        Digest::SHA256.hexdigest(canonical_json(obj))
+      end
+    end
+  end
+end

data/lib/kumi/pack.rb

@@ -0,0 +1,15 @@
+# frozen_string_literal: true
+
+module Kumi
+  module Pack
+    module_function
+
+    def build(schema:, out_dir:, targets: %w[ruby], include_ir: false)
+      Builder.build(schema: schema, out_dir: out_dir, targets: targets, include_ir: include_ir)
+    end
+
+    def print(schema:, targets: %w[ruby], include_ir: false)
+      Builder.print(schema: schema, targets: targets, include_ir: include_ir)
+    end
+  end
+end

data/lib/kumi/registry_v2/loader.rb

@@ -0,0 +1,81 @@
+# frozen_string_literal: true
+
+require "yaml"
+
+module Kumi
+  module RegistryV2
+    module Loader
+      module_function
+
+      # { "core.mul" => Function(id: "core.mul", kind: :elementwise, params: [...]) }
+      def load_functions(dir, func_struct)
+        files = Dir.glob(File.join(dir, "**", "*.y{a,}ml")).sort
+        funcs = files.flat_map { |p| (YAML.load_file(p) || {}).fetch("functions", []) }
+        # funcs.each_with_object({}) do |h, acc|
+        #   acc[h.fetch("id").to_s] = {
+        #     kind: h.fetch("kind").to_s.to_sym,
+        #     aliases: Array(h["aliases"]).map!(&:to_s),
+        #     params: h.fetch("params")
+        #   }
+        # end
+        funcs.each_with_object({}) do |h, acc|
+          f = func_struct.new(
+            id: h.fetch("id").to_s,
+            kind: h.fetch("kind").to_s.to_sym,
+            aliases: Array(h["aliases"]).map!(&:to_s),
+            params: h.fetch("params"),
+            dtype: h["dtype"],
+            expand: h["expand"],
+            options: symbolize_keys(h["options"] || {}),
+            folding_class_method: h["folding_class_method"],
+            reduction_strategy: h["reduction_strategy"]&.to_sym
+          )
+          raise "duplicate function id `#{f.id}`" if acc.key?(f.id)
+
+          acc[f.id] = f
+        end
+      end
+
+      def symbolize_keys(h)
+        h.each_with_object({}) { |(k, v), out| out[k.to_sym] = v }
+      end
+
+      def deep_symbolize_keys(obj)
+        case obj
+        when Hash
+          obj.each_with_object({}) { |(k, v), out| out[k.to_sym] = deep_symbolize_keys(v) }
+        when Array
+          obj.map { |v| deep_symbolize_keys(v) }
+        else
+          obj
+        end
+      end
+
+      # { ["core.mul", :ruby] => Kernel }
+      def load_kernels(root, kernel_struct)
+        targets = Dir.glob(File.join(root, "*")).select { |p| File.directory?(p) }.map { |p| File.basename(p).to_sym }
+        out = {}
+        targets.each do |t|
+          Dir.glob(File.join(root, t.to_s, "**", "*.y{a,}ml")).sort.each do |p|
+            (YAML.load_file(p) || {}).fetch("kernels", []).each do |h|
+              k = kernel_struct.new(
+                id: h.fetch("id"),
+                fn_id: h.fetch("fn").to_s,
+                target: t,
+                impl: h["impl"],
+                identity: h["identity"],
+                inline: h["inline"],
+                fold_inline: h["fold_inline"]
+              )
+              key = [k.fn_id, t]
+              raise "duplicate kernel for #{key}" if out.key?(key)
+
+              out[key] = k
+            end
+          end
+        end
+        out
+      end
+    end
+  end
+end

data/lib/kumi/registry_v2.rb

@@ -0,0 +1,118 @@
+# frozen_string_literal: true
+
+require "json"
+require "digest"
+require_relative "registry_v2/loader" # Assuming loader is in a sub-file
+
+module Kumi
+  module RegistryV2
+    DEFAULT_FUNCTIONS_DIR = File.expand_path("../../data/functions", __dir__)
+    DEFAULT_KERNELS_DIR = File.expand_path("../../data/kernels", __dir__)
+    SELECT_ID = "__select__"
+
+    # --- NEW: Define the Function struct ---
+    Function = Struct.new(:id, :kind, :dtype, :aliases, :params, :options, :expand, :folding_class_method, :reduction_strategy,
+                          keyword_init: true) do
+      def reduce? = kind == :reduce
+      def select? = id == SELECT_ID
+      def elementwise? = kind == :elementwise
+
+      def param_names
+        @param_names ||= params.map { |p| p["name"].to_sym }
+      end
+
+      def dtype_rule
+        @dtype_rule ||= Core::Functions::TypeRules.compile_dtype_rule(dtype, param_names)
+      end
+    end
+
+    Kernel = Struct.new(:id, :fn_id, :target, :impl, :identity, :inline, :fold_inline, keyword_init: true)
+
+    class Instance
+      def initialize(functions_by_id, kernels_by_key)
+        @functions = functions_by_id                         # "core.mul" => Function<...>
+        @alias     = build_alias(@functions)                 # "count" => "agg.count"
+        @kernels   = kernels_by_key                          # [fn_id, target_sym] => Kernel
+        @by_id     = @kernels.values.to_h { |k| [k.id, k] }
+      end
+
+      # -------- functions --------
+      def function(id)
+        @functions.fetch(resolve_function(id))
+      end
+
+      def resolve_function(id)
+        s = id.to_s
+        return s if @functions.key?(s)
+
+        @alias.fetch(s) do
+          raise "unknown function #{id}"
+        end
+      end
+
+      def function_kind(id)        = function(id).kind
+      def function_reduce?(id)     = function(id).reduce?
+      def function_elementwise?(id) = function(id).elementwise?
+      def function_select?(id)      = resolve_function(id) == SELECT_ID
+
+      # -------- kernels (no changes here) --------
+      def kernel_for(id, target:)
+        fid = resolve_function(id)
+        @kernels[[fid, target.to_sym]] or raise "no kernel for #{fid} on #{target}"
+      end
+
+      def kernel_id_for(id, target:)
+        fid = resolve_function(id)
+        k = @kernels[[fid, target.to_sym]] or raise "no kernel for #{fid} on #{target}"
+        k.id
+      end
+
+      def kernel_identity_for(id, dtype:, target:)
+        fid = resolve_function(id)
+        k = @kernels[[fid, target.to_sym]] or raise "no kernel for #{fid} on #{target}"
+
+        map = k.identity or raise "no identity for #{fid} on #{target}"
+
+        identity = map[dtype.to_s] || map["any"]
+
+        return identity if identity
+
+        raise "no identity for dtype #{dtype} on #{fid}"
+      end
+
+      def impl_for(kernel_id)
+        (@by_id[kernel_id] or raise "unknown kernel #{kernel_id}").impl
+      end
+
+      def registry_ref
+        # ... (implementation is fine, but needs to access struct attributes) ...
+        stable = {
+          kernels: @by_id.values.map do |k|
+            { "id" => k.id, "fn" => k.fn_id, "target" => k.target.to_s, "impl" => k.impl }
+          end.sort_by { _1["id"] },
+          functions: @functions.transform_values do |f|
+            { "kind" => f.kind.to_s, "aliases" => f.aliases, "params" => f.params }
+          end
+        }
+        "sha256:#{Digest::SHA256.hexdigest(JSON.generate(stable))}"
+      end
+
+      private
+
+      def build_alias(functions)
+        functions.values.each_with_object({}) do |func, acc|
+          func.aliases.each { |al| acc[al] = func.id }
+        end
+      end
+    end
+
+    module_function
+
+    def load(functions_dir: DEFAULT_FUNCTIONS_DIR, kernels_root: DEFAULT_KERNELS_DIR)
+      # Pass the new struct to the loader.
+      fn_map = Loader.load_functions(functions_dir, Function)
+      kn_map = Loader.load_kernels(kernels_root, Kernel)
+      Instance.new(fn_map, kn_map)
+    end
+  end
+end