PyInit

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
-	vars  map[string]Global
-	fns   map[string]Function
-	stubs map[string]Function
-	pyfns map[string]PyFunction
-	Prog  Program
+	mod    llvm.Module
+	vars   map[string]Global
+	fns    map[string]Function
+	stubs  map[string]Function
+	pyfns  map[string]PyFunction
+	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.
-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))
+// PyInit initializes Python for a main package.
+func (p Package) PyInit() bool {
+	if fn := p.FuncOf("main"); fn != nil {
+		b := fn.NewBuilder()
+		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
 }