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;