<commit_before>package org.iq80.leveldb.util;

import com.google.common.base.Preconditions;
import com.google.common.primitives.Ints;
import org.iq80.leveldb.impl.InternalKey;
import org.iq80.leveldb.impl.MemTable.MemTableIterator;
import org.iq80.leveldb.impl.SeekingIterator;

import java.util.Comparator;
import java.util.Iterator;
import java.util.List;
import java.util.Map.Entry;
import java.util.NoSuchElementException;

public final class DbIterator extends AbstractSeekingIterator<InternalKey, Slice>
{
    private final MemTableIterator memTableIterator;
    private final MemTableIterator immutableMemTableIterator;
    private final List<InternalTableIterator> level0Files;
    private final List<LevelIterator> levels;

    private final Comparator<InternalKey> comparator;

    private final ComparableIterator[] heap;
    private int heapSize = 0;

    public DbIterator(MemTableIterator memTableIterator,
            MemTableIterator immutableMemTableIterator,
            List<InternalTableIterator> level0Files,
            List<LevelIterator> levels,
            Comparator<InternalKey> comparator)
    {
        this.memTableIterator = memTableIterator;
        this.immutableMemTableIterator = immutableMemTableIterator;
        this.level0Files = level0Files;
        this.levels = levels;
        this.comparator = comparator;

        this.heap = new ComparableIterator[3 + level0Files.size() + levels.size()];
        resetPriorityQueue();
    }

    @Override
    protected void seekToFirstInternal()
    {
        if (memTableIterator != null) {
            memTableIterator.seekToFirst();
        }
        if (immutableMemTableIterator != null) {
            immutableMemTableIterator.seekToFirst();
        }
        for (InternalTableIterator level0File : level0Files) {
            level0File.seekToFirst();
        }
        for (LevelIterator level : levels) {
            level.seekToFirst();
        }
        resetPriorityQueue();
    }

    @Override
    protected void seekInternal(InternalKey targetKey)
    {
        if (memTableIterator != null) {
            memTableIterator.seek(targetKey);
        }
        if (immutableMemTableIterator != null) {
            immutableMemTableIterator.seek(targetKey);
        }
        for (InternalTableIterator level0File : level0Files) {
            level0File.seek(targetKey);
        }
        for (LevelIterator level : levels) {
            level.seek(targetKey);
        }
        resetPriorityQueue();
    }

    @Override
    protected Entry<InternalKey, Slice> getNextElement()
    {
        if (heapSize == 0) {
            return null;
        }

        ComparableIterator smallest = heap[0];
        Entry<InternalKey, Slice> result = smallest.next();

        // if the smallest iterator has more elements, put it back in the heap,
        // otherwise use the last element in the queue
        ComparableIterator replacementElement;
        if (smallest.hasNext()) {
            replacementElement = smallest;
        }
        else {
            heapSize--;
            replacementElement = heap[heapSize];
            heap[heapSize] = null;
        }

        if (replacementElement != null) {
            heap[0] = replacementElement;
            heapSiftDown(0);
        }

        return result;
    }

