kumi 0.0.10 → 0.0.11

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

@@ -4,135 +4,152 @@ module Kumi
   module Core
     module Analyzer
       module Passes
-        # RESPONSIBILITY: Collect field metadata from input declarations and validate consistency
-        # DEPENDENCIES: :definitions
-        # PRODUCES: :inputs - Hash mapping field names to {type:, domain:} metadata
-        # INTERFACE: new(schema, state).run(errors)
+        # Emits per-node metadata:
+        #   :type, :domain
+        #   :container     => :scalar | :field | :array
+        #   :access_mode   => :field | :element          # how THIS node is read once reached
+        #   :enter_via     => :hash | :array           # how we HOP from parent to THIS node
+        #   :consume_alias => true|false               # inline array hop (alias is not a hash key)
+        #   :children      => { name => node_meta }    # optional
+        #
+        # Invariants:
+        # - Any nested array (child depth ≥ 1) must declare its element (i.e., have children).
+        # - Depth-0 inputs always: enter_via :hash, consume_alias false, access_mode :field.
         class InputCollector < PassBase
           def run(errors)
             input_meta = {}
 
-            schema.inputs.each do |field_decl|
-              unless field_decl.is_a?(Kumi::Syntax::InputDeclaration)
-                report_error(errors, "Expected InputDeclaration node, got #{field_decl.class}", location: field_decl.loc)
-                next
-              end
-
-              name = field_decl.name
-              existing = input_meta[name]
-
-              if existing
-                # Check for compatibility and merge
-                merged_meta = merge_field_metadata(existing, field_decl, errors)
-                input_meta[name] = merged_meta if merged_meta
-              else
-                # New field - collect its metadata
-                input_meta[name] = collect_field_metadata(field_decl, errors)
-              end
+            schema.inputs.each do |decl|
+              name = decl.name
+              input_meta[name] = collect_field_metadata(decl, errors, depth: 0, name: name)
             end
 
-            state.with(:inputs, freeze_nested_hash(input_meta))
+            input_meta.each_value(&:deep_freeze!)
+            state.with(:input_metadata, input_meta.freeze)
           end
 
           private
 
-          def collect_field_metadata(field_decl, errors)
-            validate_domain_type(field_decl, errors) if field_decl.domain
-
-            metadata = {
-              type: field_decl.type,
-              domain: field_decl.domain,
-              access_mode: field_decl.access_mode
-            }
-
-            # Process children if present
-            if field_decl.children && !field_decl.children.empty?
-              children_meta = {}
-              field_decl.children.each do |child_decl|
-                unless child_decl.is_a?(Kumi::Syntax::InputDeclaration)
-                  report_error(errors, "Expected InputDeclaration node in children, got #{child_decl.class}", location: child_decl.loc)
-                  next
-                end
-                children_meta[child_decl.name] = collect_field_metadata(child_decl, errors)
+          # ---------- builders ----------
+
+          def collect_field_metadata(decl, errors, depth:, name:)
+            children = nil
+            if decl.children&.any?
+              children = {}
+              decl.children.each do |child|
+                children[child.name] = collect_field_metadata(child, errors, depth: depth + 1, name: child.name)
               end
-              metadata[:children] = children_meta
             end
 
-            metadata
+            access_mode = decl.access_mode || :field
+
+            meta = Structs::InputMeta.new(
+              type: decl.type,
+              domain: decl.domain,
+              container: kind_from_type(decl.type),
+              access_mode: access_mode,
+              enter_via: :hash,
+              consume_alias: false,
+              children: children
+            )
+            stamp_edges_from!(meta, errors, parent_depth: depth)
+            validate_access_modes!(meta, errors, parent_depth: depth)
+            meta
           end
 
-          def merge_field_metadata(existing, field_decl, errors)
-            name = field_decl.name
-
-            # Check for type compatibility
-            if existing[:type] != field_decl.type && field_decl.type && existing[:type]
-              report_error(errors,
-                           "Field :#{name} declared with conflicting types: #{existing[:type]} vs #{field_decl.type}",
-                           location: field_decl.loc)
+          # ---------- edge stamping + validation ----------
+          #
+          # Sets child[:enter_via], child[:consume_alias], child[:access_mode] defaults,
+          # and validates nested arrays declare their element.
+          #
+          # Rules:
+          # - Common: any ARRAY child at child-depth ≥ 1 must have children (no bare nested array).
+          # - Parent :object → any child:
+          #     child.enter_via = :hash; child.consume_alias = false; child.access_mode ||= :field
+          # - Parent :array:
+          #     * If exactly one child:
+          #         - child.container ∈ {:scalar, :array} → via :array, consume_alias true, access_mode :element
+          #         - child.container == :field         → via :hash,  consume_alias false, access_mode :field
+          #     * Else (element object): every child → via :hash, consume_alias false, access_mode :field
+          def stamp_edges_from!(parent_meta, errors, parent_depth:)
+            kids = parent_meta.children || {}
+            return if kids.empty?
+
+            # Validate nested arrays anywhere below root
+            kids.each do |kname, child|
+              next unless child.container == :array
+
+              if !child.children || child.children.empty?
+                report_error(errors, "Nested array at :#{kname} must declare its element", location: nil)
+              end
             end
 
