keychron: let the M6 set its DPI stage count

src/drivers/keychron/m6-hid.test.ts

@@ -205,6 +205,38 @@ test("rejects stages past the count the mouse reports", async () => {
   await assert.rejects(m6.setDpiStageValue(4, 800), /between 1 and 3/);
 });
 
+test("the stage count grows and shrinks without disturbing the stored DPI", async () => {
+  const fake = new FakeM6Device();
+  const m6 = client(fake);
+  assert.equal(await m6.setDpiStageCount(5), 5);
+  // Growing reveals slots the mouse was already holding, it does not invent them.
+  assert.deepEqual((await m6.readStatus()).dpiStages, [400, 800, 1600, 3200, 5000]);
+  assert.equal(lastSent(fake, 0x40)?.[14], 5);
+  assert.equal(await m6.setDpiStageCount(2), 2);
+  assert.deepEqual((await m6.readStatus()).dpiStages, [400, 800]);
+  // Shrinking keeps the hidden slots intact for the next time they are shown.
+  assert.equal(await m6.setDpiStageCount(5), 5);
+  assert.deepEqual((await m6.readStatus()).dpiStages, [400, 800, 1600, 3200, 5000]);
+});
+
+test("shrinking below the active stage pulls it back into range", async () => {
+  const fake = new FakeM6Device();
+  const m6 = client(fake);
+  await m6.setActiveDpiStage(2);
+  assert.equal((await m6.readStatus()).activeDpiStage, 2);
+  await m6.setDpiStageCount(1);
+  const status = await m6.readStatus();
+  assert.equal(status.activeDpiStage, 0);
+  assert.equal(status.dpi, 400);
+});
+
+test("rejects a stage count the mouse cannot hold", async () => {
+  const m6 = client(new FakeM6Device());
+  for (const count of [0, 6, 2.5, Number.NaN]) {
+    await assert.rejects(m6.setDpiStageCount(count), /between 1 and 5/);
+  }
+});
+
 test("writes the polling rate as an index into the mouse's table", async () => {
   const fake = new FakeM6Device();
   assert.equal(await client(fake).setPollingRate(4000), 4000);

src/drivers/keychron/m6-hid.ts

@@ -184,7 +184,7 @@ export class KeychronM6HidClient {
         statusNote: "Lift-off: Low is 0.7 mm, Medium is 1 mm, High is 2 mm.",
         dpiStageEditor: {
           maxStages: DPI_STAGE_COUNT,
-          countEditable: false,
+          countEditable: true,
           minDpi: DPI_MIN,
           maxDpi: DPI_MAX,
           stepDpi: DPI_STEP,
@@ -248,6 +248,24 @@ export class KeychronM6HidClient {
     return confirmed;
   }
 
+  async setDpiStageCount(count: number): Promise<number> {
+    if (!Number.isInteger(count) || count < 1 || count > DPI_STAGE_COUNT) {
+      throw new Error(`The Keychron M6 holds between 1 and ${DPI_STAGE_COUNT} DPI stages.`);
+    }
+    const settings = await this.readSettings();
+    settings.stageCount = count;
+    // All five hardware slots keep their DPI; the count only decides how many
+    // the DPI button cycles through. The active stage rides along on the same
+    // packet, so shrinking past it would write a stage the mouse cannot hold.
+    if (settings.activeDpiStage >= count) settings.activeDpiStage = count - 1;
+    await this.writeSettings(this.dpiSettingsPacket(settings));
+    const confirmed = (await this.readSettings()).stageCount;
+    if (confirmed !== count) {
+      throw new Error(`The Keychron M6 kept ${confirmed} DPI stages instead of ${count}.`);
+    }
+    return confirmed;
+  }
+
   async setPollingRate(rateHz: number): Promise<number> {
     const settings = await this.readSettings();
     const pollingIndex = settings.pollingTable.findIndex((value) => POLLING_RATES[value] === rateHz);