From 9ac0450255c70b5a4cf8fbbc7ea4c2e24ebe45a2 Mon Sep 17 00:00:00 2001
From: xushiwei <x@goplus.org>
Date: Sun, 12 May 2024 18:27:23 +0800
Subject: [PATCH] llgo/ssa: LoadPyModSyms

---
 .gitignore          |   1 +
 cl/compile.go       |  10 +++---
 py/_pyg/module.c    |  25 +++++++++++++
 py/llgo.cfg         |   6 ++++
 py/llgo_autogen.lla | Bin 523 -> 1822 bytes
 py/module.go        |   3 ++
 py/module.ll        |  84 ++++++++++++++++++++++++++++++++++++++++++++
 ssa/decl.go         |  27 ++++++++------
 ssa/package.go      |  40 +++++++++++++++++----
 ssa/ssa_test.go     |   4 +--
 10 files changed, 176 insertions(+), 24 deletions(-)
 create mode 100644 py/_pyg/module.c
 create mode 100644 py/llgo.cfg
 create mode 100644 py/module.ll

diff --git a/.gitignore b/.gitignore
index d3ecf298..41c277a2 100644
--- a/.gitignore
+++ b/.gitignore
@@ -18,6 +18,7 @@ _go/
 _runtime/
 _tinygo/
 build.dir/
+.vscode/
 
 # Test binary, built with `go test -c`
 *.test
diff --git a/cl/compile.go b/cl/compile.go
index 998958ec..68a8f2ad 100644
--- a/cl/compile.go
+++ b/cl/compile.go
@@ -211,7 +211,7 @@ var (
 	argvTy = types.NewPointer(types.NewPointer(types.Typ[types.Int8]))
 )
 
