llgo/ssa: b.NewPyModVar, b.ImportPyMod, PyObjectPtr, PyObjectPtrPtr

README.md

@@ -65,8 +65,8 @@ The `_demo` directory contains our demos (it start with `_` to prevent the `go`
 And the `py/_demo` directory contains python related demos:
 
 * [hellopy](py/_demo/hellopy/hello.go): link Python to Go and say `Hello world`
-* [clpy](py/_demo/clpy/cleval.go): compile Python code and eval.
+* [clpy](py/_demo/clpy/cleval.go): compile Python code and eval
-* [callpy](py/_demo/callpy/call.go): call Python standard library function `math.sqrt`.
+* [callpy](py/_demo/callpy/call.go): call Python standard library function `math.sqrt`
 
 ### How to run demos
 

cl/cltest/cltest.go

@@ -147,6 +147,13 @@ func TestCompileEx(t *testing.T, src any, fname, expected string) {
 		}
 		return rt
 	})
+	prog.SetPython(func() *types.Package {
+		rt, err := imp.Import(llssa.PkgPython)
+		if err != nil {
+			t.Fatal("load python failed:", err)
+		}
+		return rt
+	})
 
 	ret, err := cl.NewPackage(prog, foo, files)
 	if err != nil {

cl/compile.go

@@ -347,14 +347,15 @@ func (p *context) compileBlock(b llssa.Builder, block *ssa.BasicBlock, n int, do
 	if pyModInit {
 		jump := block.Instrs[last].(*ssa.Jump)
 		jumpTo := p.jumpTo(jump)
-		modName := pysymPrefix + p.pyMod
+		modPath := p.pyMod
+		modName := pysymPrefix + modPath
 		modPtr := p.pkg.NewPyModVar(modName).Expr
 		mod := b.Load(modPtr)
 		cond := b.BinOp(token.NEQ, mod, b.Prog.Null(mod.Type))
 		newBlk := p.fn.MakeBlock()
 		b.If(cond, jumpTo, newBlk)
 		b.SetBlock(newBlk)
-		// TODO(xsw): pyModInit
+		b.Store(modPtr, b.ImportPyMod(modPath))
 		b.Jump(jumpTo)
 	}
 	return ret

internal/build/build.go

@@ -120,18 +120,31 @@ func Do(args []string, conf *Config) {
 		cl.SetDebug(cl.DbgFlagAll)
 	}
 
+	var needRt bool
 	var rt []*packages.Package
 	prog := llssa.NewProgram(nil)
+	load := func() []*packages.Package {
+		if rt == nil {
+			var err error
+			rt, err = packages.Load(cfg, llssa.PkgRuntime, llssa.PkgPython)
+			check(err)
+		}
+		return rt
+	}
 	prog.SetRuntime(func() *types.Package {
-		rt, err = packages.Load(cfg, llssa.PkgRuntime)
+		needRt = true
-		check(err)
+		rt := load()
 		return rt[0].Types
 	})
+	prog.SetPython(func() *types.Package {
+		rt := load()
+		return rt[1].Types
+	})
 
 	pkgs := buildAllPkgs(prog, initial, mode, verbose)
 
 	var runtimeFiles []string
-	if rt != nil {
+	if needRt {
 		runtimeFiles = allLinkFiles(rt)
 	}
 	if mode != ModeBuild {

internal/llgen/llgen.go

@@ -30,6 +30,7 @@ import (
 	"golang.org/x/tools/go/ssa/ssautil"
 
 	llssa "github.com/goplus/llgo/ssa"
+	cpackages "golang.org/x/tools/go/packages"
 )
 
 func Init() {
@@ -71,11 +72,10 @@ func Gen(pkgPath, inFile string, src any) string {
 	}
 
 	prog := llssa.NewProgram(nil)
-	prog.SetRuntime(func() *types.Package {
+	initRtAndPy(prog, &cpackages.Config{
-		rt, err := imp.Import(llssa.PkgRuntime)
+		Mode: loadSyntax | cpackages.NeedDeps,
-		check(err)
-		return rt
 	})
+
 	ret, err := cl.NewPackage(prog, ssaPkg, files)
 	check(err)
 

internal/llgen/llgenf.go

@@ -38,6 +38,27 @@ const (
 	loadSyntax  = loadTypes | packages.NeedSyntax | packages.NeedTypesInfo
 )
 
+func initRtAndPy(prog llssa.Program, cfg *packages.Config) {
+	var pkgRtAndPy []*packages.Package
+	load := func() []*packages.Package {
+		if pkgRtAndPy == nil {
+			var err error
+			pkgRtAndPy, err = packages.Load(cfg, llssa.PkgRuntime, llssa.PkgPython)
+			check(err)
+		}
+		return pkgRtAndPy
+	}
+
+	prog.SetRuntime(func() *types.Package {
+		rt := load()
+		return rt[0].Types
+	})
+	prog.SetPython(func() *types.Package {
+		rt := load()
+		return rt[1].Types
+	})
+}
+
 func GenFrom(fileOrPkg string) string {
 	cfg := &packages.Config{
 		Mode: loadSyntax | packages.NeedDeps,
@@ -52,11 +73,7 @@ func GenFrom(fileOrPkg string) string {
 	ssaPkg.Build()
 
 	prog := llssa.NewProgram(nil)
-	prog.SetRuntime(func() *types.Package {
+	initRtAndPy(prog, cfg)
-		rt, err := packages.Load(cfg, llssa.PkgRuntime)
-		check(err)
-		return rt[0].Types
-	})
 
 	if Verbose {
 		ssaPkg.WriteTo(os.Stderr)

ssa/package.go

@@ -26,6 +26,7 @@ import (
 )
 
 const (
+	PkgPython  = "github.com/goplus/llgo/py"
 	PkgRuntime = "github.com/goplus/llgo/internal/runtime"
 )
 
@@ -103,6 +104,9 @@ type aProgram struct {
 	rt    *types.Package
 	rtget func() *types.Package
 
+	py    *types.Package
+	pyget func() *types.Package
+
 	target *Target
 	td     llvm.TargetData
 	// tm  llvm.TargetMachine
@@ -130,8 +134,13 @@ type aProgram struct {
 	uintptrTy Type
 	intTy     Type
 	f64Ty     Type
+	pyObjPtr  Type
+	pyObjPPtr Type
+
+	pyImpTy *types.Signature
 
 	needRuntime bool
+	needPyInit  bool
 }
 
 // A Program presents a program.
@@ -157,6 +166,17 @@ func NewProgram(target *Target) Program {
 	return &aProgram{ctx: ctx, gocvt: newGoTypes(), target: target, td: td}
 }
 
+// SetPython sets the Python package.
+// Its type can be *types.Package or func() *types.Package.
+func (p Program) SetPython(py any) {
+	switch v := py.(type) {
+	case *types.Package:
+		p.py = v
+	case func() *types.Package:
+		p.pyget = v
+	}
+}
+
 // SetRuntime sets the runtime.
 // Its type can be *types.Package or func() *types.Package.
 func (p Program) SetRuntime(runtime any) {
@@ -181,12 +201,25 @@ func (p Program) runtime() *types.Package {
 	return p.rt
 }
 
+func (p Program) python() *types.Package {
+	if p.py == nil {
+		p.py = p.pyget()
+	}
+	return p.py
+}
+
 func (p Program) rtNamed(name string) *types.Named {
 	t := p.runtime().Scope().Lookup(name).Type().(*types.Named)
 	t, _ = p.gocvt.cvtNamed(t)
 	return t
 }
 
+func (p Program) pyNamed(name string) *types.Named {
+	// TODO(xsw): does python type need to convert?
+	t := p.python().Scope().Lookup(name).Type().(*types.Named)
+	return t
+}
+
 func (p Program) rtType(name string) Type {
 	return p.rawType(p.rtNamed(name))
 }
@@ -231,6 +264,23 @@ func (p Program) NewPackage(name, pkgPath string) Package {
 	return &aPackage{mod, gbls, fns, stubs, p}
 }
 
+// PyObjectPtrPtr returns the **py.Object type.
+func (p Program) PyObjectPtrPtr() Type {
+	if p.pyObjPPtr == nil {
+		p.pyObjPPtr = p.Pointer(p.PyObjectPtr())
+	}
+	return p.pyObjPPtr
+}
+
+// PyObjectPtr returns the *py.Object type.
+func (p Program) PyObjectPtr() Type {
+	if p.pyObjPtr == nil {
+		objPtr := types.NewPointer(p.pyNamed("Object"))
+		p.pyObjPtr = p.rawType(objPtr)
+	}
+	return p.pyObjPtr
+}
+
 // Void returns void type.
 func (p Program) Void() Type {
 	if p.voidTy == nil {
@@ -327,13 +377,6 @@ func (p Package) NewConst(name string, val constant.Value) NamedConst {
 }
 */
 
-// NewPyModVar creates a new global variable for a Python module.
-func (p Package) NewPyModVar(name string) Global {
-	ret := p.NewVar(name, types.NewPointer(types.Typ[types.Int]), InC)
-	ret.impl.SetLinkage(llvm.LinkOnceAnyLinkage)
-	return ret
-}
-
 // NewVar creates a new global variable.
 func (p Package) NewVar(name string, typ types.Type, bg Background) Global {
 	t := p.Prog.Type(typ, bg)
@@ -375,6 +418,11 @@ func (p Package) rtFunc(fnName string) Expr {
 	return p.NewFunc(name, sig, InGo).Expr
 }
 
+func (p Package) cpyFunc(fullName string, sig *types.Signature) Expr {
+	p.Prog.needPyInit = true
+	return p.NewFunc(fullName, sig, InC).Expr
+}
+
 func (p Package) closureStub(b Builder, t *types.Struct, v Expr) Expr {
 	name := v.impl.Name()
 	prog := b.Prog
@@ -461,3 +509,33 @@ func (p *Package) WriteFile(file string) (err error) {
 */
 
 // -----------------------------------------------------------------------------
+
+func (p Program) tyImportPyModule() *types.Signature {
+	if p.pyImpTy == nil {
+		charPtr := types.NewPointer(types.Typ[types.Int8])
+		objPtr := p.PyObjectPtr().raw.Type
+		params := types.NewTuple(types.NewParam(token.NoPos, nil, "", charPtr))
+		results := types.NewTuple(types.NewParam(token.NoPos, nil, "", objPtr))
+		p.pyImpTy = types.NewSignatureType(nil, nil, nil, params, results, false)
+	}
+	return p.pyImpTy
+}
+
+// NewPyModVar creates a new global variable for a Python module.
+func (p Package) NewPyModVar(name string) Global {
+	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)
+	return g
+}
+
+// ImportPyMod imports a Python module.
+func (b Builder) ImportPyMod(path string) Expr {
+	pkg := b.Func.Pkg
+	fnImp := pkg.cpyFunc("PyImport_ImportModule", b.Prog.tyImportPyModule())
+	return b.Call(fnImp, b.CStr(path))
+}
+
+// -----------------------------------------------------------------------------