Add mount function calls in typescript

Author: aasoni

crates/bindings-typescript/src/server/runtime.ts

@@ -10,6 +10,7 @@ import {
 import {
   RawModuleDef,
   ViewResultHeader,
+  type RawReducerDefV10,
   type RawTableDefV10,
   type Typespace,
 } from '../lib/autogen/types';
@@ -27,6 +28,7 @@ import {
   type UniqueIndex,
 } from '../lib/indexes';
 import { callProcedure } from './procedures';
+import type { Reducers } from './reducers';
 import {
   type AuthCtx,
   type JsonObject,
@@ -42,7 +44,7 @@ import type { DbView } from './db_view';
 import { getErrorConstructor, SenderError } from './errors';
 import { Range, type Bound } from './range';
 import { makeRandom, type Random } from './rng';
-import type { SchemaInner } from './schema';
+import type { MountedDispatchInfo, SchemaInner } from './schema';
 
 const { freeze } = Object;
 
@@ -212,13 +214,17 @@ export const ReducerCtxImpl = class ReducerCtx<
     me: InstanceType<typeof this>,
     sender: Identity,
     timestamp: Timestamp,
-    connectionId: ConnectionId | null
+    connectionId: ConnectionId | null,
+    dbView?: DbView<any>
   ) {
     me.sender = sender;
     me.timestamp = timestamp;
     me.connectionId = connectionId;
     me.#uuidCounter = undefined;
     me.#senderAuth = undefined;
+    if (dbView !== undefined) {
+      me.db = dbView;
+    }
   }
 
   get databaseIdentity() {
@@ -272,21 +278,55 @@ export const callUserFunction = function __spacetimedb_end_short_backtrace<
   return fn(...args);
 };
 
+type FlatMountDispatch = {
+  reducerFns: Reducers;
+  reducerDefs: RawReducerDefV10[];
+  tables: Array<{ accessorName: string; tableDef: RawTableDefV10 }>;
+  typespace: Typespace;
+  dbView_: DbView<any> | undefined;
+};
+
+function flattenMountDispatches(
+  dispatches: MountedDispatchInfo[]
+): FlatMountDispatch[] {
+  const result: FlatMountDispatch[] = [];
+  for (const d of dispatches) {
+    result.push({
+      reducerFns: d.reducerFns,
+      reducerDefs: d.reducerDefs,
+      tables: d.tables,
+      typespace: d.typespace,
+      dbView_: undefined,
+    });
+    result.push(...flattenMountDispatches(d.subDispatches));
+  }
+  return result;
+}
+
 export const makeHooks = (schema: SchemaInner): ModuleHooks =>
   new ModuleHooksImpl(schema);
 
 class ModuleHooksImpl implements ModuleHooks {
   #schema: SchemaInner;
   #dbView_: DbView<any> | undefined;
   #reducerArgsDeserializers;
-  /** Cache the `ReducerCtx` object to avoid allocating anew for ever reducer call. */
+  #consumerReducerCount: number;
+  #flatMounts: FlatMountDispatch[];
+  /** Cache the `ReducerCtx` object to avoid allocating anew for every reducer call. */
   #reducerCtx_: InstanceType<typeof ReducerCtxImpl> | undefined;
 
   constructor(schema: SchemaInner) {
     this.#schema = schema;
-    this.#reducerArgsDeserializers = schema.moduleDef.reducers.map(
+    this.#consumerReducerCount = schema.reducers.length;
+    this.#flatMounts = flattenMountDispatches(schema.mountedDispatchInfos);
+
+    const consumerDeserializers = schema.moduleDef.reducers.map(
       ({ params }) => ProductType.makeDeserializer(params, schema.typespace)
     );
+    const mountedDeserializers = this.#flatMounts.flatMap(({ reducerDefs, typespace }) =>
+      reducerDefs.map(({ params }) => ProductType.makeDeserializer(params, typespace))
+    );
+    this.#reducerArgsDeserializers = [...consumerDeserializers, ...mountedDeserializers];
   }
 
   get #dbView() {
@@ -300,6 +340,18 @@ class ModuleHooksImpl implements ModuleHooks {
     ));
   }
 
