kumi 0.0.17 → 0.0.19

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

@@ -0,0 +1,229 @@
+# frozen_string_literal: true
+
+module Kumi
+  module Core
+    module Analyzer
+      module Passes
+        # Semantic NAST Pass (SNAST)
+        # - Rewrites intrinsic control and reductions into first-class nodes.
+        # - Attaches semantic stamps to every node: meta[:stamp] = { axes:, dtype: }.
+        # - Uses side tables for types/scopes; no meta[:plan].
+        #
+        # Reduction rule (default sugar):
+        #   If not explicitly annotated, a reduction over arguments reduces the LAST axis:
+        #     a = lub_by_prefix(arg_axes_list)
+        #     over = [a.last]; out_axes = a[0...-1]
+        #
+        # Inputs (state):
+        #   :nast_module          => Kumi::Core::NAST::Module (topologically ordered)
+        #   :metadata_table       => Hash[node_key => { result_scope:, result_type:, arg_scopes?: ... }]
+        #   :declaration_table    => Hash[name => { result_scope:, result_type: }]
+        #   :input_table          => [{path_fqn:, axes:, dtype:}] or Hash[path_fqn] => { axes:, dtype: }
+        #
+        # Output (state):
+        #   :snast_module         => Kumi::Core::NAST::Module (with NAST::Select / NAST::Reduce nodes)
+        #
+        # TODO: If downstream never keys by node ids, consider removing dependence on node.id.
+        # TODO: Use Error helpers with provenance
+        class SNASTPass < PassBase
+          def run(errors)
+            @nast_module       = get_state(:nast_module,       required: true)
+            @metadata_table    = get_state(:metadata_table,    required: true)
+            @declaration_table = get_state(:declaration_table, required: true)
+            @input_table       = get_state(:input_table,       required: true)
+            @index_table       = get_state(:index_table,       required: true)
+            @registry          = get_state(:registry,          required: true)
+            @errors = errors
+
+            debug "Building SNAST from #{@nast_module.decls.size} declarations"
+            snast_module = @nast_module.accept(self)
+            state.with(:snast_module, snast_module.freeze)
+          end
+
+          # ---------- Visitor entry points ----------
+
+          def visit_module(mod)
+            # decls is expected to be a Hash[name => Declaration]
+            mod.class.new(decls: mod.decls.transform_values { |d| d.accept(self) })
+          end
+
+          def visit_declaration(d)
+            meta = @declaration_table.fetch(d.name)
+            body = d.body.accept(self)
+            out  = d.class.new(id: d.id, name: d.name, body:, loc: d.loc, meta: { kind: d.kind })
+            stamp!(out, meta[:result_scope], meta[:result_type])
+          end
+
+          # ---------- Leaves ----------
+
+          def visit_const(n)
+            meta = meta_for(n)
+            out = n.class.new(id: n.id, value: n.value, loc: n.loc)
+            stamp!(out, [], meta[:type])
+          end
+
+          def visit_input_ref(n)
+            ent = lookup_input(n.path_fqn)
+            out = n.class.new(id: n.id, path: n.path, loc: n.loc)
+            stamp!(out, ent[:axes], ent[:dtype])
+          end
+
+          def visit_index_ref(n)
+            m = meta_for(n)
+            out = n.class.new(id: n.id, name: n.name, input_fqn: n.input_fqn, loc: n.loc)
+            stamp!(out, m[:scope], m[:type])
+          end
+
+          def visit_ref(n)
+            m = meta_for(n)
+            out = n.class.new(id: n.id, name: n.name, loc: n.loc)
+            stamp!(out, m[:result_scope], m[:result_type])
+          end
+
+          def visit_tuple(n)
+            args = n.args.map { _1.accept(self) }
+            m    = meta_for(n)
+            out  = n.class.new(id: n.id, args:, loc: n.loc)
+            stamp!(out, m[:result_scope], m[:result_type])
+          end
+
+          def visit_hash(n)
+            pairs = n.pairs.map { _1.accept(self) }
+            m    = meta_for(n)
+            out  = n.class.new(id: n.id, pairs:, loc: n.loc)
+            stamp!(out, m[:scope], m[:type])
+          end
+
+          def visit_pair(n)
+            value = n.value.accept(self)
+            m = meta_for(n)
+            out = n.class.new(id: n.id, key: n.key, value:)
+            stamp!(out, m[:scope], m[:type])
+          end
+
+          def visit_call(n)
+            if @registry.function_select?(n.fn)
+              c = n.args[0].accept(self)
+              t = n.args[1].accept(self)
+              f = n.args[2].accept(self)
+              target_axes = lub_by_prefix([axes_of(t), axes_of(f)])
+              target_axes = axes_of(c) if target_axes.empty?
+              unless prefix?(axes_of(c), target_axes)
+                raise Kumi::Core::Errors::SemanticError,
+                      "select mask axes #{axes_of(c).inspect} must prefix #{target_axes.inspect} at: #{n.loc}"
+              end
+
+              out = NAST::Select.new(id: n.id, cond: c, on_true: t, on_false: f, loc: n.loc, meta: n.meta.dup)
+              return stamp!(out, target_axes, dtype_of(t))
+            end
+
+            if @registry.function_reduce?(n.fn)
+              raise "Reducers should only have one arg" if n.args.size != 1 # TODO: -> sugar to collapse variadics?
+
+              arg_node = n.args.first
+              visited_arg = arg_node.accept(self)
+              arg_meta = visited_arg[:meta]
+              arg_type = arg_meta[:stamp][:dtype]
+
+              if Kumi::Core::Types.collection?(arg_type)
+                # --- Path for FOLD (Scalar or Vectorized) ---w
+                # The argument is semantically a tuple. Create a Fold node.
+
+                # We still need to visit the child node to build the SNAST tree
+
+                fold_node = NAST::Fold.new(
+                  id: n.id,
+                  fn: @registry.resolve_function(n.fn),
+                  arg: visited_arg, # The arg is the tuple/reference to the tuple
+                  loc: n.loc,
+                  meta: n.meta.dup
+                )
+
+                # The output type is the reduced scalar type (e.g., :integer for max).
+                # The axes are PRESERVED because a fold is an element-wise operation
+                # on the container of tuples.
+                result_meta = meta_for(n)
+                return stamp!(fold_node, result_meta[:result_scope], result_meta[:result_type])
+              else
+                # --- Path for REDUCE (Vectorized Arrays) ---
+                in_axes = axes_of(visited_arg)
+
+                if in_axes.empty?
+                  raise Kumi::Core::Errors::SemanticError,
+                        "reduce function called on a non-collection scalar: #{arg_type}"
+                end
+
+                result_meta = meta_for(n)
+                out_axes = Array(result_meta[:result_scope])
+
+                raise Kumi::Core::Errors::SemanticError, "reduce: out axes must prefix arg axes" unless prefix?(out_axes, in_axes)
+
+                over_axes = in_axes.drop(out_axes.length)
+                reduce_node = NAST::Reduce.new(
+                  id: n.id,
+                  fn: @registry.resolve_function(n.fn),
+                  over: over_axes,
+                  arg: visited_arg,
+                  loc: n.loc,
+                  meta: n.meta.dup
+                )
+                return stamp!(reduce_node, out_axes, result_meta[:result_type])
+              end
+            end
+
+            # regular elementwise
+            args = n.args.map { _1.accept(self) }
+            m    = meta_for(n)
+            out  = n.class.new(id: n.id, fn: @registry.resolve_function(n.fn), args:, opts: n.opts, loc: n.loc)
+            stamp!(out, m[:result_scope], m[:result_type])
+          end
+
+          # ---------- Helpers ----------
+
+          def stamp!(node, axes, dtype)
+            node.meta[:stamp] = { axes: Array(axes), dtype: dtype }.freeze
+            node
+          end
+
+          def meta_for(node) = @metadata_table.fetch(node_key(node))
+          def axes_of(n)  = Array(n.meta[:stamp]&.dig(:axes))
+          def dtype_of(n) = n.meta[:stamp]&.dig(:dtype)
+
+          # Least upper bound by prefix. All entries must be a prefix of the longest.
+          def lub_by_prefix(list)
+            return [] if list.empty?
+
+            cand = list.max_by(&:length) || []
+            list.each do |ax|
+              raise Kumi::Core::Errors::SemanticError, "prefix mismatch: #{ax.inspect} vs #{cand.inspect}" unless prefix?(ax, cand)
+            end
+            cand
+          end
+
+          def prefix?(pre, full)
+            pre.each_with_index.all? { |tok, i| full[i] == tok }
+          end
+
+          # Default reduce sugar: over last axis of the LUB of argument axes.
+          # Returns { over:, out_axes: }.
+          def reduce_last_axis(args_axes_list)
+            a = lub_by_prefix(args_axes_list)
+            raise Kumi::Core::Errors::SemanticError, "cannot reduce scalar" if a.empty?
+
+            { over: [a.last], out_axes: a[0...-1] }
+          end
+
+          def lookup_input(fqn)
+            if @input_table.respond_to?(:find)
+              @input_table.find { |x| x[:path_fqn] == fqn } || raise("Input not found for #{fqn}")
+            else
+              @input_table.fetch(fqn) { raise("Input not found for #{fqn}") }
+            end
+          end
+
+          def node_key(n) = "#{n.class}_#{n.id}"
+        end
+      end
+    end
+  end
+end

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

