react.tsx

enum HookType {
  STATE,
  EFFECT,
}

type EffectHook = {
  type: HookType.EFFECT;
  cleanup?: () => void;
  dependencies?: any[];
};

type StateHook = {
  type: HookType.STATE;
  state: any;
};

type Hook = StateHook | EffectHook;

enum NodeType {
  COMPONENT,
  HOST,
  TEXT,
}

type ComponentElement = {
  kind: NodeType.COMPONENT;
  component: (props: Props) => ReactElement;
  props: Props;
};

type HostElement = {
  kind: NodeType.HOST;
  tag: string;
  props: Props;
};

type TextElement = {
  kind: NodeType.TEXT;
  content: string;
};

type ReactElement = ComponentElement | HostElement | TextElement;

type BaseNode = {
  parent: ReactNode | null;
  descendants: ReactNode[];
};

type ComponentNode = ComponentElement & BaseNode & { hooks: Hook[] };
type HostNode = HostElement & BaseNode & { dom: any };
type TextNode = TextElement & { parent: ReactNode; dom: any };

type ReactNode = ComponentNode | HostNode | TextNode;

type Props = {
  children?: ReactElement[];
  [key: string]: any;
};

type Job = {
  node: ReactNode;
  element: ReactElement | null;
};

export function createElement(
  component: any,
  props: any,
  ...children: (ReactElement | string | number | null)[]
): ReactElement {
  const p = {
    ...(props || {}),
    children: children
      .flat()
      .map((child: ReactElement | string | number | null) => {
        if (typeof child === "string") {
          return {
            kind: NodeType.TEXT,
            content: child,
          };
        } else if (typeof child === "number") {
          return {
            kind: NodeType.TEXT,
            content: child.toString(),
          };
        } else {
          // Null and false will be passed here and filtered below.
          return child;
        }
      })
      .filter(Boolean),
  };

  if (typeof component === "function") {
    return {
      kind: NodeType.COMPONENT,
      component,
      props: p,
    };
  } else if (typeof component === "string") {
    return {
      kind: NodeType.HOST,
      tag: component,
      props: p,
    };
  }

  throw new Error("Something went wrong.");
}

// We assume something is event when it follow onSomething naming pattern.
const isEvent = (key: string): boolean => !!key.match(new RegExp("on[A-Z].*"));

// Event is transformed to browser event name by removing 'on' and lowercasing.
// 'onInput' becomes 'input'.
const eventToKeyword = (key: string): string =>
  key.replace("on", "").toLowerCase();

// camelCaseWillBecome kebab-case-with-dashes.
const camelCaseToKebab = (str: string) =>
  str.replace(/([a-z])([A-Z])/g, "$1-$2").toLowerCase();

function findClosestHostNode(node: ReactNode): HostNode {
  let current = node;

  while (current.kind !== NodeType.HOST && current.parent) {
    current = current.parent;
  }

  // Only interested in looking for host node as text node wouldn't have
  // children anyway.
  if (current.kind !== NodeType.HOST) {
    throw new Error("Couldn't find node.");
  }

  return current;
}

function createDom(element: HostElement): Element {
  const html = document.createElement(element.tag);

  Object.entries(element.props).forEach(([key, value]) => {
    if (key === "children" || key === "ref") {
      // Skip.
    } else if (isEvent(key)) {
      // We want to translate 'onClick' to 'click'.
      html.addEventListener(eventToKeyword(key), value);
    } else {
      html.setAttribute(camelCaseToKebab(key), value);
    }
  });

  return html;
}

function createTextNode(text: string): Text {
  return document.createTextNode(text);
}

function updateTextNode(current: TextNode, text: string) {
  current.dom.nodeValue = text;
}

function updateDom(current: HostNode, expected: HostElement) {
  const html = current.dom as HTMLElement;

  // First we iterate over current props, which means we might either update
  // them or decide to remove them.
  Object.keys(current.props).forEach((key) => {
    if (key === "children" || key === "ref") {
      // Skip.
    } else if (isEvent(key)) {
      html.removeEventListener(eventToKeyword(key), current.props[key]);
    } else {
      // Prop will be removed.
      if (!expected.props[key]) {
        html.removeAttribute(key);
      }

      // Prop will be updated.
      if (expected.props[key]) {
        html.setAttribute(camelCaseToKebab(key), expected.props[key] as string);
      }
    }
  });

  // Now, iterating over new props we will overwrite the ones we have in
  // expected object.
  Object.keys(expected.props).forEach((key) => {
    if (key === "children" || key === "ref") {
      // Skip.
    } else if (isEvent(key)) {
      html.addEventListener(eventToKeyword(key), expected.props[key]);
    } else {
      // Prop will be added.
      if (!current.props[key]) {
        html.setAttribute(camelCaseToKebab(key), expected.props[key] as string);
      }
    }
  });
}

function removeDom(node: ReactNode) {
  if (node.kind === NodeType.HOST || node.kind === NodeType.TEXT) {
    node.dom.parentNode?.removeChild(node.dom);
  } else {
    node.descendants.forEach((child) => {
      removeDom(child);
    });
  }
}

