Embed Katex as Wasm

See gohugoio#11927

internal/ext/deleteme.go

@@ -0,0 +1,69 @@
+package ext
+
+import (
+	"context"
+	"fmt"
+	"os"
+	"sync"
+
+	"github.com/tetratelabs/wazero"
+	"github.com/tetratelabs/wazero/api"
+	"github.com/tetratelabs/wazero/experimental"
+)
+
+func deleteMe(name string, wasm []byte, needsQuickJSProvider bool) func(ctx context.Context, r wazero.Runtime, inouts []*inOut) (func() error, error) {
+	return func(ctx context.Context, r wazero.Runtime, inouts []*inOut) (func() error, error) {
+		compiledModule, err := r.CompileModule(ctx, wasm)
+		if err != nil {
+			return nil, err
+		}
+
+		var compiledQuickJS wazero.CompiledModule
+		if needsQuickJSProvider {
+			compiledQuickJS, err = r.CompileModule(ctx, quickjsWasm)
+			if err != nil {
+				return nil, err
+			}
+		}
+
+		return func() error {
+			var wg sync.WaitGroup
+			for i, c := range inouts {
+				name := fmt.Sprintf("%s_%d", name, i)
+				c := c
+
+				wg.Add(1)
+				go func() {
+					defer wg.Done()
+					configBase := wazero.NewModuleConfig().WithStderr(os.Stderr).WithStartFunctions()
+					config := configBase.WithName(name)
+					if needsQuickJSProvider {
+						// TODO1 else, configure with stdout etc.
+						qmod, err := r.InstantiateModule(ctx, compiledQuickJS, configBase.WithName("").WithStdout(c.stdout).WithStdin(c.stdin))
+						if err != nil {
+							panic(err)
+						}
+						ctx = experimental.WithImportResolver(ctx,
+							func(name string) api.Module {
+								if name == "javy_quickjs_provider_v2" {
+									return qmod
+								}
+								return nil
+							})
+
+					}
+
+					mod, err := r.InstantiateModule(ctx, compiledModule, config)
+					if err != nil {
+						panic(err)
+					}
+					if _, err := mod.ExportedFunction("_start").Call(ctx); err != nil {
+						panic(err)
+					}
+				}()
+			}
+			wg.Wait()
+			return nil
+		}, nil
+	}
+}

internal/ext/katex.go

