App.tsx

import React, { useEffect, useState } from 'react';

import { SyncProvider, useSync } from './SyncContext';
import { connectSyncedStore } from './store';
import type { SyncedStoreConnection } from './store';


export function App() {
  const [connection, setConnection] = useState<
    SyncedStoreConnection | undefined
  >(undefined);

  useEffect(() => {
    console.log('Reconnecting to synced store...');
    const { rootStore, provider } = connectSyncedStore();
    setConnection({ rootStore, provider });
    return () => provider.disconnect();
  }, []);

  if (connection) {
    return (
      <SyncProvider
        rootStore={connection.rootStore}
        provider={connection.provider}
      >
        <Main />
      </SyncProvider>
    );
  } else {
    return null;
  }
}


function Main() {
  const { state, actions, selectors } = useSync();
  
  // ... render based on 'state'. It is reactive, based on SyncedStore.
  // ... reuse logic from selectors (they are just functions to encapsulate knowledge of store shape, just a nice code organization thing, not essential for memoization etc; we rely on SyncedStore for that
  // ... make changes by calling functions under 'actions'

}

store.ts

import * as Y from "yjs";
import { WebsocketProvider } from "y-websocket";
import { getYjsValue, syncedStore } from "@syncedstore/core";
import type { UUID } from "./types";

import type { AwarenessState } from "./SyncContext";

class BaseSelectors {
  constructor(
    public localAwarenessState: AwarenessState | undefined,
    public awarenessStates: AwarenessState[],
    public store: RootStore
  ) {
    this.localAwarenessState = localAwarenessState;
    this.awarenessStates = awarenessStates;
    this.store = store;
  }
}

// Todos --------------------------------
export type TodoId = UUID;

export type Todo = {
  todoId: TodoId;
  title: string;
  completed: boolean;
};

type TodosShape = {
  [index: TodoId]: Todo;
};

// TODO: rename 'state' to 'slice' since this is namespaced?
export class TodosActions {
  state: TodosShape;

  constructor(state: TodosShape) {
    this.state = state;
  }

  add(todo: Todo) {
    this.state[todo.todoId] = todo;
  }

  toggle(todoId: TodoId) {
    const todo = this.state[todoId];
    if (todo) {
      todo.completed = !todo.completed;
    }
  }

  delete(todoId: TodoId) {
    delete this.state[todoId];
  }
}

// TODO namespace and provide this.slice instead of this.store?
class TodosSelectors extends BaseSelectors {
  incompleteTodos(): Todo[] {
    return this.store.todos.filter((todo: Todo) => todo.completed);
  }
}

const initialTodos = {}

// Root / all --------------------------------

type StoreShape = {
  todos: TodosShape;
};

const sliceNames = [
  'todos',
] as const;

const createRootStore = () => {
  return syncedStore<StoreShape>({
    todos: {},
  });
};
export type RootStore = ReturnType<typeof createRootStore>;

export class Actions {
  store: RootStore;

  todos: TodosActions;

  constructor(store: RootStore) {
    this.store = store;

    // TODO fix the Partial<...>? for the map types to avoid need for assertion
    this.todos = new TodosActions(store.todos as TodosShape);
  }
}

// TODO rename store to state in selectors
export class Selectors extends BaseSelectors {
  todos: TodosSelectors;

  constructor(
    localAwarenessState: AwarenessState,
    awarenessStates: AwarenessState[],
    store: RootStore,
  ) {
    const params = [localAwarenessState, awarenessStates, store] as const;

    super(...params);

    this.todos = new TodosSelectors(...params);
  }
}

const initialStore = {
  todos: initialTodos,
};

const docId = 'aaaa123';
const WEBSOCKET_ENDPOINT = 'ws://localhost:1234';

const isArray = (x: any): x is Array<any> => Array.isArray(x);

