Proof-of-concept r3f editor

backend.ts

import { Application, Router, RouterContext } from 'https://deno.land/x/[email protected]/mod.ts';
import {
  JsxAttributeLike,
  Project,
  SyntaxKind,
  SourceFile,
  ts,
  Type,
  JsxSelfClosingElement,
  JsxOpeningElement,
} from 'https://deno.land/x/[email protected]/mod.ts';
import * as posixPath from 'https://deno.land/[email protected]/path/posix.ts';

const app = new Application();
const router = new Router();
const project = new Project({
  // TODO: Dynamically find ts config instead of hardcoding
  tsConfigFilePath: '/Users/mdougall/projects/FAZE/tsconfig.json',
});
const watchers = new Map<string, Deno.FsWatcher>();

app.use((ctx, next) => {
  ctx.response.headers.set('Access-Control-Allow-Origin', '*');
  return next();
});

function getParam(context: RouterContext<any, any>, key: string) {
  const path = context.request.url.searchParams.get(key);
  if (!path) {
    throw new Error('invariant');
  }

  return path;
}

function getJsxAttributeValue(prop: JsxAttributeLike) {
  const expression = prop.getChildrenOfKind(SyntaxKind.JsxExpression)[0]?.getExpression();
  let result = 'undefined';

  if (expression) {
    result = expression
      .getFullText()
      .replace(/\n| /g, '')
      .replace(/(,)(}|])/g, (match) => match.replace(',', ''));
  }

  const stringLiteral = prop.getChildrenOfKind(SyntaxKind.StringLiteral)[0];
  if (stringLiteral) {
    result = stringLiteral.getText();
  }

  try {
    return JSON.parse(result);
  } catch (_e) {
    return result;
  }
}

function unrollType(type: Type, _nested = false): any {
  const args = type.getTypeArguments();
  if (args.length) {
    return args.map((arg) => unrollType(arg, true)) as string[];
  }

  if (type.isUnion()) {
    // From the boolean hack above it will return two
    const types = type
      .getUnionTypes()
      .filter((type) => !type.isUndefined())
      .map((n) => {
        const value = unrollType(n, true);

        if (n.isTuple()) {
          return {
            kind: 'tuple',
            value,
          };
        }

        if (n.isArray()) {
          return {
            kind: 'array',
            value,
          };
        }

        return {
          kind: 'type',
          value,
        };
      });

    if (types.length === 1) {
      return types[0];
    }

    return {
      kind: 'union',
      type: types.map((type) => type.value),
    };
  }

  if (_nested) {
    return type.getText();
  }

  return {
    kind: 'type',
    value: type.getText(),
  };
}

function getJsxPropTypes(sourceFile: SourceFile, elementName: string) {
  if (/[a-z]/.exec(elementName[0])) {
    return {};
  }

  const propTypes: Record<string, { name: string; required: boolean; type: any }> = {};
  const symbol = sourceFile.getLocalOrThrow(elementName);
  try {
    const componentFunction = symbol
      .getDeclarations()[0]
      .asKindOrThrow(SyntaxKind.ImportSpecifier)
      .getType()
      .getSymbolOrThrow()
      .getDeclarations()[0]
      .asKindOrThrow(SyntaxKind.FunctionDeclaration);

    const [props] = componentFunction.getParameters();
    props
      .getType()
      .getProperties()
      .forEach((prop) => {
        propTypes[prop.getName()] = {
          name: prop.getName(),
          required: !prop.isOptional(),
          type: unrollType(prop.getTypeAtLocation(componentFunction)),
        };
      });

    return propTypes;
  } catch (_) {
    return propTypes;
  }
}

function getNodeTransforms(sourceFile: SourceFile, elementName: string) {
  if (/[a-z]/.exec(elementName[0])) {
    return {
      translate: true,
      scale: true,
      rotate: true,
    };
  }

  const types = getJsxPropTypes(sourceFile, elementName);

  return {
    translate: !!types.position,
    scale: !!types.scale,
    rotate: !!types.rotation,
  };
}

function isSceneObject(node: ts.Node): node is ts.JsxSelfClosingElement | ts.JsxElement {
  if (ts.isJsxSelfClosingElement(node) && !node.tagName.getText().includes('aterial')) {
    return true;
  }

  if (ts.isJsxElement(node) && !node.openingElement.tagName.getText().includes('aterial')) {
    return true;
  }

  return false;
}

