From 87fa0e0e3ff6081bce4da1381df5a8cffbae8b58 Mon Sep 17 00:00:00 2001
From: Abdur-Rahman Mansoor <armansoor@missingno.dev>
Date: Fri, 12 Jul 2024 13:33:26 -0400
Subject: [PATCH] refactor build_prp_list and remove free_prp_list

---
 drivers/nvme/nvme.asm | 81 +++++++++++++------------------------------
 1 file changed, 24 insertions(+), 57 deletions(-)

diff --git a/drivers/nvme/nvme.asm b/drivers/nvme/nvme.asm
index fa60298..fd3f591 100644
--- a/drivers/nvme/nvme.asm
+++ b/drivers/nvme/nvme.asm
@@ -454,27 +454,7 @@ proc get_log_page stdcall, pci:dword, dptr:dword, lid:byte
 
 endp
 
-proc free_prp_list stdcall, prp_list_ptr:dword, nprps:dword
-
-	push 	edi esi ebx
-	mov 	esi, [prp_list_ptr]
-	invoke  KernelAlloc, PAGE_SIZE
-	test 	eax, eax
-	jz 	.exit
-	mov 	edi, eax
-	invoke  MapPage, edi, esi, PG_SW+PG_NOCACHE
-	mov 	ebx, dword [edi + PAGE_SIZE / 8]
-	invoke  FreePage, esi
-	sub 	[nprps], PAGE_SIZE / 8
-	
-.exit:
-	xor 	eax, eax
-	pop 	ebx esi edi
-	ret
-
-endp
-
-proc build_prp_list stdcall, nprps:dword
+proc build_prp_list stdcall, nprps:dword, buf_physical:dword, prp_list_ptr:dword
 
 	push 	esi ebx edi
 	sub 	esp, 16
@@ -497,11 +477,16 @@ proc build_prp_list stdcall, nprps:dword
 	; its important to cache the value
 	mov 	dword [esp + 12], 0
 
+	; the pointer to our PRP list (virtual), needed so
+	; that the caller can free the PRP list afterwards
+	mov 	ebx, [prp_list_ptr]
+	mov 	dword [ebx], 0
+
 	xor 	edi, edi
 	xor 	esi, esi
 	mov 	ecx, [nprps]
-	mov 	ebx, ecx
-	shl 	ecx, 3
+	shl 	ecx, 3 ; multiply by 8 since each PRP pointer is a QWORD
+
 	; we'll store consecutive PRP list buffers here, for example
 	; given 2 PRP lists, we allocate 2 continuous pages
 	invoke  KernelAlloc, ecx ; store pointers to the PRP entries here
@@ -509,22 +494,20 @@ proc build_prp_list stdcall, nprps:dword
 	jz 	.err
 	mov 	dword [esp + 12], eax
 	mov 	edi, eax
+	mov 	dword [ebx], eax
+	
+	; note we assume buf_physical is page-aligned
+	mov 	esi, [buf_physical]
 
 .build_prp_list:
-	mov 	ecx, ebx
-	cmp 	ecx, PAGE_SIZE / 8
-	jne 	@f
-	mov 	ecx, PAGE_SIZE / 8
+	; ensure we don't cross a page boundary
+	mov 	ebx, [nprps]
+	cmp 	ebx, PAGE_SIZE / 8
+	jb 	@f
+	mov 	ebx, PAGE_SIZE / 8
+	sub 	[nprps], ebx
 
 @@:
-	imul 	ecx, PAGE_SIZE
-
-	; allocate a massive buffer for the PRPs themselves
-	invoke  KernelAlloc, ecx
-	test 	eax, eax
-	jz 	.err
-	invoke  GetPhysAddr
-	mov 	esi, eax
 	xor 	ecx, ecx
 	cmp 	dword [esp], 0
 	jz 	@f
@@ -534,6 +517,7 @@ proc build_prp_list stdcall, nprps:dword
 	mov 	eax, dword [esp + 8]
 	mov 	dword [eax + PAGE_SIZE - 8], esi
 	mov 	dword [eax + PAGE_SIZE - 4], 0
+	jmp 	.loop
 
 @@:
 	mov 	dword [esp], esi
@@ -548,7 +532,7 @@ proc build_prp_list stdcall, nprps:dword
 	jne 	.loop
 	; check if we we need to build another PRP list
 	mov 	dword [esp + 8], edi
-	mov 	ebx, [nprps]
+	add 	edi, PAGE_SIZE
 	cmp 	ebx, PAGE_SIZE / 8
 	je 	.build_prp_list
 	mov 	eax, dword [esp]
@@ -558,20 +542,6 @@ proc build_prp_list stdcall, nprps:dword
 	ret
 
 .err:
-	mov 	esi, dword [esp]
-	test 	esi, esi
-	jz 	@f
-	sub 	[nprps], PAGE_SIZE / 8
-	add 	edi, PAGE_SIZE
-	stdcall free_prp_list, esi, [esp + 4]
-
-@@:
-	mov 	eax, dword [esp + 12]
-	test 	eax, eax
-	jz 	@f
-	invoke  KernelFree, eax
-
-@@:
 	add 	esp, 16
 	pop 	edi ebx esi
 	xor 	eax, eax
@@ -639,7 +609,10 @@ proc alloc_dptr stdcall, ns:dword, prps_ptr:dword, numsectors:dword, nprps_ptr:d
 	mov 	esi, [nprps_ptr]
 	dec 	eax
 	mov 	dword [esi], eax
-	stdcall build_prp_list, eax
+	mov 	edx, dword [edi]
+	and 	edx, PAGE_SIZE - 1
+	add 	edx, PAGE_SIZE
+	stdcall build_prp_list, eax, edx
 	test 	eax, eax
 	jz 	.err
 	DEBUGF  DBG_INFO, "Successfully allocated PRP list at: %x\n", eax
@@ -647,12 +620,6 @@ proc alloc_dptr stdcall, ns:dword, prps_ptr:dword, numsectors:dword, nprps_ptr:d
 	jmp 	.success
 	
 .err:
-	mov 	eax, dword [edi]
-	test 	eax, eax
-	jz 	@f
-	invoke  FreePage, eax
-
-@@:
 	xor 	eax, eax
 	pop 	edi esi
 	ret