@@ -14,6 +14,7 @@ module Kumi
             definitions = get_state(:declarations, required: false) || {}
 
             order = compute_topological_order(dependency_graph, definitions, errors)
+
             state.with(:evaluation_order, order)
           end
 

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

@@ -11,6 +11,7 @@ module Kumi
         class TypeChecker < VisitorPass
           def run(errors)
             functions_required = Set.new
+            @registry = state[:registry]
 
             visit_nodes_of_type(Kumi::Syntax::CallExpression, errors: errors) do |node, _decl, errs|
               validate_function_call(node, errs)
@@ -31,6 +32,18 @@ module Kumi
           end
 
           def get_function_signature(node, errors)
+            func = begin
+              @registry.function(node.fn_name)
+            rescue StandardError
+              nil
+            end
+            if func
+              return {
+                arity: func.params.size,
+                param_types: func.params.map { |p| p["type"] ? Kumi::Core::Types.parse(p["type"]) : nil }
+              }
+            end
+
             Kumi::Registry.signature(node.fn_name)
           rescue Kumi::Errors::UnknownFunction
             # Use old format for backward compatibility, but node.loc provides better location

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

@@ -22,20 +22,18 @@ module Kumi
               decl = definitions[name]
               next unless decl
 
-              begin
-                # Check if this declaration is marked as vectorized
-                if broadcast_meta[:vectorized_operations]&.key?(name)
-                  # Infer the element type and wrap in array
-                  element_type = infer_vectorized_element_type(decl.expression, types, broadcast_meta)
-                  types[name] = decl.is_a?(Kumi::Syntax::TraitDeclaration) ? { array: :boolean } : { array: element_type }
-                else
-                  # Normal type inference
-                  inferred_type = infer_expression_type(decl.expression, types, broadcast_meta, name)
-                  types[name] = inferred_type
-                end
-              rescue StandardError => e
-                report_type_error(errors, "Type inference failed: #{e.message}", location: decl&.loc)
+              # Check if this declaration is marked as vectorized
+              if broadcast_meta[:vectorized_operations]&.key?(name)
+                # Infer the element type and wrap in array
+                element_type = infer_vectorized_element_type(decl.expression, types, broadcast_meta)
+                types[name] = decl.is_a?(Kumi::Syntax::TraitDeclaration) ? { array: :boolean } : { array: element_type }
+              else
+                # Normal type inference
+                inferred_type = infer_expression_type(decl.expression, types, broadcast_meta, name)
+                types[name] = inferred_type
               end
