judge/runner/expander.go

package runner

import (
	"fmt"
	"os"
	"path/filepath"
	"strings"

	"github.com/Mond1c/judge/dsl"
)

func expandPattern(pattern *dsl.Pattern) ([]*dsl.Test, error) {
	if pattern.IsDirMode() {
		return expandDirPattern(pattern)
	}
	return expandGlobPattern(pattern)
}

type patternCase struct {
	name       string
	inputPath  string
	outputPath string
	dir        string
}

func expandGlobPattern(pattern *dsl.Pattern) ([]*dsl.Test, error) {
	inputIsGlob := strings.Contains(pattern.InputGlob, "*")
	outputIsGlob := strings.Contains(pattern.OutputGlob, "*")

	if pattern.InputGlob == "" && pattern.OutputGlob == "" {
		return nil, fmt.Errorf("pattern needs at least one of input/output/dirs")
	}
	if !inputIsGlob && !outputIsGlob {
		return nil, fmt.Errorf("pattern needs at least one glob field (input or output must contain *)")
	}

	var cases []patternCase

	switch {
	case inputIsGlob && outputIsGlob:
		inputFiles, err := filepath.Glob(pattern.InputGlob)
		if err != nil {
			return nil, fmt.Errorf("invalid input glob %q: %w", pattern.InputGlob, err)
		}
		if len(inputFiles) == 0 {
			return nil, fmt.Errorf("no files matched input glob %q", pattern.InputGlob)
		}
		inputPrefix, inputSuffix := splitGlob(pattern.InputGlob)
		outputPrefix, outputSuffix := splitGlob(pattern.OutputGlob)
		for _, inputPath := range inputFiles {
			wildcard := extractWildcard(inputPath, inputPrefix, inputSuffix)
			outputPath := outputPrefix + wildcard + outputSuffix
			cases = append(cases, patternCase{
				name:       wildcard,
				inputPath:  inputPath,
				outputPath: outputPath,
			})
		}

	case inputIsGlob && !outputIsGlob:
		inputFiles, err := filepath.Glob(pattern.InputGlob)
		if err != nil {
			return nil, fmt.Errorf("invalid input glob %q: %w", pattern.InputGlob, err)
		}
		if len(inputFiles) == 0 {
			return nil, fmt.Errorf("no files matched input glob %q", pattern.InputGlob)
		}
		inputPrefix, inputSuffix := splitGlob(pattern.InputGlob)
		for _, inputPath := range inputFiles {
			wildcard := extractWildcard(inputPath, inputPrefix, inputSuffix)
			cases = append(cases, patternCase{
				name:       wildcard,
				inputPath:  inputPath,
				outputPath: pattern.OutputGlob,
			})
		}

	case !inputIsGlob && outputIsGlob:
		outputFiles, err := filepath.Glob(pattern.OutputGlob)
		if err != nil {
			return nil, fmt.Errorf("invalid output glob %q: %w", pattern.OutputGlob, err)
		}
		if len(outputFiles) == 0 {
			return nil, fmt.Errorf("no files matched output glob %q", pattern.OutputGlob)
		}
		outputPrefix, outputSuffix := splitGlob(pattern.OutputGlob)
		for _, outputPath := range outputFiles {
			wildcard := extractWildcard(outputPath, outputPrefix, outputSuffix)
			cases = append(cases, patternCase{
				name:       wildcard,
				inputPath:  pattern.InputGlob,
				outputPath: outputPath,
			})
		}
	}

	return buildTests(cases, pattern.Args)
}

func expandDirPattern(pattern *dsl.Pattern) ([]*dsl.Test, error) {
	dirs, err := filepath.Glob(pattern.DirsGlob)
	if err != nil {
		return nil, fmt.Errorf("invalid dirs glob %q: %w", pattern.DirsGlob, err)
	}
	if len(dirs) == 0 {
		return nil, fmt.Errorf("no directories matched %q", pattern.DirsGlob)
	}

	var cases []patternCase
	for _, dir := range dirs {
		info, err := os.Stat(dir)
		if err != nil || !info.IsDir() {
			continue
		}
		cases = append(cases, patternCase{
			name:       filepath.Base(dir),
			inputPath:  filepath.Join(dir, pattern.InputFile),
			outputPath: filepath.Join(dir, pattern.OutputFile),
			dir:        dir,
		})
	}
	return buildTests(cases, pattern.Args)
}

func buildTests(cases []patternCase, argTemplate []string) ([]*dsl.Test, error) {
	useInputAsFile := argsContain(argTemplate, "{input_path}")
	useOutputAsFile := argsContain(argTemplate, "{output_path}")

	var tests []*dsl.Test
	for _, c := range cases {
		inputContent, err := os.ReadFile(c.inputPath)
		if err != nil {
			return nil, fmt.Errorf("read input %q: %w", c.inputPath, err)
		}
		outputContent, err := os.ReadFile(c.outputPath)
		if err != nil {
			return nil, fmt.Errorf("read output %q: %w", c.outputPath, err)
		}

		t := &dsl.Test{
			Name:     fmt.Sprintf("pattern:%s", c.name),
			Env:      map[string]string{},
			InFiles:  map[string]string{},
			OutFiles: map[string]string{},
			Stderr:   dsl.NoMatcher{},
		}

		inputName := filepath.Base(c.inputPath)
		outputName := filepath.Base(c.outputPath)

		if useInputAsFile {
			t.InFiles[inputName] = string(inputContent)
		} else {
			s := string(inputContent)
			t.Stdin = &s
		}

		if useOutputAsFile {
			t.OutFiles[outputName] = string(outputContent)
			t.Stdout = dsl.NoMatcher{}
		} else {
			t.Stdout = dsl.ExactMatcher{Value: string(outputContent)}
		}

		if len(argTemplate) > 0 {
			t.Args = substituteArgs(argTemplate, map[string]string{
				"{input_path}":  inputName,
				"{output_path}": outputName,
				"{name}":        c.name,
				"{dir}":         c.dir,
			})
		}

		tests = append(tests, t)
	}
	return tests, nil
}

func argsContain(args []string, placeholder string) bool {
	for _, a := range args {
		if strings.Contains(a, placeholder) {
			return true
		}
	}
	return false
}

func substituteArgs(template []string, vars map[string]string) []string {
	out := make([]string, len(template))
	for i, a := range template {
		for k, v := range vars {
			a = strings.ReplaceAll(a, k, v)
		}
		out[i] = a
	}
	return out
}

func splitGlob(pattern string) (prefix, suffix string) {
	before, after, found := strings.Cut(pattern, "*")
	if !found {
		return pattern, ""
	}
	return before, after
}

func extractWildcard(path, prefix, suffix string) string {
	s := strings.TrimPrefix(path, prefix)
	s = strings.TrimSuffix(s, suffix)
	return s
}