From 200fe07473c1ca6696bd168538334824c94342c9 Mon Sep 17 00:00:00 2001
From: Aofei Sheng <aofei@aofeisheng.com>
Date: Mon, 12 Aug 2024 13:43:44 +0800
Subject: [PATCH] cmptest: add support for comparison with `llgo.expect` files

Fixes #671
---
 cmd/internal/run/run.go   | 10 +++++++---
 internal/build/build.go   | 13 +++++++------
 internal/build/cmptest.go | 37 +++++++++++++++++++++++++++++++++++--
 3 files changed, 49 insertions(+), 11 deletions(-)

diff --git a/cmd/internal/run/run.go b/cmd/internal/run/run.go
index 4d314cde..ec79238a 100644
--- a/cmd/internal/run/run.go
+++ b/cmd/internal/run/run.go
@@ -37,8 +37,8 @@ var Cmd = &base.Command{
 
 // llgo cmptest
 var CmpTestCmd = &base.Command{
-	UsageLine: "llgo cmptest [build flags] package [arguments...]",
-	Short:     "Compile programs by llgo and go, run them and do comparative tests (stdout/stderr/exitcode)",
+	UsageLine: "llgo cmptest [-genexpect] [build flags] package [arguments...]",
+	Short:     "Compile and run with llgo, compare result (stdout/stderr/exitcode) with go or llgo.expect; generate llgo.expect file if -genexpect is specified",
 }
 
 func init() {
@@ -55,9 +55,13 @@ func runCmpTest(cmd *base.Command, args []string) {
 }
 
 func runCmdEx(cmd *base.Command, args []string, mode build.Mode) {
+	conf := build.NewDefaultConf(mode)
+	if mode == build.ModeCmpTest && len(args) > 0 && args[0] == "-genexpect" {
+		conf.GenExpect = true
+		args = args[1:]
+	}
 	args, runArgs, err := parseRunArgs(args)
 	check(err)
-	conf := build.NewDefaultConf(mode)
 	conf.RunArgs = runArgs
 	build.Do(args, conf)
 }
diff --git a/internal/build/build.go b/internal/build/build.go
index aa59250a..2c056199 100644
--- a/internal/build/build.go
+++ b/internal/build/build.go
@@ -64,11 +64,12 @@ func needLLFile(mode Mode) bool {
 }
 
 type Config struct {
-	BinPath string
-	AppExt  string   // ".exe" on Windows, empty on Unix
-	OutFile string   // only valid for ModeBuild when len(pkgs) == 1
-	RunArgs []string // only valid for ModeRun
-	Mode    Mode
+	BinPath   string
+	AppExt    string   // ".exe" on Windows, empty on Unix
+	OutFile   string   // only valid for ModeBuild when len(pkgs) == 1
+	RunArgs   []string // only valid for ModeRun
+	GenExpect bool     // only valid for ModeCmpTest
+	Mode      Mode
 }
 
 func NewDefaultConf(mode Mode) *Config {
@@ -462,7 +463,7 @@ func linkMainPkg(ctx *context, pkg *packages.Package, pkgs []*aPackage, llFiles
 			os.Exit(s.ExitCode())
 		}
 	case ModeCmpTest:
-		cmpTest("", pkgPath, app, conf.RunArgs)
+		cmpTest(filepath.Dir(pkg.GoFiles[0]), pkgPath, app, conf.GenExpect, conf.RunArgs)
 	}
 	return
 }
diff --git a/internal/build/cmptest.go b/internal/build/cmptest.go
index d7b8a521..fc5cea0f 100644
--- a/internal/build/cmptest.go
+++ b/internal/build/cmptest.go
@@ -24,18 +24,51 @@ import (
 	"log"
 	"os"
 	"os/exec"
+	"path/filepath"
 )
 
-func cmpTest(dir, pkgPath, llApp string, runArgs []string) {
-	var goOut, goErr bytes.Buffer
+func cmpTest(dir, pkgPath, llApp string, genExpect bool, runArgs []string) {
 	var llgoOut, llgoErr bytes.Buffer
 	var llgoRunErr = runApp(runArgs, dir, &llgoOut, &llgoErr, llApp)
+
+	llgoExpect := formatExpect(llgoOut.Bytes(), llgoErr.Bytes(), llgoRunErr)
+	llgoExpectFile := filepath.Join(dir, "llgo.expect")
+	if genExpect {
+		if _, err := os.Stat(llgoExpectFile); !errors.Is(err, os.ErrNotExist) {
+			fatal(fmt.Errorf("llgo.expect file already exists: %s", llgoExpectFile))
+		}
+		if err := os.WriteFile(llgoExpectFile, llgoExpect, 0644); err != nil {
+			fatal(err)
+		}
+		return
+	}
+	if b, err := os.ReadFile(llgoExpectFile); err == nil {
+		checkEqual("llgo.expect", llgoExpect, b)
+		return
+	} else if !errors.Is(err, os.ErrNotExist) {
+		fatal(err)
+	}
+
+	var goOut, goErr bytes.Buffer
 	var goRunErr = runApp(runArgs, dir, &goOut, &goErr, "go", "run", pkgPath)
+
 	checkEqual("output", llgoOut.Bytes(), goOut.Bytes())
 	checkEqual("stderr", llgoErr.Bytes(), goErr.Bytes())
 	checkEqualRunErr(llgoRunErr, goRunErr)
 }
 
+func formatExpect(stdout, stderr []byte, runErr error) []byte {
+	var exitCode int
+	if runErr != nil {
+		if ee, ok := runErr.(*exec.ExitError); ok {
+			exitCode = ee.ExitCode()
+		} else { // This should never happen, but just in case.
+			exitCode = 255
+		}
+	}
+	return []byte(fmt.Sprintf("#stdout\n%s\n#stderr\n%s\n#exit %d\n", stdout, stderr, exitCode))
+}
+
 func checkEqualRunErr(llgoRunErr, goRunErr error) {
 	if llgoRunErr == goRunErr {
 		return