build: refactor main module generation

internal/build/build.go

@@ -782,11 +782,14 @@ func linkMainPkg(ctx *context, pkg *packages.Package, pkgs []*aPackage, outputPa
 		}
 	})
 	// Generate main module file (needed for global variables even in library modes)
-	entryObjFile, err := genMainModuleFile(ctx, llssa.PkgRuntime, pkg, needRuntime, needPyInit)
+	entryPkg, err := genMainModule(ctx, llssa.PkgRuntime, pkg, needRuntime, needPyInit)
+	if err != nil {
+		return err
+	}
+	entryObjFile, err := exportObject(ctx, entryPkg.PkgPath, entryPkg.ExportFile, []byte(entryPkg.LPkg.String()))
 	if err != nil {
 		return err
 	}
-	// defer os.Remove(entryLLFile)
 	objFiles = append(objFiles, entryObjFile)
 
 	// Compile extra files from target configuration
@@ -908,118 +911,6 @@ func needStart(ctx *context) bool {
 	}
 }
 
-func genMainModuleFile(ctx *context, rtPkgPath string, pkg *packages.Package, needRuntime, needPyInit bool) (path string, err error) {
-	var (
-		pyInitDecl string
-		pyInit     string
-		rtInitDecl string
-		rtInit     string
-	)
-	mainPkgPath := pkg.PkgPath
-	if needRuntime {
-		rtInit = "call void @\"" + rtPkgPath + ".init\"()"
-		rtInitDecl = "declare void @\"" + rtPkgPath + ".init\"()"
-	}
-	if needPyInit {
-		pyInit = "call void @Py_Initialize()"
-		pyInitDecl = "declare void @Py_Initialize()"
-	}
-	declSizeT := "%size_t = type i64"
-	if is32Bits(ctx.buildConf.Goarch) {
-		declSizeT = "%size_t = type i32"
-	}
-	stdioDecl := ""
-	stdioNobuf := ""
-	if IsStdioNobuf() {
-		stdioDecl = `
-@stdout = external global ptr
-@stderr = external global ptr
-@__stdout = external global ptr
-@__stderr = external global ptr
-declare i32 @setvbuf(ptr, ptr, i32, %size_t)
-	`
-		stdioNobuf = `
-; Set stdout with no buffer
-%stdout_is_null = icmp eq ptr @stdout, null
-%stdout_ptr = select i1 %stdout_is_null, ptr @__stdout, ptr @stdout
-call i32 @setvbuf(ptr %stdout_ptr, ptr null, i32 2, %size_t 0)
-; Set stderr with no buffer
-%stderr_ptr = select i1 %stdout_is_null, ptr @__stderr, ptr @stderr
-call i32 @setvbuf(ptr %stderr_ptr, ptr null, i32 2, %size_t 0)
-	`
-	}
-	// TODO(lijie): workaround for libc-free
-	// Remove main/_start when -buildmode and libc are ready
-	startDefine := `
-define weak void @_start() {
-  ; argc = 0
-  %argc = add i32 0, 0
-  ; argv = null
-  %argv = inttoptr i64 0 to i8**
-  call i32 @main(i32 %argc, i8** %argv)
-  ret void
-}
-`
-	mainDefine := "define i32 @main(i32 noundef %0, ptr nocapture noundef readnone %1) local_unnamed_addr"
-	if !needStart(ctx) && isWasmTarget(ctx.buildConf.Goos) {
-		mainDefine = "define hidden noundef i32 @__main_argc_argv(i32 noundef %0, ptr nocapture noundef readnone %1) local_unnamed_addr"
-	}
-	if !needStart(ctx) {
-		startDefine = ""
-	}
-
-	var mainCode string
-	// For library modes (c-archive, c-shared), only generate global variables
-	if ctx.buildConf.BuildMode != BuildModeExe {
-		mainCode = `; ModuleID = 'main'
-source_filename = "main"
-@__llgo_argc = global i32 0, align 4
-@__llgo_argv = global ptr null, align 8
-`
-	} else {
-		// For executable mode, generate full main function
-		mainCode = fmt.Sprintf(`; ModuleID = 'main'
-source_filename = "main"
-%s
-@__llgo_argc = global i32 0, align 4
-@__llgo_argv = global ptr null, align 8
-%s
-%s
-%s
-declare void @"%s.init"()
-declare void @"%s.main"()
-define weak void @runtime.init() {
-  ret void
-}
-
-; TODO(lijie): workaround for syscall patch
-define weak void @"syscall.init"() {
-  ret void
-}
-
-%s
-
-%s {
-_llgo_0:
-  store i32 %%0, ptr @__llgo_argc, align 4
-  store ptr %%1, ptr @__llgo_argv, align 8
-  %s
-  %s
-  %s
-  call void @runtime.init()
-  call void @"%s.init"()
-  call void @"%s.main"()
-  ret i32 0
-}
-`, declSizeT, stdioDecl,
-			pyInitDecl, rtInitDecl, mainPkgPath, mainPkgPath,
-			startDefine, mainDefine, stdioNobuf,
-			pyInit, rtInit, mainPkgPath, mainPkgPath)
-	}
-
-	return exportObject(ctx, pkg.PkgPath+".main", pkg.ExportFile+"-main", []byte(mainCode))
-}
-
 func is32Bits(goarch string) bool {
 	return goarch == "386" || goarch == "arm" || goarch == "mips" || goarch == "wasm"
 }

