kumi 0.0.12 → 0.0.14

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

@@ -0,0 +1,98 @@
+# frozen_string_literal: true
+
+require_relative "errors"
+
+module Kumi
+  module Core
+    module Functions
+      # Represents a single dimension in a signature with NEP 20 support.
+      #
+      # A dimension can be:
+      # - Named dimension (symbol): :i, :j, :n
+      # - Fixed-size dimension (integer): 2, 3, 10
+      # - With modifiers:
+      #   - flexible (?): can be omitted if not present in all operands
+      #   - broadcastable (|1): can broadcast against size-1 dimensions
+      #
+      # Examples:
+      #   Dimension.new(:i)           # named dimension 'i'
+      #   Dimension.new(3)            # fixed-size dimension of size 3
+      #   Dimension.new(:n, flexible: true)      # dimension 'n' that can be omitted
+      #   Dimension.new(:i, broadcastable: true) # dimension 'i' that can broadcast
+      class Dimension
+        attr_reader :name, :flexible, :broadcastable
+
+        def initialize(name, flexible: false, broadcastable: false)
+          @name = name
+          @flexible = flexible
+          @broadcastable = broadcastable
+
+          validate!
+          freeze
+        end
+
+        def fixed_size?
+          @name.is_a?(Integer)
+        end
+
+        def named?
+          @name.is_a?(Symbol)
+        end
+
+        def flexible?
+          @flexible
+        end
+
+        def broadcastable?
+          @broadcastable
+        end
+
+        def size
+          fixed_size? ? @name : nil
+        end
+
+        def ==(other)
+          other.is_a?(Dimension) &&
+            name == other.name &&
+            flexible == other.flexible &&
+            broadcastable == other.broadcastable
+        end
+
+        def eql?(other)
+          self == other
+        end
+
+        def hash
+          [name, flexible, broadcastable].hash
+        end
+
+        def to_s
+          str = name.to_s
+          str += "?" if flexible?
+          str += "|1" if broadcastable?
+          str
+        end
+
+        def inspect
+          "#<Dimension #{self}>"
+        end
+
+        private
+
+        def validate!
+          unless name.is_a?(Symbol) || name.is_a?(Integer)
+            raise SignatureError, "dimension name must be a symbol or integer, got: #{name.inspect}"
+          end
+
+          raise SignatureError, "fixed-size dimension must be positive, got: #{name}" if name.is_a?(Integer) && name <= 0
+
+          raise SignatureError, "dimension cannot be both flexible and broadcastable" if flexible? && broadcastable?
+
+          return unless fixed_size? && flexible?
+
+          raise SignatureError, "fixed-size dimension cannot be flexible"
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,20 @@
+module Kumi::Core::Functions
+  DType = Struct.new(:name, keyword_init: true)
+  module DTypes
+    BOOL = DType.new(name: :bool)
+    INT = DType.new(name: :int)
+    FLOAT = DType.new(name: :float)
+    STRING = DType.new(name: :string)
+    DATETIME = DType.new(name: :datetime)
+    ANY = DType.new(name: :any)
+  end
+
+  module Promotion
+    # super simple table; extend later
+    TABLE = {
+      %i[int int] => :int,  %i[int float] => :float, %i[float int] => :float, %i[float float] => :float,
+      %i[bool int] => :int, %i[bool float] => :float, %i[bool bool] => :bool
+    }
+    def self.promote(a, b) = TABLE[[a, b]] || TABLE[[b, a]] || :any
+  end
+end

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

@@ -0,0 +1,11 @@
+# frozen_string_literal: true
+
+module Kumi
+  module Core
+    module Functions
+      class SignatureError < StandardError; end
+      class SignatureParseError < SignatureError; end
+      class SignatureMatchError < SignatureError; end
+    end
+  end
+end

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

