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 globWithAffixes(pattern string) ([]string, string, string, error) {
	files, err := filepath.Glob(pattern)
	if err != nil {
		return nil, "", "", fmt.Errorf("invalid glob %q: %w", pattern, err)
	}
	if len(files) == 0 {
		return nil, "", "", fmt.Errorf("no files matched glob %q", pattern)
	}
	prefix, suffix := splitGlob(pattern)
	return files, prefix, suffix, nil
}

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

	// TODO: i know that this is copypaste, but i do not want make clousers or ifs inside cycle for now
	switch {
	case inputIsGlob && outputIsGlob:
		inputFiles, inputPrefix, inputSuffix, err := globWithAffixes(pattern.InputGlob)
		if err != nil {
			return nil, err
		}
		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, inputPrefix, inputSuffix, err := globWithAffixes(pattern.InputGlob)
		if err != nil {
			return nil, err
		}
		for _, inputPath := range inputFiles {
			wildcard := extractWildcard(inputPath, inputPrefix, inputSuffix)
			cases = append(cases, patternCase{
				name:       wildcard,
				inputPath:  inputPath,
				outputPath: pattern.OutputGlob,
			})
		}

	case !inputIsGlob && outputIsGlob:
		outputFiles, outputPrefix, outputSuffix, err := globWithAffixes(pattern.OutputGlob)
		if err != nil {
			return nil, err
		}
		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 buildTest(c *patternCase, argTemplate []string, useInputAsFile, useOutputAsFile bool) (*dsl.Test, error) {
	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.SetInputFile(inputName, inputContent)
	} else {
		t.SetStdin(inputContent)
	}

	if useOutputAsFile {
		t.SetOutputFile(outputName, outputContent)
	} else {
		t.SetStdout(outputContent)
	}

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

	return t, nil
}

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 {
		t, err := buildTest(&c, argTemplate, useInputAsFile, useOutputAsFile)
		if err != nil {
			return nil, err
		}
		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
}