Merge pull request #1396 from goplus/pr-1395

build: fix ldflags rewrites and prep caching

Author: xushiwei

cl/_testgo/rewrite/dep/dep.go

@@ -0,0 +1,11 @@
+package dep
+
+import "fmt"
+
+var VarName = "dep-default"
+var VarPlain string
+
+func PrintVar() {
+	fmt.Printf("dep.VarName: %s\n", VarName)
+	fmt.Printf("dep.VarPlain: %s\n", VarPlain)
+}

cl/_testgo/rewrite/main.go

@@ -0,0 +1,23 @@
+package main
+
+import (
+	"fmt"
+	"runtime"
+
+	dep "github.com/goplus/llgo/cl/_testgo/rewrite/dep"
+)
+
+var VarName = "main-default"
+var VarPlain string
+
+func printLine(label, value string) {
+	fmt.Printf("%s: %s\n", label, value)
+}
+
+func main() {
+	printLine("main.VarName", VarName)
+	printLine("main.VarPlain", VarPlain)
+	dep.PrintVar()
+	printLine("runtime.GOROOT()", runtime.GOROOT())
+	printLine("runtime.Version()", runtime.Version())
+}

cl/_testgo/rewrite/out.ll

@@ -0,0 +1 @@
+;

cl/compile.go

@@ -127,6 +127,58 @@ type context struct {
 	cgoArgs    []llssa.Expr
 	cgoRet     llssa.Expr
 	cgoSymbols []string
+	rewrites   map[string]string
+}
+
+func (p *context) rewriteValue(name string) (string, bool) {
+	if p.rewrites == nil {
+		return "", false
+	}
+	dot := strings.LastIndex(name, ".")
+	if dot < 0 || dot == len(name)-1 {
+		return "", false
+	}
+	varName := name[dot+1:]
+	val, ok := p.rewrites[varName]
+	return val, ok
+}
+
+// isStringPtrType checks if typ is a pointer to the basic string type (*string).
+// This is used to validate that -ldflags -X can only rewrite variables of type *string,
+// not derived string types like "type T string".
+func (p *context) isStringPtrType(typ types.Type) bool {
+	ptr, ok := typ.(*types.Pointer)
+	if !ok {
+		return false
+	}
+	basic, ok := ptr.Elem().(*types.Basic)
+	return ok && basic.Kind() == types.String
+}
+
+func (p *context) globalFullName(g *ssa.Global) string {
+	name, _, _ := p.varName(g.Pkg.Pkg, g)
+	return name
+}
+
+func (p *context) rewriteInitStore(store *ssa.Store, g *ssa.Global) (string, bool) {
+	if p.rewrites == nil {
+		return "", false
+	}
+	fn := store.Block().Parent()
+	if fn == nil || fn.Synthetic != "package initializer" {
+		return "", false
+	}
+	if _, ok := store.Val.(*ssa.Const); !ok {
+		return "", false
+	}
+	if !p.isStringPtrType(g.Type()) {
+		return "", false
+	}
+	value, ok := p.rewriteValue(p.globalFullName(g))
+	if !ok {
+		return "", false
+	}
+	return value, true
 }
 
 type pkgState byte
@@ -176,7 +228,16 @@ func (p *context) compileGlobal(pkg llssa.Package, gbl *ssa.Global) {
 		log.Println("==> NewVar", name, typ)
 	}
 	g := pkg.NewVar(name, typ, llssa.Background(vtype))
-	if define {
+	if value, ok := p.rewriteValue(name); ok {
+		if p.isStringPtrType(gbl.Type()) {
+			g.Init(pkg.ConstString(value))
+		} else {
+			log.Printf("warning: ignoring rewrite for non-string variable %s (type: %v)", name, gbl.Type())
+			if define {
+				g.InitNil()
+			}
+		}
+	} else if define {
 		g.InitNil()
 	}
 }
@@ -816,6 +877,13 @@ func (p *context) compileInstr(b llssa.Builder, instr ssa.Instruction) {
 				return
 			}
 		}
+		if p.rewrites != nil {
+			if g, ok := va.(*ssa.Global); ok {
+				if _, ok := p.rewriteInitStore(v, g); ok {
+					return
+				}
+			}
+		}
 		ptr := p.compileValue(b, va)
 		val := p.compileValue(b, v.Val)
 		b.Store(ptr, val)