@@ -0,0 +1,45 @@
+# frozen_string_literal: true
+
+module Kumi
+  module Core
+    module Functions
+      KernelHandle = Struct.new(:kind, :callable, :null_policy, :options, keyword_init: true)
+
+      module KernelAdapter
+        module_function
+
+        def build_for(function, backend_entry)
+          impl = backend_entry.impl
+          mod = ruby_module_for(function)
+
+          raise Kumi::Core::Errors::CompilationError, "Missing Ruby kernel #{impl} for #{function.name}" unless mod.respond_to?(impl)
+
+          kind = function.class_sym
+          KernelHandle.new(
+            kind: kind,
+            callable: mod.method(impl),
+            null_policy: function.null_policy,
+            options: function.options || {}
+          )
+        end
+
+        def ruby_module_for(function)
+          case function.domain.to_sym
+          when :core
+            function.class_sym == :aggregate ? Kumi::Kernels::Ruby::AggregateCore : Kumi::Kernels::Ruby::ScalarCore
+          when :string
+            Kumi::Kernels::Ruby::StringScalar
+          when :datetime
+            Kumi::Kernels::Ruby::DatetimeScalar
+          when :struct
+            Kumi::Kernels::Ruby::VectorStruct
+          when :mask
+            Kumi::Kernels::Ruby::MaskScalar
+          else
+            raise Kumi::Core::Errors::CompilationError, "Unknown domain #{function.domain}"
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,119 @@
+require "yaml"
+module Kumi::Core::Functions
+  Function = Data.define(:name, :domain, :opset, :class_sym, :signatures,
+                         :type_vars, :dtypes, :null_policy, :algebra, :vectorization,
+                         :options, :traits, :shape_fn, :kernels)
+
+  KernelEntry = Data.define(:backend, :impl, :priority, :when_)
+
+  class Loader
+    def self.load_file(path)
+      doc = YAML.load_file(path)
+      functions = doc.map { |h| build_function(h) }
+      validate!(functions)
+      functions.freeze
+    end
+
+    def self.build_function(h)
+      sigs = Array(h.fetch("signature")).map { |s| parse_signature(s) }
+      kernels = Array(h.fetch("kernels", [])).map do |k|
+        KernelEntry.new(backend: k.fetch("backend").to_sym,
+                        impl: k.fetch("impl").to_sym,
+                        priority: k.fetch("priority", 0).to_i,
+                        when_: k["when"]&.transform_keys!(&:to_sym))
+      end
+      Function.new(
+        name: h.fetch("name"),
+        domain: h.fetch("domain"),
+        opset: h.fetch("opset").to_i,
+        class_sym: h.fetch("class").to_sym,
+        signatures: sigs,
+        type_vars: h["type_vars"] || {},
+        dtypes: h["dtypes"] || {},
+        null_policy: (h["null_policy"] || "propagate").to_sym,
+        algebra: (h["algebra"] || {}).transform_keys!(&:to_sym),
+        vectorization: (h["vectorization"] || {}).transform_keys!(&:to_sym),
+        options: h["options"] || {},
+        traits: (h["traits"] || {}).transform_keys!(&:to_sym),
+        shape_fn: h["shape_fn"],
+        kernels: kernels
+      )
+    end
+
+    # " (i),(j)->(i,j)@product " → Signature
+    def self.parse_signature(s)
+      lhs, rhs = s.split("->").map(&:strip)
+      out_axes, policy = rhs.split("@").map(&:strip)
+      
+      # Parse input shapes by splitting on commas between parentheses
+      in_shapes = parse_input_shapes(lhs)
+      
+      Signature.new(in_shapes: in_shapes,
+                    out_shape: parse_axes(out_axes),
+                    join_policy: policy&.to_sym)
+    end
+
+    # Parse "(i),(j)" or "(i,j),(k,l)" properly
+    def self.parse_input_shapes(lhs)
+      # Find all parenthesized groups
+      shapes = []
+      current_pos = 0
+      
+      while current_pos < lhs.length
+        # Find the next opening parenthesis
+        start_paren = lhs.index('(', current_pos)
+        break unless start_paren
+        
+        # Find the matching closing parenthesis
+        paren_count = 0
+        end_paren = start_paren
+        
+        (start_paren..lhs.length-1).each do |i|
+          case lhs[i]
+          when '('
+            paren_count += 1
+          when ')'
+            paren_count -= 1
+            if paren_count == 0
+              end_paren = i
+              break
+            end
+          end
+        end
+        
+        # Extract the shape
+        shape_str = lhs[start_paren..end_paren]
+        shapes << parse_axes(shape_str)
+        current_pos = end_paren + 1
+      end
+      
+      shapes
+    end
+
+    def self.parse_axes(txt)
+      txt = txt.strip
+      return [] if txt == "()" || txt.empty?
+
+      inner = txt.sub("(", "").sub(")", "")
+      if inner.empty?
+        []
+      else
+        inner.split(",").map do |a|
+          dim_name = a.strip.to_sym
+          Dimension.new(dim_name)  # Convert to Dimension objects
+        end
+      end
+    end
+
+    def self.validate!(fns)
+      names = {}
+      fns.each do |f|
+        key = [f.domain, f.name, f.opset]
+        raise "duplicate function #{key}" if names[key]
+
+        names[key] = true
+        raise "no kernels for #{f.name}" if f.kernels.empty?
+      end
+    end
+  end
+end

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