    private void resetPriorityQueue()
    {
        int i = 0;
        if (memTableIterator != null && memTableIterator.hasNext()) {
            heapAdd(new ComparableIterator(memTableIterator, comparator, i++, memTableIterator.next()));
        }
        if (immutableMemTableIterator != null && immutableMemTableIterator.hasNext()) {
            heapAdd(new ComparableIterator(immutableMemTableIterator, comparator, i++, immutableMemTableIterator.next()));
        }
        for (InternalTableIterator level0File : level0Files) {
            if (level0File.hasNext()) {
                heapAdd(new ComparableIterator(level0File, comparator, i++, level0File.next()));
            }
        }
        for (LevelIterator level : levels) {
            if (level.hasNext()) {
                heapAdd(new ComparableIterator(level, comparator, i++, level.next()));
            }
        }<commit_msg>When the priority queue is reset, also reset the heap. After a seek, the iterator was still accessing key/values before the seek.<commit_after>package org.iq80.leveldb.util;

import com.google.common.base.Preconditions;
import com.google.common.primitives.Ints;
import org.iq80.leveldb.impl.InternalKey;
import org.iq80.leveldb.impl.MemTable.MemTableIterator;
import org.iq80.leveldb.impl.SeekingIterator;

import java.util.Comparator;
import java.util.Iterator;
import java.util.List;
import java.util.Map.Entry;
import java.util.NoSuchElementException;

public final class DbIterator extends AbstractSeekingIterator<InternalKey, Slice>
{
    private final MemTableIterator memTableIterator;
    private final MemTableIterator immutableMemTableIterator;
    private final List<InternalTableIterator> level0Files;
    private final List<LevelIterator> levels;

    private final Comparator<InternalKey> comparator;

    private final ComparableIterator[] heap;
    private int heapSize = 0;

    public DbIterator(MemTableIterator memTableIterator,
            MemTableIterator immutableMemTableIterator,
            List<InternalTableIterator> level0Files,
            List<LevelIterator> levels,
            Comparator<InternalKey> comparator)
    {
        this.memTableIterator = memTableIterator;
        this.immutableMemTableIterator = immutableMemTableIterator;
        this.level0Files = level0Files;
        this.levels = levels;
        this.comparator = comparator;

        this.heap = new ComparableIterator[3 + level0Files.size() + levels.size()];
        resetPriorityQueue();
    }

    @Override
    protected void seekToFirstInternal()
    {
        if (memTableIterator != null) {
            memTableIterator.seekToFirst();
        }
        if (immutableMemTableIterator != null) {
            immutableMemTableIterator.seekToFirst();
        }
        for (InternalTableIterator level0File : level0Files) {
            level0File.seekToFirst();
        }
        for (LevelIterator level : levels) {
            level.seekToFirst();
        }
        resetPriorityQueue();
    }

    @Override
    protected void seekInternal(InternalKey targetKey)
    {
        if (memTableIterator != null) {
            memTableIterator.seek(targetKey);
        }
        if (immutableMemTableIterator != null) {
            immutableMemTableIterator.seek(targetKey);
        }
        for (InternalTableIterator level0File : level0Files) {
            level0File.seek(targetKey);
        }
        for (LevelIterator level : levels) {
            level.seek(targetKey);
        }
        resetPriorityQueue();
    }

    @Override
    protected Entry<InternalKey, Slice> getNextElement()
    {
        if (heapSize == 0) {
            return null;
        }

        ComparableIterator smallest = heap[0];
        Entry<InternalKey, Slice> result = smallest.next();

        // if the smallest iterator has more elements, put it back in the heap,
        // otherwise use the last element in the queue
        ComparableIterator replacementElement;
        if (smallest.hasNext()) {
            replacementElement = smallest;
        }
        else {
            heapSize--;
            replacementElement = heap[heapSize];
            heap[heapSize] = null;
        }

        if (replacementElement != null) {
            heap[0] = replacementElement;
            heapSiftDown(0);
        }

        return result;
    }

    private void resetPriorityQueue()
    {
        int i = 0;
        heapSize = 0;
        if (memTableIterator != null && memTableIterator.hasNext()) {
            heapAdd(new ComparableIterator(memTableIterator, comparator, i++, memTableIterator.next()));
        }
        if (immutableMemTableIterator != null && immutableMemTableIterator.hasNext()) {
            heapAdd(new ComparableIterator(immutableMemTableIterator, comparator, i++, immutableMemTableIterator.next()));
        }
        for (InternalTableIterator level0File : level0Files) {
            if (level0File.hasNext()) {
                heapAdd(new ComparableIterator(level0File, comparator, i++, level0File.next()));
            }
        }
        for (LevelIterator level : levels) {
            if (level.hasNext()) {
                heapAdd(new ComparableIterator(level, comparator, i++, level.next()));
            }
        }
{'fraction_non_alphanumeric': 0.08955748068368063, 'fraction_numerical': 0.007375321938656053, 'mean_word_length': 2.484094616639478, 'pattern_counts': {'":': 0, '<': 21, '<?xml version=': 0, '>': 21, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4890510', 'n_tokens_mistral': 2258, 'n_tokens_neox': 2204, 'n_words': 493}