@@ -980,12 +1048,22 @@ type Patches = map[string]Patch
 
 // NewPackage compiles a Go package to LLVM IR package.
 func NewPackage(prog llssa.Program, pkg *ssa.Package, files []*ast.File) (ret llssa.Package, err error) {
-	ret, _, err = NewPackageEx(prog, nil, pkg, files)
+	ret, _, err = NewPackageEx(prog, nil, nil, pkg, files)
 	return
 }
 
 // NewPackageEx compiles a Go package to LLVM IR package.
-func NewPackageEx(prog llssa.Program, patches Patches, pkg *ssa.Package, files []*ast.File) (ret llssa.Package, externs []string, err error) {
+//
+// Parameters:
+//   - prog: target LLVM SSA program context
+//   - patches: optional package patches applied during compilation
+//   - rewrites: per-package string initializers rewritten at compile time
+//   - pkg: SSA package to compile
+//   - files: parsed AST files that belong to the package
+//
+// The rewrites map uses short variable names (without package qualifier) and
+// only affects string-typed globals defined in the current package.
+func NewPackageEx(prog llssa.Program, patches Patches, rewrites map[string]string, pkg *ssa.Package, files []*ast.File) (ret llssa.Package, externs []string, err error) {
 	pkgProg := pkg.Prog
 	pkgTypes := pkg.Pkg
 	oldTypes := pkgTypes
@@ -1018,6 +1096,7 @@ func NewPackageEx(prog llssa.Program, patches Patches, pkg *ssa.Package, files [
 			types.Unsafe: {kind: PkgDeclOnly}, // TODO(xsw): PkgNoInit or PkgDeclOnly?
 		},
 		cgoSymbols: make([]string, 0, 128),
+		rewrites:   rewrites,
 	}
 	ctx.initPyModule()
 	ctx.initFiles(pkgPath, files, pkgName == "C")

cl/rewrite_internal_test.go

@@ -0,0 +1,169 @@
+//go:build !llgo
+// +build !llgo
+
+package cl
+
+import (
+	"go/ast"
+	"go/parser"
+	"go/token"
+	"go/types"
+	"runtime"
+	"strings"
+	"testing"
+
+	gpackages "github.com/goplus/gogen/packages"
+	llssa "github.com/goplus/llgo/ssa"
+	"github.com/goplus/llgo/ssa/ssatest"
+	"golang.org/x/tools/go/ssa"
+	"golang.org/x/tools/go/ssa/ssautil"
+)
+
+func init() {
+	llssa.Initialize(llssa.InitAll | llssa.InitNative)
+}
+
+func compileWithRewrites(t *testing.T, src string, rewrites map[string]string) string {
+	t.Helper()
+	fset := token.NewFileSet()
+	file, err := parser.ParseFile(fset, "rewrite.go", src, 0)
+	if err != nil {
+		t.Fatalf("parse failed: %v", err)
+	}
+	importer := gpackages.NewImporter(fset)
+	mode := ssa.SanityCheckFunctions | ssa.InstantiateGenerics
+	pkg, _, err := ssautil.BuildPackage(&types.Config{Importer: importer}, fset,
+		types.NewPackage(file.Name.Name, file.Name.Name), []*ast.File{file}, mode)
+	if err != nil {
+		t.Fatalf("build package failed: %v", err)
+	}
+	prog := ssatest.NewProgramEx(t, nil, importer)
+	prog.TypeSizes(types.SizesFor("gc", runtime.GOARCH))
+	ret, _, err := NewPackageEx(prog, nil, rewrites, pkg, []*ast.File{file})
+	if err != nil {
+		t.Fatalf("NewPackageEx failed: %v", err)
+	}
+	return ret.String()
+}
+
+func TestRewriteGlobalStrings(t *testing.T) {
+	const src = `package rewritepkg
+var VarInit = "original_value"
+var VarPlain string
+func Use() string { return VarInit + VarPlain }
+`
+	ir := compileWithRewrites(t, src, map[string]string{
+		"VarInit":  "rewrite_init",
+		"VarPlain": "rewrite_plain",
+	})
+	if strings.Contains(ir, "original_value") {
+		t.Fatalf("original initializer still present:\n%s", ir)
+	}
+	for _, want := range []string{`c"rewrite_init"`, `c"rewrite_plain"`} {
+		if !strings.Contains(ir, want) {
+			t.Fatalf("missing %s in IR:\n%s", want, ir)
+		}
+	}
+}
+
+func TestRewriteSkipsNonConstStores(t *testing.T) {
+	const src = `package rewritepkg
+import "strings"
+var VarInit = strings.ToUpper("original_value")
+var VarPlain string
+func Use() string { return VarInit + VarPlain }
+`
+	ir := compileWithRewrites(t, src, map[string]string{
+		"VarInit":  "rewrite_init",
+		"VarPlain": "rewrite_plain",
+	})
+	if !strings.Contains(ir, `c"rewrite_init"`) {
+		t.Fatalf("expected rewrite_init constant to remain:\n%s", ir)
+	}
+	if !strings.Contains(ir, "strings.ToUpper") {
+		t.Fatalf("expected call to strings.ToUpper in IR:\n%s", ir)
+	}
+}
+
+func TestRewriteValueNoDot(t *testing.T) {
+	ctx := &context{rewrites: map[string]string{"VarInit": "rewrite_init"}}
+	if _, ok := ctx.rewriteValue("VarInit"); ok {
+		t.Fatalf("rewriteValue should skip names without package prefix")
+	}
+	if _, ok := ctx.rewriteValue("pkg."); ok {
+		t.Fatalf("rewriteValue should skip trailing dot names")
+	}
+}
+
+func TestIsStringPtrTypeDefault(t *testing.T) {
+	ctx := &context{}
+	if ctx.isStringPtrType(types.NewPointer(types.Typ[types.Int])) {
+		t.Fatalf("expected non-string pointer to return false")
+	}
+}
+
+func TestIsStringPtrTypeBranches(t *testing.T) {
+	ctx := &context{}
+	if ctx.isStringPtrType(types.NewSlice(types.Typ[types.String])) {
+		t.Fatalf("slice should trigger default branch and return false")
+	}
+	if ctx.isStringPtrType(nil) {
+		t.Fatalf("nil type should return false")
+	}
+	if !ctx.isStringPtrType(types.NewPointer(types.Typ[types.String])) {
+		t.Fatalf("*string should return true")
+	}
+}
+
+func TestRewriteIgnoredInNonInitStore(t *testing.T) {
+	const src = `package rewritepkg
+var VarInit = "original_value"
+func Override() { VarInit = "override_value" }
+`
+	ir := compileWithRewrites(t, src, map[string]string{"VarInit": "rewrite_init"})
+	if !strings.Contains(ir, `c"override_value"`) {
+		t.Fatalf("override store should retain original literal:\n%s", ir)
+	}
+	if !strings.Contains(ir, `c"rewrite_init"`) {
+		t.Fatalf("global initializer should still be rewritten:\n%s", ir)
+	}
+}
+
+func TestRewriteMissingEntry(t *testing.T) {
+	const src = `package rewritepkg
+var VarInit = "original_value"
+var VarOther = "other_value"
+`
+	ir := compileWithRewrites(t, src, map[string]string{"VarInit": "rewrite_init"})
+	if !strings.Contains(ir, `c"other_value"`) {
+		t.Fatalf("VarOther should keep original initializer:\n%s", ir)
+	}
+	if !strings.Contains(ir, `c"rewrite_init"`) {
+		t.Fatalf("VarInit should still be rewritten:\n%s", ir)
+	}
+}
+
+func TestRewriteIgnoresNonStringVar(t *testing.T) {
+	const src = `package rewritepkg
+type wrapper struct{ v int }
+var VarStruct = wrapper{v: 1}
+`
+	ir := compileWithRewrites(t, src, map[string]string{"VarStruct": "rewrite_struct"})
+	if strings.Contains(ir, `c"rewrite_struct"`) {
+		t.Fatalf("non-string variables must not be rewritten:\n%s", ir)
+	}
+}
+
+func TestRewriteIgnoresStringAlias(t *testing.T) {
+	const src = `package rewritepkg
+type T string
+var VarAlias T = "original_value"
+`
+	ir := compileWithRewrites(t, src, map[string]string{"VarAlias": "rewrite_alias"})
+	if strings.Contains(ir, `c"rewrite_alias"`) {
+		t.Fatalf("string alias types must not be rewritten:\n%s", ir)
+	}
+	if !strings.Contains(ir, `c"original_value"`) {
+		t.Fatalf("original value should remain for alias type:\n%s", ir)
+	}
+}