-            # Check for domain compatibility
-            if existing[:domain] != field_decl.domain && field_decl.domain && existing[:domain]
-              report_error(errors,
-                           "Field :#{name} declared with conflicting domains: #{existing[:domain].inspect} vs #{field_decl.domain.inspect}",
-                           location: field_decl.loc)
-            end
+            case parent_meta.container
+            when :object
+              kids.each_value do |child|
+                child.enter_via = :hash
+                child.consume_alias = false
+                child.access_mode = :field
+              end
 
-            # Validate domain type if provided
-            validate_domain_type(field_decl, errors) if field_decl.domain
-
-            # Merge metadata (later declarations override nil values)
-            merged = {
-              type: field_decl.type || existing[:type],
-              domain: field_decl.domain || existing[:domain],
-              access_mode: field_decl.access_mode || existing[:access_mode]
-            }
-
-            # Merge children if present
-            if field_decl.children && !field_decl.children.empty?
-              existing_children = existing[:children] || {}
-              new_children = {}
-
-              field_decl.children.each do |child_decl|
-                unless child_decl.is_a?(Kumi::Syntax::InputDeclaration)
-                  report_error(errors, "Expected InputDeclaration node in children, got #{child_decl.class}", location: child_decl.loc)
-                  next
+            when :array
+              # Array parents MUST explicitly declare their access mode
+              access_mode = parent_meta.access_mode
+              raise "Array must explicitly declare access_mode (:field or :element)" unless access_mode
+
+              case access_mode
+              when :field
+                # Array of objects: all children are fields accessed via hash
+                kids.each_value do |child|
+                  child.enter_via = :hash
+                  child.consume_alias = false
+                  child.access_mode = :field
                 end
 
-                child_name = child_decl.name
-                new_children[child_name] = if existing_children[child_name]
-                                             merge_field_metadata(existing_children[child_name], child_decl, errors)
-                                           else
-                                             collect_field_metadata(child_decl, errors)
-                                           end
-              end
+              when :element
+                _name, only = kids.first
+                only.enter_via = :array
+                only.consume_alias = true
+                only.access_mode = :element
 
-              merged[:children] = new_children
-            elsif existing[:children]
-              merged[:children] = existing[:children]
+              else
+                raise "Invalid access_mode :#{access_mode} for array (must be :field or :element)"
+              end
             end
-
-            merged
           end
 
-          def freeze_nested_hash(hash)
-            hash.each_value do |value|
-              freeze_nested_hash(value) if value.is_a?(Hash)
-            end
-            hash.freeze
-          end
+          # Enforce access_mode semantics are only used in valid contexts.
+          def validate_access_modes!(parent_meta, errors, parent_depth:)
+            kids = parent_meta.children || {}
+            return if kids.empty?
+
+            kids.each do |kname, child|
+              mode = child.access_mode
+              next unless mode
 
-          def validate_domain_type(field_decl, errors)
-            domain = field_decl.domain
-            return if valid_domain_type?(domain)
+              unless %i[field element].include?(mode)
+                report_error(errors, "Invalid access_mode for :#{kname}: #{mode.inspect}", location: nil)
+                next
+              end
 
-            report_error(errors,
-                         "Field :#{field_decl.name} has invalid domain constraint: #{domain.inspect}. Domain must be a Range, Array, or Proc",
-                         location: field_decl.loc)
+              if mode == :element
+                if parent_meta.container == :array
+                  single = (kids.size == 1)
+                  unless single && %i[scalar array].include?(child.container)
+                    report_error(errors, "access_mode :element only valid for single scalar/array element (at :#{kname})", location: nil)
+                  end
+                else
+                  report_error(errors, "access_mode :element only valid under array parent (at :#{kname})", location: nil)
+                end
+              end
+            end
           end
 
-          def valid_domain_type?(domain)
-            domain.is_a?(Range) || domain.is_a?(Array) || domain.is_a?(Proc)
+          def kind_from_type(t)
+            return :array if t == :array
+            return :field if t == :field
+
+            :scalar
           end
         end
       end

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

