2 files changed, 44 insertions, 19 deletions

diff --git a/src/peripherals/spi.spec.ts b/src/peripherals/spi.spec.ts
index 9e6dac5..da429fe 100644
--- a/src/peripherals/spi.spec.ts
+++ b/src/peripherals/spi.spec.ts
@@ -96,26 +96,26 @@ describe('SPI', () => {
     expect(spi.isMaster).toBe(true);
   });
 
-  it('should call the `onTransfer` callback when initiating an SPI trasfer by writing to SPDR', () => {
+  it('should call the `onByteTransfer` callback when initiating an SPI trasfer by writing to SPDR', () => {
     const cpu = new CPU(new Uint16Array(1024));
     const spi = new AVRSPI(cpu, spiConfig, FREQ_16MHZ);
-    spi.onTransfer = jest.fn();
+    spi.onByte = jest.fn();
 
     cpu.writeData(SPCR, SPE | MSTR);
     cpu.writeData(SPDR, 0x8f);
 
-    expect(spi.onTransfer).toHaveBeenCalledWith(0x8f);
+    expect(spi.onByte).toHaveBeenCalledWith(0x8f);
   });
 
   it('should ignore SPDR writes when the SPE bit in SPCR is clear', () => {
     const cpu = new CPU(new Uint16Array(1024));
     const spi = new AVRSPI(cpu, spiConfig, FREQ_16MHZ);
-    spi.onTransfer = jest.fn();
+    spi.onByte = jest.fn();
 
     cpu.writeData(SPCR, MSTR);
     cpu.writeData(SPDR, 0x8f);
 
-    expect(spi.onTransfer).not.toHaveBeenCalled();
+    expect(spi.onByte).not.toHaveBeenCalled();
   });
 
   it('should transmit a byte successfully (integration)', () => {
@@ -155,9 +155,9 @@ describe('SPI', () => {
 
     let byteReceivedFromAsmCode: number | null = null;
 
-    spi.onTransfer = (value) => {
+    spi.onByte = (value) => {
       byteReceivedFromAsmCode = value;
-      return 0x5b; // we copy this byte to
+      cpu.addClockEvent(() => spi.completeTransfer(0x5b), spi.transferCycles);
     };
 
     const runner = new TestProgramRunner(cpu, () => 0);
@@ -228,7 +228,9 @@ describe('SPI', () => {
   it('should should only update SPDR when tranfer finishes (double buffering)', () => {
     const cpu = new CPU(new Uint16Array(1024));
     const spi = new AVRSPI(cpu, spiConfig, FREQ_16MHZ);
-    spi.onTransfer = jest.fn(() => 0x88);
+    spi.onByte = () => {
+      cpu.addClockEvent(() => spi.completeTransfer(0x88), spi.transferCycles);
+    };
 
     cpu.writeData(SPCR, SPE | MSTR);
     cpu.writeData(SPDR, 0x8f);
diff --git a/src/peripherals/spi.ts b/src/peripherals/spi.ts
index 7b753ab..3167c07 100644
--- a/src/peripherals/spi.ts
+++ b/src/peripherals/spi.ts
@@ -31,15 +31,28 @@ export const spiConfig: SPIConfig = {
   SPDR: 0x4e,
 };
 
-export type SPITransferCallback = (value: u8) => u8;
+export type SPITransferCallback = (value: u8) => number;
+export type SPIByteTransferCallback = (value: u8) => void;
 
 const bitsPerByte = 8;
 
 export class AVRSPI {
-  public onTransfer: SPITransferCallback | null = null;
+  /** @deprecated Use onByte() instead */
+  public onTransfer: SPITransferCallback = () => 0;
+
+  /**
+   * SPI byte transfer callback. Invoked whenever the user code starts an SPI transaction.
+   * You can override this with your own SPI handler logic.
+   *
+   * The callback receives a argument: the byte sent over the SPI MOSI line.
+   * It should call `completeTransfer()` within `transferCycles` CPU cycles.
+   */
+  public onByte: SPIByteTransferCallback = (value) => {
+    const valueIn = this.onTransfer(value);
+    this.cpu.addClockEvent(() => this.completeTransfer(valueIn), this.transferCycles);
+  };
 
   private transmissionActive = false;
-  private receivedByte: u8 = 0;
 
   // Interrupts
   private SPI: AVRInterruptConfig = {
@@ -68,14 +81,8 @@ export class AVRSPI {
       cpu.data[SPSR] &= ~SPSR_WCOL;
       this.cpu.clearInterrupt(this.SPI);
 
-      this.receivedByte = this.onTransfer?.(value) ?? 0;
-      const cyclesToComplete = this.clockDivider * bitsPerByte;
       this.transmissionActive = true;
-      this.cpu.addClockEvent(() => {
-        this.cpu.data[SPDR] = this.receivedByte;
-        this.cpu.setInterruptFlag(this.SPI);
-        this.transmissionActive = false;
-      }, cyclesToComplete);
+      this.onByte(value);
       return true;
     };
     cpu.writeHooks[SPCR] = (value: u8) => {
@@ -89,7 +96,18 @@ export class AVRSPI {
 
   reset() {
     this.transmissionActive = false;
-    this.receivedByte = 0;
+  }
+
+  /**
+   * Completes an SPI transaction. Call this method only from the `onByte` callback.
+   *
+   * @param receivedByte Byte read from the SPI MISO line.
+   */
+  completeTransfer(receivedByte: number) {
+    const { SPDR } = this.config;
+    this.cpu.data[SPDR] = receivedByte;
+    this.cpu.setInterruptFlag(this.SPI);
+    this.transmissionActive = false;
   }
 
   get isMaster() {
@@ -128,6 +146,11 @@ export class AVRSPI {
     throw new Error('Invalid divider value!');
   }
 
+  /** Number of cycles to complete a single byte SPI transaction */
+  get transferCycles() {
+    return this.clockDivider * bitsPerByte;
+  }
+
   /**
    * The SPI freqeuncy is only relevant to Master mode.
    * In slave mode, the frequency can be as high as F(osc) / 4.