<commit_before>package bt.torrent;

import bt.BtException;
import bt.data.DataStatus;
import bt.data.IChunkDescriptor;
import bt.net.IPeerConnection;
import bt.protocol.InvalidMessageException;
import bt.protocol.Request;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Predicate;

public class PieceManager implements IPieceManager {

    private static final Logger LOGGER = LoggerFactory.getLogger(PieceManager.class);

    private List<IChunkDescriptor> chunks;
    private AtomicInteger completePieces;

    /**
     * Indicates if there is at least one verified chunk in the local torrent files.
     */
    private volatile boolean haveAnyData;
    private byte[] bitfield;

    private PieceSelectorHelper pieceSelectorHelper;
    private Assignments assignments;
    private Predicate<Integer> selectionValidator;

    public PieceManager(PieceSelector selector, List<IChunkDescriptor> chunks) {

        this.chunks = chunks;
        completePieces = new AtomicInteger();

        bitfield = createBitfield(chunks);
        for (byte b : bitfield) {
            if (b != 0) {
                haveAnyData = true;
                break;
            }
        }

        pieceSelectorHelper = new PieceSelectorHelper(new PieceStats(chunks.size()), selector,
                pieceIndex -> !checkPieceCompleted(pieceIndex), bitfield.length);

        assignments = new Assignments();
        selectionValidator = pieceIndex -> !assignments.getAssignee(pieceIndex).isPresent();
    }

    /**
     * Creates a standard bittorrent bitfield, where n-th bit
     * (couting from high position to low) indicates the availability of n-th piece.
     */
    private byte[] createBitfield(List<IChunkDescriptor> chunks) {

        int chunkCount = chunks.size();
        byte[] bitfield = new byte[(int) Math.ceil(chunkCount / 8d)];
        int completePieces = 0;
        int bitfieldIndex = 0;
        while (chunkCount > 0) {
            int b = 0, offset = bitfieldIndex * 8;
            int k = chunkCount < 8? chunkCount : 8;
            for (int i = 0; i < k; i++) {
                IChunkDescriptor chunk = chunks.get(offset + i);
                if (chunk.getStatus() == DataStatus.VERIFIED) {
                    b += 0b1 << (7 - i);
                    completePieces++;
                }
            }
            bitfield[bitfieldIndex] = (byte) b;
            bitfieldIndex++;
            chunkCount -= 8;
        }
        this.completePieces.addAndGet(completePieces);
        return bitfield;
    }

    @Override
    public boolean haveAnyData() {
        return haveAnyData;
    }

    @Override
    public byte[] getBitfield() {
        return Arrays.copyOf(bitfield, bitfield.length);
    }

    // TODO: Check if peer has everything (i.e. is a seeder) and store him separately
    // this should help to improve performance of the next piece selection algorithm
    @Override
    public void peerHasBitfield(IPeerConnection peer, byte[] peerBitfield) {

        if (peerBitfield.length == bitfield.length) {
            byte[] bs = Arrays.copyOf(peerBitfield, peerBitfield.length);
            pieceSelectorHelper.addPeerBitfield(peer, bs);
        } else {
            throw new BtException("bitfield has wrong size: " + peerBitfield.length);
        }
    }

    @Override
    public void peerHasPiece(IPeerConnection peer, Integer pieceIndex) {
        validatePieceIndex(pieceIndex);
        pieceSelectorHelper.addPeerPiece(peer, pieceIndex);
    }

    @Override
    public boolean checkPieceCompleted(Integer pieceIndex) {
        return checkPieceCompleted(pieceIndex, false);
    }

    @Override
    public boolean checkPieceVerified(Integer pieceIndex) {
        return checkPieceCompleted(pieceIndex, true);
    }

    private boolean checkPieceCompleted(Integer pieceIndex, boolean verifiedOnly) {
        boolean completed = false;
        try {
            if (getBit(bitfield, pieceIndex) == 1) {
                completed = true;

            } else {
                IChunkDescriptor chunk = chunks.get(pieceIndex);
                switch (chunk.getStatus()) {
                    case VERIFIED: {
                        setBit(bitfield, pieceIndex);
                        haveAnyData = true;
                        completePieces.incrementAndGet();
                        completed = true;
                        break;
                    }
                    case COMPLETE: {
                        // chunk is completed but hasn't been verified yet
                        completed = !verifiedOnly;
                        break;
                    }
                }
            }
        } catch (Exception e) {
            LOGGER.error("Failed to check if chunk is complete {piece index: " + pieceIndex + "}", e);
        }

        if (completed && assignments.getAssignee(pieceIndex).isPresent()) {
            assignments.removeAssignee(pieceIndex);
        }
        return completed;
    }

    @Override
    public boolean mightSelectPieceForPeer(IPeerConnection peer) {
        Optional<Integer> piece = pieceSelectorHelper.selectPieceForPeer(peer, selectionValidator);
        return piece.isPresent() && !assignments.getAssignee(piece.get()).isPresent();
    }