let _currentNode: ReactNode | null;
let _hookIndex = 0;

const _tasks: Job[] = [];
const _effects: (() => void)[] = [];
let _updating = false;

function update(node: ReactNode, element: ReactElement | null) {
  const previousNode = _currentNode;
  const previousIndex = _hookIndex;

  let elements: ReactElement[] = [];

  // First we find out what to work on. For component we have to call render
  // function to find out what what React elements it will return this time.
  if (node.kind === NodeType.COMPONENT) {
    _currentNode = node;
    _hookIndex = 0;
    elements = [node.component(node.props)];
  } else if (element && "props" in element && node.kind === NodeType.HOST) {
    elements = element.props.children || [];
  } else if (
    (element && element.kind === NodeType.TEXT) ||
    node.kind === NodeType.TEXT
  ) {
    // Operations for the text node were done in the parent node.
    return;
  }

  const length = Math.max(node.descendants.length, elements.length);
  const pairs: [
    left: ReactNode | undefined,
    right: ReactElement | undefined
  ][] = [];

  for (let i = 0; i < length; i++) {
    pairs.push([node.descendants[i], elements[i]]);
  }

  for (const [current, expected] of pairs) {
    const shouldUpdate =
      // Both are defined and...
      current &&
      expected &&
      //...it was and still is component and previous and current functional
      // component function is the same one (which means we are rerendering) or...
      ((current.kind === NodeType.COMPONENT &&
        expected.kind === NodeType.COMPONENT &&
        current.component === expected.component) ||
        // ... it was and still is host node and tag has not changed (which means
        // we just need to update props) or...
        (current.kind === NodeType.HOST &&
          expected.kind === NodeType.HOST &&
          current.tag === expected.tag) ||
        // ...it was and is text node (hich means we update text content).
        (current.kind === NodeType.TEXT && expected.kind === NodeType.TEXT));
    const shouldReplace = current && expected;
    const shouldAdd = !current && expected !== undefined;
    const shouldRemove = current !== undefined && !expected;

    if (shouldUpdate) {
      updateNode(current!, expected!);
    } else if (shouldReplace) {
      replaceNode(node, current!, expected!);
    } else if (shouldAdd) {
      addNode(node, expected!);
    } else if (shouldRemove) {
      removeNode(node, current!);
    }
  }

  _currentNode = previousNode;
  _hookIndex = previousIndex;
}

function addNode(node: ComponentNode | HostNode, expected: ReactElement) {
  let newNode: ReactNode;

  if (expected.kind === NodeType.COMPONENT) {
    newNode = {
      ...expected,
      parent: node,
      descendants: [],
      hooks: [],
    };
  } else if (expected.kind === NodeType.HOST) {
    const nodeConstruction: any = {
      ...expected,
      parent: node,
      descendants: [],
    };

    const dom = createDom(expected);
    const hostNodeParent = findClosestHostNode(node);
    hostNodeParent.dom.appendChild(dom);

    nodeConstruction.dom = dom;
    newNode = nodeConstruction;
  } else if (expected.kind === NodeType.TEXT) {
    const nodeConstruction: any = {
      ...expected,
      parent: node,
      descendants: [],
    };

    const dom = createTextNode(expected.content);

    const hostNodeParent = findClosestHostNode(node);
    hostNodeParent.dom.appendChild(dom);

    nodeConstruction.dom = dom;
    newNode = nodeConstruction;
  } else {
    throw new Error("Unknown node type.");
  }

  node.descendants.push(newNode);
  update(newNode, expected);
}

function updateNode(current: ReactNode, expected: ReactElement) {
  if (current.kind === NodeType.HOST && expected.kind === NodeType.HOST) {
    updateDom(current, expected);
  } else if (
    // Text value changed.
    current.kind === NodeType.TEXT &&
    expected.kind === NodeType.TEXT &&
    current.content !== expected.content
  ) {
    current.content = expected.content;
    updateTextNode(current, expected.content);
  }

  // Props can be replaced.
  if ("props" in current && "props" in expected) {
    current.props = expected.props;
  }

  update(current, expected);
}

