find byte stride and entries per sector

boot0.s

@@ -20,6 +20,8 @@
 ; bp - 0x08: partition array starting LBA
 ; bp - 0x0a: number of GPT partitions (saturated to 16-bits)
 ; bp - 0x0c: sector stride when loading GPT entries
+; bp - 0x0e: byte stride when loading GPT entries
+; bp - 0x10: entries per sector when loding GPT entries
 
 ; BIOS puts our boot sector at 0000:7c00
 org 0x7c00
@@ -114,29 +116,55 @@ bits 16
   mov bx, [di + 0x52]
   or bx, bx
   jz .gpt_n_partitions_loaded
-  mov ax, 0xff ; Number of partitions overflows 16 bits
+  ; Number of partitions overflows 16 bits, so we just concern ourselves with the first 65535.
+  ; That's an awful lot of partitions anyway.
+  mov ax, 0xffff
 .gpt_n_partitions_loaded:
   push ax
 
   ; Load GPT entry size
-  mov ax, [di + 0x54]
-  mov bx, [di + 0x56]
-  mov cx, ax
-  mov dx, bx
-  ; Find number of sectors per entry
-  shr ax, 9
-  shr bx, 9
-  or bx, bx
-  jnz .panic ; Sectors per entry overflows 16 bits
-  mov bx, ax
-  cmp bx, 1
-  jnz .gpt_sector_stride_loaded
-  inc bx ; Sector stride must be at least one or we'll load the same sector forever!
-.gpt_sector_stride_loaded:
-  push bx
+  mov eax, [di + 0x54] ; Operand size override otherwise this is going to be painful
+  mov ebx, eax
+  ; Assert that the entry size is 128 * 2^n for some integer n>=0. This is required for a valid GPT
+  ; and has the nice properties that:
+  ; - If each entry is larger than a sector (512 bytes), they'll be sector-aligned.
+  ; - If each entry is smaller than a sector, an integer number of them will fit into a sector.
+  or eax, eax  ; Test size != 0 because 128 * 2^n != 0
+  jz .panic
+  and eax, 127 ; Test size is a multiple of 128
+  jnz .panic
+  mov eax, ebx
+  shr eax, 7   ; Divide by 128
+  mov ecx, eax
+  dec eax
+  and eax, ecx ; Test size/128 is a power of 2. popcnt gives an exception for some reason...
+  jnz .panic
+  mov eax, ebx
 
-  ; ax stores sectors per entry
-  ; FIXME: if sectors per entry > 1, then entry size must be a multiple of sector size
+  ; Find the "sector stride", which is the number of sectors we increment by each time we want to
+  ; load a new entry.
+  shr eax, 9      ; Divide by sector size to get sectors per entry
+  cmp eax, 0xffff ; Make sure sectors per entry fits in 16 bits
+  ja .panic
+  or ax, ax
+  jnz .gpt_sector_stride_loaded
+  ; Sector stride must be at least one or we'll load the same sector each time!
+  inc ax
+.gpt_sector_stride_loaded:
+  push ax
+
+  cmp ebx, 512
+  jb .gpt_find_entries_per_sector
+  push word 0 ; Arbitrary byte stride since there's only one entry per sector
+  push word 1 ; 1 entry per sector, since an entry is larger than a sector
+  jmp .gpt_found_entries_per_sector
+.gpt_find_entries_per_sector:
+  push bx     ; Store byte stride = entry length in this case
+  xor dx, dx
+  mov ax, 512
+  div bx      ; Find entries per sector
+  push ax
+.gpt_found_entries_per_sector:
 
   mov al, [0x7e00]
   mov byte fs:[0x0000], al