From bdb398e6aee0569e8b10be44e69ac720d861d7e1 Mon Sep 17 00:00:00 2001
From: Clark McCauley <clarkmccauley@gmail.com>
Date: Mon, 7 Oct 2024 09:46:37 -0600
Subject: [PATCH] Exposed extism module and added methods to get exported
 functions (#79)

---
 extism.go  | 21 ++++++++-------------
 module.go  | 30 ++++++++++++++++++++++++++++++
 runtime.go |  4 ++--
 3 files changed, 40 insertions(+), 15 deletions(-)
 create mode 100644 module.go

diff --git a/extism.go b/extism.go
index 6245355..5648076 100644
--- a/extism.go
+++ b/extism.go
@@ -24,11 +24,6 @@ import (
 	"github.com/tetratelabs/wazero/sys"
 )
 
-type module struct {
-	module api.Module
-	wasm   []byte
-}
-
 type PluginCtxKey string
 type InputOffsetKey string
 
@@ -119,8 +114,8 @@ func (l LogLevel) String() string {
 // Plugin is used to call WASM functions
 type Plugin struct {
 	Runtime *Runtime
-	Modules map[string]module
-	Main    module
+	Modules map[string]Module
+	Main    Module
 	Timeout time.Duration
 	Config  map[string]string
 	// NOTE: maybe we can have some nice methods for getting/setting vars
@@ -424,7 +419,7 @@ func NewPlugin(
 		return nil, fmt.Errorf("manifest can't be empty")
 	}
 
-	modules := map[string]module{}
+	modules := map[string]Module{}
 
 	// NOTE: this is only necessary for guest modules because
 	// host modules have the same access privileges as the host itself
@@ -500,7 +495,7 @@ func NewPlugin(
 			return nil, err
 		}
 
-		modules[data.Name] = module{module: m, wasm: data.Data}
+		modules[data.Name] = Module{inner: m}
 	}
 
 	i := 0
@@ -514,7 +509,7 @@ func NewPlugin(
 		varMax = int64(manifest.Memory.MaxVarBytes)
 	}
 	for _, m := range modules {
-		if m.module.Name() == "main" {
+		if m.inner.Name() == "main" {
 			p := &Plugin{
 				Runtime:              &c,
 				Modules:              modules,
@@ -619,9 +614,9 @@ func (plugin *Plugin) GetErrorWithContext(ctx context.Context) string {
 	return string(mem)
 }
 
-// FunctionExists returns true when the named function is present in the plugin's main module
+// FunctionExists returns true when the named function is present in the plugin's main Module
 func (plugin *Plugin) FunctionExists(name string) bool {
-	return plugin.Main.module.ExportedFunction(name) != nil
+	return plugin.Main.inner.ExportedFunction(name) != nil
 }
 
 // Call a function by name with the given input, returning the output
@@ -646,7 +641,7 @@ func (plugin *Plugin) CallWithContext(ctx context.Context, name string, data []b
 
 	ctx = context.WithValue(ctx, InputOffsetKey("inputOffset"), intputOffset)
 
-	var f = plugin.Main.module.ExportedFunction(name)
+	var f = plugin.Main.inner.ExportedFunction(name)
 
 	if f == nil {
 		return 1, []byte{}, fmt.Errorf("unknown function: %s", name)
diff --git a/module.go b/module.go
new file mode 100644
index 0000000..c5d928d
--- /dev/null
+++ b/module.go
@@ -0,0 +1,30 @@
+package extism
+
+import "github.com/tetratelabs/wazero/api"
+
+// Module is a wrapper around a wazero module. It allows us to provide
+// our own API and stability guarantees despite any changes that wazero
+// may choose to make.
+type Module struct {
+	inner api.Module
+}
+
+// ExportedFunctions returns a map of functions exported from the module
+// keyed by the function name.
+func (m *Module) ExportedFunctions() map[string]FunctionDefinition {
+	v := make(map[string]FunctionDefinition)
+	for name, def := range m.inner.ExportedFunctionDefinitions() {
+		v[name] = FunctionDefinition{inner: def}
+	}
+	return v
+}
+
+// FunctionDefinition represents a function defined in a module. It provides
+// a wrapper around the underlying wazero function definition.
+type FunctionDefinition struct {
+	inner api.FunctionDefinition
+}
+
+func (f *FunctionDefinition) Name() string {
+	return f.inner.Name()
+}
diff --git a/runtime.go b/runtime.go
index 9fa8c00..d079003 100644
--- a/runtime.go
+++ b/runtime.go
@@ -25,12 +25,12 @@ type guestRuntime struct {
 func detectGuestRuntime(ctx context.Context, p *Plugin) guestRuntime {
 	m := p.Main
 
-	runtime, ok := haskellRuntime(ctx, p, m.module)
+	runtime, ok := haskellRuntime(ctx, p, m.inner)
 	if ok {
 		return runtime
 	}
 
-	runtime, ok = wasiRuntime(ctx, p, m.module)
+	runtime, ok = wasiRuntime(ctx, p, m.inner)
 	if ok {
 		return runtime
 	}