internal/build/main_module.go

@@ -0,0 +1,221 @@
+//go:build !llgo
+// +build !llgo
+
+/*
+ * Copyright (c) 2024 The GoPlus Authors (goplus.org). All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package build
+
+import (
+	"go/token"
+	"go/types"
+	"runtime"
+
+	"github.com/goplus/llgo/internal/packages"
+	llvm "github.com/goplus/llvm"
+
+	llssa "github.com/goplus/llgo/ssa"
+)
+
+func genMainModule(ctx *context, rtPkgPath string, pkg *packages.Package, needRuntime, needPyInit bool) (Package, error) {
+	prog := ctx.prog
+	mainPkg := prog.NewPackage("", pkg.PkgPath+".main")
+
+	argcVar := mainPkg.NewVarEx("__llgo_argc", prog.Pointer(prog.Int32()))
+	argcVar.Init(prog.Zero(prog.Int32()))
+
+	charPtrType := prog.Pointer(prog.Byte())
+	argvValueType := prog.Pointer(charPtrType)
+	argvVar := mainPkg.NewVarEx("__llgo_argv", prog.Pointer(argvValueType))
+	argvVar.InitNil()
+
+	exportFile := pkg.ExportFile
+	if exportFile == "" {
+		exportFile = pkg.PkgPath
+	}
+	mainAPkg := &aPackage{
+		Package: &packages.Package{
+			PkgPath:    pkg.PkgPath + ".main",
+			ExportFile: exportFile + "-main",
+		},
+		LPkg: mainPkg,
+	}
+
+	if ctx.buildConf.BuildMode != BuildModeExe {
+		return mainAPkg, nil
+	}
+
+	runtimeStub := defineWeakNoArgStub(mainPkg, "runtime.init")
+	defineWeakNoArgStub(mainPkg, "syscall.init")
+
+	var pyInit llssa.Function
+	if needPyInit {
+		pyInit = declareNoArgFunc(mainPkg, "Py_Initialize")
+	}
+
+	var rtInit llssa.Function
+	if needRuntime {
+		rtInit = declareNoArgFunc(mainPkg, rtPkgPath+".init")
+	}
+
+	mainInit := declareNoArgFunc(mainPkg, pkg.PkgPath+".init")
+	mainMain := declareNoArgFunc(mainPkg, pkg.PkgPath+".main")
+
+	entryFn := defineEntryFunction(ctx, mainPkg, argcVar, argvVar, argvValueType, runtimeStub, mainInit, mainMain, pyInit, rtInit)
+
+	if needStart(ctx) {
+		defineStart(mainPkg, entryFn, argvValueType)
+	}
+
+	return mainAPkg, nil
+}
+
+func defineEntryFunction(ctx *context, pkg llssa.Package, argcVar, argvVar llssa.Global, argvType llssa.Type, runtimeStub, mainInit, mainMain llssa.Function, pyInit, rtInit llssa.Function) llssa.Function {
+	prog := pkg.Prog
+	entryName := "main"
+	if !needStart(ctx) && isWasmTarget(ctx.buildConf.Goos) {
+		entryName = "__main_argc_argv"
+	}
+	sig := newEntrySignature(argvType.RawType())
+	fn := pkg.NewFunc(entryName, sig, llssa.InC)
+	fnVal := pkg.Module().NamedFunction(entryName)
+	if entryName != "main" {
+		fnVal.SetVisibility(llvm.HiddenVisibility)
+		fnVal.SetUnnamedAddr(true)
+	}
+	b := fn.MakeBody(1)
+	b.Store(argcVar.Expr, fn.Param(0))
+	b.Store(argvVar.Expr, fn.Param(1))
+	if IsStdioNobuf() {
+		emitStdioNobuf(b, pkg, ctx.buildConf.Goarch)
+	}
+	if pyInit != nil {
+		b.Call(pyInit.Expr)
+	}
+	if rtInit != nil {
+		b.Call(rtInit.Expr)
+	}
+	b.Call(runtimeStub.Expr)
+	b.Call(mainInit.Expr)
+	b.Call(mainMain.Expr)
+	b.Return(prog.IntVal(0, prog.Int32()))
+	return fn
+}
+
+func defineStart(pkg llssa.Package, entry llssa.Function, argvType llssa.Type) {
+	fn := pkg.NewFunc("_start", llssa.NoArgsNoRet, llssa.InC)
+	pkg.Module().NamedFunction("_start").SetLinkage(llvm.WeakAnyLinkage)
+	b := fn.MakeBody(1)
+	prog := pkg.Prog
+	b.Call(entry.Expr, prog.IntVal(0, prog.Int32()), prog.Nil(argvType))
+	b.Return()
+}
+
+func declareNoArgFunc(pkg llssa.Package, name string) llssa.Function {
+	return pkg.NewFunc(name, llssa.NoArgsNoRet, llssa.InC)
+}
+
+func defineWeakNoArgStub(pkg llssa.Package, name string) llssa.Function {
+	fn := pkg.NewFunc(name, llssa.NoArgsNoRet, llssa.InC)
+	pkg.Module().NamedFunction(name).SetLinkage(llvm.WeakAnyLinkage)
+	b := fn.MakeBody(1)
+	b.Return()
+	return fn
+}
+
+func emitStdioNobuf(b llssa.Builder, pkg llssa.Package, goarch string) {
+	prog := pkg.Prog
+	streamType := prog.VoidPtr()
+	streamPtrType := prog.Pointer(streamType)
+	stdout := declareExternalPtrGlobal(pkg, "stdout", streamType)
+	stderr := declareExternalPtrGlobal(pkg, "stderr", streamType)
+	stdoutAlt := declareExternalPtrGlobal(pkg, "__stdout", streamType)
+	stderrAlt := declareExternalPtrGlobal(pkg, "__stderr", streamType)
+	sizeType := sizeTypeForArch(prog, goarch)
+	charPtrType := prog.Pointer(prog.Byte())
+	setvbuf := declareSetvbuf(pkg, streamPtrType, charPtrType, prog.Int32(), sizeType)
+
+	stdoutSlot := b.AllocaT(streamPtrType)
+	b.Store(stdoutSlot, stdout)
+	stderrSlot := b.AllocaT(streamPtrType)
+	b.Store(stderrSlot, stderr)
+	cond := b.BinOp(token.EQL, stdout, prog.Nil(streamPtrType))
+	b.IfThen(cond, func() {
+		b.Store(stdoutSlot, stdoutAlt)
+		b.Store(stderrSlot, stderrAlt)
+	})
+	stdoutPtr := b.Load(stdoutSlot)
+	stderrPtr := b.Load(stderrSlot)
+
+	mode := prog.IntVal(2, prog.Int32())
+	zeroSize := prog.Zero(sizeType)
+	nullBuf := prog.Nil(charPtrType)
+
+	b.Call(setvbuf.Expr, stdoutPtr, nullBuf, mode, zeroSize)
+	b.Call(setvbuf.Expr, stderrPtr, nullBuf, mode, zeroSize)
+}
+
+func declareExternalPtrGlobal(pkg llssa.Package, name string, valueType llssa.Type) llssa.Expr {
+	ptrType := pkg.Prog.Pointer(valueType)
+	global := pkg.NewVarEx(name, ptrType)
+	pkg.Module().NamedGlobal(name).SetLinkage(llvm.ExternalLinkage)
+	return global.Expr
+}
+
+func declareSetvbuf(pkg llssa.Package, streamPtrType, bufPtrType, intType, sizeType llssa.Type) llssa.Function {
+	sig := newSignature(
+		[]types.Type{
+			streamPtrType.RawType(),
+			bufPtrType.RawType(),
+			intType.RawType(),
+			sizeType.RawType(),
+		},
+		[]types.Type{intType.RawType()},
+	)
+	return pkg.NewFunc("setvbuf", sig, llssa.InC)
+}
+
+func tupleOf(tys ...types.Type) *types.Tuple {
+	if len(tys) == 0 {
+		return types.NewTuple()
+	}
+	vars := make([]*types.Var, len(tys))
+	for i, t := range tys {
+		vars[i] = types.NewParam(token.NoPos, nil, "", t)
+	}
+	return types.NewTuple(vars...)
+}
+
+func newSignature(params []types.Type, results []types.Type) *types.Signature {
+	return types.NewSignatureType(nil, nil, nil, tupleOf(params...), tupleOf(results...), false)
+}
+
+func newEntrySignature(argvType types.Type) *types.Signature {
+	return newSignature(
+		[]types.Type{types.Typ[types.Int32], argvType},
+		[]types.Type{types.Typ[types.Int32]},
+	)
+}
+
+func sizeTypeForArch(prog llssa.Program, arch string) llssa.Type {
+	if arch == "" {
+		arch = runtime.GOARCH
+	}
+	if is32Bits(arch) {
+		return prog.Uint32()
+	}
+	return prog.Uint64()
+}

internal/build/main_module_test.go

@@ -0,0 +1,76 @@
+//go:build !llgo
+// +build !llgo
+
+package build
+
+import (
+	"strings"
+	"testing"
+
+	"github.com/goplus/llvm"
+
+	"github.com/goplus/llgo/internal/packages"
+	llssa "github.com/goplus/llgo/ssa"
+)
+
+func init() {
+	llssa.Initialize(llssa.InitAll)
+}
+
+func TestGenMainModuleExecutable(t *testing.T) {
+	llvm.InitializeAllTargets()
+	t.Setenv(llgoStdioNobuf, "")
+	ctx := &context{
+		prog: llssa.NewProgram(nil),
+		buildConf: &Config{
+			BuildMode: BuildModeExe,
+			Goos:      "linux",
+			Goarch:    "amd64",
+		},
+	}
+	pkg := &packages.Package{PkgPath: "example.com/foo", ExportFile: "foo.a"}
+	mod, err := genMainModule(ctx, llssa.PkgRuntime, pkg, true, true)
+	if err != nil {
+		t.Fatalf("genMainModule() error = %v", err)
+	}
+	if mod.ExportFile != "foo.a-main" {
+		t.Fatalf("unexpected export file: %s", mod.ExportFile)
+	}
+	ir := mod.LPkg.String()
+	checks := []string{
+		"define i32 @main(",
+		"call void @Py_Initialize()",
+		"call void @\"example.com/foo.init\"()",
+		"define weak void @_start()",
+	}
+	for _, want := range checks {
+		if !strings.Contains(ir, want) {
+			t.Fatalf("main module IR missing %q:\n%s", want, ir)
+		}
+	}
+}
+
+func TestGenMainModuleLibrary(t *testing.T) {
+	llvm.InitializeAllTargets()
+	t.Setenv(llgoStdioNobuf, "")
+	ctx := &context{
+		prog: llssa.NewProgram(nil),
+		buildConf: &Config{
+			BuildMode: BuildModeCArchive,
+			Goos:      "linux",
+			Goarch:    "amd64",
+		},
+	}
+	pkg := &packages.Package{PkgPath: "example.com/foo", ExportFile: "foo.a"}
+	mod, err := genMainModule(ctx, llssa.PkgRuntime, pkg, false, false)
+	if err != nil {
+		t.Fatalf("genMainModule() error = %v", err)
+	}
+	ir := mod.LPkg.String()
+	if strings.Contains(ir, "define i32 @main") {
+		t.Fatalf("library mode should not emit main function:\n%s", ir)
+	}
+	if !strings.Contains(ir, "@__llgo_argc = global i32 0") {
+		t.Fatalf("library mode missing argc global:\n%s", ir)
+	}
+}