Custom aggregate functions.

2026-01-11 21:49:13 +00:00 · 2023-07-01 15:14:45 +01:00
parent fec1f8d32a
commit 3844e81404
9 changed files with 221 additions and 32 deletions
--- a/func.go
+++ b/func.go
@@ -35,48 +35,124 @@ func (c *Conn) CreateFunction(name string, nArg int, flag FunctionFlag, fn func(
 	return c.error(r)
 }

+// CreateWindowFunction defines a new aggregate or window function.
+//
+// https://www.sqlite.org/c3ref/create_function.html
+func (c *Conn) CreateWindowFunction(name string, nArg int, flag FunctionFlag, fn AggregateFunction) error {
+	call := c.api.createAggregate
+	namePtr := c.arena.string(name)
+	funcPtr := util.AddHandle(c.ctx, fn)
+	if _, ok := fn.(WindowFunction); ok {
+		call = c.api.createWindow
+	}
+	r := c.call(call,
+		uint64(c.handle), uint64(namePtr), uint64(nArg),
+		uint64(flag), uint64(funcPtr))
+	return c.error(r)
+}
+
+// ScalarFunction is the interface a scalar function should implement.
+//
+// https://www.sqlite.org/appfunc.html
+type ScalarFunction interface {
+	Func(ctx Context, arg ...Value)
+}
+
+// AggregateFunction is the interface an aggregate function should implement.
+//
+// https://www.sqlite.org/appfunc.html
+type AggregateFunction interface {
+	Step(ctx Context, arg ...Value)
+	Final(ctx Context)
+}
+
+// WindowFunction is the interface an aggregate window function should implement.
+//
+// https://www.sqlite.org/windowfunctions.html
+type WindowFunction interface {
+	AggregateFunction
+	Value(ctx Context)
+	Inverse(ctx Context, arg ...Value)
+}
+
 func exportHostFunctions(env wazero.HostModuleBuilder) wazero.HostModuleBuilder {
-	util.ExportFuncVI(env, "go_destroy", cbDestroy)
-	util.ExportFuncIIIIII(env, "go_compare", cbCompare)
-	util.ExportFuncVIII(env, "go_func", cbFunc)
-	util.ExportFuncVIII(env, "go_step", cbStep)
-	util.ExportFuncVI(env, "go_final", cbFinal)
-	util.ExportFuncVI(env, "go_value", cbValue)
-	util.ExportFuncVIII(env, "go_inverse", cbInverse)
+	util.ExportFuncVI(env, "go_destroy", callbackDestroy)
+	util.ExportFuncIIIIII(env, "go_compare", callbackCompare)
+	util.ExportFuncVIII(env, "go_func", callbackFunc)
+	util.ExportFuncVIII(env, "go_step", callbackStep)
+	util.ExportFuncVI(env, "go_final", callbackFinal)
+	util.ExportFuncVI(env, "go_value", callbackValue)
+	util.ExportFuncVIII(env, "go_inverse", callbackInverse)
 	return env
 }

-func cbDestroy(ctx context.Context, mod api.Module, pApp uint32) {
+func callbackDestroy(ctx context.Context, mod api.Module, pApp uint32) {
 	util.DelHandle(ctx, pApp)
 }

-func cbCompare(ctx context.Context, mod api.Module, pApp, nKey1, pKey1, nKey2, pKey2 uint32) uint32 {
+func callbackCompare(ctx context.Context, mod api.Module, pApp, nKey1, pKey1, nKey2, pKey2 uint32) uint32 {
 	fn := util.GetHandle(ctx, pApp).(func(a, b []byte) int)
 	return uint32(fn(util.View(mod, pKey1, uint64(nKey1)), util.View(mod, pKey2, uint64(nKey2))))
 }

-func cbFunc(ctx context.Context, mod api.Module, pCtx, nArg, pArg uint32) {
+func callbackFunc(ctx context.Context, mod api.Module, pCtx, nArg, pArg uint32) {
 	module := ctx.Value(moduleKey{}).(*module)
 	pApp := uint32(module.call(module.api.userData, uint64(pCtx)))
+	fn := util.GetHandle(ctx, pApp).(func(ctx Context, arg ...Value))
+	fn(Context{
+		module: module,
+		handle: pCtx,
+	}, callbackArgs(module, nArg, pArg)...)
+}
+
+func callbackStep(ctx context.Context, mod api.Module, pCtx, nArg, pArg uint32) {
+	module := ctx.Value(moduleKey{}).(*module)
+	pApp := uint32(module.call(module.api.userData, uint64(pCtx)))
+	fn := util.GetHandle(ctx, pApp).(AggregateFunction)
+	fn.Step(Context{
+		module: module,
+		handle: pCtx,
+	}, callbackArgs(module, nArg, pArg)...)
+}
+
+func callbackFinal(ctx context.Context, mod api.Module, pCtx uint32) {
+	module := ctx.Value(moduleKey{}).(*module)
+	pApp := uint32(module.call(module.api.userData, uint64(pCtx)))
+	fn := util.GetHandle(ctx, pApp).(AggregateFunction)
+	fn.Final(Context{
+		module: module,
+		handle: pCtx,
+		final:  true,
+	})
+}
+
+func callbackValue(ctx context.Context, mod api.Module, pCtx uint32) {
+	module := ctx.Value(moduleKey{}).(*module)
+	pApp := uint32(module.call(module.api.userData, uint64(pCtx)))
+	fn := util.GetHandle(ctx, pApp).(WindowFunction)
+	fn.Value(Context{
+		module: module,
+		handle: pCtx,
+	})
+}
+
+func callbackInverse(ctx context.Context, mod api.Module, pCtx, nArg, pArg uint32) {
+	module := ctx.Value(moduleKey{}).(*module)
+	pApp := uint32(module.call(module.api.userData, uint64(pCtx)))
+	fn := util.GetHandle(ctx, pApp).(WindowFunction)
+	fn.Inverse(Context{
+		module: module,
+		handle: pCtx,
+	}, callbackArgs(module, nArg, pArg)...)
+}
+
+func callbackArgs(module *module, nArg, pArg uint32) []Value {
 	args := make([]Value, nArg)
 	for i := range args {
 		args[i] = Value{
-			handle: util.ReadUint32(mod, pArg+ptrlen*uint32(i)),
 			module: module,
+			handle: util.ReadUint32(module.mod, pArg+ptrlen*uint32(i)),
 		}
 	}
-	context := Context{
-		handle: pCtx,
-		module: module,
-	}
-	fn := util.GetHandle(ctx, pApp).(func(ctx Context, arg ...Value))
-	fn(context, args...)
+	return args
 }
-
-func cbStep(ctx context.Context, mod api.Module, pCtx, nArg, pArg uint32) {}
-
-func cbFinal(ctx context.Context, mod api.Module, pCtx uint32) {}
-
-func cbValue(ctx context.Context, mod api.Module, pCtx uint32) {}
-
-func cbInverse(ctx context.Context, mod api.Module, pCtx, nArg, pArg uint32) {}