@@ -8,8 +8,9 @@ import (
 var katexWasm []byte
 
 var katexOptions = Options{
-	CompileModule: compileFunc("katex", katexWasm, true),
-	PoolSize:      8,
+	Runtime:  Binary{Name: "javy_quickjs_provider_v2", Data: quickjsWasm},
+	Main:     Binary{Name: "renderkatex", Data: katexWasm},
+	PoolSize: 8,
 }
 
 // StartKatex starts a new dispatcher for the Katex module.

internal/ext/warpc.go

@@ -15,6 +15,7 @@ import (
 	"time"
 
 	"github.com/gohugoio/hugo/common/hugio"
+	"golang.org/x/sync/errgroup"
 
 	"github.com/tetratelabs/wazero"
 	"github.com/tetratelabs/wazero/api"
@@ -216,15 +217,60 @@ func (call *call[Q, R]) done() {
 	}
 }
 
+// Binary represents a WebAssembly binary.
+type Binary struct {
+	// The name of the binary.
+	// For 	quickjs, this must match the instance import name, "javy_quickjs_provider_v2".
+	// For the main module, we only use this for caching.
+	Name string
+
+	// THe wasm binary.
+	Data []byte
+}
+
 type Options struct {
 	Ctx context.Context
 
-	CompileModule       func(ctx context.Context, r wazero.Runtime, io []*inOut) (func() error, error)
+	// E.g. quickjs wasm. May be omitted if not needed.
+	Runtime Binary
+
+	// The main module to instantiate.
+	Main Binary
+
 	CompilationCacheDir string
 	PoolSize            int
 }
 
+type CompileModuleContext struct {
+	Opts    Options
+	Runtime wazero.Runtime
+}
+
+type CompiledModule struct {
+	// Runtime (e.g. QuickJS) may be nil if not needed (e.g. embedded in Module).
+	Runtime wazero.CompiledModule
+
+	// If Runtime is not nil, this should be the name of the instance.
+	RuntimeName string
+
+	// The main module to instantiate.
+	// This will be insantiated multiple times in a pool,
+	// so it does not need a name.
+	Module wazero.CompiledModule
+}
+
 func Start[Q, R IDGetter](opts Options) (Dispatcher[Q, R], error) {
+	if opts.Main.Data == nil {
+		return nil, errors.New("Main.Data must be set")
+	}
+	if opts.Main.Name == "" {
+		return nil, errors.New("Main.Name must be set")
+	}
+
+	if opts.Runtime.Data != nil && opts.Runtime.Name == "" {
+		return nil, errors.New("Runtime.Name must be set")
+	}
+
 	if opts.PoolSize == 0 {
 		opts.PoolSize = 1
 	}
@@ -242,9 +288,7 @@ func newDispatcher[Q, R IDGetter](opts Options) (*dispatcherPool[Q, R], error) {
 	if opts.Ctx == nil {
 		opts.Ctx = context.Background()
 	}
-	if opts.CompileModule == nil {
-		return nil, errors.New("InstansiateModule is required")
-	}
+
 	ctx := opts.Ctx
 
 	runtimeConfig := wazero.NewRuntimeConfig()
@@ -280,14 +324,63 @@ func newDispatcher[Q, R IDGetter](opts Options) (*dispatcherPool[Q, R], error) {
 		}
 	}
 
-	run, err := opts.CompileModule(ctx, r, inOuts)
+	var (
+		runtimeModule wazero.CompiledModule
+		mainModule    wazero.CompiledModule
+		err           error
+	)
+
+	if opts.Runtime.Data != nil {
+		runtimeModule, err = r.CompileModule(ctx, opts.Runtime.Data)
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	mainModule, err = r.CompileModule(ctx, opts.Main.Data)
 	if err != nil {
 		return nil, err
 	}
 
+	run := func() error {
+		g, ctx := errgroup.WithContext(ctx)
+		for _, c := range inOuts {
+			c := c
+			g.Go(func() error {
+				configBase := wazero.NewModuleConfig().WithStdout(c.stdout).WithStdin(c.stdin).WithStartFunctions()
+				if opts.Runtime.Data != nil {
+					// This needs to be anonymous, it will be resolved in the import resolver below.
+					runtimeInstance, err := r.InstantiateModule(ctx, runtimeModule, configBase.WithName(""))
+					if err != nil {
+						return err
+					}
+					ctx = experimental.WithImportResolver(ctx,
+						func(name string) api.Module {
+							if name == opts.Runtime.Name {
+								return runtimeInstance
+							}
+							return nil
+						},
+					)
+				}
+
+				mainInstance, err := r.InstantiateModule(ctx, mainModule, configBase.WithName("").WithStdout(c.stdout).WithStdin(c.stdin))
+				if err != nil {
+					return err
+				}
+				if _, err := mainInstance.ExportedFunction("_start").Call(ctx); err != nil {
+					return err
+				}
+
+				return nil
+			})
+		}
+		return g.Wait()
+	}
+
 	done := make(chan struct{})
 	go func() {
-		// This will block until stdin is closed.
+		// This will block until stdin is closed or it encounters an error.
 		err := run()
 		if err != nil {
 			panic(err)
@@ -336,6 +429,30 @@ func printStackTrace(w io.Writer) {
 	fmt.Fprintf(w, "%s", buf)
 }
 
+// TODO1 delete me.
+func compileFuncTake2(name string, wasm []byte, needsQuickJSProvider bool) func(ctx CompileModuleContext) (CompiledModule, error) {
+	return func(ctx CompileModuleContext) (CompiledModule, error) {
+		compiledModule, err := ctx.Runtime.CompileModule(ctx.Opts.Ctx, wasm)
+		if err != nil {
+			return CompiledModule{}, err
+		}
+
+		var compiledQuickJS wazero.CompiledModule
+		if needsQuickJSProvider {
+			compiledQuickJS, err = ctx.Runtime.CompileModule(ctx.Opts.Ctx, quickjsWasm)
+			if err != nil {
+				return CompiledModule{}, err
+			}
+		}
+
+		return CompiledModule{
+			Runtime:     compiledQuickJS,
+			RuntimeName: "javy_quickjs_provider_v2",
+			Module:      compiledModule,
+		}, nil
+	}
+}
+
 func compileFunc(name string, wasm []byte, needsQuickJSProvider bool) func(ctx context.Context, r wazero.Runtime, inouts []*inOut) (func() error, error) {
 	return func(ctx context.Context, r wazero.Runtime, inouts []*inOut) (func() error, error) {
 		compiledModule, err := r.CompileModule(ctx, wasm)