Merge pull request #149 from xushiwei/q

PyInit

ssa/decl.go

@@ -298,7 +298,11 @@ func (p Package) NewPyFunc(name string, sig *types.Signature) PyFunction {
 	if v, ok := p.pyfns[name]; ok {
 		return v
 	}
-	obj := p.NewVar(name, p.Prog.PyObjectPtrPtr().RawType(), InC)
+	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)
 	ty := &aType{obj.ll, rawType{sig}, vkPyFunc}
 	expr := Expr{obj.impl, ty}
 	ret := &aPyFunction{expr, obj}

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)
+)
+
 // -----------------------------------------------------------------------------

ssa/ssa_test.go

@@ -25,6 +25,28 @@ import (
 	"github.com/goplus/llvm"
 )
 
+func TestSetBlock(t *testing.T) {
+	defer func() {
+		if r := recover(); r == nil {
+			t.Log("SetBlock: no error?")
+		}
+	}()
+	fn := &aFunction{}
+	b := &aBuilder{Func: fn}
+	b.SetBlock(&aBasicBlock{})
+}
+
+func TestSetBlockEx(t *testing.T) {
+	defer func() {
+		if r := recover(); r == nil {
+			t.Log("SetBlockEx: no error?")
+		}
+	}()
+	fn := &aFunction{}
+	b := &aBuilder{Func: fn}
+	b.SetBlockEx(&aBasicBlock{fn: fn}, -1)
+}
+
 func TestSetPython(t *testing.T) {
 	prog := NewProgram(nil)
 	typ := types.NewPackage("foo", "foo")
@@ -125,6 +147,10 @@ func TestPyFunc(t *testing.T) {
 	if pkg.NewPyFunc("a", sig) != a {
 		t.Fatal("NewPyFunc(a) failed")
 	}
+	foo := pkg.NewPyModVar("foo")
+	if pkg.NewPyModVar("foo") != foo {
+		t.Fatal("NewPyModVar(foo) failed")
+	}
 }
 
 func TestVar(t *testing.T) {

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
 }