@oobleck/fluid-backend/src/model/notebook.ts

Fluid Framework backend for nteract RTC

import { combineLatest, fromEvent, merge, Observable, Subject, Subscription } from "rxjs";
import { filter, map, scan } from "rxjs/operators";
import {
  createGroupOp,
  createInsertSegmentOp,
  createRemoveRangeOp,
  DataObject,
  DataObjectFactory,
  ISharedMap,
  MergeTreeDeltaType,
  SharedMap,
  SharedObjectSequence,
  SharedSequence,
  SubSequence
} from "@fluid-experimental/fluid-framework";
import { ISequencedDocumentMessage } from "@fluidframework/protocol-definitions";
import { IFluidHandle } from "@fluidframework/core-interfaces";
import {
  CellEvent,
  CellInput,
  ICellInsertedEvent,
  ICellMovedEvent,
  ICellRemovedEvent,
  ICellReplacedEvent,
  MetadataEntryDef,
  NotebookContentInput,
  NotebookEvent
} from "../schema";
import { AugmentedCellEdit, CellModel, INotebookModel } from "./types";
import { CodeCellDDS } from "./codeCell";
import { TextCellDDS } from "./textCell";

enum PropertyKey {
  Cells = "cells",
  CellOrder = "cellOrder",
  Metadata = "metadata"
}

const fromDocumentMessage = <T extends SharedMap | SharedSequence<string>>(target: T) =>
  fromEvent<[ISequencedDocumentMessage, boolean, T]>(target, "op");

export class NotebookDDS extends DataObject<{}, NotebookContentInput> implements INotebookModel {
  public static DataObjectName = "notebook-model";
  private cellMap!: SharedMap;
  private cellOrder!: SharedObjectSequence<string>;
  private metadataMap!: ISharedMap;
  private readonly subscriptions: Subscription[] = [];

  public static readonly Factory = new DataObjectFactory(
    NotebookDDS.DataObjectName,
    NotebookDDS,
    [SharedMap.getFactory(), SharedObjectSequence.getFactory()],
    {},
    [CodeCellDDS.Factory.registryEntry, TextCellDDS.Factory.registryEntry]
  );

  //#region INotebookModel
  cellEdits$ = new Subject<AugmentedCellEdit>();

  cellEvents$ = new Subject<CellEvent>();

  events$!: Observable<NotebookEvent>;

  get metadata(): MetadataEntryDef[] {
    const result: MetadataEntryDef[] = [];
    this.metadataMap?.forEach((value, key) => result.push({ key, value }));
    return result;
  }

  async getCells(): Promise<CellModel[]> {
    const componentPromises: Promise<CellModel>[] =
      this.cellOrder.getItems(0).map((cellId) => {
        const handle = this.cellMap.get(cellId);
        return handle.get();
      });
    return Promise.all(componentPromises);
  }

  async getCell(id: string): Promise<CellModel | undefined> {
    const cellHandle = this.cellMap.get(id);
    if (cellHandle) {
      const cellComponent = await cellHandle.get();
      return cellComponent;
    }
    return undefined;
  }

  async insertCell(input: CellInput, insertAt: number): Promise<CellModel> {
    const component = await this.createCellComponent(input);
    this.cellMap.set(component.id, component.handle);
    this.cellOrder.insert(insertAt, [component.id]);
    return component;
  }

  deleteCell(id: string): void {
    const removeIndex = this.cellOrder.getItems(0).indexOf(id);
    if (removeIndex !== -1) {
      this.cellOrder.remove(removeIndex, removeIndex + 1);
    }
    this.cellMap.delete(id);
  }

  async replaceCell(id: string, cell: CellInput): Promise<CellModel> {
    const component = await this.createCellComponent(cell);
    this.cellMap.set(component.id, component.handle);

    const removeIndex = this.cellOrder.getItems(0).indexOf(id);
    const removeOp = createRemoveRangeOp(removeIndex, removeIndex + 1);
    const insertOp = createInsertSegmentOp(removeIndex, new SubSequence([component.id]));
    this.cellOrder.groupOperation(createGroupOp(removeOp, insertOp));

    this.cellMap.delete(id);

    return component;
  }

  moveCell(id: string, insertAt: number): void {
    const removeIndex = this.cellOrder.getItems(0).indexOf(id);
    const removeOp = createRemoveRangeOp(removeIndex, removeIndex + 1);
    const insertOp = createInsertSegmentOp(insertAt, new SubSequence([id]));
    this.cellOrder.groupOperation(createGroupOp(removeOp, insertOp));
  }

  updateMetadata(key: string, value: unknown): void {
    this.metadataMap.set(key, value);
  }

  async clearCellOutputs(id?: string) {
    if (id) {
      const cell = await this.getCell(id);
      if (cell?.cellType === "CodeCell") {
        cell.clearOutputs();
      }
    } else {
      const cells = await this.getCells();
      cells.forEach((cell) => {
        if (cell.cellType === "CodeCell") {
          cell.clearOutputs();
        }
      });
    }
    return true;
  }
  //#endregion

  //#region DataObject
  protected async initializingFirstTime(input?: NotebookContentInput): Promise<void> {
    const cellMap = SharedMap.create(this.runtime);
    const cellOrder = SharedObjectSequence.create<string>(this.runtime);
    const metadataMap = SharedMap.create(this.runtime);

    if (input) {
      const cellIds = [];
      for (const cell of input.cells) {
        const { id, handle } = await this.createCellComponent(cell);
        cellIds.push(id);
        cellMap.set(id, handle);
      }
      cellOrder.insert(0, cellIds);

      input.metadata?.forEach(({ key, value }) => {
        metadataMap.set(key, value);
      });
    }

    this.root
      .set(PropertyKey.Cells, cellMap.handle)
      .set(PropertyKey.CellOrder, cellOrder.handle)
      .set(PropertyKey.Metadata, metadataMap.handle);
  }

