From 690389e6d1efef1dee2bf2cfa9c33fed50ffc06d Mon Sep 17 00:00:00 2001
From: Riccardo Mori <patacca@autistici.org>
Date: Thu, 7 Mar 2024 10:33:52 +0100
Subject: [PATCH] Consider the first chunk that we encounter as valid

---
 include/quokka/Layout.h |  4 +++-
 src/Layout.cpp          | 19 ++++++++++++++++---
 2 files changed, 19 insertions(+), 4 deletions(-)

diff --git a/include/quokka/Layout.h b/include/quokka/Layout.h
index 6bed67c..c489340 100644
--- a/include/quokka/Layout.h
+++ b/include/quokka/Layout.h
@@ -271,8 +271,10 @@ class HeadIterator {
   /**
    * Retrieve the next chunk address if any are found
    * @param address Address to start from
+   * @param skip_current Do not consider the chunk starting at the current
+   * address
    */
-  void SetNextChunk(ea_t address);
+  void SetNextChunk(ea_t address, bool skip_current = true);
 
   /**
    * Compute the next address and state
diff --git a/src/Layout.cpp b/src/Layout.cpp
index d9db010..df5e887 100644
--- a/src/Layout.cpp
+++ b/src/Layout.cpp
@@ -61,11 +61,24 @@ void HeadIterator::InitAddresses(ea_t address) {
   this->next_unk_addr = next_unknown(address, this->max_ea);
   this->next_head_addr = next_head(address, this->max_ea);
 
-  this->SetNextChunk(address);
+  this->SetNextChunk(address, false);
 }
 
-void HeadIterator::SetNextChunk(ea_t address) {
-  func_t* func = get_next_fchunk(address);
+void HeadIterator::SetNextChunk(ea_t address, bool skip_current /* = true*/) {
+  func_t* func;
+
+  // Check if there is a chunk starting at provided address
+  if (!skip_current) {
+    func = get_fchunk(address);
+    if (func != nullptr && func->start_ea == address) {
+      this->next_chunk_addr = address;
+      this->next_func_chunk = func;
+      return;
+    }
+  }
+
+  // Find the next chunk, not considering the one at the current address
+  func = get_next_fchunk(address);
   if (func != nullptr) {
     this->next_chunk_addr = func->start_ea;
     this->next_func_chunk = func;