+              # rescue StandardError => e
+              # report_type_error(errors, "Type inference failed: #{e.message}", location: decl&.loc)
             end
 
             state.with(:inferred_types, types)
@@ -108,9 +106,9 @@ module Kumi
                 expected_type = param_types[i] || param_types.last
                 next if expected_type.nil?
 
-                unless Types.compatible?(arg_type, expected_type)
-                  # Could add warning here in future, but for now just infer best type
-                end
+                # unless Types.compatible?(arg_type, expected_type)
+                # Could add warning here in future, but for now just infer best type
+                # end
               end
             end
 

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

@@ -14,10 +14,11 @@ module Kumi
           Atom        = Kumi::Core::AtomUnsatSolver::Atom
 
           def run(errors)
+            return state
             definitions = get_state(:declarations)
             @input_meta = get_state(:input_metadata) || {}
             @definitions = definitions
-            @evaluator = ConstantEvaluator.new(definitions)
+            @evaluator = UnsatConstantEvaluator.new(definitions)
 
             each_decl do |decl|
               if decl.expression.is_a?(CascadeExpression)

data/lib/kumi/core/analyzer/plans.rb

@@ -46,6 +46,18 @@ module Kumi
             !total_reduction?
           end
         end
+
+        InputPlan = Struct.new(
+          :source_path,     # Array<Symbol>
+          :axes,            # Array<Symbol> (logical site axes)
+          :dtype,           # Symbol/String
+          :key_policy,      # Symbol
+          :missing_policy,  # Symbol
+          :navigation_steps, # Array<Hash>   [{axis:, kind:, key:, alias:, loop_idx:, ...}]
+          :path_fqn, # String        "a.b.c"
+          :open_axis,
+          keyword_init: true
+        )
       end
     end
   end