+  #getMountDbView(mountIdx: number): DbView<any> {
+    const m = this.#flatMounts[mountIdx];
+    return (m.dbView_ ??= freeze(
+      Object.fromEntries(
+        m.tables.map(({ accessorName, tableDef }) => [
+          accessorName,
+          makeTableView(m.typespace, tableDef),
+        ])
+      )
+    ));
+  }
+
   get #reducerCtx() {
     return (this.#reducerCtx_ ??= new ReducerCtxImpl(
       Identity.zero(),
@@ -333,19 +385,42 @@ class ModuleHooksImpl implements ModuleHooks {
     timestamp: bigint,
     argsBuf: DataView
   ): void {
-    const moduleCtx = this.#schema;
     const deserializeArgs = this.#reducerArgsDeserializers[reducerId];
     BINARY_READER.reset(argsBuf);
     const args = deserializeArgs(BINARY_READER);
     const senderIdentity = new Identity(sender);
+
+    let fn: ((...args: any[]) => any) | undefined;
+    let dbView: DbView<any>;
+
+    if (reducerId < this.#consumerReducerCount) {
+      fn = this.#schema.reducers[reducerId];
+      dbView = this.#dbView;
+    } else {
+      let offset = this.#consumerReducerCount;
+      for (let i = 0; i < this.#flatMounts.length; i++) {
+        const m = this.#flatMounts[i];
+        if (reducerId < offset + m.reducerFns.length) {
+          fn = m.reducerFns[reducerId - offset];
+          dbView = this.#getMountDbView(i);
+          break;
+        }
+        offset += m.reducerFns.length;
+      }
+      if (fn === undefined) {
+        throw new RangeError(`unknown reducerId ${reducerId}`);
+      }
+    }
+
     const ctx = this.#reducerCtx;
     ReducerCtxImpl.reset(
       ctx,
       senderIdentity,
       new Timestamp(timestamp),
-      ConnectionId.nullIfZero(new ConnectionId(connId))
+      ConnectionId.nullIfZero(new ConnectionId(connId)),
+      dbView!
     );
-    callUserFunction(moduleCtx.reducers[reducerId], ctx, args);
+    callUserFunction(fn, ctx, args);
   }
 
   __call_view__(

crates/bindings-typescript/src/server/schema.ts

@@ -1,6 +1,11 @@
 import { moduleHooks, type ModuleDefaultExport } from 'spacetime:sys@2.0';
 import { CaseConversionPolicy, Lifecycle } from '../lib/autogen/types';
-import type { RawModuleDefV10 } from '../lib/autogen/types';
+import type {
+  RawModuleDefV10,
+  RawReducerDefV10,
+  RawTableDefV10,
+  Typespace,
+} from '../lib/autogen/types';
 import {
   type ParamsAsObject,
   type ParamsObj,
@@ -44,6 +49,14 @@ import {
 } from './views';
 import type { UntypedTableDef } from '../lib/table';
 
+export type MountedDispatchInfo = {
+  reducerFns: Reducers;
+  reducerDefs: RawReducerDefV10[];
+  typespace: Typespace;
+  tables: Array<{ accessorName: string; tableDef: RawTableDefV10 }>;
+  subDispatches: MountedDispatchInfo[];
+};
+
 export class SchemaInner<
   S extends UntypedSchemaDef = UntypedSchemaDef,
 > extends ModuleContext {
@@ -65,6 +78,7 @@ export class SchemaInner<
     string
   > = new Map();
   pendingSchedules: PendingSchedule[] = [];
+  mountedDispatchInfos: MountedDispatchInfo[] = [];
 
   constructor(getSchemaType: (ctx: SchemaInner<S>) => S) {
     super();
@@ -162,6 +176,10 @@ export class Schema<S extends UntypedSchemaDef> implements ModuleDefaultExport {
     return this.#ctx.typespace;
   }
 
+  get mountedDispatchInfos(): MountedDispatchInfo[] {
+    return this.#ctx.mountedDispatchInfos;
+  }
+
   /** Internal: register exports and materialize the RawModuleDefV10 for upload. */
   buildRawModuleDefV10(
     exports: object,
@@ -174,6 +192,32 @@ export class Schema<S extends UntypedSchemaDef> implements ModuleDefaultExport {
     return this.#ctx.rawModuleDefV10();
   }
 
+  /**
+   * @internal – called by schema() when processing a mounted namespace entry.
+   * Registers the library's exports and returns both the serialized module def
+   * and the runtime dispatch info needed by ModuleHooksImpl for __call_reducer__.
+   */
+  buildMountForDispatch(
+    exports: object
+  ): { rawDef: RawModuleDefV10; dispatch: MountedDispatchInfo } {
+    const rawDef = this.buildRawModuleDefV10(exports, {
+      ignoreNonModuleExports: true,
+    });
+    return {
+      rawDef,
+      dispatch: {
+        reducerFns: [...this.#ctx.reducers],
+        reducerDefs: [...this.#ctx.moduleDef.reducers],
+        typespace: this.#ctx.moduleDef.typespace,
+        tables: Object.values(this.#ctx.schemaType.tables).map(t => ({
+          accessorName: t.accessorName,
+          tableDef: t.tableDef,
+        })),
+        subDispatches: [...this.#ctx.mountedDispatchInfos],
+      },
+    };
+  }
+
   /**
    * Defines a SpacetimeDB reducer function.
    *
@@ -617,12 +661,9 @@ export function schema<const H extends Record<string, SchemaEntry>>(
         );
       }
       if (isMountedModuleNamespace(entry)) {
-        ctx.addMount({
-          namespace: accName,
-          module: entry.default.buildRawModuleDefV10(entry, {
-            ignoreNonModuleExports: true,
-          }),
-        });
+        const { rawDef, dispatch } = entry.default.buildMountForDispatch(entry);
+        ctx.addMount({ namespace: accName, module: rawDef });
+        ctx.mountedDispatchInfos.push(dispatch);
         continue;
       }
       if (!isUntypedTableSchema(entry)) {

crates/bindings-typescript/tests/schema_mounts.test.ts

@@ -114,4 +114,70 @@ describe('schema mounts', () => {
       })
     ).toThrow(/looks like a default import/);
   });
+
+  it('populates mountedDispatchInfos with reducer fns and table metadata', () => {
+    const sessions = table(
+      { name: 'sessions' },
+      { id: t.u64().primaryKey().autoInc() }
+    );
+
+    const authSchema = schema({ sessions });
+    const cleanExpiredSessions = authSchema.reducer(() => {});
+    const authLib = { default: authSchema, cleanExpiredSessions };
+
+    const players = table({ name: 'players' }, { id: t.u32().primaryKey() });
+    const consumer = schema({ players, myauth: authLib });
+
+    const infos = consumer.mountedDispatchInfos;
+    expect(infos).toHaveLength(1);
+
+    const info = infos[0];
+    expect(info.reducerFns).toHaveLength(1);
+    expect(info.reducerDefs).toHaveLength(1);
+    expect(info.reducerDefs[0].sourceName).toBe('cleanExpiredSessions');
+    expect(info.tables).toHaveLength(1);
+    expect(info.tables[0].accessorName).toBe('sessions');
+    expect(info.subDispatches).toHaveLength(0);
+  });
+
+  it('flattens nested mount dispatches depth-first', () => {
+    // baz library: 1 reducer
+    const bazTable = table({ name: 'baz_items' }, { id: t.u32().primaryKey() });
+    const bazSchema = schema({ bazTable });
+    const bazReducer = bazSchema.reducer(() => {});
+    const bazLib = { default: bazSchema, bazReducer };
+
+    // auth library: 1 own reducer, mounts baz
+    const sessions = table(
+      { name: 'sessions' },
+      { id: t.u64().primaryKey().autoInc() }
+    );
+    const authSchema = schema({ sessions, baz: bazLib });
+    const authReducer = authSchema.reducer(() => {});
+    const authLib = { default: authSchema, authReducer };
+
+    // consumer: 1 own reducer, mounts auth
+    const players = table({ name: 'players' }, { id: t.u32().primaryKey() });
+    const consumer = schema({ players, myauth: authLib });
+    const consumerReducer = consumer.reducer(() => {});
+
+    // Verify depth-first structure:
+    // consumer.mountedDispatchInfos[0] = myauth (authReducer)
+    // consumer.mountedDispatchInfos[0].subDispatches[0] = myauth.baz (bazReducer)
+    const infos = consumer.mountedDispatchInfos;
+    expect(infos).toHaveLength(1);
+
+    const authInfo = infos[0];
+    expect(authInfo.reducerFns).toHaveLength(1);
+    expect(authInfo.reducerDefs[0].sourceName).toBe('authReducer');
+    expect(authInfo.subDispatches).toHaveLength(1);
+
+    const bazInfo = authInfo.subDispatches[0];
+    expect(bazInfo.reducerFns).toHaveLength(1);
+    expect(bazInfo.reducerDefs[0].sourceName).toBe('bazReducer');
+    expect(bazInfo.subDispatches).toHaveLength(0);
+
+    // Unused variable check
+    void consumerReducer;
+  });
 });