2 files changed, 41 insertions, 18 deletions

diff --git a/src/peripherals/timer.spec.ts b/src/peripherals/timer.spec.ts
index 40db1f6..43768d2 100644
--- a/src/peripherals/timer.spec.ts
+++ b/src/peripherals/timer.spec.ts
@@ -83,7 +83,7 @@ describe('timer', () => {
     cpu.writeData(TCCR0B, CS01 | CS00); // Set prescaler to 64
     cpu.cycles = 1;
     cpu.tick();
-    cpu.cycles = 64;
+    cpu.cycles = 1 + 64;
     cpu.tick();
     const tcnt = cpu.readData(TCNT0);
     expect(tcnt).toEqual(1);
@@ -353,11 +353,11 @@ describe('timer', () => {
     cpu.cycles = 1;
     cpu.tick();
 
-    cpu.cycles = 511;
+    cpu.cycles = 1 + 511;
     cpu.tick();
     expect(cpu.readData(TCNT2)).toEqual(1);
 
-    cpu.cycles = 512;
+    cpu.cycles = 1 + 512;
     cpu.tick();
     expect(cpu.readData(TCNT2)).toEqual(2);
   });
@@ -366,7 +366,7 @@ describe('timer', () => {
     const { program, instructionCount } = asmProgram(`
       LDI r16, 0x1      ; TCCR0B = 1 << CS00
       OUT 0x25, r16
-      LDI r16, 0x0      ; TCNT0 <- 0x30
+      LDI r16, 0x0      ; TCNT0 <- 0
       OUT 0x26, r16
       NOP
       LDS r1, 0x46      ; r1 <- TCNT0 (2 cycles)
@@ -396,6 +396,22 @@ describe('timer', () => {
     expect(cpu.readData(R17)).toEqual(2);
   });
 
+  it('should not keep counting for one more instruction when the timer is disabled (issue #72)', () => {
+    const { program, instructionCount } = asmProgram(`
+      EOR r1, r1      ; r1 = 0;
+      LDI r16, 0x1    ; TCCR2B = 1 << CS20;
+      STS 0xb1, r16   ; Should start counting after this instruction,
+      STS 0xb1, r1    ; and stop counting *after* this one.
+      NOP
+      LDS r17, 0xb2   ; TCNT2 should equal 2 at this point (not counting the NOP)
+  `);
+    const cpu = new CPU(program);
+    new AVRTimer(cpu, timer2Config);
+    const runner = new TestProgramRunner(cpu);
+    runner.runInstructions(instructionCount);
+    expect(cpu.readData(R17)).toEqual(2);
+  });
+
   describe('Phase-correct PWM mode', () => {
     it('should count up to TOP, down to 0, and then set TOV flag', () => {
       const { program, instructionCount } = asmProgram(`
diff --git a/src/peripherals/timer.ts b/src/peripherals/timer.ts
index b8a6668..188832d 100644
--- a/src/peripherals/timer.ts
+++ b/src/peripherals/timer.ts
@@ -249,6 +249,7 @@ export class AVRTimer {
   private compA: CompBitsValue;
   private compB: CompBitsValue;
   private tcntUpdated = false;
+  private updateDivider = false;
   private countingUp = true;
   private divider = 0;
 
@@ -326,10 +327,9 @@ export class AVRTimer {
     };
     cpu.writeHooks[config.TCCRB] = (value) => {
       this.cpu.data[config.TCCRB] = value;
-      this.divider = this.config.dividers[this.CS];
-      this.reschedule();
-      this.tcntUpdated = true;
-      this.cpu.updateClockEvent(this.count, 0);
+      this.updateDivider = true;
+      this.cpu.clearClockEvent(this.count);
+      this.cpu.addClockEvent(this.count, 0);
       this.updateWGMConfig();
       return true;
     };
@@ -357,6 +357,7 @@ export class AVRTimer {
     this.tcntNext = 0;
     this.tcntUpdated = false;
     this.countingUp = false;
+    this.updateDivider = true;
   }
 
   get TCCRA() {
@@ -399,8 +400,9 @@ export class AVRTimer {
   }
 
   count = (reschedule = true) => {
-    const { divider, lastCycle } = this;
-    const delta = this.cpu.cycles - lastCycle;
+    const { divider, lastCycle, cpu } = this;
+    const { cycles } = cpu;
+    const delta = cycles - lastCycle;
     if (divider && delta >= divider) {
       const counterDelta = Math.floor(delta / divider);
       this.lastCycle += counterDelta * divider;
@@ -417,23 +419,28 @@ export class AVRTimer {
         this.timerUpdated();
       }
       if ((timerMode === TimerMode.Normal || timerMode === TimerMode.FastPWM) && val > newVal) {
-        this.cpu.setInterruptFlag(this.OVF);
+        cpu.setInterruptFlag(this.OVF);
       }
     }
     if (this.tcntUpdated) {
       this.tcnt = this.tcntNext;
       this.tcntUpdated = false;
     }
-    if (reschedule) {
-      this.cpu.addClockEvent(this.count, this.lastCycle + divider - this.cpu.cycles);
+    if (this.updateDivider) {
+      const newDivider = this.config.dividers[this.CS];
+      this.lastCycle = newDivider ? this.cpu.cycles : 0;
+      this.updateDivider = false;
+      this.divider = newDivider;
+      if (newDivider) {
+        cpu.addClockEvent(this.count, this.lastCycle + newDivider - cpu.cycles);
+      }
+      return;
+    }
+    if (reschedule && divider) {
+      cpu.addClockEvent(this.count, this.lastCycle + divider - cpu.cycles);
     }
   };
 
-  private reschedule() {
-    this.cpu.clearClockEvent(this.count);
-    this.cpu.addClockEvent(this.count, this.lastCycle + this.divider - this.cpu.cycles);
-  }
-
   private phasePwmCount(value: u16, delta: u8) {
     while (delta > 0) {
       if (this.countingUp) {