prerun improvements

Author: dbentley-vsql

mysql-test/suite/villagesql/extension/r/extension_simple_type_usage.result

@@ -68,6 +68,12 @@ FROM test_bytearray t1, test_bytearray t2
 WHERE t1.id = 1 AND t2.id = 2;
 concatenated
 'hello   world123'
+# Test ba_call_index (exercises typed prerun/postrun + state-passing VDF)
+SELECT vsql_simple.ba_call_index() as idx FROM test_bytearray ORDER BY id;
+idx
+1
+2
+3
 # Clean up
 DROP TABLE test_bytearray;
 # Uninstall the extension

mysql-test/suite/villagesql/extension/t/extension_simple_type_usage.test

@@ -44,6 +44,9 @@ SELECT CONCAT("'", vsql_simple.ba_concat(t1.data, t2.data), "'") as concatenated
 FROM test_bytearray t1, test_bytearray t2
 WHERE t1.id = 1 AND t2.id = 2;
 
+--echo # Test ba_call_index (exercises typed prerun/postrun + state-passing VDF)
+SELECT vsql_simple.ba_call_index() as idx FROM test_bytearray ORDER BY id;
+
 --echo # Clean up
 DROP TABLE test_bytearray;
 

villagesql/examples/vsql-simple/src/extension.cc

@@ -112,6 +112,25 @@ void odd_chars(CustomArg in, CustomResult out) {
   out.set_length(kBytearrayLen);
 }
 
