PyInit

cl/_testpy/callpy/out.ll

@@ -26,6 +26,7 @@ _llgo_2:                                          ; preds = %_llgo_1, %_llgo_0
 
 define void @main(i32 %0, ptr %1) {
 _llgo_0:
+  call void @Py_Initialize()
   store i32 %0, ptr @__llgo_argc, align 4
   store ptr %1, ptr @__llgo_argv, align 8
   call void @"github.com/goplus/llgo/internal/runtime.init"()
@@ -57,3 +58,5 @@ declare double @PyFloat_AsDouble()
 declare i32 @printf(ptr, ...)
 
 declare ptr @PyBytes_AsString()
+
+declare void @Py_Initialize()

cl/_testpy/math/out.ll

@@ -1,7 +1,7 @@
 ; ModuleID = 'math'
 source_filename = "math"
 
-@__llgo_py.math.sqrt = external global ptr
+@__llgo_py.math.sqrt = linkonce global ptr null
 @"math.init$guard" = global ptr null
 @__llgo_py.math = linkonce global ptr null
 @0 = private unnamed_addr constant [5 x i8] c"math\00", align 1

cl/cltest/cltest.go

@@ -159,6 +159,11 @@ func TestCompileEx(t *testing.T, src any, fname, expected string) {
 	if err != nil {
 		t.Fatal("cl.NewPackage failed:", err)
 	}
+
+	if prog.NeedPyInit() { // call PyInit if needed
+		ret.PyInit()
+	}
+
 	if v := ret.String(); v != expected {
 		t.Fatalf("\n==> got:\n%s\n==> expected:\n%s\n", v, expected)
 	}

cl/compile.go

@@ -380,8 +380,7 @@ const (
 )
 
 func callRuntimeInit(b llssa.Builder, pkg llssa.Package) {
-	sig := types.NewSignatureType(nil, nil, nil, nil, nil, false)
+	fn := pkg.NewFunc(RuntimeInit, llssa.NoArgsNoRet, llssa.InC) // don't need to convert runtime.init
-	fn := pkg.NewFunc(RuntimeInit, sig, llssa.InC) // don't need to convert runtime.init
 	b.Call(fn.Expr)
 }
 

cl/compile_test.go

@@ -33,7 +33,7 @@ func TestFromTestpy(t *testing.T) {
 }
 
 func TestFromTestlibc(t *testing.T) {
-	cltest.FromDir(t, "", "./_testlibc", false)
+	cltest.FromDir(t, "", "./_testlibc", true)
 }
 
 func TestFromTestrt(t *testing.T) {

cl/import.go

@@ -214,7 +214,7 @@ func (p *context) initLink(line string, prefix int, f func(inPkgName string) (fu
 			} else {
 				panic(line + ": no specified call convention. eg. //go:linkname Printf C.printf")
 			}
-		} else {
+		} else if c := inPkgName[0]; c >= 'A' && c <= 'Z' {
 			fmt.Fprintln(os.Stderr, "==>", line)
 			fmt.Fprintf(os.Stderr, "llgo: linkname %s not found and ignored\n", inPkgName)
 		}

internal/build/build.go

@@ -160,12 +160,22 @@ func Do(args []string, conf *Config) {
 	}
 }
 
-func setNeedRuntime(pkg *packages.Package) {
+func setNeedRuntimeOrPyInit(pkg *packages.Package, needRuntime, needPyInit bool) {
-	pkg.ID = "" // just use pkg.Module to mark it needs runtime
+	v := []byte{'0', '0'}
+	if needRuntime {
+		v[0] = '1'
+	}
+	if needPyInit {
+		v[1] = '1'
+	}
+	pkg.ID = string(v) // just use pkg.ID to mark it needs runtime
 }
 
-func isNeedRuntime(pkg *packages.Package) bool {
+func isNeedRuntimeOrPyInit(pkg *packages.Package) (needRuntime, needPyInit bool) {
-	return pkg.ID == ""
+	if len(pkg.ID) == 2 {
+		return pkg.ID[0] == '1', pkg.ID[1] == '1'
+	}
+	return
 }
 
 func buildAllPkgs(prog llssa.Program, initial []*packages.Package, mode Mode, verbose bool) (pkgs []*aPackage) {
@@ -204,9 +214,7 @@ func buildAllPkgs(prog llssa.Program, initial []*packages.Package, mode Mode, ve
 			pkg.ExportFile = command
 		default:
 			buildPkg(prog, aPkg, mode, verbose)
-			if prog.NeedRuntime() {
+			setNeedRuntimeOrPyInit(pkg, prog.NeedRuntime(), prog.NeedPyInit())
-				setNeedRuntime(pkg)
-			}
 		}
 	}
 	return
@@ -225,27 +233,43 @@ func linkMainPkg(pkg *packages.Package, pkgs []*aPackage, runtimeFiles []string,
 	args[1] = app
 	args[2] = "-Wno-override-module"
 	needRuntime := false
+	needPyInit := false
 	packages.Visit([]*packages.Package{pkg}, nil, func(p *packages.Package) {
 		if p.ExportFile != "" && !isRuntimePkg(p.PkgPath) { // skip packages that only contain declarations
 			args = appendLinkFiles(args, p.ExportFile)
+			need1, need2 := isNeedRuntimeOrPyInit(p)
 			if !needRuntime {
-				needRuntime = isNeedRuntime(p)
+				needRuntime = need1
+			}
+			if !needPyInit {
+				needPyInit = need2
 			}
 		}
 	})
+
+	var aPkg *aPackage
+	for _, v := range pkgs {
+		if v.Package == pkg { // found this package
+			aPkg = v
+			break
+		}
+	}
+
+	dirty := false
 	if needRuntime && runtimeFiles != nil {
 		args = append(args, runtimeFiles...)
 	} else {
-		for _, aPkg := range pkgs {
+		dirty = true
-			if aPkg.Package == pkg { // make empty runtime.init if no runtime needed
+		fn := aPkg.LPkg.FuncOf(cl.RuntimeInit)
-				lpkg := aPkg.LPkg
+		fn.MakeBody(1).Return()
-				lpkg.FuncOf(cl.RuntimeInit).MakeBody(1).Return()
+	}
-				if needLLFile(mode) {
+	if needPyInit {
-					os.WriteFile(pkg.ExportFile, []byte(lpkg.String()), 0644)
+		dirty = aPkg.LPkg.PyInit()
-				}
+	}
-				break
+
-			}
+	if dirty && needLLFile(mode) {
-		}
+		lpkg := aPkg.LPkg
+		os.WriteFile(pkg.ExportFile, []byte(lpkg.String()), 0644)
 	}
 
 	if verbose || mode != ModeRun {

internal/llgen/llgen.go

@@ -79,6 +79,10 @@ func Gen(pkgPath, inFile string, src any) string {
 	ret, err := cl.NewPackage(prog, ssaPkg, files)
 	check(err)
 
+	if prog.NeedPyInit() { // call PyInit if needed
+		ret.PyInit()
+	}
+
 	return ret.String()
 }
 

internal/llgen/llgenf.go

@@ -82,6 +82,10 @@ func GenFrom(fileOrPkg string) string {
 	ret, err := cl.NewPackage(prog, ssaPkg, pkg.Syntax)
 	check(err)
 
+	if prog.NeedPyInit() { // call PyInit if needed
+		ret.PyInit()
+	}
+
 	return ret.String()
 }
 

ssa/cl_test.go

@@ -23,7 +23,7 @@ import (
 )
 
 func TestFromTestpy(t *testing.T) {
-	cltest.FromDir(t, "", "../cl/_testpy", false)
+	cltest.FromDir(t, "", "../cl/_testpy", true)
 }
 
 func TestFromTestrt(t *testing.T) {

ssa/decl.go

@@ -299,6 +299,7 @@ func (p Package) NewPyFunc(name string, sig *types.Signature) PyFunction {
 		return v
 	}
 	prog := p.Prog
+	prog.needPyInit = true
 	obj := p.NewVar(name, prog.PyObjectPtrPtr().RawType(), InC)
 	obj.Init(prog.Null(obj.Type))
 	obj.impl.SetLinkage(llvm.LinkOnceAnyLinkage)

ssa/package.go

@@ -194,6 +194,11 @@ func (p Program) NeedRuntime() bool {
 	return p.needRuntime
 }
 
+// NeedPyInit returns if the current package needs Python initialization.
+func (p Program) NeedPyInit() bool {
+	return p.needPyInit
+}
+
 func (p Program) runtime() *types.Package {
 	if p.rt == nil {
 		p.rt = p.rtget()
@@ -262,8 +267,11 @@ func (p Program) NewPackage(name, pkgPath string) Package {
 	fns := make(map[string]Function)
 	stubs := make(map[string]Function)
 	pyfns := make(map[string]PyFunction)
+	pymods := make(map[string]Global)
 	p.needRuntime = false
-	return &aPackage{mod, gbls, fns, stubs, pyfns, p}
+	// Don't need reset p.needPyInit here
+	// p.needPyInit = false
+	return &aPackage{mod, gbls, fns, stubs, pyfns, pymods, p}
 }
 
 // PyObjectPtrPtr returns the **py.Object type.
@@ -363,12 +371,13 @@ func (p Program) Float64() Type {
 // initializer) and "init#%d", the nth declared init function,
 // and unspecified other things too.
 type aPackage struct {
-	mod   llvm.Module
+	mod    llvm.Module
-	vars  map[string]Global
+	vars   map[string]Global
-	fns   map[string]Function
+	fns    map[string]Function
-	stubs map[string]Function
+	stubs  map[string]Function
-	pyfns map[string]PyFunction
+	pyfns  map[string]PyFunction
-	Prog  Program
+	pymods map[string]Global
+	Prog   Program
 }
 
 type Package = *aPackage
@@ -506,23 +515,38 @@ func (p Program) tyCallOneArg() *types.Signature {
 	return p.callOneArg
 }
 
-// ImportPyMod imports a Python module.
+// PyInit initializes Python for a main package.
-func (b Builder) ImportPyMod(path string) Expr {
+func (p Package) PyInit() bool {
-	pkg := b.Func.Pkg
+	if fn := p.FuncOf("main"); fn != nil {
-	fnImp := pkg.pyFunc("PyImport_ImportModule", b.Prog.tyImportPyModule())
+		b := fn.NewBuilder()
-	return b.Call(fnImp, b.CStr(path))
+		b.SetBlockEx(fn.Block(0), AtStart).CallPyInit()
+		b.Dispose()
+		return true
+	}
+	return false
 }
 
 // NewPyModVar creates a new global variable for a Python module.
 func (p Package) NewPyModVar(name string) Global {
+	if v, ok := p.pymods[name]; ok {
+		return v
+	}
 	prog := p.Prog
 	objPtr := prog.PyObjectPtrPtr().raw.Type
 	g := p.NewVar(name, objPtr, InC)
 	g.Init(prog.Null(g.Type))
 	g.impl.SetLinkage(llvm.LinkOnceAnyLinkage)
+	p.pymods[name] = g
 	return g
 }
 
+// ImportPyMod imports a Python module.
+func (b Builder) ImportPyMod(path string) Expr {
+	pkg := b.Func.Pkg
+	fnImp := pkg.pyFunc("PyImport_ImportModule", b.Prog.tyImportPyModule())
+	return b.Call(fnImp, b.CStr(path))
+}
+
 func (b Builder) pyCall(fn Expr, args []Expr) (ret Expr) {
 	prog := b.Prog
 	pkg := b.Func.Pkg
@@ -542,4 +566,14 @@ func (b Builder) pyCall(fn Expr, args []Expr) (ret Expr) {
 	return
 }
 
+// CallPyInit calls Py_Initialize.
+func (b Builder) CallPyInit() (ret Expr) {
+	fn := b.Func.Pkg.pyFunc("Py_Initialize", NoArgsNoRet)
+	return b.Call(fn)
+}
+
+var (
+	NoArgsNoRet = types.NewSignatureType(nil, nil, nil, nil, nil, false)
+)
+
 // -----------------------------------------------------------------------------

ssa/stmt_builder.go

@@ -57,15 +57,40 @@ type aBuilder struct {
 // Builder represents a builder for creating instructions in a function.
 type Builder = *aBuilder
 
-// SetBlock sets the current block to the specified basic block.
+// Dispose disposes of the builder.
+func (b Builder) Dispose() {
+	b.impl.Dispose()
+}
+
+// SetBlock means SetBlockEx(blk, AtEnd).
 func (b Builder) SetBlock(blk BasicBlock) Builder {
-	if b.Func != blk.fn {
-		panic("mismatched function")
-	}
 	if debugInstr {
 		log.Printf("Block _llgo_%v:\n", blk.idx)
 	}
-	b.impl.SetInsertPointAtEnd(blk.impl)
+	b.SetBlockEx(blk, AtEnd)
+	return b
+}
+
+type InsertPoint int
+
+const (
+	AtEnd InsertPoint = iota
+	AtStart
+)
+
+// SetBlockEx sets blk as current basic block and pos as its insert point.
+func (b Builder) SetBlockEx(blk BasicBlock, pos InsertPoint) Builder {
+	if b.Func != blk.fn {
+		panic("mismatched function")
+	}
+	switch pos {
+	case AtEnd:
+		b.impl.SetInsertPointAtEnd(blk.impl)
+	case AtStart:
+		b.impl.SetInsertPointBefore(blk.impl.FirstInstruction())
+	default:
+		panic("SetBlockEx: invalid pos")
+	}
 	return b
 }