feat(raycicmd): add matrix_expand

Rayci will expand matrix into individual Buildkite steps, allowing for fine-grained depends_on targeting of specific matrix combinations.

Topic: matrix-expand
Labels: draft
Signed-off-by: andrew <andrew@anyscale.com>

Author: andrew-anyscale

raycicmd/matrix_expand.go

@@ -0,0 +1,379 @@
+package raycicmd
+
+import (
+	"fmt"
+	"sort"
+	"strings"
+)
+
+// matrixConfig represents a parsed matrix definition.
+type matrixConfig struct {
+	Setup map[string][]string // dimension -> values
+}
+
+// matrixInstance represents one expanded combination.
+type matrixInstance struct {
+	Values map[string]string // dimension -> selected value
+}
+
+// matrixSelector represents a dependency selector with optional matrix filter.
+type matrixSelector struct {
+	Key    string            // base step key
+	Matrix map[string]string // partial dimension constraints (nil = all)
+}
+
+// parseMatrixConfig parses the matrix field from a step.
+// Examples:
+// Named dimensions matrix
+//
+//	matrix:
+//	  setup:
+//	    python: ["3.10", "3.11"]
+//	    cuda: ["12.1.1", "12.8.1"]
+//
+// Simple array matrix (no dimension name)
+//
+//	matrix:
+//	  - "darwin"
+//	  - "Linux"
+//	  - "Windows"
+func parseMatrixConfig(v any) (*matrixConfig, error) {
+	cfg := &matrixConfig{
+		Setup: make(map[string][]string),
+	}
+
+	switch val := v.(type) {
+	case []any:
+		// Simple array: matrix: ["3.10", "3.11"]
+		values, err := anySliceToStringSlice(val)
+		if err != nil {
+			return nil, fmt.Errorf("invalid matrix array: %w", err)
+		}
+		cfg.Setup[""] = values
+
+	case map[string]any:
+		// Map with setup and optional adjustments
+		if setup, ok := val["setup"]; ok {
+			setupMap, ok := setup.(map[string]any)
+			if !ok {
+				return nil, fmt.Errorf("matrix.setup must be a map")
+			}
+			for dim, vals := range setupMap {
+				valsSlice, ok := vals.([]any)
+				if !ok {
+					return nil, fmt.Errorf("matrix.setup.%s must be an array", dim)
+				}
+				strVals, err := anySliceToStringSlice(valsSlice)
+				if err != nil {
+					return nil, fmt.Errorf("invalid values for dimension %s: %w", dim, err)
+				}
+				cfg.Setup[dim] = strVals
+			}
+		} else {
+			return nil, fmt.Errorf("matrix map must have 'setup' key")
+		}
+
+		if _, ok := val["adjustments"]; ok {
+			return nil, fmt.Errorf("matrix.adjustments is not supported")
+		}
+
+	default:
+		return nil, fmt.Errorf("matrix must be an array or map, got %T", v)
+	}
+
+	return cfg, nil
+}
+
+// expandMatrix generates all combinations from a matrix config.
+func expandMatrix(cfg *matrixConfig) []*matrixInstance {
+	if len(cfg.Setup) == 0 {
+		return nil
+	}
+
+	dims := sortedMatrixDimensions(cfg)
+	return generateMatrixCartesianProduct(cfg.Setup, dims)
+}
+
+func sortedMatrixDimensions(cfg *matrixConfig) []string {
+	dims := make([]string, 0, len(cfg.Setup))
+	for dim := range cfg.Setup {
+		dims = append(dims, dim)
+	}
+	sort.Strings(dims)
+	return dims
+}
+
+func generateMatrixCartesianProduct(setup map[string][]string, dims []string) []*matrixInstance {
+	if len(dims) == 0 {
+		return []*matrixInstance{{Values: make(map[string]string)}}
+	}
+
+	firstDim := dims[0]
+	restDims := dims[1:]
+
+	restInstances := generateMatrixCartesianProduct(setup, restDims)
+
+	var result []*matrixInstance
+	for _, val := range setup[firstDim] {
+		for _, rest := range restInstances {
+			inst := &matrixInstance{
+				Values: make(map[string]string),
+			}
+			for k, v := range rest.Values {
+				inst.Values[k] = v
+			}
+			inst.Values[firstDim] = val
+			result = append(result, inst)
+		}
+	}
+
+	return result
+}
+
+// generateMatrixInstanceKey creates the expanded key for an instance.
+// Format: {base-key}-{dim1}{val1}-{dim2}{val2} (dims sorted alphabetically)
+func generateMatrixInstanceKey(baseKey string, cfg *matrixConfig, values map[string]string) string {
+	dims := sortedMatrixDimensions(cfg)
+
+	parts := []string{baseKey}
+	for _, dim := range dims {
+		val := values[dim]
+		if dim == "" {
+			// Simple array matrix (no dimension name) - use value as-is (e.g., 3.10)
+			parts = append(parts, val)
+		} else {
+			// Named dimensions: use dimension-value format (e.g., python3.11)
+			parts = append(parts, dim+val)
+		}
+	}
+
+	return strings.Join(parts, "-")
+}
+
+// generateMatrixTags creates auto-generated tags for matrix values.
+func generateMatrixTags(values map[string]string) []string {
+	var tags []string
+	for dim, val := range values {
+		if dim == "" {
+			// Simple array: just use the value as tag (e.g., 3.10)
+			tags = append(tags, val)
+		} else {
+			// Named dimension: dimension-value format (e.g., python-3.11)
+			tags = append(tags, dim+"-"+val)
+		}
+	}
+	sort.Strings(tags)
+	return tags
+}
+
+// substituteMatrixValues replaces {{matrix.X}} placeholders in any value.
+func substituteMatrixValues(v any, values map[string]string) any {
+	switch val := v.(type) {
+	case string:
+		result := val
+		for dim, dimVal := range values {
+			if dim == "" {
+				// Simple array: use {{matrix}}
+				result = strings.ReplaceAll(result, "{{matrix}}", dimVal)
+			} else {
+				// Named dimension: use {{matrix.dim}} (e.g., {{matrix.python}})
+				result = strings.ReplaceAll(result, "{{matrix."+dim+"}}", dimVal)
+			}
+		}
+		return result
+
+	case map[string]any:
+		result := make(map[string]any)
+		for k, v := range val {
+			result[k] = substituteMatrixValues(v, values)
+		}
+		return result
+
+	case []any:
+		result := make([]any, len(val))
+		for i, v := range val {
+			result[i] = substituteMatrixValues(v, values)
+		}
+		return result
+
+	default:
+		return v
+	}
+}
+
+// parseMatrixDependsOn parses depends_on field, returns list of selectors.
+// Handles: nil, string, []string, []any, []map with {key, matrix} selector syntax
+func parseMatrixDependsOn(v any) ([]*matrixSelector, error) {
+	if v == nil {
+		return nil, nil
+	}
+
+	switch val := v.(type) {
+	case string:
+		return []*matrixSelector{{Key: val}}, nil
+
+	case []string:
+		var selectors []*matrixSelector
+		for _, key := range val {
+			selectors = append(selectors, &matrixSelector{Key: key})
+		}
+		return selectors, nil
+
+	case []any:
+		var selectors []*matrixSelector
+		for i, item := range val {
+			switch itemVal := item.(type) {
+			case string:
+				selectors = append(selectors, &matrixSelector{Key: itemVal})
+
+			case map[string]any:
+				sel, err := parseMatrixSelectorMap(itemVal)
+				if err != nil {
+					return nil, fmt.Errorf("depends_on[%d]: %w", i, err)
+				}
+				selectors = append(selectors, sel)
+
+			default:
+				return nil, fmt.Errorf("depends_on[%d]: unexpected type %T", i, item)
+			}
+		}
+		return selectors, nil
+
+	default:
+		return nil, fmt.Errorf("depends_on must be string or array, got %T", v)
+	}
+}
+
+func parseMatrixSelectorMap(m map[string]any) (*matrixSelector, error) {
+	sel := &matrixSelector{}
+
+	key, ok := m["key"]
+	if !ok {
+		return nil, fmt.Errorf("selector missing 'key' field")
+	}
+	keyStr, ok := key.(string)
+	if !ok {
+		return nil, fmt.Errorf("selector 'key' must be a string")
+	}
+	sel.Key = keyStr
+
+	if matrix, ok := m["matrix"]; ok {
+		matrixMap, ok := matrix.(map[string]any)
+		if !ok {
+			return nil, fmt.Errorf("selector 'matrix' must be a map")
+		}
+		sel.Matrix = make(map[string]string)
+		for k, v := range matrixMap {
+			s, ok := v.(string)
+			if !ok {
+				return nil, fmt.Errorf("selector 'matrix.%s' must be a string", k)
+			}
+			sel.Matrix[k] = s
+		}
+	}
+
+	return sel, nil
+}
+
+// expandMatrixSelector expands a selector to matching expanded keys.
+func expandMatrixSelector(
+	sel *matrixSelector,
+	registry map[string]*matrixConfig,
+	expandedKeys map[string][]string, // baseKey -> list of expanded keys
+) ([]string, error) {
+	cfg, ok := registry[sel.Key]
+	if !ok {
+		// Not a matrix step, just return the key as-is
+		return []string{sel.Key}, nil
+	}
+
+	if sel.Matrix == nil {
+		// No filter - return the base key (which is the meta-step)
+		return []string{sel.Key}, nil
+	}
+
+	// Validate that all selector dimensions exist in the config
+	for dim := range sel.Matrix {
+		if _, ok := cfg.Setup[dim]; !ok {
+			return nil, fmt.Errorf("selector dimension %q not found in matrix for %q", dim, sel.Key)
+		}
+	}
+
+	// Find all expanded keys that match the selector
+	allExpanded := expandedKeys[sel.Key]
+	var matches []string
+
+	instances := expandMatrix(cfg)
+	for i, inst := range instances {
+		if matchesMatrixSelector(inst, sel) {
+			if i < len(allExpanded) {
+				matches = append(matches, allExpanded[i])
+			}
+		}
+	}
+
+	if len(matches) == 0 {
+		return nil, fmt.Errorf("no matches for selector {key: %q, matrix: %v}", sel.Key, sel.Matrix)
+	}
+
+	return matches, nil
+}
+
+func matchesMatrixSelector(inst *matrixInstance, sel *matrixSelector) bool {
+	for dim, val := range sel.Matrix {
+		if inst.Values[dim] != val {
+			return false
+		}
+	}
+	return true
+}
+
+// hasMatrixPlaceholder checks if a string contains any {{matrix...}} placeholder.
+func hasMatrixPlaceholder(s string) bool {
+	return strings.Contains(s, "{{matrix")
+}
+
+// anySliceToStringSlice converts []any to []string.
+func anySliceToStringSlice(arr []any) ([]string, error) {
+	result := make([]string, len(arr))
+	for i, v := range arr {
+		s, ok := v.(string)
+		if !ok {
+			return nil, fmt.Errorf("element %d is not a string: %T", i, v)
+		}
+		result[i] = s
+	}
+	return result, nil
+}
+
+// deepCopyStepMap creates a deep copy of a step map.
+func deepCopyStepMap(step map[string]any) map[string]any {
+	return deepCopyAnyMap(step)
+}
+
+func deepCopyAnyMap(m map[string]any) map[string]any {
+	result := make(map[string]any)
+	for k, v := range m {
+		result[k] = deepCopyAnyValue(v)
+	}
+	return result
+}
+
+func deepCopyAnyValue(v any) any {
+	switch val := v.(type) {
+	case map[string]any:
+		return deepCopyAnyMap(val)
+	case []any:
+		return deepCopyAnySlice(val)
+	default:
+		return v
+	}
+}
+
+func deepCopyAnySlice(s []any) []any {
+	result := make([]any, len(s))
+	for i, v := range s {
+		result[i] = deepCopyAnyValue(v)
+	}
+	return result
+}