@@ -0,0 +1,68 @@
+module Kumi::Core::Functions
+  class RegistryV2
+    def initialize(functions:)
+      @by_name = functions.group_by(&:name).transform_values { |v| v.sort_by(&:opset).freeze }.freeze
+    end
+
+    # Factory method to load from YAML configuration
+    def self.load_from_file(path = nil)
+      path ||= File.join(__dir__, "../../../..", "config", "functions.yaml")
+      functions = Loader.load_file(path)
+      new(functions: functions)
+    end
+
+    def fetch(name, opset: nil)
+      list = @by_name[name] or raise KeyError, "unknown function #{name}"
+      opset ? list.find { |f| f.opset == opset } : list.last
+    end
+
+    # Get function signatures for NEP-20 signature resolution
+    # This bridges RegistryV2 with our existing SignatureResolver
+    def get_function_signatures(name, opset: nil)
+      begin
+        fn = fetch(name, opset: opset)
+        # Convert Signature objects to string representations for NEP-20 parser
+        fn.signatures.map(&:to_signature_string)
+      rescue KeyError
+        []  # Function not found in RegistryV2 - fall back to legacy registry
+      end
+    end
+
+    # Enhanced signature resolution using NEP-20 resolver
+    def choose_signature(fn, arg_shapes)
+      # Use our NEP-20 SignatureResolver for proper dimension handling
+      sig_strings = fn.signatures.map(&:to_signature_string)
+      parsed_sigs = sig_strings.map { |s| SignatureParser.parse(s) }
+      
+      plan = SignatureResolver.choose(signatures: parsed_sigs, arg_shapes: arg_shapes)
+      
+      # Return both the original function signature and the resolution plan
+      {
+        function: fn,
+        signature: plan[:signature],
+        plan: plan
+      }
+    rescue SignatureMatchError => e
+      raise ArgumentError, "no matching signature for #{fn.name} with shapes #{arg_shapes.inspect}: #{e.message}"
+    end
+
+    def resolve_kernel(fn, backend:, conditions: {})
+      ks = fn.kernels.select { |k| k.backend == backend.to_sym }
+      ks = ks.select { |k| conditions.all? { |ck, cv| k.when_&.fetch(ck, cv) == cv } } unless conditions.empty?
+      ks.max_by(&:priority) or raise "no kernel for #{fn.name} backend=#{backend}"
+    end
+
+    # Introspection methods
+    def all_function_names
+      @by_name.keys
+    end
+
+    def function_exists?(name)
+      @by_name.key?(name)
+    end
+
+    def all_functions
+      @by_name.values.flatten
+    end
+  end
+end

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

