From e2e1e5049096c90f2d1e514b9771829ce2f08fb8 Mon Sep 17 00:00:00 2001
From: Abdur-Rahman Mansoor <armansoor@missingno.dev>
Date: Thu, 23 May 2024 19:15:34 -0400
Subject: [PATCH] implement more stuff

---
 drivers/nvme/nvme.asm | 58 +++++++++++++++++++++++++++++++++----------
 drivers/nvme/nvme.inc | 10 ++++++--
 2 files changed, 53 insertions(+), 15 deletions(-)

diff --git a/drivers/nvme/nvme.asm b/drivers/nvme/nvme.asm
index d1e4ad6..b0b3607 100644
--- a/drivers/nvme/nvme.asm
+++ b/drivers/nvme/nvme.asm
@@ -242,31 +242,26 @@ proc nvme_init stdcall, pci:dword
 		PDEBUGF  DBG_INFO, "PCI(%u.%u.%u): NVMe contiguous queues required\n", byte [pci + pcidev.bus], byte [pci + pcidev.devfn]
 
 .cqr_not_req:
+		; For some reason, bit 7 (No I/O command set supported) is also set to 1 despite bit 0 (NVM command set)
+	    ; being set to 1.. so I am not sure if bit 7 should be checked at all.. investigate later.
 		mov      ebx, dword [eax + NVME_REG_MAP.CAP + 4]
-		mov      ecx, ebx
-		test     ebx, CAP_CSS_NVM_CMDSET
+		test 	 ebx, CAP_CSS_NVM_CMDSET
 		jz       .exit_fail
 
 		if __DEBUG__
+			mov 	 ecx, ebx
 			and      ebx, CAP_MPSMIN
 			and      ecx, CAP_MPSMAX
 			shr      ebx, 16
 			shr      ecx, 16
 			PDEBUGF  DBG_INFO, "PCI(%u.%u.%u): NVMe memory page size minimum: %u\n", byte [pci + pcidev.bus], byte [pci + pcidev.devfn], ebx
 			PDEBUGF  DBG_INFO, "PCI(%u.%u.%u): NVMe memory page size maximum: %u\n", byte [pci + pcidev.bus], byte [pci + pcidev.devfn], ecx
-
-			mov      ebx, dword [eax + NVME_REG_MAP.CC]
-			mov      ecx, ebx
-			and      ebx, CC_IOSQES
-			and      ecx, CC_IOCQES
-			shl      ebx, 16
-			shl      ecx, 16
-
-			; TODO: Change entry sizes to their appropriate values
-			PDEBUGF  DBG_INFO, "PCI(%u.%u.%u): NVMe I/O submission queue entry size: %u\n", byte [pci + pcidev.bus], byte [pci + pcidev.devfn], ebx
-			PDEBUGF  DBG_INFO, "PCI(%u.%u.%u): NVMe I/O completion queue entry size: %u\n", byte [pci + pcidev.bus], byte [pci + pcidev.devfn], ecx
 		end if
 
+		; Reset controller
+		mov 	ebx, dword [eax + NVME_REG_MAP.CC]
+		stdcall nvme_controller_reset, [pci] 
+
         xor     eax, eax
         inc     eax
 		pop     ecx ebx
@@ -280,6 +275,43 @@ proc nvme_init stdcall, pci:dword
 
 endp
 
+proc nvme_controller_reset stdcall, pci:dword
+
+		push 	ebx
+		PDEBUGF DBG_INFO, "PCI(%u.%u.%u): Resetting NVMe controller\n", byte [pci + pcidev.bus], byte [pci + pcidev.devfn]
+		mov 	eax, dword [pci + pcidev.mmio_ptr]
+	    mov 	ebx, dword [eax + NVME_REG_MAP.CSTS]
+
+; Wait for controller to be brought to idle state, CSTS.RDY should be cleared to 0 when this happens
+.wait:
+		test    ebx, CSTS_RDY
+		jnz 	.wait
+		mov 	ebx, dword [eax + NVME_REG_MAP.CC]
+		PDEBUGF DBG_INFO, "PCI(%u.%u.%u): NVMe.CC = 0x%x\n", byte [pci + pcidev.bus], byte [pci + pcidev.devfn], ebx
+		and 	ebx, 0xfffffffe
+		mov 	dword [eax + NVME_REG_MAP.CC], ebx
+		stdcall nvme_wait, [pci]
+		pop 	ebx
+		ret
+
+endp
+
+proc nvme_wait stdcall, pci:dword
+
+		push 	esi
+		mov 	eax, dword [pci + pcidev.mmio_ptr]
+		mov 	eax, dword [eax + NVME_REG_MAP.CAP]
+		and 	eax, CAP_TO
+		shr 	eax, 24
+		DEBUGF  DBG_INFO, "CSTS.TO = %u\n", eax
+		mov 	esi, eax
+		;imul 	esi, 500
+		;invoke  Sleep
+		pop 	esi
+		ret
+
+endp
+
 proc nvme_cleanup
 
 		DEBUGF  DBG_INFO, "(NVMe): Cleaning up...\n"
diff --git a/drivers/nvme/nvme.inc b/drivers/nvme/nvme.inc
index 9472425..86ba5d7 100644
--- a/drivers/nvme/nvme.inc
+++ b/drivers/nvme/nvme.inc
@@ -89,7 +89,6 @@ CAP_TO = 0xff000000
 CAP_DSTRD = 1 or (1 shl 1) or (1 shl 2) or (1 shl 3)
 CAP_NSSRS = 1 shl 4
 CAP_CSS_NVM_CMDSET = 1 shl 5
-CAP_CSS_OMCS = 1 shl 11
 CAP_CSS_NOIO = 1 shl 12
 CAP_BPS = 1 shl 14
 CAP_CPS_COSCOP = 1 shl 15
@@ -104,7 +103,7 @@ CAP_CRMS_CRWMS = 1 shl 28
 CAP_CRMS_CRIMS = 1 shl 29
 
 ; Controller Configuration Bits
-CC_ENABLE = 1
+CC_EN = 1
 CC_CSS = (1 shl 4) or (1 shl 5) or (1 shl 6)
 CC_MPS = (1 shl 7) or (1 shl 8) or (1 shl 9) or (1 shl 10)
 CC_AMS = (1 shl 11) or (1 shl 12) or (1 shl 13)
@@ -157,6 +156,13 @@ CQ_STATUS_SC_MADIE_URE   = 0x81 ; Unrecovered Read Error
 CQ_STATUS_SC_MADIE_ACDEN = 0x86 ; Access Denied
 CQ_STATUS_SC_MADIE_DOULB = 0x87 ; Deallocated or Unwritten Logical Block
 
+; Controller Status (CSTS) Values
+CSTS_RDY = 1
+CSTS_CFS = 1 shl 1
+CSTS_SHST = (1 shl 2) or (1 shl 3)
+CSTS_NSSRO = 1 shl 4
+CSTS_PP = 1 shl 5
+
 struct NVME_REG_MAP 
 	CAP      rq 1 ; Controller Capabilities
 	VS       rd 1 ; Version