function transformSouceFileSync(sourceFile: SourceFile) {
  const parsedPath = posixPath.parse(sourceFile.getFilePath());
  const destination = posixPath.join('./.tmp', parsedPath.base);
  const transformedSource = sourceFile.copy(destination, { overwrite: true });

  transformedSource.transform((traversal) => {
    const node = traversal.visitChildren();
    const pos = node.pos - 30; // TODO: Why is the 30 offset needed? What is affecting it?

    if (!isSceneObject(node)) {
      return node;
    }

    if (ts.isJsxSelfClosingElement(node)) {
      const transform = getNodeTransforms(transformedSource, node.tagName.getText());

      return traversal.factory.createJsxElement(
        traversal.factory.createJsxOpeningElement(
          traversal.factory.createIdentifier('group'),
          [],
          traversal.factory.createJsxAttributes(
            [
              node.attributes.properties.find((x) => x.name?.getText() === 'key'),
              traversal.factory.createJsxAttribute(
                traversal.factory.createIdentifier('userData'),
                traversal.factory.createJsxExpression(
                  undefined,
                  traversal.factory.createObjectLiteralExpression([
                    traversal.factory.createPropertyAssignment(
                      '__r3fEditor',
                      traversal.factory.createObjectLiteralExpression([
                        traversal.factory.createPropertyAssignment(
                          'pos',
                          traversal.factory.createNumericLiteral(pos)
                        ),
                        traversal.factory.createPropertyAssignment(
                          'translate',
                          transform.translate
                            ? traversal.factory.createTrue()
                            : traversal.factory.createFalse()
                        ),
                        traversal.factory.createPropertyAssignment(
                          'rotate',
                          transform.rotate
                            ? traversal.factory.createTrue()
                            : traversal.factory.createFalse()
                        ),
                        traversal.factory.createPropertyAssignment(
                          'scale',
                          transform.scale
                            ? traversal.factory.createTrue()
                            : traversal.factory.createFalse()
                        ),
                      ])
                    ),
                  ])
                )
              ),
            ].filter(Boolean) as ts.JsxAttributeLike[]
          )
        ),
        [node],
        traversal.factory.createJsxClosingElement(traversal.factory.createIdentifier('group'))
      );
    }

    if (ts.isJsxElement(node)) {
      const transform = getNodeTransforms(transformedSource, node.openingElement.tagName.getText());

      return traversal.factory.updateJsxElement(
        node,
        traversal.factory.createJsxOpeningElement(
          traversal.factory.createIdentifier('group'),
          [],
          traversal.factory.createJsxAttributes([
            traversal.factory.createJsxAttribute(
              traversal.factory.createIdentifier('userData'),
              traversal.factory.createJsxExpression(
                undefined,
                traversal.factory.createObjectLiteralExpression([
                  traversal.factory.createPropertyAssignment(
                    '__r3fEditor',
                    traversal.factory.createObjectLiteralExpression([
                      traversal.factory.createPropertyAssignment(
                        'pos',
                        traversal.factory.createNumericLiteral(pos)
                      ),
                      traversal.factory.createPropertyAssignment(
                        'translate',
                        transform.translate
                          ? traversal.factory.createTrue()
                          : traversal.factory.createFalse()
                      ),
                      traversal.factory.createPropertyAssignment(
                        'rotate',
                        transform.rotate
                          ? traversal.factory.createTrue()
                          : traversal.factory.createFalse()
                      ),
                      traversal.factory.createPropertyAssignment(
                        'scale',
                        transform.scale
                          ? traversal.factory.createTrue()
                          : traversal.factory.createFalse()
                      ),
                    ])
                  ),
                ])
              )
            ),
          ])
        ),
        [node],
        traversal.factory.createJsxClosingElement(traversal.factory.createIdentifier('group'))
      );
    }

    return node;
  });

  // TODO: Why isn't this needed anymore?
  // transformedSource.getImportDeclarations().forEach((decl) => {
  //   if (decl.isModuleSpecifierRelative()) {
  //     const previousModuleSpecifier = decl.getModuleSpecifierValue();
  //     decl.setModuleSpecifier('../' + previousModuleSpecifier);
  //   }
  // });

  transformedSource.saveSync();
}