-func (p *context) compileFuncDecl(pkg llssa.Package, f *ssa.Function, call bool) (llssa.Function, llssa.PyFunction, int) {
+func (p *context) compileFuncDecl(pkg llssa.Package, f *ssa.Function, call bool) (llssa.Function, llssa.PyObject, int) {
 	pkgTypes, name, ftype := p.funcName(f, true)
 	if ftype != goFunc {
 		if ftype == pyFunc {
@@ -284,14 +284,14 @@ func (p *context) compileFuncDecl(pkg llssa.Package, f *ssa.Function, call bool)
 
 // funcOf returns a function by name and set ftype = goFunc, cFunc, etc.
 // or returns nil and set ftype = llgoCstr, llgoAlloca, llgoUnreachable, etc.
-func (p *context) funcOf(fn *ssa.Function) (aFn llssa.Function, pyFn llssa.PyFunction, ftype int) {
+func (p *context) funcOf(fn *ssa.Function) (aFn llssa.Function, pyFn llssa.PyObject, ftype int) {
 	pkgTypes, name, ftype := p.funcName(fn, false)
 	switch ftype {
 	case pyFunc:
 		if kind, mod := pkgKindByScope(pkgTypes.Scope()); kind == PkgPyModule {
 			pkg := p.pkg
 			fnName := pysymPrefix + mod + "." + name
-			if pyFn = pkg.PyFuncOf(fnName); pyFn == nil {
+			if pyFn = pkg.PyObjOf(fnName); pyFn == nil {
 				pyFn = pkg.NewPyFunc(fnName, fn.Signature, true)
 				return
 			}
@@ -361,7 +361,7 @@ func (p *context) compileBlock(b llssa.Builder, block *ssa.BasicBlock, n int, do
 		jumpTo := p.jumpTo(jump)
 		modPath := p.pyMod
 		modName := pysymPrefix + modPath
-		modPtr := p.pkg.NewPyModVar(modName).Expr
+		modPtr := p.pkg.NewPyModVar(modName, true).Expr
 		mod := b.Load(modPtr)
 		cond := b.BinOp(token.NEQ, mod, b.Prog.Null(mod.Type))
 		newBlk := p.fn.MakeBlock()
@@ -763,7 +763,7 @@ func (p *context) compileInstr(b llssa.Builder, instr ssa.Instruction) {
 	}
 }
 
-func (p *context) compileFunction(v *ssa.Function) (goFn llssa.Function, pyFn llssa.PyFunction, kind int) {
+func (p *context) compileFunction(v *ssa.Function) (goFn llssa.Function, pyFn llssa.PyObject, kind int) {
 	// v.Pkg == nil: means auto generated function?
 	if v.Pkg == p.goPkg || v.Pkg == nil {
 		// function in this package
diff --git a/py/_pyg/module.c b/py/_pyg/module.c
new file mode 100644
index 00000000..7ce91821
--- /dev/null
+++ b/py/_pyg/module.c
@@ -0,0 +1,25 @@
+#include <stdlib.h>
+#include <stdarg.h>
+
+// example:
+// llgoLoadPyModSyms(mod, "name1", &func1, "name2", &func2, NULL)
+
+typedef struct PyObject PyObject;
+
+PyObject* PyObject_GetAttrString(PyObject* mod, const char* attrName);
+
+void llgoLoadPyModSyms(PyObject* mod, ...) {
+    va_list ap;
+    va_start(ap, mod);
+    for (;;) {
+        const char* name = va_arg(ap, const char*);
+        if (name == NULL) {
+            break;
+        }
+        PyObject** pfunc = va_arg(ap, PyObject**);
+        if (*pfunc == NULL) {
+            *pfunc = PyObject_GetAttrString(mod, name);
+        }
+    }
+    va_end(ap);
+}
diff --git a/py/llgo.cfg b/py/llgo.cfg
new file mode 100644
index 00000000..09270308
--- /dev/null
+++ b/py/llgo.cfg
@@ -0,0 +1,6 @@
+{
+    "cl": [
+        "clang -emit-llvm -S -o module.ll -c _pyg/module.c",
+        "rm llgo_autogen.lla; zip llgo_autogen.lla llgo_autogen.ll module.ll",
+    ]
+}
diff --git a/py/llgo_autogen.lla b/py/llgo_autogen.lla
index 79b52437fa407fb23c5da0a55c2f952b2d5d7ecc..ab226c5eee6f15bb2cb9eddb5bc7a0613c4f92e3 100644
GIT binary patch
delta 1339
zcmZ|Pe>l?#90&04X0uiq4vW&7o{+RnF`XZWp&TO;v7*Pc`8iXxLzbo7LYEVg-#2T1
zR2ZR2#CFZm&Gf^$(KIB4L`^>!&UOFXbI)_n=lMME=l#d~-`5pL8KRyGVyQuN0RY$l
z%xBjGcGBHvQXv4)f&c(GU<ky-QW!KUmPYgUg9B>gc}`&y>(5Bs1b`rmAOQF}eH{=S
z`@&lH{GWUdUf(OKM<<e$uLFr_2z6;dYG{Sv4w8Mop~jc*zRsPsC9E!EjlZ~Ll=gRj
zYyz&Cy^LpR;xuu%TSTRe_a31`ic`Q;Z%tqyt9Lq!tm@kH!9z^Tj}z|d)kn`FME9)x
z<13Fi<Mqy#1`LH16bFxEHC0zS6KqUIp{fUQu99VqCPB*a1?T=YKWYdM+SY1t^}~A@
zN~66z{KT=2{@a<0^XFEKmx{$<8``n@k2fu~hqCe+87X{9O%nELa;<+-Dn(i`l0&$e
zV~*A7j4E_-BO$AgM{b4xaIB8Xq#N(WcT=T0%EReNt7cU_i(CEflfqvQz7i@1N=DL_
z$}X{g9kp2!yL>`Ap9iNNF?wIXH9d(4{2au60=U_sH6NI)wL~)Iy7rcFTWr3ND7WIp
ztg_H)lH+1pPXZfGE=Zp#$RC0%4@gLpmGA+5uxz?r!Y1098|_dU6bKg$(bt?W9`fX@
z_TR1CGWDo~=(jsRPs{Zeh!)OzA6fGyd2e@{m*?SaM=S@=CC=K^GUSEYT_YGYe(Qp!
z(?|Dt_0a_PZcvV!R<5DFbmOCk0t%W{Tay++8I^^UT(y@e@W18%xN@|4C^u?IyUg<f
zUFFg0fIy_LTCuWBOSIA7j)U3m0rz_j$z2Lp_6sdRo>`w^!+<DPEz#BDJCD<e&qwQD
zdAGI8ZR!~d$}=X~vBeIMcSWgGZB=H)Wk0LYY@Mp-wMG=>SeQ{|(7@~QH#C+arl`@e
zz<N@oVQ3w<8bmPVp~Cp5&Ysk?Oc@rTYAYT&?sToHpMC4zXdyI5cI}O8v@P98G~ttI
z9RpMSIenC8x=h~yCR6}Yt4ASrOtFo)3mB`TVo47@>)kM@W6T+FtJqaC)xgQzZWlVz
z*)@Pv;&vT2jl58*a<hzJ=|zuKIsJaqqw$&CvF^!=*$ot>YU5ly9fU11TYjuwVy^EM
zLkX8#P5QME3#fH+J#Qu00;Nx#i-HQQv><hnj04r8PTy+#P;9lK!g-R0#L<c8?TFPE
z13OQi2y%ME#@~LHXC@d=*gS89F0w1%wngu3(*Z);jJE;wOdt#11)^41ru)-SGtQGz
zbPxWtsT7R~Y9V4>o{I=Y@Cc_~OPtgxPh#+WJ~5oL+o!q(KwPItIYB<Yj&Ac3W_bPP
z%crI{&y7j!=V;0&+S>)(HPu~*ooKu;-h5#V`C|Hy3`I&b@1oS}{+r=XBa&&t_8CI+
zoG^w9JJ`g6K^tLs0nC(vV@3ogiFkBoCA;s=3#M!KndVprRmN)Q?X=)1yD4*mqn7Mp
zIbWST7dVli)R4^&JpCz}KL+)0JT6fmS7X!ml?z=nA1999i`!o~F>idGU|)L$hP_!w
zghLfEzTygZ@br!3khO#Ln+jVdEm8)_;ziL_B7N{&A4b+{fO$Lfb4R&1t;bpdLFZz5
z7z_qbL(hU+%^qIzDsFC26f0eP-j5q%ilGGAWqyaY?VaS{+2o1n>}(_DtN3N%5-gVk
v3OEG<qtyOSH2rny0AL+1fS1ZykAL|F{l~Ze*yclos{cI!uODwcDPP-PI~!lB

delta 32
hcmbQo*UhqF5#!`EwonmPHU=PI1j0}ty?_zK0|0a`1o{8~

diff --git a/py/module.go b/py/module.go
index 59e669c5..17f19b14 100644
--- a/py/module.go
+++ b/py/module.go
@@ -50,4 +50,7 @@ func Import(name *Object) *Object
 // llgo:link (*Object).ModuleGetDict C.PyModule_GetDict
 func (m *Object) ModuleGetDict() *Object { return nil }
 
+// llgo:link (*Object).ModuleLoadSyms C.llgoLoadPyModSyms
+func (m *Object) ModuleLoadSyms(__llgo_va_list ...any) {}
+
 // -----------------------------------------------------------------------------
diff --git a/py/module.ll b/py/module.ll
new file mode 100644
index 00000000..390aaa64
--- /dev/null
+++ b/py/module.ll
@@ -0,0 +1,84 @@
+; ModuleID = '_pyg/module.c'
+source_filename = "_pyg/module.c"
+target datalayout = "e-m:o-i64:64-i128:128-n32:64-S128"
+target triple = "arm64-apple-macosx13.0.0"
+
+%struct.PyObject = type opaque
+
+; Function Attrs: noinline nounwind optnone ssp uwtable(sync)
+define void @llgoLoadPyModSyms(%struct.PyObject* noundef %0, ...) #0 {
+  %2 = alloca %struct.PyObject*, align 8
+  %3 = alloca i8*, align 8
+  %4 = alloca i8*, align 8
+  %5 = alloca i8*, align 8
+  %6 = alloca %struct.PyObject**, align 8
+  %7 = alloca %struct.PyObject**, align 8
+  store %struct.PyObject* %0, %struct.PyObject** %2, align 8
+  %8 = bitcast i8** %3 to i8*
+  call void @llvm.va_start(i8* %8)
+  br label %9
+
+9:                                                ; preds = %26, %1
+  %10 = va_arg i8** %3, i8*
+  store i8* %10, i8** %5, align 8
+  %11 = load i8*, i8** %5, align 8
+  store i8* %11, i8** %4, align 8
+  %12 = load i8*, i8** %4, align 8
+  %13 = icmp eq i8* %12, null
+  br i1 %13, label %14, label %15
+
+14:                                               ; preds = %9
+  br label %27
+
+15:                                               ; preds = %9
+  %16 = va_arg i8** %3, %struct.PyObject**
+  store %struct.PyObject** %16, %struct.PyObject*** %7, align 8
+  %17 = load %struct.PyObject**, %struct.PyObject*** %7, align 8
+  store %struct.PyObject** %17, %struct.PyObject*** %6, align 8
+  %18 = load %struct.PyObject**, %struct.PyObject*** %6, align 8
+  %19 = load %struct.PyObject*, %struct.PyObject** %18, align 8
+  %20 = icmp eq %struct.PyObject* %19, null
+  br i1 %20, label %21, label %26
+
+21:                                               ; preds = %15
+  %22 = load %struct.PyObject*, %struct.PyObject** %2, align 8
+  %23 = load i8*, i8** %4, align 8
+  %24 = call %struct.PyObject* @PyObject_GetAttrString(%struct.PyObject* noundef %22, i8* noundef %23)
+  %25 = load %struct.PyObject**, %struct.PyObject*** %6, align 8
+  store %struct.PyObject* %24, %struct.PyObject** %25, align 8
+  br label %26
+
+26:                                               ; preds = %21, %15
+  br label %9
+
+27:                                               ; preds = %14
+  %28 = bitcast i8** %3 to i8*
+  call void @llvm.va_end(i8* %28)
+  ret void
+}
+
+; Function Attrs: nocallback nofree nosync nounwind willreturn
+declare void @llvm.va_start(i8*) #1
+
+declare %struct.PyObject* @PyObject_GetAttrString(%struct.PyObject* noundef, i8* noundef) #2
+
+; Function Attrs: nocallback nofree nosync nounwind willreturn
+declare void @llvm.va_end(i8*) #1
+
+attributes #0 = { noinline nounwind optnone ssp uwtable(sync) "frame-pointer"="non-leaf" "min-legal-vector-width"="0" "no-trapping-math"="true" "probe-stack"="__chkstk_darwin" "stack-protector-buffer-size"="8" "target-cpu"="apple-m1" "target-features"="+aes,+crc,+crypto,+dotprod,+fp-armv8,+fp16fml,+fullfp16,+lse,+neon,+ras,+rcpc,+rdm,+sha2,+sha3,+sm4,+v8.5a,+zcm,+zcz" }
+attributes #1 = { nocallback nofree nosync nounwind willreturn }
+attributes #2 = { "frame-pointer"="non-leaf" "no-trapping-math"="true" "probe-stack"="__chkstk_darwin" "stack-protector-buffer-size"="8" "target-cpu"="apple-m1" "target-features"="+aes,+crc,+crypto,+dotprod,+fp-armv8,+fp16fml,+fullfp16,+lse,+neon,+ras,+rcpc,+rdm,+sha2,+sha3,+sm4,+v8.5a,+zcm,+zcz" }
+
+!llvm.module.flags = !{!0, !1, !2, !3, !4, !5, !6, !7, !8}
+!llvm.ident = !{!9}
+
+!0 = !{i32 2, !"SDK Version", [2 x i32] [i32 13, i32 3]}
+!1 = !{i32 1, !"wchar_size", i32 4}
+!2 = !{i32 8, !"branch-target-enforcement", i32 0}
+!3 = !{i32 8, !"sign-return-address", i32 0}
+!4 = !{i32 8, !"sign-return-address-all", i32 0}
+!5 = !{i32 8, !"sign-return-address-with-bkey", i32 0}
+!6 = !{i32 7, !"PIC Level", i32 2}
+!7 = !{i32 7, !"uwtable", i32 1}
+!8 = !{i32 7, !"frame-pointer", i32 1}
+!9 = !{!"Apple clang version 14.0.3 (clang-1403.0.22.14.1)"}
diff --git a/ssa/decl.go b/ssa/decl.go
index f5d1b02c..925ea0db 100644
--- a/ssa/decl.go
+++ b/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
 }
 
 // -----------------------------------------------------------------------------
diff --git a/ssa/package.go b/ssa/package.go
index be863bce..5893c298 100644
--- a/ssa/package.go
+++ b/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
diff --git a/ssa/ssa_test.go b/ssa/ssa_test.go
index 1afddd0d..b2bfe311 100644
--- a/ssa/ssa_test.go
+++ b/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")
 	}
 }