+// BA_CALL_INDEX: return the 1-based index of this call within the current
+// statement. Demonstrates the typed prerun + State& pattern: prerun
+// allocates a counter via emplace_state<T> and each row reads-and-
+// increments. No postrun needed — the SDK's installed deleter frees the
+// CallCounter after the last row.
+
+struct CallCounter {
+  long long n = 0;
+};
+
+void ba_call_index_prerun(PrerunArgs, PrerunResult out) {
+  out.emplace_state<CallCounter>();
+}
+
+void ba_call_index(CallCounter &state, IntResult out) {
+  state.n++;
+  out.set(state.n);
+}
+
 // BA_CONCAT: concatenate two bytearrays (returns STRING with 16 bytes)
 void ba_concat(CustomArg a, CustomArg b, StringResult out) {
   if (a.is_null() || b.is_null()) {
@@ -153,4 +172,8 @@ VEF_GENERATE_ENTRY_POINTS(make_extension()
                                         .returns(STRING)
                                         .param(BYTEARRAY)
                                         .param(BYTEARRAY)
+                                        .build())
+                              .func(make_func<&ba_call_index>("ba_call_index")
+                                        .returns(INT)
+                                        .prerun<&ba_call_index_prerun>()
                                         .build()))

villagesql/sdk/include/villagesql/abi/types.h

@@ -519,6 +519,16 @@ typedef struct {
   // Extension-allocated state. Set this to pass data to vdf and postrun.
   // Caller initializes to NULL.
   void *user_data;
+
+  // Optional SDK-managed deleter for user_data. If non-NULL, the server
+  // calls user_data_deleter(user_data) after the postrun returns (or after
+  // the last vdf call, if no postrun was registered). Used by the C++ SDK's
+  // PrerunResult::emplace_state<T> to install typed_delete<T> so extensions
+  // do not have to write a postrun just to free state.
+  //
+  // Raw C ABI users typically leave this NULL and free user_data themselves
+  // in their postrun.
+  void (*user_data_deleter)(void *user_data);
 } vef_prerun_result_t;
 
 typedef void (*vef_prerun_func_t)(vef_context_t *ctx, vef_prerun_args_t *args,

villagesql/sdk/include/villagesql/extension.h

@@ -114,19 +114,32 @@
 // PRERUN/POSTRUN FUNCTIONS
 // ------------------------
 //
-// For prerun/postrun functions (per-statement setup/teardown):
+// Prerun (per-statement setup) and postrun (per-statement teardown) hooks
+// use typed wrappers — see <villagesql/vsql/pre_post_run.h>. Raw ABI
+// signatures are not accepted.
+//
+//   void my_prerun(vsql::PrerunArgs args, vsql::PrerunResult out) {
+//     // validate args.type_at(i), call out.error(...) on failure,
+//     // out.emplace_state<MyState>(...) for per-statement state, etc.
+//   }
+//   void my_postrun(vsql::PostrunArgs args) {
+//     // optional; only needed if you used set_user_data with a non-C++
+//     // allocator. emplace_state<T> is auto-freed by the SDK.
+//   }
 //
 //   make_func<&my_impl>("my_func")
 //     .returns(STRING)
-//     .prerun<&my_prerun>()   // Called before first row
-//     .postrun<&my_postrun>() // Called after last row
+//     .prerun<&my_prerun>()
+//     .postrun<&my_postrun>()
 //     .build()
 //
-// Note: Prerun and postrun functions can be a cumbersome API. The func builder
-// already handles simple cases (e.g., type checking for functions with fixed
-// args and allocating fixed buffer sizes). We want to cover more cases. If
-// you find that you need to use prerun or postrun functions, please come talk
-// to us so we can understand your use case.
+// If your VDF wants the per-statement state, declare it as the first
+// parameter: void my_impl(MyState& s, IntArg a, IntResult out).
+//
+// The builder also handles common cases without you writing a hook (e.g.
+// arg-count/type validation for fixed-arity functions, fixed-size result
+// buffers). Reach for prerun only when you need custom validation, a
+// dynamically-sized result buffer, or per-statement state.
 //
 //
 // AGGREGATE FUNCTIONS

villagesql/sdk/include/villagesql/vsql/func_builder.h

@@ -47,6 +47,7 @@
 #include <villagesql/abi/types.h>
 #include <villagesql/vsql/func_types.h>
 #include <villagesql/vsql/maybe_params.h>
+#include <villagesql/vsql/pre_post_run.h>
 #include <villagesql/vsql/type_params_cache.h>
 
 namespace vsql {
@@ -341,6 +342,51 @@ struct AggResultWithOutputWrapper {
   }
 };
 
+// =============================================================================
+// Prerun / Postrun Adapters
+// =============================================================================
+
+// Thunks that convert a typed user prerun/postrun (void(PrerunArgs,
+// PrerunResult) / void(PostrunArgs)) into the raw vef_*_func_t ABI shape.
+// Used by FuncBuilder::prerun/postrun when the user's hook is typed.
+
+template <auto Hook>
+void typed_prerun_thunk(vef_context_t *, vef_prerun_args_t *args,
+                        vef_prerun_result_t *result) {
+  Hook(PrerunArgs(args), PrerunResult(result));
+}
+
+template <auto Hook>
+void typed_postrun_thunk(vef_context_t *, vef_postrun_args_t *args,
+                         vef_postrun_result_t *) {
+  Hook(PostrunArgs(args));
+}
+
+// Predicates used by FuncBuilder::prerun/postrun to decide which thunk
+// (if any) to install. Each is true exactly when the hook's signature
+// matches the typed shape for its slot.
+
+template <auto Hook>
+constexpr bool is_typed_prerun() {
+  using Params = typename FuncParamTypes<decltype(Hook)>::type;
+  if constexpr (std::tuple_size_v<Params> != 2) {
+    return false;
+  } else {
+    return std::is_same_v<std::tuple_element_t<0, Params>, PrerunArgs> &&
+           std::is_same_v<std::tuple_element_t<1, Params>, PrerunResult>;
+  }
+}
+
+template <auto Hook>
+constexpr bool is_typed_postrun() {
+  using Params = typename FuncParamTypes<decltype(Hook)>::type;
+  if constexpr (std::tuple_size_v<Params> != 1) {
+    return false;
+  } else {
+    return std::is_same_v<std::tuple_element_t<0, Params>, PostrunArgs>;
+  }
+}
+
 // =============================================================================
 // Wrapper Template
 // =============================================================================
@@ -383,6 +429,80 @@ struct Wrapper {
   }
 };
 
+// Wrapper for VDFs whose first parameter is `State&` — typed per-statement
+// state allocated in prerun via PrerunResult::emplace_state<State>. The
+// wrapper dereferences args->user_data as State* and forwards a reference.
+//
+// Param tuple shape: <State&, TypedArg..., ResultWrapper>. NumParams is the
+// SQL argument count (not counting state or result), so typed args live at
+// indices [1, NumParams].
+template <auto Func, typename State, size_t NumParams>
+struct WrapperTypedState {
+  static void invoke(vef_context_t *ctx, vef_vdf_args_t *args,
+                     vef_vdf_result_t *result) {
+    invoke_impl(ctx, args, result, std::make_index_sequence<NumParams>{});
+  }
+
+ private:
+  template <size_t... Is>
+  static void invoke_impl(vef_context_t *ctx, vef_vdf_args_t *args,
+                          vef_vdf_result_t *result,
+                          std::index_sequence<Is...>) {
+    using Params = typename FuncParamTypes<decltype(Func)>::type;
+    State &state = *static_cast<State *>(args->user_data);
+    std::array<vef_invalue_t, NumParams> vals{
+        get_invalue(ctx, args, static_cast<unsigned int>(Is))...};
+    Func(state, make_arg<std::tuple_element_t<1 + Is, Params>>(&vals[Is])...,
+         make_result<std::tuple_element_t<1 + NumParams, Params>>(result));
+  }
+
+  template <typename T>
+  static T make_arg(vef_invalue_t *v) {
+    return T(v);
+  }
+  template <typename T>
+  static T make_result(vef_vdf_result_t *r) {
+    return T(r);
+  }
+};
+
+// Wrapper for VDFs whose first parameter is `void*` — raw escape hatch for
+// extensions that manage state with custom allocators, polymorphic state,
+// or anything that doesn't fit emplace_state<T>. The wrapper forwards
+// args->user_data straight through.
+//
+// Param tuple shape: <void*, TypedArg..., ResultWrapper>. Same indexing as
+// WrapperTypedState.
+template <auto Func, size_t NumParams>
+struct WrapperVoidState {
+  static void invoke(vef_context_t *ctx, vef_vdf_args_t *args,
+                     vef_vdf_result_t *result) {
+    invoke_impl(ctx, args, result, std::make_index_sequence<NumParams>{});
+  }
+
+ private:
+  template <size_t... Is>
+  static void invoke_impl(vef_context_t *ctx, vef_vdf_args_t *args,
+                          vef_vdf_result_t *result,
+                          std::index_sequence<Is...>) {
+    using Params = typename FuncParamTypes<decltype(Func)>::type;
+    std::array<vef_invalue_t, NumParams> vals{
+        get_invalue(ctx, args, static_cast<unsigned int>(Is))...};
+    Func(args->user_data,
+         make_arg<std::tuple_element_t<1 + Is, Params>>(&vals[Is])...,
+         make_result<std::tuple_element_t<1 + NumParams, Params>>(result));
+  }
+
+  template <typename T>
+  static T make_arg(vef_invalue_t *v) {
+    return T(v);
+  }
+  template <typename T>
+  static T make_result(vef_vdf_result_t *r) {
+    return T(r);
+  }
+};
+
 // =============================================================================
 // Type Operation VDF Wrappers
 // =============================================================================
@@ -962,15 +1082,21 @@ class FuncBuilder {
     return *this;
   }
 
-  template <vef_prerun_func_t Hook>
+  template <auto Hook>
   constexpr FuncBuilder<Func, NumParams> &prerun() {
-    prerun_ = Hook;
+    static_assert(detail::is_typed_prerun<Hook>(),
+                  "prerun<Hook>(): Hook must be void(PrerunArgs, "
+                  "PrerunResult). Raw ABI signatures are not accepted.");
+    prerun_ = &detail::typed_prerun_thunk<Hook>;
     return *this;
   }
 
-  template <vef_postrun_func_t Hook>
+  template <auto Hook>
   constexpr FuncBuilder<Func, NumParams> &postrun() {
-    postrun_ = Hook;
+    static_assert(detail::is_typed_postrun<Hook>(),
+                  "postrun<Hook>(): Hook must be void(PostrunArgs). Raw ABI "
+                  "signatures are not accepted.");
+    postrun_ = &detail::typed_postrun_thunk<Hook>;
     return *this;
   }
 
@@ -982,7 +1108,28 @@ class FuncBuilder {
     using UniquePTuple = typename detail::unique_params_types<AllParams>::type;
 
     detail::FuncWithMetadata meta{};
-    meta.f = &detail::Wrapper<Func, NumParams>::invoke;
+    // Dispatch on the first parameter type to decide whether the function
+    // wants state passed in. Three shapes:
+    //   void(TypedArgs..., ResultWrapper)             — stateless
+    //   void(State&, TypedArgs..., ResultWrapper)     — typed state, mutable
+    //   void(const State&, TypedArgs..., ResultWrapper) — typed state, const
+    //   void(void*, TypedArgs..., ResultWrapper)      — raw state
+    if constexpr (std::tuple_size_v<AllParams> >= 2) {
+      using First = std::tuple_element_t<0, AllParams>;
+      if constexpr (std::is_same_v<First, void *>) {
+        meta.f = &detail::WrapperVoidState<Func, NumParams>::invoke;
+      } else if constexpr (std::is_lvalue_reference_v<First>) {
+        // Both `State&` and `const State&` route to WrapperTypedState; the
+        // implicit conversion from `State&` to `const State&` happens at the
+        // call site if the user opted for read-only.
+        using State = std::remove_const_t<std::remove_reference_t<First>>;
+        meta.f = &detail::WrapperTypedState<Func, State, NumParams>::invoke;
+      } else {
+        meta.f = &detail::Wrapper<Func, NumParams>::invoke;
+      }
+    } else {
+      meta.f = &detail::Wrapper<Func, NumParams>::invoke;
+    }
     meta.prerun = prerun_;
     meta.postrun = postrun_;
     meta.return_type = detail::to_vef_type(return_type_);