dev, kernel: AHCI support and HeFS improvements, alongside scheduler allocation tree fix.

Signed-off-by: Amlal <amlal@nekernel.org>

7 files changed, 266 insertions, 146 deletions

diff --git a/dev/kernel/FSKit/HeFS.h b/dev/kernel/FSKit/HeFS.h
index c18f05d4..fd462b17 100644
--- a/dev/kernel/FSKit/HeFS.h
+++ b/dev/kernel/FSKit/HeFS.h
@@ -39,6 +39,7 @@
 struct HEFS_BOOT_NODE;

 struct HEFS_INDEX_NODE;

 struct HEFS_INDEX_NODE_DIRECTORY;

+struct HEFS_JOURNAL_NODE;

 

 enum {

   kHeFSHardDrive         = 0xC0,  // Hard Drive

@@ -162,6 +163,7 @@ struct PACKED ALIGN(8) HEFS_INDEX_NODE final {
 };

 

 enum {

+  kHeFSInvalidColor = 0,

   kHeFSRed = 100,

   kHeFSBlack,

   kHeFSColorCount,

diff --git a/dev/kernel/HALKit/AMD64/HalPagingMgrAMD64.cc b/dev/kernel/HALKit/AMD64/HalPagingMgrAMD64.cc
index 8e80c4d8..d7a8baf6 100644
--- a/dev/kernel/HALKit/AMD64/HalPagingMgrAMD64.cc
+++ b/dev/kernel/HALKit/AMD64/HalPagingMgrAMD64.cc
@@ -90,7 +90,7 @@ UIntPtr hal_get_phys_address(VoidPtr virt) {
 
   mmi_page_status((Detail::PTE*) pte);
 
-  return pte->PhysicalAddress;
+  return (pte->PhysicalAddress << 12) | (kVMAddr & 0xFFF);
 }
 
 /***********************************************************************************/
@@ -125,16 +125,21 @@ EXTERN_C Int32 mm_map_page(VoidPtr virtual_address, VoidPtr physical_address, UI
   UInt64*      pt  = reinterpret_cast<UInt64*>(pde & ~kPageMask);
   Detail::PTE* pte = (Detail::PTE*) pt[(kVMAddr >> 12) & kMask9];
 
-  pte->Present         = !!(flags & kMMFlagsPresent);
-  pte->Wr              = !!(flags & kMMFlagsWr);
-  pte->User            = !!(flags & kMMFlagsUser);
-  pte->Nx              = !!(flags & kMMFlagsNX);
-  pte->Pcd             = !!(flags & kMMFlagsPCD);
-  pte->Pwt             = !!(flags & kMMFlagsPwt);
-  pte->PhysicalAddress = ((UIntPtr) (physical_address));
+  pte->Present = !!(flags & kMMFlagsPresent);
+  pte->Wr      = !!(flags & kMMFlagsWr);
+  pte->User    = !!(flags & kMMFlagsUser);
+  pte->Nx      = !!(flags & kMMFlagsNX);
+  pte->Pcd     = !!(flags & kMMFlagsPCD);
+  pte->Pwt     = !!(flags & kMMFlagsPwt);
+
+  pte->PhysicalAddress = ((UIntPtr) (physical_address)) >> 12;
 
   hal_invl_tlb(virtual_address);
 
+  asm volatile("clflush (%0)" : : "r"(virtual_address) : "memory");
+
+  asm volatile("mfence" ::: "memory");
+
   mmi_page_status(pte);
 
   return kErrorSuccess;
diff --git a/dev/kernel/HALKit/AMD64/Storage/AHCI+Generic.cc b/dev/kernel/HALKit/AMD64/Storage/AHCI+Generic.cc
index aa9ab2cb..a28fa0ab 100644
--- a/dev/kernel/HALKit/AMD64/Storage/AHCI+Generic.cc
+++ b/dev/kernel/HALKit/AMD64/Storage/AHCI+Generic.cc
@@ -24,6 +24,7 @@
 #include <KernelKit/ProcessScheduler.h>
 #include <NewKit/Utils.h>
 #include <StorageKit/AHCI.h>
+#include <StorageKit/DMA.h>
 #include <modules/AHCI/AHCI.h>
 #include <modules/ATA/ATA.h>
 
@@ -54,10 +55,11 @@ using namespace Kernel;
 
 STATIC PCI::Device kSATADev;
 STATIC HbaMemRef   kSATAHba;
-STATIC Lba         kSATASectorCount      = 0UL;
-STATIC UInt16      kSATAIndex            = 0U;
-STATIC Char        kCurrentDiskModel[50] = {"GENERIC SATA"};
-STATIC UInt16      kSATAPortsImplemented = 0U;
+STATIC Lba         kSATASectorCount                                 = 0UL;
+STATIC UInt16      kSATAIndex                                       = 0U;
+STATIC Char        kCurrentDiskModel[50]                            = {"GENERIC SATA"};
+STATIC UInt16      kSATAPortsImplemented                            = 0U;
+STATIC             ALIGN(4096) UInt8 kIdentifyData[kAHCISectorSize] = {0};
 
 template <BOOL Write, BOOL CommandOrCTRL, BOOL Identify>
 STATIC Void drv_std_input_output_ahci(UInt64 lba, UInt8* buffer, SizeT sector_sz,
@@ -67,42 +69,37 @@ STATIC Int32 drv_find_cmd_slot_ahci(HbaPort* port) noexcept;
 
 STATIC Void drv_compute_disk_ahci() noexcept;
 
+/// @brief Identify device and read LBA info, Disk OEM vendor...
 STATIC Void drv_compute_disk_ahci() noexcept {
   kSATASectorCount = 0UL;
 
-  /// Normally 512 bytes, but add an additional 512 bytes to make 1 KIB.
-  const UInt16 kSzIdent = 512;
+  rt_set_memory(kIdentifyData, 0, kAHCISectorSize);
 
-  /// Push it to the stack
-  UInt8* identify_data = new UInt8[kSzIdent];
+  drv_std_input_output_ahci<NO, YES, YES>(0, kIdentifyData, kAHCISectorSize, kAHCISectorSize);
 
-  /// Send AHCI command for identification.
-  drv_std_input_output_ahci<NO, YES, YES>(0, identify_data, kAHCISectorSize, kSzIdent);
+  // --> Reinterpret the 512-byte buffer as an array of 256 UInt16 words
+  UInt16* identify_words = reinterpret_cast<UInt16*>(kIdentifyData);
 
   /// Extract 48-bit LBA.
-  UInt64 lba48_sectors = 0;
-  lba48_sectors |= (UInt64) identify_data[100];
-  lba48_sectors |= (UInt64) identify_data[101] << 16;
-  lba48_sectors |= (UInt64) identify_data[102] << 32;
+  UInt64 lba48_sectors = 0UL;
+  lba48_sectors |= (UInt64) identify_words[100];
+  lba48_sectors |= (UInt64) identify_words[101] << 16;
+  lba48_sectors |= (UInt64) identify_words[102] << 32;
 
-  /// Now verify if lba48
   if (lba48_sectors == 0)
-    kSATASectorCount = (identify_data[61] << 16) | identify_data[60];
+    kSATASectorCount = (identify_words[61] << 16) | identify_words[60];
   else
     kSATASectorCount = lba48_sectors;
 
   for (Int32 i = 0; i < 20; i++) {
-    kCurrentDiskModel[i * 2]     = (identify_data[27 + i] >> 8) & 0xFF;
-    kCurrentDiskModel[i * 2 + 1] = identify_data[27 + i] & 0xFF;
+    kCurrentDiskModel[i * 2]     = (identify_words[27 + i] >> 8) & 0xFF;
+    kCurrentDiskModel[i * 2 + 1] = identify_words[27 + i] & 0xFF;
   }
 
   kCurrentDiskModel[40] = '\0';
 
   (Void)(kout << "SATA Sector Count: " << hex_number(kSATASectorCount) << kendl);
   (Void)(kout << "SATA Disk Model: " << kCurrentDiskModel << kendl);
-
-  delete[] identify_data;
-  identify_data = nullptr;
 }
 
 /// @brief Finds a command slot for a HBA port.
@@ -127,81 +124,108 @@ STATIC Int32 drv_find_cmd_slot_ahci(HbaPort* port) noexcept {
 template <BOOL Write, BOOL CommandOrCTRL, BOOL Identify>
 STATIC Void drv_std_input_output_ahci(UInt64 lba, UInt8* buffer, SizeT sector_sz,
                                       SizeT size_buffer) noexcept {
-  UIntPtr slot = 0UL;
-
-  slot = drv_find_cmd_slot_ahci(&kSATAHba->Ports[kSATAIndex]);
+  NE_UNUSED(sector_sz);
 
-  if (slot == ~0UL) {
+  if (!buffer || size_buffer == 0) {
+    kout << "Invalid buffer for AHCI I/O.\r";
     err_global_get() = kErrorDisk;
     return;
   }
 
-  if (size_buffer > mib_cast(4) || sector_sz > kAHCISectorSize) return;
+  UIntPtr slot = drv_find_cmd_slot_ahci(&kSATAHba->Ports[kSATAIndex]);
 
-  if (!Write) {
-    // Zero-memory the buffer field.
-    rt_set_memory(buffer, 0, size_buffer);
+  if (slot == ~0UL) {
+    kout << "No free command slot!\r";
+    err_global_get() = kErrorDisk;
+    return;
   }
 
-  /// prepare command header.
   volatile HbaCmdHeader* command_header =
-      ((volatile HbaCmdHeader*) (((UInt64) kSATAHba->Ports[kSATAIndex].Clb)));
-
-  /// Offset to specific command slot.
+      (volatile HbaCmdHeader*) ((UInt64) kSATAHba->Ports[kSATAIndex].Clb);
   command_header += slot;
 
-  /// check for command header.
   MUST_PASS(command_header);
 
-  command_header->Struc.Cfl   = sizeof(FisRegH2D) / sizeof(UInt32);
-  command_header->Struc.Write = Write;
-  command_header->Prdtl       = 8;
+  // Clear old command table memory
+  volatile HbaCmdTbl* command_table =
+      (volatile HbaCmdTbl*) (((UInt64) command_header->Ctbau << 32) | command_header->Ctba);
+  rt_set_memory((VoidPtr) command_table, 0, sizeof(HbaCmdTbl));
 
-  auto ctba_phys     = ((UInt64) command_header->Ctbau << 32) | command_header->Ctba;
-  auto command_table = reinterpret_cast<volatile HbaCmdTbl*>(ctba_phys);
+  VoidPtr ptr = Kernel::rtl_dma_alloc(size_buffer, 4096);
 
-  MUST_PASS(command_table);
-
-  UIntPtr buffer_phys     = HAL::hal_get_phys_address(buffer);
+  // Build the PRDT
   SizeT   bytes_remaining = size_buffer;
+  SizeT   prdt_index      = 0;
+  UIntPtr buffer_phys     = (UIntPtr) ptr;
+
+  while (bytes_remaining > 0 && prdt_index < 8) {
+    SizeT chunk_size = bytes_remaining;
+    if (chunk_size > 8 * 1024)  // AHCI recommends ~8 KiB per PRDT
+      chunk_size = 8 * 1024;
+
+    command_table->Prdt[prdt_index].Dba  = (UInt32) (buffer_phys & 0xFFFFFFFF);
+    command_table->Prdt[prdt_index].Dbau = (UInt32) (buffer_phys >> 32);
+    command_table->Prdt[prdt_index].Dbc  = (UInt32) (chunk_size - 1);
+    command_table->Prdt[prdt_index].Ie   = NO;
+
+    buffer_phys += chunk_size;
+    bytes_remaining -= chunk_size;
+    ++prdt_index;
+  }
 
-  command_table->Prdt[0].Dba  = (UInt32) (buffer_phys & 0xFFFFFFFF);
-  command_table->Prdt[0].Dbau = (UInt32) (buffer_phys >> 32);
-  command_table->Prdt[0].Dbc  = bytes_remaining - 1;
-  command_table->Prdt[0].Ie   = NO;
+  if (bytes_remaining > 0) {
+    kout << "Warning: AHCI PRDT overflow, cannot map full buffer.\r";
+    err_global_get() = kErrorDisk;
+    return;
+  }
+
+  command_header->Prdtl       = prdt_index;
+  command_header->Struc.Cfl   = sizeof(FisRegH2D) / sizeof(UInt32);
+  command_header->Struc.Write = Write;
 
   volatile FisRegH2D* h2d_fis = (volatile FisRegH2D*) (&command_table->Cfis[0]);
 
   h2d_fis->FisType   = kFISTypeRegH2D;
   h2d_fis->CmdOrCtrl = CommandOrCTRL;
   h2d_fis->Command =
-      (Identify ? (kAHCICmdIdentify) : (Write ? kAHCICmdWriteDmaEx : kAHCICmdReadDmaEx));
+      (Identify ? kAHCICmdIdentify : (Write ? kAHCICmdWriteDmaEx : kAHCICmdReadDmaEx));
 
-  h2d_fis->Lba0 = (lba) & 0xFF;
+  h2d_fis->Lba0 = (lba >> 0) & 0xFF;
   h2d_fis->Lba1 = (lba >> 8) & 0xFF;
   h2d_fis->Lba2 = (lba >> 16) & 0xFF;
-
-  h2d_fis->Device = kSATALBAMode;
-
   h2d_fis->Lba3 = (lba >> 24) & 0xFF;
   h2d_fis->Lba4 = (lba >> 32) & 0xFF;
   h2d_fis->Lba5 = (lba >> 40) & 0xFF;
 
-  h2d_fis->CountLow  = (size_buffer) & 0xFF;
-  h2d_fis->CountHigh = (size_buffer >> 8) & 0xFF;
+  if (Identify) {
+    h2d_fis->Device    = 0;  // DO NOT set LBAMODE flag
+    h2d_fis->CountLow  = 1;  // IDENTIFY always transfers 1 sector
+    h2d_fis->CountHigh = 0;
+  } else {
+    h2d_fis->Device    = kSATALBAMode;
+    h2d_fis->CountLow  = (size_buffer / kAHCISectorSize) & 0xFF;
+    h2d_fis->CountHigh = ((size_buffer / kAHCISectorSize) >> 8) & 0xFF;
+  }
+
+  rtl_dma_flush(ptr, size_buffer);
 
+  // Issue command
   kSATAHba->Ports[kSATAIndex].Ci = (1 << slot);
 
-  for (Int32 i = 0; i < 1000000; ++i) {
+  while (YES) {
     if (!(kSATAHba->Ports[kSATAIndex].Ci & (1 << slot))) break;
   }
 
+  rtl_dma_flush(ptr, size_buffer);
+
   if (kSATAHba->Is & kHBAErrTaskFile) {
+    kout << "AHCI Task File Error during I/O.\r";
     err_global_get() = kErrorDiskIsCorrupted;
     return;
+  } else {
+    rt_copy_memory(ptr, buffer, size_buffer);
+    err_global_get() = kErrorSuccess;
   }
-
-  err_global_get() = kErrorSuccess;
 }
 
 /***
@@ -251,6 +275,47 @@ STATIC BOOL ahci_enable_and_probe() {
   return YES;
 }
 
+STATIC Bool drv_init_command_structures_ahci() {
+  // Allocate 4KiB for Command List (32 headers)
+  VoidPtr clb_mem = Kernel::rtl_dma_alloc(4096, 1024);
+  if (!clb_mem) {
+    kout << "Failed to allocate CLB memory!\r";
+    return NO;
+  }
+
+  UIntPtr clb_phys = HAL::hal_get_phys_address(clb_mem);
+
+  kSATAHba->Ports[kSATAIndex].Clb  = (UInt32)(clb_phys & 0xFFFFFFFF);
+  kSATAHba->Ports[kSATAIndex].Clbu = (UInt32)(clb_phys >> 32);
+
+  // Clear it
+  rt_set_memory(clb_mem, 0, 4096);
+
+  // For each command slot (up to 32)
+  volatile HbaCmdHeader* header = (volatile HbaCmdHeader*) clb_mem;
+
+  for (Int32 i = 0; i < 32; ++i) {
+    // Allocate 4KiB for Command Table
+    VoidPtr ct_mem = Kernel::rtl_dma_alloc(4096, 128);
+    if (!ct_mem) {
+      kout << "Failed to allocate CTB memory for slot " << hex_number(i);
+      kout << "!\r";
+      return NO;
+    }
+
+    UIntPtr ct_phys = HAL::hal_get_phys_address(ct_mem);
+
+    header[i].Ctba  = (UInt32)(ct_phys & 0xFFFFFFFF);
+    header[i].Ctbau = (UInt32)(ct_phys >> 32);
+
+    // Clear the command table
+    rt_set_memory((VoidPtr) ct_mem, 0, 4096);
+  }
+
+  return YES;
+}
+
+
 /// @brief Initializes an AHCI disk.
 /// @param pi the amount of ports that have been detected.
 /// @param atapi reference value, tells whether we should detect ATAPI instead of SATA.
@@ -295,6 +360,11 @@ STATIC Bool drv_std_init_ahci(UInt16& pi, BOOL& atapi) {
           kSATAIndex = ahci_index;
           kSATAHba   = mem_ahci;
 
+          if (!drv_init_command_structures_ahci()) {
+            err_global_get() = kErrorDisk;
+            return NO;
+          }
+
           goto success_hba_fetch;
         }
 
diff --git a/dev/kernel/StorageKit/DMA.h b/dev/kernel/StorageKit/DMA.h
new file mode 100644
index 00000000..953683eb
--- /dev/null
+++ b/dev/kernel/StorageKit/DMA.h
@@ -0,0 +1,58 @@
+/*
+ * Copyright (c) 2025 Amlal El Mahrouss. All rights reserved.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ * THE SOFTWARE.
+ */
+
+#pragma once
+
+#include <KernelKit/DebugOutput.h>
+#include <NewKit/Defines.h>
+
+#ifdef __NE_AMD64__
+#define DMA_POOL_START (0x1000000)
+#define DMA_POOL_SIZE (0x1000000)
+
+namespace Kernel {
+inline UInt8* kDmaPoolPtr = (UInt8*) DMA_POOL_START;
+
+inline VoidPtr rtl_dma_alloc(SizeT size, SizeT align) {
+  UIntPtr addr = (UIntPtr) kDmaPoolPtr;
+
+  addr = (addr + (align - 1)) & ~(align - 1);  // Align up
+
+  if (addr + size >= DMA_POOL_START + DMA_POOL_SIZE) {
+    kout << "DMA Pool exhausted!\r";
+
+    return nullptr;
+  }
+
+  kDmaPoolPtr = (UInt8*) (addr + size);
+  return (VoidPtr) addr;
+}
+
+inline Void rtl_dma_flush(Void* ptr, SizeT size_buffer) {
+  for (SizeT i = 0; i < size_buffer; ++i) {
+    asm volatile("clflush (%0)" : : "r"((UInt8*) ptr + i) : "memory");
+  }
+  asm volatile("mfence" ::: "memory");
+}
+}  // namespace Kernel
+#else
+#endif
\ No newline at end of file
diff --git a/dev/kernel/src/BitMapMgr.cc b/dev/kernel/src/BitMapMgr.cc
index 382d0be5..cc5797b3 100644
--- a/dev/kernel/src/BitMapMgr.cc
+++ b/dev/kernel/src/BitMapMgr.cc
@@ -62,8 +62,6 @@ namespace HAL {
 
         if (user) flags |= kMMFlagsUser;
 
-        flags |= HAL::kMMFlagsPCD;
-
         return flags;
       }
 
diff --git a/dev/kernel/src/FS/HeFS.cc b/dev/kernel/src/FS/HeFS.cc
index d7cce9d6..eb6e6fc1 100644
--- a/dev/kernel/src/FS/HeFS.cc
+++ b/dev/kernel/src/FS/HeFS.cc
@@ -79,8 +79,7 @@ namespace Detail {
     NE_UNUSED(node);
 
     if (!dir || !node) {
-      ke_panic(RUNTIME_CHECK_FILESYSTEM,
-               "Error: Invalid directory node/boot_node in RB-Tree traversal.");
+      return;
     }
 
     if (dir->fChild != 0) {
@@ -92,7 +91,7 @@ namespace Detail {
     } else if (dir->fPrev != 0) {
       start = dir->fPrev;
     } else {
-      start = node->fStartIND;
+      start = 0;
     }
   }
 
@@ -218,7 +217,7 @@ namespace Detail {
       auto start = root->fStartIND;
       auto end   = root->fEndIND;
 
-      auto hop_watch = 0;
+      auto hop_watch = 0UL;
 
       while (YES) {
         if (start > end) {
@@ -226,16 +225,11 @@ namespace Detail {
           break;
         }
 
-        if (hop_watch > 100) {
+        if (hop_watch++ > (root->fEndIND / sizeof(HEFS_INDEX_NODE_DIRECTORY))) {
           kout << "Error: Hop watch exceeded, filesystem is stalling.\r";
           break;
         }
 
-        if (start == 0) {
-          ++hop_watch;
-          start = root->fStartIND;
-        }
-
         mnt->fPacket.fPacketLba     = start;
         mnt->fPacket.fPacketSize    = sizeof(HEFS_INDEX_NODE_DIRECTORY);
         mnt->fPacket.fPacketContent = dir;
@@ -316,11 +310,11 @@ namespace Detail {
           return NO;
         }
 
-        if (tmpdir->fDeleted || !tmpdir->fCreated) {
-          dir->fChild = tmpdir->fChild;
+        if (tmpdir->fDeleted) {
+          dir->fChild  = tmpdir->fChild;
           dir->fParent = tmpdir->fParent;
-          dir->fNext = tmpdir->fNext;
-          dir->fPrev = tmpdir->fPrev;
+          dir->fNext   = tmpdir->fNext;
+          dir->fPrev   = tmpdir->fPrev;
 
           mnt->fPacket.fPacketLba     = start;
           mnt->fPacket.fPacketSize    = sizeof(HEFS_INDEX_NODE_DIRECTORY);
@@ -331,7 +325,9 @@ namespace Detail {
           return YES;
         }
 
+        auto old_start = start;
         hefsi_traverse_tree(tmpdir, root, start);
+        if (start == 0 || start == old_start) break;
       }
 
       return YES;
@@ -349,7 +345,7 @@ namespace Detail {
       auto start = root->fStartIND;
       auto end   = root->fEndIND;
 
-      auto hop_watch = 0;
+      auto hop_watch = 0UL;
 
       while (YES) {
         if (start > end) {
@@ -357,16 +353,11 @@ namespace Detail {
           break;
         }
 
-        if (hop_watch > 100) {
+        if (hop_watch++ > (root->fEndIND / sizeof(HEFS_INDEX_NODE_DIRECTORY))) {
           kout << "Error: Hop watch exceeded, filesystem is stalling.\r";
           break;
         }
 
-        if (start == 0) {
-          ++hop_watch;
-          start = root->fStartIND;
-        }
-
         mnt->fPacket.fPacketLba     = start;
         mnt->fPacket.fPacketSize    = sizeof(HEFS_INDEX_NODE_DIRECTORY);
         mnt->fPacket.fPacketContent = dir;
@@ -390,7 +381,9 @@ namespace Detail {
           }
         }
 
+        auto old_start = start;
         hefsi_traverse_tree(dir, root, start);
+        if (start == 0 || start == old_start) break;
       }
 
       delete dir;
@@ -427,7 +420,7 @@ namespace Detail {
 
       auto start_cnt = 0UL;
 
-      auto hop_watch = 0;
+      auto hop_watch = 0UL;
 
       while (YES) {
         if (start > end) {
@@ -435,16 +428,11 @@ namespace Detail {
           break;
         }
 
-        if (hop_watch > 100) {
+        if (hop_watch++ > (root->fEndIND / sizeof(HEFS_INDEX_NODE_DIRECTORY))) {
           kout << "Error: Hop watch exceeded, filesystem is stalling.\r";
           break;
         }
 
-        if (start == 0) {
-          ++hop_watch;
-          start = root->fStartIND;
-        }
-
         mnt->fPacket.fPacketLba     = start;
         mnt->fPacket.fPacketSize    = sizeof(HEFS_INDEX_NODE_DIRECTORY);
         mnt->fPacket.fPacketContent = dir;
@@ -507,7 +495,9 @@ namespace Detail {
           }
         }
 
+        auto old_start = start;
         hefsi_traverse_tree(dir, root, start);
+        if (start == 0 || start == old_start) break;
       }
 
       delete dir;
@@ -539,19 +529,16 @@ namespace Detail {
 
       auto start = root->fStartIND;
 
-      auto hop_watch = 0;
+      auto hop_watch = 0UL;
 
       while (YES) {
-        if (hop_watch > 100) {
+        if (start > root->fEndIND) break;
+
+        if (hop_watch++ > (root->fEndIND / sizeof(HEFS_INDEX_NODE_DIRECTORY))) {
           kout << "Error: Hop watch exceeded, filesystem is stalling.\r";
           break;
         }
 
-        if (start == 0) {
-          ++hop_watch;
-          start = root->fStartIND;
-        }
-
         mnt->fPacket.fPacketLba     = start;
         mnt->fPacket.fPacketSize    = sizeof(HEFS_INDEX_NODE_DIRECTORY);
         mnt->fPacket.fPacketContent = dir;
@@ -590,7 +577,11 @@ namespace Detail {
           }
         }
 
+        auto old_start = start;
+
         hefsi_traverse_tree(dir, root, start);
+
+        if (start == 0 || start == old_start) break;
       }
 
       delete dir;
@@ -614,9 +605,10 @@ namespace Detail {
 
       auto start = root->fStartIND;
 
+      SizeT hop_watch = 0UL;
+
       while (YES) {
-        if (start > root->fEndIND)
-          break;
+        if (start > root->fEndIND) break;
 
         mnt->fPacket.fPacketLba     = start;
         mnt->fPacket.fPacketSize    = sizeof(HEFS_INDEX_NODE_DIRECTORY);
@@ -624,6 +616,11 @@ namespace Detail {
 
         mnt->fInput(mnt->fPacket);
 
+        if (hop_watch++ > (root->fSectorCount)) {  // <-- NEW (stall protection)
+          kout << "Warning: Traversal stalled during balancing.\r";
+          break;  // Exit cleanly instead of hanging forever
+        }
+
         if (!mnt->fPacket.fPacketGood) {
           delete dir;
           dir = nullptr;
@@ -700,14 +697,21 @@ namespace Detail {
             return NO;
           }
 
+          auto old_start = start;
           hefsi_traverse_tree(dir, root, start);
 
+          if (start == 0 || start == old_start) break;
+
           continue;
         } else {
           if (dir_parent->fNext == start) {
             hefsi_rotate_left(dir, start, mnt);
+
+            auto old_start = start;
             hefsi_traverse_tree(dir, root, start);
 
+            if (start == 0 || start == old_start) break;
+
             continue;
           }
 
@@ -729,12 +733,12 @@ namespace Detail {
           }
 
           hefsi_rotate_right(dir, start, mnt);
-          hefsi_traverse_tree(dir, root, start);
-
-          continue;
         }
 
+        auto old_start = start;
         hefsi_traverse_tree(dir, root, start);
+        
+        if (start == 0 || start == old_start) break;
       }
 
       delete dir;
@@ -852,8 +856,6 @@ _Output Bool HeFileSystemParser::Format(_Input _Output DriveTrait* drive, _Input
   drive->fPacket.fPacketSize    = sizeof(HEFS_BOOT_NODE);
   drive->fPacket.fPacketContent = root;
 
-  HEFS_BOOT_NODE root_cpy = *root;
-
   drive->fOutput(drive->fPacket);
 
   if (!drive->fPacket.fPacketGood) {
@@ -865,20 +867,20 @@ _Output Bool HeFileSystemParser::Format(_Input _Output DriveTrait* drive, _Input
     return NO;
   }
 
-  SizeT cnt = kHeFSBlockCount;
+  /// @note this allocates 4 ind at format.
+  SizeT cnt = 4UL;
 
-  Lba next = root_cpy.fStartIND;
-  Lba prev = next;
+  Lba next = drive->fLbaStart + sizeof(HEFS_BOOT_NODE) + sizeof(HEFS_BOOT_NODE);
 
   HEFS_INDEX_NODE_DIRECTORY* index_node = new HEFS_INDEX_NODE_DIRECTORY();
-  
+
   // Pre-allocate index node directory tree
   for (SizeT i = 0; i < cnt; ++i) {
     rt_set_memory(index_node, 0, sizeof(HEFS_INDEX_NODE_DIRECTORY));
     wrt_copy_memory((VoidPtr) u"/$", index_node->fName, wrt_string_len(u"/$"));
 
-    index_node->fFlags   = flags;
-    index_node->fKind    = kHeFSFileKindDirectory;
+    index_node->fFlags = flags;
+    index_node->fKind  = kHeFSFileKindDirectory;
 
     index_node->fDeleted = kHeFSTimeMax;
 
@@ -890,23 +892,25 @@ _Output Bool HeFileSystemParser::Format(_Input _Output DriveTrait* drive, _Input
     index_node->fUID   = 0;
     index_node->fGID   = 0;
     index_node->fMode  = 0;
-    index_node->fColor = kHeFSBlack;
 
-    index_node->fParent = next;
+    index_node->fColor = kHeFSBlack;
     index_node->fChild  = 0;
-    index_node->fNext = next + sizeof(HEFS_INDEX_NODE_DIRECTORY);
-    index_node->fPrev = prev;
+    index_node->fParent = 0;
+    index_node->fNext   = 0;
+    index_node->fPrev   = 0;
 
-    drive->fPacket.fPacketLba  = next;
-    drive->fPacket.fPacketSize = sizeof(HEFS_INDEX_NODE_DIRECTORY);
+    drive->fPacket.fPacketLba     = next;
+    drive->fPacket.fPacketSize    = sizeof(HEFS_INDEX_NODE_DIRECTORY);
     drive->fPacket.fPacketContent = index_node;
 
-    prev = next;
     next += sizeof(HEFS_INDEX_NODE_DIRECTORY);
 
     drive->fOutput(drive->fPacket);
   }
 
+  delete index_node;
+  index_node = nullptr;
+
   // Create the directories, something UNIX inspired but more explicit and forward looking.
 
   this->CreateDirectory(drive, kHeFSEncodingUTF16, u"/boot");
@@ -921,7 +925,7 @@ _Output Bool HeFileSystemParser::Format(_Input _Output DriveTrait* drive, _Input
   this->CreateFile(drive, kHeFSEncodingBinary, u"/", u".hefs");
 
   delete root;
-  root     = nullptr;
+  root = nullptr;
 
   if (drive->fPacket.fPacketGood) return YES;
 
@@ -967,37 +971,18 @@ _Output Bool HeFileSystemParser::CreateDirectory(_Input DriveTrait* drive, _Inpu
 
   Detail::hefsi_balance_filesystem(root, drive);
 
-  auto dirent = Detail::hefs_fetch_index_node_directory(root, drive, dir);
-
-  if (dirent) {
-    kout << "Error: Directory already exists.\r";
-
-    delete dirent;
-    dirent = nullptr;
-
-    delete node;
-    delete root;
-
-    return NO;
-  }
-
-  dirent = new HEFS_INDEX_NODE_DIRECTORY();
+  auto dirent = new HEFS_INDEX_NODE_DIRECTORY();
 
   rt_set_memory(dirent, 0, sizeof(HEFS_INDEX_NODE_DIRECTORY));
 
   wrt_copy_memory((VoidPtr) dir, dirent->fName, wrt_string_len(dir));
 
   dirent->fAccessed   = 0;
-  dirent->fCreated    = 0;  /// TODO: Add the current time.
+  dirent->fCreated    = kHeFSTimeMax;  /// TODO: Add the current time.
   dirent->fDeleted    = 0;
   dirent->fModified   = 0;
   dirent->fEntryCount = 0;
 
-  dirent->fParent = root->fStartIND;  // No parent (it's the real root)
-  dirent->fChild  = root->fEndIND;    // No children yet
-  dirent->fNext   = 0;                // No next
-  dirent->fPrev   = 0;                // No previous
-
   dirent->fKind     = kHeFSFileKindDirectory;
   dirent->fFlags    = flags;
   dirent->fChecksum = 0;
diff --git a/dev/kernel/src/UserProcessScheduler.cc b/dev/kernel/src/UserProcessScheduler.cc
index 5d983f34..96938d10 100644
--- a/dev/kernel/src/UserProcessScheduler.cc
+++ b/dev/kernel/src/UserProcessScheduler.cc
@@ -108,11 +108,13 @@ STATIC USER_PROCESS::USER_HEAP_TREE* sched_try_go_upper_heap_tree(
   tree = tree->MemoryParent;
 
   if (tree) {
-    tree = tree->MemoryParent;
+    auto tree_tmp = tree->MemoryNext;
 
-    if (!tree) {
-      return nullptr;
+    if (!tree_tmp) {
+      return tree;
     }
+
+    return tree_tmp;
   }
 
   return tree;