data/lib/kumi/core/analyzer/state_serde.rb

@@ -26,7 +26,7 @@ module Kumi
         end
 
         def load_json(json_str)
-          h = JSON.parse(json_str)  # Don't symbolize keys - let decode_json_safe handle it
+          h = JSON.parse(json_str) # Don't symbolize keys - let decode_json_safe handle it
           ::Kumi::Core::Analyzer::AnalysisState.new(decode_json_safe(h))
         end
 
@@ -49,8 +49,8 @@ module Kumi
             # Check for special encoding markers first (before key transformation)
             if    x.key?("$sym") then x["$sym"].to_sym
             elsif x.key?("$set") then Set.new(x["$set"].map { |item| decode_json_safe(item) })
-            elsif x.key?("$ir")  then x["$ir"]  # Keep as string inspection for JSON round-trip
-            else 
+            elsif x.key?("$ir")  then x["$ir"] # Keep as string inspection for JSON round-trip
+            else
               # Regular hash - transform keys to symbols and recursively decode values
               x.transform_keys(&:to_sym).transform_values { |value| decode_json_safe(value) }
             end
@@ -61,4 +61,4 @@ module Kumi
       end
     end
   end
-end
+end

data/lib/kumi/core/analyzer/structs/access_plan.rb

@@ -6,8 +6,8 @@ module Kumi
       module Structs
         # One plan for a specific path and mode (path:mode)
         AccessPlan = Struct.new(:path, :containers, :leaf, :scope, :depth, :mode,
-                                :on_missing, :key_policy, :operations, keyword_init: true) do
-          def initialize(path:, containers:, leaf:, scope:, depth:, mode:, on_missing:, key_policy:, operations:)
+                                :on_missing, :key_policy, :operations, :chain, keyword_init: true) do
+          def initialize(path:, containers:, leaf:, scope:, depth:, mode:, on_missing:, key_policy:, operations:, chain:)
             super
             freeze
           end

data/lib/kumi/core/analyzer/unsat_constant_evaluator.rb