    @Override
    public Optional<Integer> selectPieceForPeer(IPeerConnection peer) {
        Optional<Integer> assignedPiece = assignments.getAssignedPiece(peer);<commit_msg>Update for PieceManager (failing tests)<commit_after>package bt.torrent;

import bt.BtException;
import bt.data.DataStatus;
import bt.data.IChunkDescriptor;
import bt.net.IPeerConnection;
import bt.protocol.InvalidMessageException;
import bt.protocol.Request;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Predicate;

public class PieceManager implements IPieceManager {

    private static final Logger LOGGER = LoggerFactory.getLogger(PieceManager.class);

    private List<IChunkDescriptor> chunks;
    private AtomicInteger completePieces;

    /**
     * Indicates if there is at least one verified chunk in the local torrent files.
     */
    private volatile boolean haveAnyData;
    private byte[] bitfield;

    private PieceSelectorHelper pieceSelectorHelper;
    private Assignments assignments;
    private Predicate<Integer> selectionValidator;

    public PieceManager(PieceSelector selector, List<IChunkDescriptor> chunks) {

        this.chunks = chunks;
        completePieces = new AtomicInteger();

        bitfield = createBitfield(chunks);
        for (byte b : bitfield) {
            if (b != 0) {
                haveAnyData = true;
                break;
            }
        }

        pieceSelectorHelper = new PieceSelectorHelper(new PieceStats(chunks.size()), selector,
                pieceIndex -> !checkPieceCompleted(pieceIndex), bitfield.length);

        assignments = new Assignments();
        selectionValidator = pieceIndex -> !assignments.getAssignee(pieceIndex).isPresent();
    }

    /**
     * Creates a standard bittorrent bitfield, where n-th bit
     * (couting from high position to low) indicates the availability of n-th piece.
     */
    private byte[] createBitfield(List<IChunkDescriptor> chunks) {

        int chunkCount = chunks.size();
        byte[] bitfield = new byte[(int) Math.ceil(chunkCount / 8d)];
        int completePieces = 0;
        int bitfieldIndex = 0;
        while (chunkCount > 0) {
            int b = 0, offset = bitfieldIndex * 8;
            int k = chunkCount < 8? chunkCount : 8;
            for (int i = 0; i < k; i++) {
                IChunkDescriptor chunk = chunks.get(offset + i);
                if (chunk.getStatus() == DataStatus.VERIFIED) {
                    b += 0b1 << (7 - i);
                    completePieces++;
                }
            }
            bitfield[bitfieldIndex] = (byte) b;
            bitfieldIndex++;
            chunkCount -= 8;
        }
        this.completePieces.addAndGet(completePieces);
        return bitfield;
    }

    @Override
    public boolean haveAnyData() {
        return haveAnyData;
    }

    @Override
    public byte[] getBitfield() {
        return Arrays.copyOf(bitfield, bitfield.length);
    }

    // TODO: Check if peer has everything (i.e. is a seeder) and store him separately
    // this should help to improve performance of the next piece selection algorithm
    @Override
    public void peerHasBitfield(IPeerConnection peer, byte[] peerBitfield) {

        if (peerBitfield.length == bitfield.length) {
            byte[] bs = Arrays.copyOf(peerBitfield, peerBitfield.length);
            pieceSelectorHelper.addPeerBitfield(peer, bs);
        } else {
            throw new BtException("bitfield has wrong size: " + peerBitfield.length);
        }
    }

    @Override
    public void peerHasPiece(IPeerConnection peer, Integer pieceIndex) {
        validatePieceIndex(pieceIndex);
        pieceSelectorHelper.addPeerPiece(peer, pieceIndex);
    }

    @Override
    public boolean checkPieceCompleted(Integer pieceIndex) {
        return checkPieceCompleted(pieceIndex, false);
    }

    @Override
    public boolean checkPieceVerified(Integer pieceIndex) {
        return checkPieceCompleted(pieceIndex, true);
    }

    private boolean checkPieceCompleted(Integer pieceIndex, boolean verifiedOnly) {
        boolean completed = false;
        try {
            if (getBit(bitfield, pieceIndex) == 1) {
                completed = true;

            } else {
                IChunkDescriptor chunk = chunks.get(pieceIndex);
                switch (chunk.getStatus()) {
                    case VERIFIED: {
                        setBit(bitfield, pieceIndex);
                        haveAnyData = true;
                        completePieces.incrementAndGet();
                        completed = true;
                        break;
                    }
                    case COMPLETE: {
                        // chunk is completed but hasn't been verified yet
                        completed = !verifiedOnly;
                        break;
                    }
                }
            }
        } catch (Exception e) {
            LOGGER.error("Failed to check if chunk is complete {piece index: " + pieceIndex + "}", e);
        }

        if (completed && assignments.getAssignee(pieceIndex).isPresent()) {
            assignments.removeAssignee(pieceIndex);
        }
        return completed;
    }

    @Override
    public boolean mightSelectPieceForPeer(IPeerConnection peer) {
        if (peer.isClosed() || assignments.getAssignedPiece(peer).isPresent()) {
            return false;
        }
        Optional<Integer> piece = pieceSelectorHelper.selectPieceForPeer(peer, selectionValidator);
        return piece.isPresent() && !assignments.getAssignee(piece.get()).isPresent();
    }

    @Override
    public Optional<Integer> selectPieceForPeer(IPeerConnection peer) {
        Optional<Integer> assignedPiece = assignments.getAssignedPiece(peer);
{'fraction_non_alphanumeric': 0.08863064781892291, 'fraction_numerical': 0.002948573410805654, 'mean_word_length': 2.2393258426966294, 'pattern_counts': {'":': 0, '<': 25, '<?xml version=': 0, '>': 23, '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': '401874', 'n_tokens_mistral': 3117, 'n_tokens_neox': 2974, 'n_words': 762}