From 0b16efacd32356b24d3e63b104c3e81cfd70838c Mon Sep 17 00:00:00 2001
From: Abdur-Rahman Mansoor <armansoor@missingno.dev>
Date: Sat, 22 Jun 2024 15:18:43 -0400
Subject: [PATCH] separate `determine_active_nsids` and `is_active_namespace`

---
 drivers/nvme/nvme.asm | 108 +++++++++++++++++++++++-------------------
 1 file changed, 58 insertions(+), 50 deletions(-)

diff --git a/drivers/nvme/nvme.asm b/drivers/nvme/nvme.asm
index bd5c970..3ea3713 100644
--- a/drivers/nvme/nvme.asm
+++ b/drivers/nvme/nvme.asm
@@ -165,73 +165,81 @@ proc create_namespace stdcall, pci:dword, cid:word
 
 endp
 
-; See page 248 of the NVMe 1.4 specification for reference
-proc determine_active_nsids stdcall, pci:dword
-
-		push 	esi edi ebx
-		xor 	esi, esi
-		xor 	edi, edi
+proc is_active_namespace stdcall, pci:dword, nsid:dword
+		
+		push 	esi edi
 		invoke  KernelAlloc, 0x1000
 		test 	eax, eax
-		jz 		.fail
+		jnz 	@f
+		pop 	edi esi
+		ret
+
+@@:
 		mov 	esi, eax
 		invoke  GetPhysAddr
-		stdcall nvme_identify, [pci], 0, eax, CNS_IDCS
-		mov 	eax, dword [esi + IDENTC.nn]
-		mov 	edx, [pci]
-		mov 	dword [edx + pcidev.nn], eax
-		test 	eax, eax
-		jz 		.fail
-		invoke  KernelAlloc, 0x1000
-		test 	eax, eax
-		jz 		.fail
-		mov 	edi, eax
-		invoke  GetPhysAddr
-		mov 	ebx, eax
+		stdcall nvme_identify, [pci], [nsid], eax, CNS_IDNS
 		xor 	ecx, ecx
 
-.loop:
-		cmp 	ecx, dword [esi + IDENTC.nn]
-		jg 		.success
-		stdcall nvme_identify, [pci], ecx, ebx, CNS_IDNS ; identify N'th namespace
-		push 	ecx
-		xor 	ecx, ecx
-
-; We want to check if entire struct is zeroed out, if so, it is not an active NSID,
-; otherwise it is
-.inner_loop:
-		mov 	eax, dword [edi + ecx * 4]
+@@:
+		mov 	eax, dword [esi + ecx * 4]
 		test 	eax, eax
 		jnz 	.is_active_nsid
 		inc 	ecx
 		cmp 	ecx, 0x1000 / 4
-		jne 	.inner_loop
-		pop 	ecx
-		inc 	ecx
-		jmp 	.loop
-
-.is_active_nsid:
-		pop 	ecx ; pop active NSID value into ecx
-		mov 	eax, [pci]
-		mov 	eax, dword [eax + pcidev.nsids]
-		mov 	dword [eax + ecx * 4], ecx
-		inc 	ecx
-		jmp 	.loop
-
-.fail:
-		invoke  KernelFree, esi ; assumes KernelFree can check if its nullptr (hopefully)
-		invoke  KernelFree, edi
-		pop 	ebx edi esi
+		jne 	@b
+		
+.not_active_nsid:
+		invoke  KernelFree, esi
+		pop 	edi esi
 		xor 	eax, eax
 		ret
 
-.success:
+.is_active_nsid:
 		invoke  KernelFree, esi
-		invoke  KernelFree, edi
-		pop 	ebx edi esi
+		pop 	edi esi
 		xor 	eax, eax
 		inc 	eax
 		ret
+
+endp
+
+; See page 248 of the NVMe 1.4 specification for reference
+proc determine_active_nsids stdcall, pci:dword
+
+		push 	esi edi
+		invoke  KernelAlloc, 0x1000
+		test 	eax, eax
+		jnz 	@f
+		pop 	edi esi
+		ret
+
+@@:
+		mov 	esi, eax
+		invoke  GetPhysAddr
+		stdcall nvme_identify, [pci], 0, eax, CNS_IDCS
+		mov 	edi, [pci]
+		mov 	eax, dword [esi + IDENTC.nn]
+		mov 	dword [edi + pcidev.nn], eax
+		mov 	edi, [edi + pcidev.nsids]
+		xor 	ecx, ecx
+
+.loop:
+		cmp 	ecx, dword [esi + IDENTC.nn]
+		je 		.ret
+		stdcall is_active_namespace, [pci], ecx
+		test 	eax, eax
+		jz 		.not_active_namespace
+		mov 	dword [edi + ecx * 4], ecx
+
+.not_active_namespace:
+		inc 	ecx
+		jmp 	.loop
+
+.ret:
+		invoke  KernelFree, esi
+		pop 	edi esi
+		xor 	eax, eax
+		ret
 		
 endp