@@ -0,0 +1,70 @@
+# frozen_string_literal: true
+
+require_relative "dimension"
+
+module Kumi
+  module Core
+    module Functions
+      # Shape utilities with NEP 20 support. A "shape" is an Array<Dimension> of dimension objects.
+      # [] == scalar, [Dimension.new(:i)] == vector along :i, etc.
+      module Shape
+        module_function
+
+        def scalar?(shape) = shape.empty?
+
+        def equal?(a, b) = a.map(&:name) == b.map(&:name)
+
+        # NEP 20 broadcast rules:
+        # - scalar can broadcast to any expected shape
+        # - fixed-size dimensions must match exactly
+        # - broadcastable dimensions with |1 modifier can broadcast against size-1
+        # - flexible dimensions with ? can be omitted if not present in all operands
+        def broadcastable?(got:, expected:)
+          return true if scalar?(got)
+          return false if got.length != expected.length
+
+          got.zip(expected).all? do |got_dim, exp_dim|
+            broadcastable_dimension?(got: got_dim, expected: exp_dim)
+          end
+        end
+
+        def broadcastable_dimension?(got:, expected:)
+          # Same name and modifiers
+          return true if got == expected
+
+          # Same name, different modifiers - check compatibility
+          if got.name == expected.name
+            # Fixed-size dimensions must match exactly
+            if got.fixed_size? || expected.fixed_size?
+              return got.size == expected.size if got.fixed_size? && expected.fixed_size?
+              return false # one fixed, one not - incompatible
+            end
+
+            # Both named dimensions with same name are compatible
+            return true
+          end
+
+          # Different names - only broadcastable with |1 modifier
+          expected.broadcastable? && scalar?([got])
+        end
+
+        # Check if a dimension can be omitted (NEP 20 flexible dimensions)
+        def flexible?(dim)
+          dim.is_a?(Dimension) && dim.flexible?
+        end
+
+        # Check if a dimension can broadcast (NEP 20 broadcastable dimensions)
+        def broadcastable_dimension?(dim)
+          dim.is_a?(Dimension) && dim.broadcastable?
+        end
+
+        # Convenience: find dimensions in set a that are not in set b
+        def dimensions_minus(a, b)
+          a_names = a.map(&:name).to_set
+          b_names = b.map(&:name).to_set
+          (a_names - b_names).to_a
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,122 @@
+# frozen_string_literal: true
+
+module Kumi
+  module Core
+    module Functions
+      # Signature is a small immutable value object describing a function's
+      # vectorization contract with NEP 20 support.
+      #
+      # in_shapes: Array<Array<Dimension>>   e.g., [[Dimension.new(:i)], [Dimension.new(:i)]]
+      # out_shape: Array<Dimension>          e.g., [Dimension.new(:i)], [Dimension.new(:i), Dimension.new(:j)], or []
+      # join_policy: nil | :zip | :product
+      # raw: original string, for diagnostics (optional)
+      class Signature
+        attr_reader :in_shapes, :out_shape, :join_policy, :raw
+
+        def initialize(in_shapes:, out_shape:, join_policy: nil, raw: nil)
+          @in_shapes   = deep_dup(in_shapes).freeze
+          @out_shape   = out_shape.dup.freeze
+          @join_policy = join_policy&.to_sym
+          @raw         = raw
+          validate!
+          freeze
+        end
+
+        def arity = @in_shapes.length
+
+        # Dimensions that appear in any input but not in output (i.e., reduced/dropped).
+        def dropped_axes
+          input_names = @in_shapes.flatten.map(&:name)
+          output_names = @out_shape.map(&:name)
+          (input_names - output_names).uniq.freeze
+        end
+
+        # True if any axis from inputs is dropped (common in aggregates).
+        def reduction?
+          !dropped_axes.empty?
+        end
+
+        def to_h
+          {
+            in_shapes: in_shapes.map(&:dup),
+            out_shape: out_shape.dup,
+            join_policy: join_policy,
+            raw: raw
+          }
+        end
+
+        def inspect
+          "#<Signature #{format_signature}#{" @#{join_policy}" if join_policy}>"
+        end
+
+        def format_signature
+          lhs = in_shapes.map { |dims| "(#{dims.map(&:to_s).join(',')})" }.join(",")
+          rhs = "(#{out_shape.map(&:to_s).join(',')})"
+          "#{lhs}->#{rhs}"
+        end
+
+        # Convert back to string representation for NEP-20 parser compatibility
+        def to_signature_string
+          sig_str = format_signature
+          join_policy ? "#{sig_str}@#{join_policy}" : sig_str
+        end
+
+        private
+
+        def validate!
+          unless @in_shapes.is_a?(Array) && @in_shapes.all? { |s| s.is_a?(Array) }
+            raise SignatureError, "in_shapes must be an array of dimension arrays"
+          end
+
+          @in_shapes.each_with_index do |dims, idx|
+            validate_dimensions!(dims, where: "in_shapes[#{idx}]")
+          end
+
+          validate_dimensions!(@out_shape, where: "out_shape")
+
+          unless [nil, :zip, :product].include?(@join_policy)
+            raise SignatureError, "join_policy must be nil, :zip, or :product; got #{@join_policy.inspect}"
+          end
+
+          # Validate NEP 20 constraints
+          validate_nep20_constraints!
+        end
+
+        def validate_dimensions!(dims, where: "shape")
+          unless dims.is_a?(Array) && dims.all? { |d| d.is_a?(Dimension) }
+            raise SignatureError, "#{where}: must be an array of Dimension objects, got: #{dims.inspect}"
+          end
+
+          # Check for duplicate dimension names within a single argument
+          names = dims.map(&:name)
+          duplicates = names.group_by { |n| n }.select { |_, v| v.size > 1 }.keys
+          raise SignatureError, "#{where}: duplicate dimension names #{duplicates.inspect}" unless duplicates.empty?
+
+          true
+        end
+
+        def validate_nep20_constraints!
+          # Broadcastable dimensions should only appear in inputs, not outputs
+          @out_shape.each do |dim|
+            raise SignatureError, "output dimension #{dim} cannot be broadcastable" if dim.broadcastable?
+          end
+
+          # Fixed-size dimensions in outputs must match corresponding input dimensions
+          all_input_dims = @in_shapes.flatten
+          @out_shape.each do |out_dim|
+            next unless out_dim.fixed_size?
+
+            matching_inputs = all_input_dims.select { |in_dim| in_dim.name == out_dim.name }
+            matching_inputs.each do |in_dim|
+              if in_dim.fixed_size? && in_dim.size != out_dim.size
+                raise SignatureError, "fixed-size dimension #{out_dim.name} has inconsistent sizes: #{in_dim.size} vs #{out_dim.size}"
+              end
+            end
+          end
+        end
+
+        def deep_dup(arr) = arr.map { |x| x.is_a?(Array) ? x.dup : x }
+      end
+    end
+  end
+end

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