@@ -0,0 +1,59 @@
+# frozen_string_literal: true
+
+module Kumi
+  module Core
+    module Analyzer
+      class UnsatConstantEvaluator
+        include Syntax
+
+        def initialize(definitions)
+          @definitions = definitions
+          @memo = {}
+        end
+
+        OPERATORS = {
+          add: :+,
+          subtract: :-,
+          multiply: :*,
+          divide: :/
+        }.freeze
+
+        def evaluate(node, visited = Set.new)
+          return :unknown unless node
+          return @memo[node] if @memo.key?(node)
+          return node.value if node.is_a?(Literal)
+
+          result = case node
+                   when DeclarationReference then evaluate_binding(node, visited)
+                   when CallExpression then evaluate_call_expression(node, visited)
+                   else :unknown
+                   end
+
+          @memo[node] = result unless result == :unknown
+          result
+        end
+
+        private
+
+        def evaluate_binding(node, visited)
+          return :unknown if visited.include?(node.name)
+
+          visited << node.name
+          definition = @definitions[node.name]
+          return :unknown unless definition
+
+          evaluate(definition.expression, visited)
+        end
+
+        def evaluate_call_expression(node, visited)
+          return :unknown unless OPERATORS.key?(node.fn_name)
+
+          args = node.args.map { |arg| evaluate(arg, visited) }
+          return :unknown if args.any?(:unknown)
+
+          args.reduce(OPERATORS[node.fn_name])
+        end
+      end
+    end
+  end
+end

data/lib/kumi/core/compiler/access_codegen.rb

@@ -58,4 +58,4 @@ module Kumi
       end
     end
   end
-end
+end

data/lib/kumi/core/compiler/access_emit/base.rb

@@ -1,4 +1,5 @@
 # frozen_string_literal: true
+
 module Kumi
   module Core
     module Compiler
@@ -8,7 +9,8 @@ module Kumi
 
           # ---------- IR segmentation ----------
           def segment_ops(ops)
-            segs, cur = [], []
+            segs = []
+            cur = []
             i = 0
             while i < ops.length
               case ops[i][:type]
@@ -107,15 +109,17 @@ module Kumi
                 end
               else # :materialize, :read
                 # Important: for :materialize this is ALWAYS nil (never [])
-                return_val = 'nil'
+                return_val = "nil"
                 map_depth.positive? ? "next #{return_val}" : "return #{return_val}"
               end
             when :skip
               if mode == :materialize
-                return_val = preview_array ? '[]' : 'nil'
+                return_val = preview_array ? "[]" : "nil"
                 map_depth.positive? ? "next #{return_val}" : "return #{return_val}"
+              elsif map_depth.positive?
+                "next"
               else
-                map_depth.positive? ? "next" : (mode == :each_indexed ? "if block; return nil; else; return out; end" : "return out")
+                (mode == :each_indexed ? "if block; return nil; else; return out; end" : "return out")
               end
             else # :error
               %(raise KeyError, "Missing key '#{key}' at '#{path_key}' (#{mode})")
@@ -170,4 +174,4 @@ module Kumi
       end
     end
   end
-end
+end

data/lib/kumi/core/compiler/access_emit/each_indexed.rb

@@ -1,8 +1,11 @@
 # frozen_string_literal: true
+
 module Kumi::Core::Compiler::AccessEmit
   module EachIndexed
     extend Base
+
     module_function
+
     def build(plan)
       policy     = plan.on_missing || :error
       key_policy = plan.key_policy || :indifferent
@@ -13,7 +16,9 @@ module Kumi::Core::Compiler::AccessEmit
       code << "  out = []\n"
       code << "  node0 = data\n"
       code << "  idx_vec = []\n"
-      nodev, depth, loop_depth = "node0", 0, 0
+      nodev = "node0"
+      depth = 0
+      loop_depth = 0
 
       segs.each do |seg|
         if seg == :array
@@ -25,7 +30,9 @@ module Kumi::Core::Compiler::AccessEmit
           code << "    idx_vec[#{loop_depth}] = i#{loop_depth}\n"
           child = "node#{depth + 1}"
           code << "    #{child} = ary#{loop_depth}[i#{loop_depth}]\n"