@@ -0,0 +1,293 @@
+# frozen_string_literal: true
+
+module Kumi
+  module Core
+    module Analyzer
+      module Passes
+        # Plans join and reduce operations for declarations.
+        # Determines reduction axes, flattening requirements, and join policies.
+        #
+        # DEPENDENCIES: :broadcasts, :scope_plans, :decl_shapes, :declarations, :input_metadata
+        # PRODUCES: :join_reduce_plans
+        class JoinReducePlanningPass < PassBase
+          include Kumi::Core::Analyzer::Plans
+
+          def run(_errors)
+            broadcasts = get_state(:broadcasts, required: false) || {}
+            scope_plans = get_state(:scope_plans, required: false) || {}
+            declarations = get_state(:declarations, required: true)
+            input_metadata = get_state(:input_metadata, required: true)
+
+            plans = {}
+
+            # Process reduction operations
+            process_reductions(broadcasts, scope_plans, declarations, input_metadata, plans)
+
+            # Process join operations (for non-reduction vectorized operations)
+            process_joins(broadcasts, scope_plans, declarations, plans)
+
+            # Return new state with join/reduce plans
+            state.with(:join_reduce_plans, plans.freeze)
+          end
+
+          private
+
+          def process_reductions(broadcasts, scope_plans, declarations, input_metadata, plans)
+            reduction_ops = broadcasts[:reduction_operations] || {}
+
+            reduction_ops.each do |name, info|
+              debug_reduction(name, info) if ENV["DEBUG_JOIN_REDUCE"]
+
+              # Get the source scope from scope_plans or infer from argument
+              source_scope = get_source_scope(name, info, scope_plans, declarations, input_metadata)
+
+              # Determine reduction axis and result scope
+              axis, result_scope = determine_reduction_axis(source_scope, info, scope_plans, name)
+
+              # Check for flattening requirements
+              flatten_indices = determine_flatten_indices(info)
+
+              plan = Reduce.new(
+                function: info[:function],
+                axis: axis,
+                source_scope: source_scope,
+                result_scope: result_scope,
+                flatten_args: flatten_indices
+              )
+
+              plans[name] = plan
+
+              debug_reduction_plan(name, plan) if ENV["DEBUG_JOIN_REDUCE"]
+            end
+          end
+
+          def process_joins(broadcasts, scope_plans, declarations, plans)
+            vectorized_ops = broadcasts[:vectorized_operations] || {}
+
+            # Process vectorized operations
+            vectorized_ops.each do |name, info|
+              # Skip if already processed as reduction
+              next if plans.key?(name)
+
+              debug_join(name, info) if ENV["DEBUG_JOIN_REDUCE"]
+
+              scope_plan = scope_plans[name]
+              next unless scope_plan
+
+              # Only need join planning if multiple arguments at different scopes
+              next unless requires_join?(declarations[name], scope_plan)
+
+              plan = Join.new(
+                policy: :zip, # Default to zip for array operations
+                target_scope: scope_plan.scope
+              )
+
+              plans[name] = plan
+
+              debug_join_plan(name, plan) if ENV["DEBUG_JOIN_REDUCE"]
+            end
+
+            # Process scalar declarations that need broadcasting to vectorized scopes
+            # (These are referenced by vectorized cascades but aren't vectorized themselves)
+            scope_plans.each do |name, scope_plan|
+              # Skip if already processed
+              next if plans.key?(name)
+              
+              # Skip if no vectorized target scope
+              next unless scope_plan.scope && !scope_plan.scope.empty?
+              
+              # Skip if already vectorized (handled above)
+              next if vectorized_ops.key?(name)
+              
+              # Check if this scalar declaration needs broadcasting
+              if needs_scalar_to_vector_broadcast?(name, scope_plan, declarations, vectorized_ops)
+                debug_scalar_broadcast(name, scope_plan) if ENV["DEBUG_JOIN_REDUCE"]
+                
+                plan = Join.new(
+                  policy: :broadcast, # Use broadcast policy for scalar-to-vector
+                  target_scope: scope_plan.scope
+                )
+
+                plans[name] = plan
+
+                debug_join_plan(name, plan) if ENV["DEBUG_JOIN_REDUCE"]
+              end
+            end
+          end
+
+          def get_source_scope(name, reduction_info, scope_plans, declarations, input_metadata)
+            # Always infer from the reduction argument - this is the full dimensional scope
+            infer_scope_from_argument(reduction_info[:argument], declarations, input_metadata)
+          end
+
+          def determine_reduction_axis(source_scope, reduction_info, scope_plans, name)
+            return [[], []] if source_scope.empty?
+
+            # Check if explicit axis is specified
+            if reduction_info[:axis]
+              axis = reduction_info[:axis]
+              result_scope = compute_result_scope(source_scope, axis)
+              return [axis, result_scope]
+            end
+
+            # Check if there's a scope plan that specifies what to preserve (result_scope)
+            scope_plan = scope_plans[name]
+            if scope_plan&.scope && !scope_plan.scope.empty?
+              desired_result_scope = scope_plan.scope
+              # Compute axis by removing the desired result dimensions
+              axis = source_scope - desired_result_scope
+              return [axis, desired_result_scope]
+            end
+
+            # Default: reduce innermost dimension (partial reduction)
+            axis = [source_scope.last]
+            result_scope = source_scope[0...-1]
+
+            [axis, result_scope]
+          end
+
+          def compute_result_scope(source_scope, axis)
+            # Remove specified axis dimensions from source scope
+            case axis
+            when :all
+              []
+            when Array
+              source_scope - axis
+            when Integer
+              # Numeric axis: remove that many innermost dimensions
+              source_scope[0...-axis]
+            else
+              source_scope
+            end
+          end
+
+          def determine_flatten_indices(reduction_info)
+            # Check for explicit flatten requirements
+            flatten = reduction_info[:flatten_argument_indices] || []
+            Array(flatten)
+          end
+
+          def requires_join?(declaration, scope_plan)
+            return false unless declaration
+            return false unless scope_plan.scope && !scope_plan.scope.empty?
+
+            expr = declaration.expression
+
+            case expr
+            when Kumi::Syntax::CallExpression
+              # Multiple arguments suggest potential join requirement
+              expr.args.size > 1
+            when Kumi::Syntax::CascadeExpression
+              # Cascades with vectorized target scope need join planning
+              # to handle cross-scope conditions and results
+              true
+            else
+              false
+            end
+          end
+
+          def infer_scope_from_argument(arg, declarations, input_metadata)
+            return [] unless arg
+
+            case arg
+            when Kumi::Syntax::InputElementReference
+              dims_from_path(arg.path, input_metadata)
+            when Kumi::Syntax::DeclarationReference
+              # Look up the declaration's scope if available
+              decl = declarations[arg.name]
+              decl ? infer_scope_from_argument(decl.expression, declarations, input_metadata) : []
+            when Kumi::Syntax::CallExpression
+              # For calls, use the deepest scope from arguments
+              scopes = arg.args.map { |a| infer_scope_from_argument(a, declarations, input_metadata) }
+              scopes.max_by(&:size) || []
+            else
+              []
+            end
+          end
+
+          def dims_from_path(path, input_metadata)
+            dims = []
+            meta = input_metadata
+
+            path.each do |seg|
+              field = meta[seg] || meta[seg.to_sym] || meta[seg.to_s]
+              break unless field
+
+              dims << seg.to_sym if field[:type] == :array
+
+              meta = field[:children] || {}
+            end
+
+            dims
+          end
+
+          def needs_scalar_to_vector_broadcast?(name, scope_plan, declarations, vectorized_ops)
+            # Check if this scalar declaration is referenced by any vectorized operation
+            # that requires it to be broadcast to a vectorized scope
+            
+            # Look for vectorized operations that reference this declaration
+            vectorized_ops.each do |vec_name, vec_info|
+              vec_decl = declarations[vec_name]
+              next unless vec_decl
+              
+              # Check if this vectorized operation references our scalar declaration
+              if declaration_references?(vec_decl.expression, name)
+                return true
+              end
+            end
+            
+            false
+          end
+
+          def declaration_references?(expr, target_name)
+            case expr
+            when Kumi::Syntax::DeclarationReference
+              expr.name == target_name
+            when Kumi::Syntax::CallExpression
+              expr.args.any? { |arg| declaration_references?(arg, target_name) }
+            when Kumi::Syntax::CascadeExpression
+              expr.cases.any? do |case_expr|
+                declaration_references?(case_expr.condition, target_name) ||
+                declaration_references?(case_expr.result, target_name)
+              end
+            when Kumi::Syntax::ArrayExpression
+              expr.elements.any? { |elem| declaration_references?(elem, target_name) }
+            else
+              false
+            end
+          end
+
+          # Debug helpers
+          def debug_reduction(name, info)
+            puts "\n=== Processing reduction: #{name} ==="
+            puts "Function: #{info[:function]}"
+            puts "Argument: #{info[:argument].class}"
+          end
+
+          def debug_reduction_plan(name, plan)
+            puts "Reduction plan for #{name}:"
+            puts "  Axis: #{plan.axis.inspect}"
+            puts "  Source scope: #{plan.source_scope.inspect}"
+            puts "  Result scope: #{plan.result_scope.inspect}"
+          end
+
+          def debug_join(name, info)
+            puts "\n=== Processing join: #{name} ==="
+            puts "Source: #{info[:source]}"
+          end
+
+          def debug_join_plan(name, plan)
+            puts "Join plan for #{name}:"
+            puts "  Target scope: #{plan.target_scope.inspect}"
+            puts "  Policy: #{plan.policy}"
+          end
+
+          def debug_scalar_broadcast(name, scope_plan)
+            puts "\n=== Processing scalar broadcast: #{name} ==="
+            puts "Target scope: #{scope_plan.scope.inspect}"
+          end
+        end
+      end
+    end
+  end
+end