optimizer: make optimizeSourcePaths work with sub queries

compiler/dag/op.go

@@ -410,3 +410,34 @@ func WalkT[T any](v reflect.Value, post func(T) T) {
 		}
 	}
 }
+
+func WalkTWithError[T any](v reflect.Value, post func(T) (T, error)) error {
+	switch v.Kind() {
+	case reflect.Array, reflect.Slice:
+		for i := range v.Len() {
+			if err := WalkTWithError(v.Index(i), post); err != nil {
+				return err
+			}
+		}
+	case reflect.Interface, reflect.Pointer:
+		if err := WalkTWithError(v.Elem(), post); err != nil {
+			return err
+		}
+	case reflect.Struct:
+		for _, field := range v.Fields() {
+			if err := WalkTWithError(field, post); err != nil {
+				return err
+			}
+		}
+	}
+	if v.CanSet() {
+		if t, ok := v.Interface().(T); ok {
+			r, err := post(t)
+			if err != nil {
+				return err
+			}
+			v.Set(reflect.ValueOf(r))
+		}
+	}
+	return nil
+}

compiler/optimizer/optimizer.go

@@ -132,8 +132,7 @@ func (o *Optimizer) Optimize(main *dag.Main) error {
 	seq = replaceSortAndHeadOrTailWithTop(seq)
 	o.optimizeParallels(seq)
 	seq = mergeFilters(seq)
-	seq, err := o.optimizeSourcePaths(seq)
-	if err != nil {
+	if err := o.optimizeSourcePaths(&seq); err != nil {
 		return err
 	}
 	seq = removePassOps(seq)
@@ -198,8 +197,8 @@ func (o *Optimizer) OptimizeDeleter(main *dag.Main, replicas int) error {
 	return nil
 }
 
-func (o *Optimizer) optimizeSourcePaths(seq dag.Seq) (dag.Seq, error) {
-	return walkEntries(seq, func(seq dag.Seq) (dag.Seq, error) {
+func (o *Optimizer) optimizeSourcePaths(seq *dag.Seq) error {
+	return dag.WalkTWithError(reflect.ValueOf(seq), func(seq dag.Seq) (dag.Seq, error) {
 		if len(seq) == 0 {
 			return nil, errors.New("internal error: optimizer encountered empty sequential operator")
 		}