@@ -0,0 +1,86 @@
+# frozen_string_literal: true
+
+require_relative "dimension"
+require_relative "signature"
+
+module Kumi
+  module Core
+    module Functions
+      # Parses NEP 20 conformant signature strings like:
+      #  "(),()->()"                    # scalar operations
+      #  "(i),(i)->(i)"                 # vector operations
+      #  "(i),(j)->(i,j)@product"       # matrix operations with join policy
+      #  "(i,j)->(i)"                   # reduction of :j
+      #  "(3),(3)->(3)"                 # fixed-size 3-vectors (cross product)
+      #  "(i?),(i?)->(i?)"               # flexible dimensions
+      #  "(i|1),(i|1)->()"               # broadcastable dimensions
+      #  "(m?,n),(n,p?)->(m?,p?)"        # matmul signature
+      class SignatureParser
+        class << self
+          def parse(str)
+            raise SignatureParseError, "empty signature" if str.nil? || str.strip.empty?
+            lhs, rhs = str.split("->", 2)&.map!(&:strip)
+            raise SignatureParseError, "signature must contain '->': #{str.inspect}" unless rhs
+
+            out_spec, policy = rhs.split("@", 2)&.map!(&:strip)
+            in_shapes = parse_many(lhs)
+            out_shape = parse_axes(out_spec)
+            join_policy = policy&.to_sym
+
+            Signature.new(in_shapes: in_shapes, out_shape: out_shape, join_policy: join_policy, raw: str)
+          rescue SignatureError => e
+            raise
+          rescue StandardError => e
+            raise SignatureParseError, "invalid signature #{str.inspect}: #{e.message}"
+          end
+
+          private
+
+          def parse_many(lhs)
+            # Handle zero arguments case
+            return [] if lhs.strip.empty?
+
+            # split by commas that are *between* groups, not inside them  
+            # simpler approach: split by '),', re-add ')' where needed
+            tokens = lhs.split("),").map { |t| t.strip.end_with?(")") ? t.strip : "#{t.strip})" }
+            tokens.map { |tok| parse_axes(tok) }
+          end
+
+          def parse_axes(spec)
+            spec = spec.strip
+            raise SignatureParseError, "missing parentheses in #{spec.inspect}" unless spec.start_with?("(") && spec.end_with?(")")
+
+            inner = spec[1..-2].strip
+            return [] if inner.empty?
+
+            inner.split(",").map { |dim_str| parse_dimension(dim_str.strip) }
+          end
+
+          # Parse a single dimension with NEP 20 modifiers
+          # Examples: "i", "3", "n?", "i|1"
+          def parse_dimension(dim_str)
+            return Dimension.new(:empty) if dim_str.empty?
+
+            # Extract modifiers
+            flexible = dim_str.end_with?("?")
+            dim_str = dim_str.chomp("?") if flexible
+
+            broadcastable = dim_str.end_with?("|1")
+            dim_str = dim_str.chomp("|1") if broadcastable
+
+            # Parse name (symbol or integer)
+            name = if dim_str.match?(/^\d+$/)
+                     dim_str.to_i
+                   else
+                     dim_str.to_sym
+                   end
+
+            Dimension.new(name, flexible: flexible, broadcastable: broadcastable)
+          rescue StandardError => e
+            raise SignatureParseError, "invalid dimension #{dim_str.inspect}: #{e.message}"
+          end
+        end
+      end
+    end
+  end
+end