export const connectSyncedStore = () => {
  const rootStore = createRootStore();
  const doc = getYjsValue(rootStore) as Y.Doc;
  const provider = new WebsocketProvider(WEBSOCKET_ENDPOINT, docId, doc, {
    connect: false,
  });

  // Hacky way to initialize a document. TODO: do this server-side.
  provider.on('sync', (isSynced: boolean) => {
    // TODO: Decouple from todos
    const isEmpty = Object.keys(rootStore.todos).length === 0;

    if (isSynced && isEmpty) {
      Y.transact(doc, () => {
        sliceNames.forEach((sliceName) => {
          const initialSlice = initialStore[sliceName];
          const rootSlice = rootStore[sliceName];

          if (isArray(initialSlice)) {
            (rootSlice as Array<any>).push(...initialSlice);
          } else {
            Object.assign(rootSlice, initialSlice);
          }
        });
      });
    }
  });

  provider.connect();

  return { rootStore, provider };
};

export type SyncedStoreConnection = ReturnType<typeof connectSyncedStore>;

SyncContext.tsx

import React, {
  useEffect,
  useCallback,
  useState,
  createContext,
  useContext,
  ReactNode,
} from "react";
import { getYjsValue } from "@syncedstore/core";
import { useSyncedStore } from "@syncedstore/react";
import type * as Y from "yjs";
import type { UserId } from "./types";
import type { RootStore } from "./store";
import { Actions, Selectors } from "./store";
import type { WebsocketProvider } from "y-websocket";

export type AwarenessState = { userId: UserId | undefined };

type SyncProviderProps = {
  provider: WebsocketProvider;
  rootStore: RootStore;
  children: ReactNode;
  forceUpdate: boolean;
};

type SyncState = {
  localAwarenessState: AwarenessState;
  awarenessStates: AwarenessState[];
  provider: WebsocketProvider;
  actions: Actions;
  selectors: Selectors;
  state: RootStore; // TODO: type this
  rootStore: RootStore; // TODO: type this
};

const SyncStateContext = createContext<
  | {
      localAwarenessState: AwarenessState;
      awarenessStates: AwarenessState[];
      provider: WebsocketProvider;
      rootStore: RootStore; // TODO: type this
    }
  | undefined
>(undefined);

export function SyncProvider({
  provider,
  rootStore,
  children,
}: SyncProviderProps) {
  // TODO: Instead of only tracking userId via awareness, perhaps track it separately and then push it out to awareness?
  // E.g. so we can have local currentUserId when offline
  const [localAwarenessState, setLocalAwarenessState] =
    useState<AwarenessState>({
      userId: undefined,
    });

  const [awarenessStates, setAwarenessStates] = useState<AwarenessState[]>([]);

  const onAwarenessUpdate = useCallback(() => {
    // TODO how to type the provider.awareness?
    setLocalAwarenessState(
      provider.awareness.getLocalState() as AwarenessState
    );
    setAwarenessStates(
      Array.from(provider.awareness.getStates().values()) as AwarenessState[]
    );
  }, [provider.awareness]);

  useEffect(() => {
    provider.awareness.on("update", onAwarenessUpdate);
    onAwarenessUpdate();
    return () => provider.awareness.off("update", onAwarenessUpdate);
  }, [provider.awareness]);

  const value = {
    localAwarenessState,
    awarenessStates,
    provider,
    rootStore,
  };

  return (
    <SyncStateContext.Provider value={value}>
      {children}
    </SyncStateContext.Provider>
  );
}

export function useSync(): SyncState {
  const context = useContext(SyncStateContext);

  if (context === undefined) {
    throw new Error("useSync must be used within a SyncProvider");
  }

  // Call useSyncedStore in this stack frame to register the observable.
  const { rootStore, provider } = context;

  const state = useSyncedStore(rootStore);

  const actions = new Actions(state);

  // Also wire up selectors to the store that was initialized in this stack
  // frame so that its selections correctly register observables.
  const selectors = new Selectors(
    context.localAwarenessState,
    context.awarenessStates,
    state
  );

  // Re-render on sync... TODO: is this necessary? was left in there for debugging
  const [syncCount, setSyncCount] = useState(0);
  const onProviderSync = useCallback(() => {
    console.log("provider sync, force re-render");
    setSyncCount((value) => value + 1);
  }, []);
  useEffect(() => {
    provider.on("sync", onProviderSync);
    return () => provider.off("sync", onProviderSync);
  }, [provider]);

  return {
    ...context,
    state,
    actions,
    selectors,
  };
}