-          nodev = child; depth += 1; loop_depth += 1
+          nodev = child
+          depth += 1
+          loop_depth += 1
         else
           seg.each do |(_, key, preview)|
             code << fetch_hash_code(node_var: nodev, key: key, key_policy: key_policy,
@@ -53,4 +60,4 @@ module Kumi::Core::Compiler::AccessEmit
       code
     end
   end
-end
+end

data/lib/kumi/core/compiler/access_emit/materialize.rb

@@ -1,8 +1,11 @@
 # frozen_string_literal: true
+
 module Kumi::Core::Compiler::AccessEmit
   module Materialize
     extend Base
+
     module_function
+
     def build(plan)
       policy     = plan.on_missing || :error
       key_policy = plan.key_policy || :indifferent
@@ -10,7 +13,9 @@ module Kumi::Core::Compiler::AccessEmit
       segs       = segment_ops(plan.operations)
 
       code = +"lambda do |data|\n"
-      nodev, depth, map_depth = "node0", 0, 0
+      nodev = "node0"
+      depth = 0
+      map_depth = 0
       code << "  #{nodev} = data\n"
 
       segs.each do |seg|
@@ -19,7 +24,9 @@ module Kumi::Core::Compiler::AccessEmit
           child = "node#{depth + 1}"
           code << "  #{nodev} = #{nodev}.map do |__e#{depth}|\n"
           code << "    #{child} = __e#{depth}\n"
-          nodev = child; depth += 1; map_depth += 1
+          nodev = child
+          depth += 1
+          map_depth += 1
         else
           seg.each do |(_, key, preview)|
             code << "  "
@@ -32,8 +39,8 @@ module Kumi::Core::Compiler::AccessEmit
       end
 
       while map_depth.positive?
-        code << "  " * map_depth + "#{nodev}\n"
-        code << "  " * (map_depth - 1) + "end\n"
+        code << (("  " * map_depth) + "#{nodev}\n")
+        code << (("  " * (map_depth - 1)) + "end\n")
         nodev = "node#{depth - 1}"
         depth -= 1
         map_depth -= 1
@@ -42,4 +49,4 @@ module Kumi::Core::Compiler::AccessEmit
       code
     end
   end
-end
+end

data/lib/kumi/core/compiler/access_emit/ravel.rb

@@ -1,8 +1,11 @@
 # frozen_string_literal: true
+
 module Kumi::Core::Compiler::AccessEmit
   module Ravel
     extend Base
+
     module_function
+
     def build(plan)
       policy     = plan.on_missing || :error
       key_policy = plan.key_policy || :indifferent
@@ -11,7 +14,9 @@ module Kumi::Core::Compiler::AccessEmit
 
       code = +"lambda do |data|\n"
       code << "  out = []\n"
-      nodev, depth, loop_depth = "node0", 0, 0
+      nodev = "node0"
+      depth = 0
+      loop_depth = 0
       code << "  #{nodev} = data\n"
 
       segs.each do |seg|
@@ -23,7 +28,9 @@ module Kumi::Core::Compiler::AccessEmit
           code << "  while (i#{loop_depth} += 1) < len#{loop_depth}\n"
           child = "node#{depth + 1}"
           code << "    #{child} = ary#{loop_depth}[i#{loop_depth}]\n"
-          nodev = child; depth += 1; loop_depth += 1
+          nodev = child
+          depth += 1
+          loop_depth += 1
         else
           seg.each do |(_, key, preview)|
             code << "  "
@@ -47,4 +54,4 @@ module Kumi::Core::Compiler::AccessEmit
       code
     end
   end
-end
+end

data/lib/kumi/core/compiler/access_emit/read.rb

@@ -1,8 +1,11 @@
 # frozen_string_literal: true
+
 module Kumi::Core::Compiler::AccessEmit
   module Read
     extend Base
+
     module_function
+
     def build(plan)
       policy     = plan.on_missing || :error
       key_policy = plan.key_policy || :indifferent
@@ -29,4 +32,4 @@ module Kumi::Core::Compiler::AccessEmit
       RUBY
     end
   end
-end
+end