  protected async initializingFromExisting(): Promise<void> {
  }

  protected async hasInitialized(): Promise<void> {
    this.cellMap = await this.root.get(PropertyKey.Cells)?.get();
    this.cellOrder = await this.root.get(PropertyKey.CellOrder)?.get();
    this.metadataMap = await this.root.get(PropertyKey.Metadata)?.get();

    if (this.isInteractive()) {
      await this.setupObservables();
    }
  }

  public dispose(): void {
    this.subscriptions.forEach((s) => s.unsubscribe());
    this.subscriptions.splice(0);
    super.dispose();
  }
  //#endregion

  //#region private
  private isInteractive() {
    return !!this.context.containerRuntime.clientDetails.capabilities.interactive;
  }

  private createCellComponent(cell: CellInput) {
    if ("code" in cell) {
      return CodeCellDDS.Factory.createChildInstance(this.context, cell.code);
    } else if ("markdown" in cell) {
      return TextCellDDS.Factory.createChildInstance(this.context, cell.markdown);
    } else if ("raw" in cell) {
      return TextCellDDS.Factory.createChildInstance(this.context, cell.raw);
    }
    throw new Error("Unsupported cell input type");
  }

  private async setupObservables() {
    const snapshot$ = fromDocumentMessage(this.cellOrder).pipe(
      scan(
        (acc, [{ sequenceNumber }, l, t]) => ({ next: t.getItems(0), prev: acc.next, sequenceNumber }),
        { next: this.cellOrder.getItems(0), prev: [] as string[], sequenceNumber: 0 }
      )
    );

    const insert$ = fromDocumentMessage(this.cellOrder).pipe(
      filter(([{ contents: delta }, local]) => {
        return !local
          && delta?.type === MergeTreeDeltaType.INSERT
          && typeof delta.pos1 === "number"
          && delta.seg?.items instanceof Array;
      }),
      map(([{ sequenceNumber, contents: delta }, local, target]) => {
        const all = target.getItems(0);
        const count = target.getItemCount();
        const pos = delta.pos1;
        const after = pos > 0 ? all[pos - 1] : undefined;
        const before = pos < count - 1 ? all[pos + 1] : undefined;
        return {
          __typename: "CellInsertedEvent",
          id: delta.seg.items[0],
          pos,
          after,
          before,
          sequenceNumber
        } as ICellInsertedEvent;
      })
    );

    const move$ = fromDocumentMessage(this.cellOrder).pipe(
      filter(([{ contents: delta }, local]) => {
        return !local
          && delta?.type === MergeTreeDeltaType.GROUP
          && delta.ops instanceof Array
          && delta.ops.length === 2;
      }),
      map(([{ sequenceNumber, contents: delta }]) => {
        const [removeOp, insertOp] = delta.ops;
        if (removeOp.pos1 === insertOp.pos1) {
          return {
            __typename: "CellReplacedEvent",
            id: "",
            replacedBy: insertOp.seg.items[0],
            pos: insertOp.pos1,
            sequenceNumber
          } as ICellReplacedEvent;
        } else {
          return {
            __typename: "CellMovedEvent",
            from: removeOp.pos1,
            to: insertOp.pos1,
            sequenceNumber
          } as ICellMovedEvent;
        }
      })
    );

    const delete$ = fromDocumentMessage(this.cellOrder).pipe(
      filter(([{ contents: delta }, local]) => {
        return !local
          && delta?.type === MergeTreeDeltaType.REMOVE
          && typeof delta.pos1 === "number"
          && typeof delta.pos2 === "number";
      }),
      map(([{ sequenceNumber, contents: delta }]) => {
        return {
          __typename: "CellRemovedEvent",
          id: "",
          pos: delta.pos1, sequenceNumber
        } as ICellRemovedEvent;
      })
    );

    this.events$ = combineLatest([merge(insert$, move$, delete$), snapshot$]).pipe(
      filter(([e, s]) => e.sequenceNumber === s.sequenceNumber),
      map(([e, s]) => {
        switch (e.__typename) {
          case "CellRemovedEvent":
          case "CellReplacedEvent":
            e = { ...e, id: s.prev[e.pos] };
            break;
        }
        return e;
      })
    );

    fromDocumentMessage(this.cellMap).pipe(
      filter(([{ contents: delta }]) => delta?.type === "set"),
      map(([{ contents: delta }]) => {
        return [delta.key, delta.value?.value];
      })
    ).subscribe(async ([id]) => {
      const handle = this.cellMap.get(id);
      await this.enlistCell(handle);
    });

    fromDocumentMessage(this.cellMap).pipe(
      filter(([{ contents: delta }]) => delta?.type === "delete"),
      map(([{ contents: delta }]) => {
        return delta.key;
      })
    );

    for (const [, cellHandle] of this.cellMap.entries()) {
      await this.enlistCell(cellHandle);
    }
  }

  private async enlistCell(handleOrModel: IFluidHandle<CellModel> | CellModel) {
    const cell: CellModel = "cellType" in handleOrModel ? handleOrModel : await handleOrModel.get();
    this.subscriptions.push(
      cell.edits$.subscribe(this.cellEdits$),
      cell.events$.subscribe(this.cellEvents$)
    );
  }
  //#endregion
}