From 482f796bad9cdb46443600816b593ee3cbc65100 Mon Sep 17 00:00:00 2001
From: Aofei Sheng <aofei@aofeisheng.com>
Date: Sat, 3 Aug 2024 08:35:19 +0800
Subject: [PATCH] ssa: add `llgo:link` support to `Builder.abiMthd`

---
 cl/builtin_test.go | 65 ++++++++++++++++++++++++++++++++++++++++++++++
 cl/compile.go      | 13 ++++++++++
 ssa/abitype.go     |  3 +++
 ssa/package.go     |  6 +++++
 4 files changed, 87 insertions(+)

diff --git a/cl/builtin_test.go b/cl/builtin_test.go
index e2b0fce2..615b3274 100644
--- a/cl/builtin_test.go
+++ b/cl/builtin_test.go
@@ -282,3 +282,68 @@ func TestErrVarOf(t *testing.T) {
 	g := &ssa.Global{Pkg: ssaPkg}
 	ctx.varOf(nil, g)
 }
+
+func TestContextResolveLinkname(t *testing.T) {
+	tests := []struct {
+		name    string
+		context *context
+		input   string
+		want    string
+		panics  bool
+	}{
+		{
+			name: "Normal",
+			context: &context{
+				link: map[string]string{
+					"foo": "C.bar",
+				},
+			},
+			input: "foo",
+			want:  "bar",
+		},
+		{
+			name: "MultipleLinks",
+			context: &context{
+				link: map[string]string{
+					"foo1": "C.bar1",
+					"foo2": "C.bar2",
+				},
+			},
+			input: "foo2",
+			want:  "bar2",
+		},
+		{
+			name:    "NoLink",
+			context: &context{link: map[string]string{}},
+			input:   "foo",
+			want:    "foo",
+		},
+		{
+			name: "InvalidLink",
+			context: &context{
+				link: map[string]string{
+					"foo": "invalid.bar",
+				},
+			},
+			input:  "foo",
+			panics: true,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if tt.panics {
+				defer func() {
+					if r := recover(); r == nil {
+						t.Error("want panic")
+					}
+				}()
+			}
+			got := tt.context.resolveLinkname(tt.input)
+			if !tt.panics {
+				if got != tt.want {
+					t.Errorf("got %q, want %q", got, tt.want)
+				}
+			}
+		})
+	}
+}
diff --git a/cl/compile.go b/cl/compile.go
index 113dfc1b..ea011b05 100644
--- a/cl/compile.go
+++ b/cl/compile.go
@@ -25,6 +25,7 @@ import (
 	"log"
 	"os"
 	"sort"
+	"strings"
 
 	"github.com/goplus/llgo/cl/blocks"
 	"github.com/goplus/llgo/internal/typepatch"
@@ -812,6 +813,7 @@ func NewPackageEx(prog llssa.Program, patches Patches, pkg *ssa.Package, files [
 	ctx.initPyModule()
 	ctx.initFiles(pkgPath, files)
 	ret.SetPatch(ctx.patchType)
+	ret.SetResolveLinkname(ctx.resolveLinkname)
 
 	if hasPatch {
 		skips := ctx.skips
@@ -904,4 +906,15 @@ func (p *context) patchType(typ types.Type) types.Type {
 	return typ
 }
 
+func (p *context) resolveLinkname(name string) string {
+	if link, ok := p.link[name]; ok {
+		prefix, ltarget, _ := strings.Cut(link, ".")
+		if prefix != "C" {
+			panic("resolveLinkname: invalid link: " + link)
+		}
+		return ltarget
+	}
+	return name
+}
+
 // -----------------------------------------------------------------------------
diff --git a/ssa/abitype.go b/ssa/abitype.go
index 9cc25d56..8671eeca 100644
--- a/ssa/abitype.go
+++ b/ssa/abitype.go
@@ -154,6 +154,9 @@ func (b Builder) abiMethodOf(mPkg *types.Package, mName string, mSig *types.Sign
 
 func (b Builder) abiMthd(mPkg *types.Package, mName string, mSig *types.Signature, name, abiTyp, ifn llvm.Value) (ret Expr, tfn llvm.Value) {
 	fullName := FuncName(mPkg, mName, mSig.Recv())
+	if b.Pkg.fnlink != nil {
+		fullName = b.Pkg.fnlink(fullName)
+	}
 	tfn = b.Pkg.NewFunc(fullName, mSig, InGo).impl // TODO(xsw): use rawType to speed up
 	if ifn.IsNil() {
 		ifn = tfn
diff --git a/ssa/package.go b/ssa/package.go
index 02b0e360..d394eb87 100644
--- a/ssa/package.go
+++ b/ssa/package.go
@@ -586,6 +586,7 @@ type aPackage struct {
 	named  map[types.Type]Expr
 	afterb unsafe.Pointer
 	patch  func(types.Type) types.Type
+	fnlink func(string) string
 
 	iRoutine int
 }
@@ -658,6 +659,11 @@ func (p Package) SetPatch(fn func(types.Type) types.Type) {
 	p.patch = fn
 }
 
+// SetResolveLinkname sets a function to resolve linkname.
+func (p Package) SetResolveLinkname(fn func(string) string) {
+	p.fnlink = fn
+}
+
 // -----------------------------------------------------------------------------
 
 func (p Package) afterBuilder() Builder {