function getSourceFile(path: string) {
  const sourceFile = project.addSourceFileAtPath(path);

  if (!watchers.get(path)) {
    const watcher = Deno.watchFs(path);
    // deno-lint-ignore no-inner-declarations
    async function watch() {
      for await (const event of watcher) {
        if (event.kind === 'modify') {
          sourceFile.refreshFromFileSystemSync();
          transformSouceFileSync(sourceFile);
        }
      }
    }

    watch();
    watchers.set(path, watcher);
    transformSouceFileSync(sourceFile);
  }

  return sourceFile;
}

router.get('/scene/open', (context) => {
  const path = getParam(context, 'path');
  const sourceFile = getSourceFile(path);

  const jsxElements = sourceFile.getDescendantsOfKind(SyntaxKind.JsxElement).map((x) => ({
    name: x.getOpeningElement().getTagNameNode().getText(),
    pos: x.getPos(),
  }));

  const jsxSelfClosing = sourceFile
    .getDescendantsOfKind(SyntaxKind.JsxSelfClosingElement)
    .map((x) => ({
      name: x.getTagNameNode().getText(),
      pos: x.getPos(),
    }));

  context.response.body = { sceneObjects: jsxElements.concat(jsxSelfClosing) };
});

function getAllJsxElements(sourceFile: SourceFile) {
  const jsxElements = sourceFile.getDescendantsOfKind(SyntaxKind.JsxOpeningElement);
  const jsxSelfClosing = sourceFile.getDescendantsOfKind(SyntaxKind.JsxSelfClosingElement);

  const elements: (JsxSelfClosingElement | JsxOpeningElement)[] = [];
  elements.push(...jsxSelfClosing, ...jsxElements);

  return elements;
}

router.get('/scene/object/:pos', (context) => {
  const path = getParam(context, 'path');
  const sourceFile = getSourceFile(path);
  const pos = Number(context.params.pos);
  const sceneObject = getAllJsxElements(sourceFile).find((node) => node.getPos() === pos);
  if (!sceneObject) {
    context.response.status = 404;
    context.response.body = { message: 'Not found' };
    return;
  }

  const props = sceneObject.getAttributes().map((prop) => ({
    name: prop.getChildAtIndex(0).getText(),
    pos: prop.getPos(),
    value: getJsxAttributeValue(prop),
  }));

  const name = sceneObject.getTagNameNode().getText();
  const types = getJsxPropTypes(sourceFile, name);

  context.response.body = { name, props, types };
});

router.get('/scene/prop/:pos', (context) => {
  const path = getParam(context, 'path');
  const value = getParam(context, 'value');
  const sourceFile = getSourceFile(path);
  const pos = Number(context.params.pos);

  const attribute = sourceFile.getDescendantAtPos(pos)?.getParentIfKind(SyntaxKind.JsxAttribute);
  if (!attribute) {
    context.response.status = 404;
    context.response.body = { message: 'Not found' };
    return;
  }

  const parsed = JSON.parse(value);

  switch (typeof parsed) {
    case 'string':
      attribute.setInitializer(`"${parsed}"`);
      break;

    default:
      attribute.setInitializer(`{${value}}`);
      break;
  }

  context.response.body = { message: 'success' };
});

router.get('/scene/close', (context) => {
  const path = getParam(context, 'path');
  const sourceFile = project.getSourceFile(path);
  if (sourceFile) {
    project.removeSourceFile(sourceFile);
  }

  const watcher = watchers.get(path);
  if (watcher) {
    watcher.close();
    watchers.delete(path);
  }

  context.response.body = { message: 'success' };
});

router.get('/scene/save', async (context) => {
  const path = getParam(context, 'path');
  const sourceFile = project.getSourceFile(path);
  if (sourceFile) {
    await sourceFile.save();
  }

  context.response.body = { message: 'success' };
});

app.use(router.routes());
app.use(router.allowedMethods());

await app.listen({ port: 8000 });

frontend.tsx

