llgo/ssa: LoadPyModSyms

ssa/decl.go

@@ -285,36 +285,41 @@ func (p Function) Block(idx int) BasicBlock {
 
 // -----------------------------------------------------------------------------
 
-type aPyFunction struct {
+type aPyObject struct {
 	Expr
 	Obj Global
 }
 
-// PyFunction represents a python function.
-type PyFunction = *aPyFunction
+// PyObject represents a python object.
+type PyObject = *aPyObject
 
 // NewPyFunc creates a new python function.
-func (p Package) NewPyFunc(name string, sig *types.Signature, doInit bool) PyFunction {
-	if v, ok := p.pyfns[name]; ok {
+func (p Package) NewPyFunc(name string, sig *types.Signature, doInit bool) PyObject {
+	if v, ok := p.pyobjs[name]; ok {
 		return v
 	}
 	prog := p.Prog
-	prog.NeedPyInit = true
 	obj := p.NewVar(name, prog.PyObjectPtrPtr().RawType(), InC)
 	if doInit {
+		prog.NeedPyInit = true
 		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}
-	p.pyfns[name] = ret
+	ret := &aPyObject{expr, obj}
+	p.pyobjs[name] = ret
 	return ret
 }
 
-// PyFuncOf returns a function by name.
-func (p Package) PyFuncOf(name string) PyFunction {
-	return p.pyfns[name]
+// PyObjOf returns a python object by name.
+func (p Package) PyObjOf(name string) PyObject {
+	return p.pyobjs[name]
+}
+
+// PyObjs returns all used python objects in this project.
+func (p Package) PyObjs() map[string]PyObject {
+	return p.pyobjs
 }
 
 // -----------------------------------------------------------------------------

ssa/package.go

@@ -140,6 +140,7 @@ type aProgram struct {
 	pyImpTy    *types.Signature
 	callNoArgs *types.Signature
 	callOneArg *types.Signature
+	loadPyModS *types.Signature
 
 	NeedRuntime bool
 	NeedPyInit  bool
@@ -257,12 +258,12 @@ func (p Program) NewPackage(name, pkgPath string) Package {
 	gbls := make(map[string]Global)
 	fns := make(map[string]Function)
 	stubs := make(map[string]Function)
-	pyfns := make(map[string]PyFunction)
+	pyobjs := make(map[string]PyObject)
 	pymods := make(map[string]Global)
 	p.NeedRuntime = false
 	// Don't need reset p.needPyInit here
 	// p.needPyInit = false
-	return &aPackage{mod, gbls, fns, stubs, pyfns, pymods, p}
+	return &aPackage{mod, gbls, fns, stubs, pyobjs, pymods, p}
 }
 
 // PyObjectPtrPtr returns the **py.Object type.
@@ -366,7 +367,7 @@ type aPackage struct {
 	vars   map[string]Global
 	fns    map[string]Function
 	stubs  map[string]Function
-	pyfns  map[string]PyFunction
+	pyobjs map[string]PyObject
 	pymods map[string]Global
 	Prog   Program
 }
@@ -506,6 +507,16 @@ func (p Program) tyCallOneArg() *types.Signature {
 	return p.callOneArg
 }
 
+func (p Program) tyLoadPyModSyms() *types.Signature {
+	if p.loadPyModS == nil {
+		objPtr := p.PyObjectPtr().raw.Type
+		paramObjPtr := types.NewParam(token.NoPos, nil, "mod", objPtr)
+		params := types.NewTuple(paramObjPtr, VArg())
+		p.loadPyModS = types.NewSignatureType(nil, nil, nil, params, nil, false)
+	}
+	return p.loadPyModS
+}
+
 // PyInit initializes Python for a main package.
 func (p Package) PyInit() bool {
 	if fn := p.FuncOf("main"); fn != nil {
@@ -518,15 +529,17 @@ func (p Package) PyInit() bool {
 }
 
 // NewPyModVar creates a new global variable for a Python module.
-func (p Package) NewPyModVar(name string) Global {
+func (p Package) NewPyModVar(name string, doInit bool) 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)
+	if doInit {
+		g.Init(prog.Null(g.Type))
+		g.impl.SetLinkage(llvm.LinkOnceAnyLinkage)
+	}
 	p.pymods[name] = g
 	return g
 }
@@ -538,6 +551,21 @@ func (b Builder) ImportPyMod(path string) Expr {
 	return b.Call(fnImp, b.CStr(path))
 }
 
+// LoadPyModSyms loads python objects from specified module.
+func (b Builder) LoadPyModSyms(modName string, objs ...PyObject) Expr {
+	pkg := b.Func.Pkg
+	fnLoad := pkg.pyFunc("llgoLoadPyModSyms", b.Prog.tyLoadPyModSyms())
+	modPtr := pkg.NewPyModVar(modName, false).Expr
+	mod := b.Load(modPtr)
+	args := make([]Expr, 1, len(objs)*2+1)
+	args[0] = mod
+	for _, o := range objs {
+		args = append(args, b.CStr(o.impl.Name()))
+		args = append(args, o.Expr)
+	}
+	return b.Call(fnLoad, args...)
+}
+
 func (b Builder) pyCall(fn Expr, args []Expr) (ret Expr) {
 	prog := b.Prog
 	pkg := b.Func.Pkg

ssa/ssa_test.go

@@ -147,8 +147,8 @@ func TestPyFunc(t *testing.T) {
 	if pkg.NewPyFunc("a", sig, false) != a {
 		t.Fatal("NewPyFunc(a) failed")
 	}
-	foo := pkg.NewPyModVar("foo")
-	if pkg.NewPyModVar("foo") != foo {
+	foo := pkg.NewPyModVar("foo", false)
+	if pkg.NewPyModVar("foo", false) != foo {
 		t.Fatal("NewPyModVar(foo) failed")
 	}
 }