Merge pull request #149 from xushiwei/q

PyInit

ssa/package.go

@@ -138,7 +138,7 @@ type aProgram struct {
 	pyObjPPtr Type
 
 	pyImpTy    *types.Signature
-	callNoArg  *types.Signature
+	callNoArgs *types.Signature
 	callOneArg *types.Signature
 
 	needRuntime bool
@@ -195,6 +195,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()
@@ -263,8 +268,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.
@@ -364,12 +372,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
@@ -486,14 +495,14 @@ func (p Program) tyImportPyModule() *types.Signature {
 	return p.pyImpTy
 }
 
-func (p Program) tyCallNoArg() *types.Signature {
-	if p.callNoArg == nil {
+func (p Program) tyCallNoArgs() *types.Signature {
+	if p.callNoArgs == nil {
 		objPtr := p.PyObjectPtr().raw.Type
 		paramObjPtr := types.NewParam(token.NoPos, nil, "", objPtr)
 		params := types.NewTuple(paramObjPtr)
-		p.callNoArg = types.NewSignatureType(nil, nil, nil, params, params, false)
+		p.callNoArgs = types.NewSignatureType(nil, nil, nil, params, params, false)
 	}
-	return p.callNoArg
+	return p.callNoArgs
 }
 
 func (p Program) tyCallOneArg() *types.Signature {
@@ -507,23 +516,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
@@ -532,7 +556,7 @@ func (b Builder) pyCall(fn Expr, args []Expr) (ret Expr) {
 	n := params.Len()
 	switch n {
 	case 0:
-		call := pkg.pyFunc("PyObject_CallNoArg", prog.tyCallNoArg())
+		call := pkg.pyFunc("PyObject_CallNoArgs", prog.tyCallNoArgs())
 		ret = b.Call(call, fn)
 	case 1:
 		call := pkg.pyFunc("PyObject_CallOneArg", prog.tyCallOneArg())
@@ -543,4 +567,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)
+)
+
 // -----------------------------------------------------------------------------