import { OrbitControls, PerspectiveCamera, TransformControls } from '@react-three/drei';
import { Canvas, ThreeEvent, useThree } from '@react-three/fiber';
import { useEffect, useState } from 'react';
import { Layers, Mesh, Object3D, Vector3 } from 'three';

const V1 = new Vector3();

// TODO: Don't hardcode scene link
const path = `/Users/mdougall/projects/FAZE/packages/game/src/scenes/park.tsx`;
const layers = new Layers();
layers.enableAll();

interface EditorNodeData {
  pos: number;
  translate: boolean;
  rotate: boolean;
  scale: boolean;
  sceneObject: Object3D;
}

const findPositionedObject = (object: Object3D): Object3D => {
  let parent: Object3D | null = object.parent;

  while (parent) {
    if (parent.position.lengthSq()) {
      return parent;
    }

    parent = parent.parent;
  }

  return object;
};

const findEditorData = (object: Object3D): EditorNodeData | null => {
  let parent: Object3D | null = object.parent;

  while (parent) {
    if ('__r3fEditor' in parent.userData) {
      return {
        ...parent.userData.__r3fEditor,
        sceneObject: findPositionedObject(object),
      } as EditorNodeData;
    }

    parent = parent.parent;
  }

  return null;
};

function Selection({ children }: { children: JSX.Element }) {
  const [selected, setSelected] = useState<EditorNodeData>();
  const [propPos, setPropPos] = useState<number | undefined>();

  const onClick = async (e: ThreeEvent<MouseEvent>) => {
    if (e.delta > 1) {
      return;
    }

    if (selected && e.object === selected?.sceneObject) {
      // Ignore this event we're already selected.
      return;
    }

    const data = findEditorData(e.object);
    if (data) {
      e.stopPropagation();
      setSelected(data);

      // Begin fetching data for this.
      const res = await fetch(`http://localhost:8000/scene/object/${data.pos}?path=${path}`);
      const json = await res.json();
      const propPos = json.props?.find((prop: any) => prop.name === 'position')?.pos;
      setPropPos(propPos);
    }
  };

  useEffect(() => {
    const callback = (e: KeyboardEvent) => {
      if (e.key === 'Escape') {
        setSelected(undefined);
        setPropPos(undefined);
      }
    };

    document.addEventListener('keyup', callback);

    return () => document.removeEventListener('keyup', callback);
  }, []);

  const onUpdate = (e: any) => {
    if (!e) {
      return;
    }

    if (!selected) {
      return;
    }

    if (e.mode === 'translate' && propPos) {
      fetch(
        `http://localhost:8000/scene/prop/${propPos}?value=${JSON.stringify(
          selected.sceneObject.getWorldPosition(V1).toArray()
        )}&path=${path}`
      );
    }
  };

  return (
    <group onClick={onClick}>
      {children}
      {selected && (
        <TransformControls
          // TODO: Scale/rotate
          mode="translate"
          enabled={selected.translate}
          onMouseUp={onUpdate}
          object={selected.sceneObject}
        />
      )}
    </group>
  );
}

function ForceSceneObjectsVisible() {
  const scene = useThree((three) => three.scene);

  useEffect(() => {
    scene.traverse((object) => {
      if (object.type === 'Mesh') {
        const mesh = object as Mesh;
        if (!mesh.visible) {
          mesh.visible = true;
        }
      }
    });
  });

  return null;
}

export function Editor() {
  // TODO: Don't hardcode scene path
  const Scene =
    require(`/Users/mdougall/projects/FAZE/packages/game/src/scenes/.tmp/park.tsx`).ParkScene;

  useEffect(() => {
    const callback = (e: KeyboardEvent) => {
      if (e.keyCode === 83 && (navigator.platform.match('Mac') ? e.metaKey : e.ctrlKey)) {
        e.preventDefault();
        fetch(`http://localhost:8000/scene/save?path=${path}`);
      }
    };

    document.addEventListener('keydown', callback);

    return () => {
      document.removeEventListener('keydown', callback);
    };
  }, []);

  return (
    <>
      <Canvas>
        <PerspectiveCamera makeDefault layers={layers} position={[0, 10, -1]} />
        <OrbitControls makeDefault />
        <Selection>
          <Scene />
        </Selection>

        <ForceSceneObjectsVisible />
      </Canvas>
    </>
  );
}