3 files changed, 103 insertions, 16 deletions

diff --git a/demo/src/execute.ts b/demo/src/execute.ts
index 7b6cae0..ee5278e 100644
--- a/demo/src/execute.ts
+++ b/demo/src/execute.ts
@@ -11,6 +11,7 @@ import {
   usart0Config
 } from 'avr8js';
 import { loadHex } from './intelhex';
+import { MicroTaskScheduler } from './task-scheduler';
 
 // ATmega328p params
 const FLASH = 0x8000;
@@ -24,8 +25,8 @@ export class AVRRunner {
   readonly portD: AVRIOPort;
   readonly usart: AVRUSART;
   readonly speed = 16e6; // 16 MHZ
-
-  private stopped = false;
+  readonly workUnitCycles = 500000;
+  readonly taskScheduler = new MicroTaskScheduler();
 
   constructor(hex: string) {
     loadHex(hex, new Uint8Array(this.program.buffer));
@@ -35,28 +36,23 @@ export class AVRRunner {
     this.portC = new AVRIOPort(this.cpu, portCConfig);
     this.portD = new AVRIOPort(this.cpu, portDConfig);
     this.usart = new AVRUSART(this.cpu, usart0Config, this.speed);
+    this.taskScheduler.start();
   }
 
-  async execute(callback: (cpu: CPU) => void) {
-    this.stopped = false;
-    const workUnitCycles = 500000;
-    let nextTick = this.cpu.cycles + workUnitCycles;
-    for (;;) {
+  // CPU main loop
+  execute(callback: (cpu: CPU) => void) {
+    const cyclesToRun = this.cpu.cycles + this.workUnitCycles;
+    while (this.cpu.cycles < cyclesToRun) {
       avrInstruction(this.cpu);
       this.timer.tick();
       this.usart.tick();
-      if (this.cpu.cycles >= nextTick) {
-        callback(this.cpu);
-        await new Promise((resolve) => setTimeout(resolve, 0));
-        if (this.stopped) {
-          break;
-        }
-        nextTick += workUnitCycles;
-      }
     }
+
+    callback(this.cpu);
+    this.taskScheduler.postTask(() => this.execute(callback));
   }
 
   stop() {
-    this.stopped = true;
+    this.taskScheduler.stop();
   }
 }
diff --git a/demo/src/task-scheduler.spec.ts b/demo/src/task-scheduler.spec.ts
new file mode 100644
index 0000000..f5fb4bf
--- /dev/null
+++ b/demo/src/task-scheduler.spec.ts
@@ -0,0 +1,52 @@
+/**
+ * @jest-environment jsdom
+ */
+/// <reference lib="dom" />
+
+import { MicroTaskScheduler } from './task-scheduler';
+
+describe('task-scheduler', () => {
+  let taskScheduler: MicroTaskScheduler;
+  let task: jest.Mock;
+
+  beforeEach(() => {
+    taskScheduler = new MicroTaskScheduler();
+    task = jest.fn();
+  });
+
+  it('should execute task', async () => {
+    taskScheduler.start();
+    taskScheduler.postTask(task);
+    await new Promise((resolve) => setTimeout(resolve, 0));
+    expect(task).toHaveBeenCalledTimes(1);
+  });
+
+  it('should execute task twice when posted twice', async () => {
+    taskScheduler.start();
+    taskScheduler.postTask(task);
+    taskScheduler.postTask(task);
+    await new Promise((resolve) => setTimeout(resolve, 0));
+    expect(task).toHaveBeenCalledTimes(2);
+  });
+
+  it('should not execute task when not started', async () => {
+    taskScheduler.postTask(task);
+    await new Promise((resolve) => setTimeout(resolve, 0));
+    expect(task).not.toHaveBeenCalled();
+  });
+
+  it('should not execute task when stopped', async () => {
+    taskScheduler.start();
+    taskScheduler.stop();
+    taskScheduler.postTask(task);
+    await new Promise((resolve) => setTimeout(resolve, 0));
+    expect(task).not.toHaveBeenCalled();
+  });
+
+  it('should not register listener twice', async () => {
+    const addEventListenerSpy = jest.spyOn(window, 'addEventListener');
+    taskScheduler.start();
+    taskScheduler.start();
+    expect(addEventListenerSpy).toHaveBeenCalledTimes(1);
+  });
+});
diff --git a/demo/src/task-scheduler.ts b/demo/src/task-scheduler.ts
new file mode 100644
index 0000000..5364d92
--- /dev/null
+++ b/demo/src/task-scheduler.ts
@@ -0,0 +1,39 @@
+// Faster setTimeout(fn, 0) implementation using postMessage API
+// Based on https://dbaron.org/log/20100309-faster-timeouts
+export type IMicroTaskCallback = () => void;
+
+export class MicroTaskScheduler {
+  readonly messageName = 'zero-timeout-message';
+
+  private executionQueue: Array<IMicroTaskCallback> = [];
+  private stopped = true;
+
+  start() {
+    if (this.stopped) {
+      this.stopped = false;
+      window.addEventListener('message', this.handleMessage, true);
+    }
+  }
+
+  stop() {
+    this.stopped = true;
+    window.removeEventListener('message', this.handleMessage, true);
+  }
+
+  postTask(fn: IMicroTaskCallback) {
+    if (!this.stopped) {
+      this.executionQueue.push(fn);
+      window.postMessage(this.messageName, '*');
+    }
+  }
+
+  private handleMessage = (event: MessageEvent) => {
+    if (event.data === this.messageName) {
+      event.stopPropagation();
+      const executeJob = this.executionQueue.shift();
+      if (executeJob !== undefined) {
+        executeJob();
+      }
+    }
+  };
+}