Replicache with Suspense

model.tsx

import React, {
  createContext,
  ReactNode,
  useCallback,
  useEffect,
  useState,
} from 'react';
import { unstable_batchedUpdates } from 'react-dom';
import { ReadonlyJSONValue, ReadTransaction, Replicache } from 'replicache';

type Mutators = {};
type Cache = Replicache<Mutators>;

function getCachePromise(path: string): Promise<Cache> {
  // should return cache post initial sync
  return null as any;
}

type Result = ReadonlyJSONValue;
type Future<T> = { read: () => T; promise: Promise<T> };
type Callback<T> = (future: Future<T>) => void;
type QueryBody<T> = (tx: ReadTransaction) => Promise<T>;
type KeyedQuery<T> = { key: string; query: QueryBody<T> };
type Query<T> = QueryBody<T> | KeyedQuery<T>;
type Subscription<T> = {
  callbacks: Array<Callback<T>>;
  future: Future<T>;
};

let hasPendingCallback = false;
let callbacks: (() => void)[] = [];

function doCallback() {
  const cbs = callbacks;
  callbacks = [];
  hasPendingCallback = false;
  unstable_batchedUpdates(() => {
    for (const callback of cbs) {
      callback();
    }
  });
}

type Status = 'pending' | 'success' | 'error';
function getFuture<T>(promise: Promise<T>): Future<T> {
  let status: Status = 'pending';
  let result: T;
  let error: unknown;
  let suspender = promise.then(
    (r) => {
      status = 'success';
      result = r;
    },
    (e) => {
      status = 'error';
      error = e;
    }
  );
  return {
    promise,
    read: () => {
      if (status === 'pending') {
        throw suspender;
      } else if (status === 'error') {
        throw error;
      }
      return result;
    },
  };
}

function getResult<T>(result: T): Future<T> {
  return {
    promise: Promise.resolve(result),
    read: () => {
      return result;
    },
  };
}

function getError<T>(error: unknown): Future<T> {
  return {
    promise: Promise.reject(error),
    read: () => {
      throw error;
    },
  };
}

export class Model {
  path: string;
  cache: Promise<Cache>;
  subscriptions = new Map<Query<any>, Subscription<any>>();
  queries = new Map<string, QueryBody<any>>();

  constructor(path: string) {
    console.log('new model', path);
    this.path = path;
    this.cache = getCachePromise(path);
  }

  private getKeyedQuery<T extends Result>({ key, query }: KeyedQuery<T>) {
    const existing = this.queries.get(key);
    if (existing) return existing;
    this.queries.set(key, query);
    return query as QueryBody<T>;
  }

  private getQueryBody<T extends Result>(query: Query<T>) {
    if ('key' in query) {
      return this.getKeyedQuery(query);
    }
    return query;
  }

  private getSubscription<T extends Result>(query: Query<T>) {
    const body = this.getQueryBody(query);
    const existing = this.subscriptions.get(body);
    if (existing) {
      return existing as Subscription<T>;
    }

    let pending = true;
    let resolve: any;
    let reject: any;

    const promise = new Promise<T>((onResolve, onReject) => {
      resolve = onResolve;
      reject = onReject;
    });

    this.cache.then((cache) =>
      cache.subscribe(body, {
        onData: (value) => {
          if (pending) {
            pending = false;
            resolve(value);
          }
          this.onResolve(body, value);
        },
        onError: (error) => {
          if (pending) {
            pending = true;
            reject(error);
          }
          this.onReject(body, error);
        },
      })
    );

    const future = getFuture(promise);
    const subscription = { callbacks: [], future };
    this.subscriptions.set(body, subscription);

    return subscription;
  }

  getFuture<T extends Result>(query: Query<T>) {
    return this.getSubscription(query).future;
  }

  getPromise<T extends Result>(query: Query<T>) {
    return this.getSubscription(query).future.promise;
  }

  subscribe<T extends Result>(query: Query<T>, callback: Callback<T>) {
    const subscription = this.getSubscription(query);
    subscription.callbacks.push(callback);
    return () => this.unsubscribe(query, callback);
  }

  private unsubscribe<T extends Result>(
    query: Query<T>,
    callback: Callback<T>
  ) {
    const subscription = this.getSubscription(query);
    const index = subscription.callbacks.indexOf(callback);
    if (index > -1) subscription.callbacks.splice(index, 1);
  }

  private onResolve<T extends Result>(query: Query<T>, result: T) {
    const subscription = this.getSubscription(query);
    subscription.future = getResult(result);
    for (const callback of subscription.callbacks) {
      callback(subscription.future);
    }
  }

  private onReject<T extends Result>(query: Query<T>, error: unknown) {
    const subscription = this.getSubscription(query);
    subscription.future = getError(error);
    for (const callback of subscription.callbacks) {
      callback(subscription.future);
    }
  }

  // we proxy our mutations here so we can do client specific stuff

  async close() {
    console.log(`closing ${this.path}`);
    await this.cache.then((cache) => cache.close());
  }
}

type GetModel = (name: string, path: string) => Model;
const ModelContext = createContext<GetModel | null>(null);

// this provides named replicaches for particular paths:
// - we have one space and one deal model for example, but they can
// point at different "spaces" and "deals"
export function ModelProvider(props: { children: ReactNode }) {
  const [models] = useState(new Map<string, Model>());
  const getModel = useCallback((name: string, path: string) => {
    const existing = models.get(name);
    if (existing) {
      if (existing.path === path) return existing;
      existing.close();
    }
    const model = new Model(path);
    models.set(name, model);
    return model;
  }, []);
  return (
    <ModelContext.Provider value={getModel}>
      {props.children}
    </ModelContext.Provider>
  );
}

// returns a named model for a particular path
export function useModel(name: string, path: string) {
  const getModel = React.useContext(ModelContext);
  if (getModel) return getModel(name, path);
  throw new Error('no model context');
}

// this is the equivalent of useSubscribe - we create convenience hooks so we
// can do useDealQuery or orSpaceQuery
export function useModelQuery<T extends Result>(model: Model, query: Query<T>) {
  const [future, setFuture] = useState(() => model.getFuture(query));

  useEffect(() => {
    return model.subscribe(query, (future) => {
      callbacks.push(() => setFuture(future));
      if (!hasPendingCallback) {
        Promise.resolve().then(doCallback);
        hasPendingCallback = true;
      }
    });
  }, [model, query]);

  return future.read();
}

// example usage

type Foo = { name: string };

async function queryFoos(tx: ReadTransaction) {
  return (await tx.scan().toArray()) as Array<Foo>;
}

function Component() {
  const model = useModel('space', 'spaces/1234');
  const foos = useModelQuery(model, queryFoos);
  return (
    <>
      {foos.map((foo) => (
        <div>{foo.name}</div>
      ))}
    </>
  );
}