function replaceNode(
  node: ComponentNode | HostNode,
  current: ReactNode,
  expected: ReactElement
) {
  let newNode: ReactNode;
  if (expected.kind === NodeType.COMPONENT) {
    newNode = {
      ...expected,
      parent: node,
      descendants: [],
      hooks: [],
    };

    removeDom(current);
  } else if (expected.kind === NodeType.HOST) {
    const firstParentWithHostNode = findClosestHostNode(node);

    const nodeConstruction: any = {
      ...expected,
      parent: node,
      descendants: [],
    };

    const dom = createDom(expected);
    if (current.kind === NodeType.HOST || current.kind === NodeType.TEXT) {
      firstParentWithHostNode.dom.replaceChild(dom, current.dom);
    } else {
      removeDom(current);
      firstParentWithHostNode.dom.appendChild(dom);
    }
    nodeConstruction.dom = dom;

    newNode = nodeConstruction;
  } else if (expected.kind === NodeType.TEXT) {
    const firstParentWithHostNode = findClosestHostNode(node);
    const nodeConstruction: any = {
      ...expected,
      parent: node,
    };

    const dom = createTextNode(expected.content);
    if (current.kind === NodeType.TEXT) {
      throw new Error("Update should have happened on this node.");
    } else if (current.kind === NodeType.HOST) {
      firstParentWithHostNode.dom.replaceChild(dom, current.dom);
      nodeConstruction.dom = dom;
    } else {
      removeDom(current);
      firstParentWithHostNode.dom.appendChild(dom);
      nodeConstruction.dom = dom;
    }

    newNode = nodeConstruction;
  } else {
    throw new Error("Couldn't resolve node kind.");
  }

  node.descendants[node.descendants.indexOf(current)] = newNode;
  update(newNode, expected);
}

function removeNode(node: ComponentNode | HostNode, current: ReactNode) {
  const indexOfCurrent = node.descendants.indexOf(current);
  removeDom(current);
  node.descendants.splice(indexOfCurrent, 1);
}

function useState<T>(initial: T): [T, (next: T | ((current: T) => T)) => void] {
  const c = _currentNode;
  const i = _hookIndex;

  if (!c || c.kind !== NodeType.COMPONENT) {
    throw new Error("Executing useState for non-function element.");
  }

  if (c.hooks[i] === undefined) {
    c.hooks[i] = {
      type: HookType.STATE,
      state: initial,
    };
  }

  const hook = c.hooks[i];
  // Just for type correctness but also good place to spot implementation bugs.
  if (hook.type !== HookType.STATE) {
    throw new Error("Something went wrong.");
  }

  const setState = (next: T | ((current: T) => T)) => {
    // https://github.com/microsoft/TypeScript/issues/37663#issuecomment-856866935
    // In case of a different iframe, window or realm, next won't be instance
    // of the same Function and will be saved instead of treated as callback.
    if (next instanceof Function) {
      hook.state = next(hook.state);
    } else {
      hook.state = next;
    }

    runUpdateLoop(c, null);
  };

  _hookIndex += 1;
  return [hook.state, setState];
}

function useEffect(
  callback: () => void | (() => void),
  dependencies?: any[]
): void {
  const c = _currentNode;
  const i = _hookIndex;

  if (!c || c.kind !== NodeType.COMPONENT) {
    throw new Error("Executing useEffect for non-function element.");
  }
  _effects.push(() => {
    if (c.hooks[i] === undefined) {
      // INITIALIZE
      const hook: EffectHook = {
        type: HookType.EFFECT,
        cleanup: undefined,
        dependencies,
      };
      c.hooks[i] = hook;
      const cleanup = callback();
      hook.cleanup = cleanup ? cleanup : undefined;
    } else if (dependencies) {
      // COMPARE DEPENDENCIES
      const hook = c.hooks[i];
      if (hook.type !== HookType.EFFECT || hook.dependencies === undefined) {
        throw new Error("Something went wrong.");
      }

      let shouldRun = false;
      for (let j = 0; j < dependencies.length; j++) {
        if (dependencies[j] !== hook.dependencies[j]) {
          shouldRun = true;
        }
      }

      if (shouldRun) {
        const cleanup = callback();
        c.hooks[i] = {
          type: HookType.EFFECT,
          cleanup: cleanup ? cleanup : undefined,
          dependencies,
        };
      }
    } else if (!dependencies) {
      // RUN ALWAYS
      const cleanup = callback();
      c.hooks[i] = {
        type: HookType.EFFECT,
        cleanup: cleanup ? cleanup : undefined,
        dependencies,
      };
    }
  });

  _hookIndex += 1;
}

function render(element: ReactElement, dom: HTMLElement) {
  const rootNode: HostNode = {
    kind: NodeType.HOST,
    tag: dom.tagName.toLowerCase(),
    props: {
      children: [element],
    },
    parent: null,
    descendants: [],
    dom,
  };

  runUpdateLoop(rootNode, createElement(rootNode.tag, rootNode.props, element));
}

function runUpdateLoop(node: ReactNode, element: ReactElement | null) {
  _tasks.push({ node, element });

  if (_updating) {
    return;
  }

  _updating = true;

  let current: Job | undefined;
  while ((current = _tasks.shift())) {
    update(current.node, current.element);

    // Run all effects queued for this update.
    let effect: (() => void) | undefined;
    while ((effect = _effects.shift())) {
      effect();
    }
  }

  _updating = false;
}

const c = createElement;

const App = () => {
  return (
    <div>
      {[1, 2, 3].map((i) => (
        <div key={i}>{i}</div>
      ))}
      <div>
        test
        {2}
      </div>
      <div>{["test"]}</div>
      <div>test</div>
      <div>{false && "test"}</div>
    </div>
  );
};

render(<App />, document.getElementById("app")!);