BZ2 : l'outsider entre DEFLATE et LZMA

#!/usr/bin/python3
#-*- encoding: Utf-8 -*-

from typing import List, Dict, Set, Sequence, Union, Optional
from binascii import crc32
from re import sub, DOTALL
from io import BytesIO

"""
    Variante de notre classe LecteurDeBits qui lit les valeurs encodées dans
    chaque octet d'un flux de bits de gauche à droite (MSB), plutôt que de
    droite à gauche (LSB). Taillée pour BZip2.
"""

class LecteurDeBitsMSB:
    
    def __init__(self, entree : bytes):
        
        self.octets = BytesIO(entree)
        
        # On initialise nos attributs.
        
        self.bits_non_lus = 0
        self.taille_bits_non_lus = 0
        
    """
        Lire un entier composé d'un certain nombre de bits,
        en commençant par les bits les plus à gauche
    """
    
    def lire_bits(self, nombre_bits : int) -> int:

        while self.taille_bits_non_lus < nombre_bits:
            
            prochain_octet = self.octets.read(1)
            
            """
                Si notre lecture du BytesIO a retourné une chaîne vide, alors
                ça veut dire qu'il n'y a plus d'octets de disponibles, on va
                déclencher une erreur
            """
            
            if not prochain_octet:
                raise EOFError
            
            """
                On n'oublie pas de convertir notre chaîne d'un octet en
                entier en ne prenant explicitement que le premier octet: "[0]"
                
                On l'ajoute à nos bits non lus
            """
            
            self.bits_non_lus <<= 8
            self.bits_non_lus |= prochain_octet[0]
            self.taille_bits_non_lus += 8
        
        # Les bits à lire sont les bits encore non lus les
        # plus à gauche (MSB = Most significant bits)
        
        bits_lus = self.bits_non_lus >> (self.taille_bits_non_lus - nombre_bits)
        
        self.taille_bits_non_lus -= nombre_bits
        
        masque_bits_non_lus = (1 << self.taille_bits_non_lus) - 1
        self.bits_non_lus &= masque_bits_non_lus
        
        return bits_lus
    
    """
        Lire simplement des octets
    """
    
    def lire_octets(self, nombre_octets : int) -> bytes:
        
        self.aligner_bits_sur_octet()
        
        octets_lus = self.octets.read(nombre_octets)
        
        if len(octets_lus) < nombre_octets:
            raise EOFError # Pas assez d'octets, déclencher une erreur, EOF = End of file = Fin de fichier
        
        return octets_lus
    
    """
        Se rendre au début de l'octet suivant
    """
    
    def aligner_bits_sur_octet(self):
        
        self.bits_non_lus = 0
        self.taille_bits_non_lus = 0
    

# Notre décodeur de Huffman, repris de DEFLATE.

class DecodeurHuffman:
    
    """
        La méthode "__init__" (en jargon technique, on l'appelle « notre
        constructeur » prendra comme premier argument une association (un
        dictionnaire Python) entre des valeurs (des entiers), et leurs tailles 
        de codes respectives d'autres entiers). À partir de ça, les codes
        Huffman seront calculés automatiquement.
    """
    
    def __init__(self, valeur_vers_taille_de_code : Dict[int, int], lecteur_de_bits : LecteurDeBitsMSB):
        
        # Chaîne de code = chaîne représentant un code écrit en chiffres binaires
        
        self.chaine_de_code_vers_valeur : Dict[str, int] = {}
        
        self.lecteur_de_bits : LecteurDeBitsMSB = lecteur_de_bits
        
        derniere_taille_de_code : int = None
        dernier_code : int = None
        
        """
            Comme expliqué plus haut dans le tutoriel, on va numéroter nos
            codes selon leurs tailles de codes respectives, puis selon leurs
            valeurs.
            
            Du coup, il va falloir passer sur les valeurs dans l'ordre, en
            prenent d'abord en compte la taille de code.
            
            Cette petite fonction, que l'on va passer à la fonction de tri de
            Python ("sorted"), devrait nous être très utile pour trier d'abord
            par tailles de codes, et ensuite par valeurs.
        """
        
        def tri_valeurs(valeur_et_taille_de_code):
            valeur, taille_de_code = valeur_et_taille_de_code # On prend la tuple...
            return (taille_de_code, valeur) # ... et on la change de sens pour dire qu'on trie par tailles de code d'abord
        
        for valeur, taille_de_code in sorted(valeur_vers_taille_de_code.items(), key = tri_valeurs):
            
            if not taille_de_code:
                continue # Taille de code de 0 ! On passe !
            
            if dernier_code is None:
                code = 0
                
            else:
                if taille_de_code < derniere_taille_de_code:
                    # Vu qu'on passe sur les tailles de code dans l'ordre, de
                    # la plus petite à la plus grande, ça n'arrive très
                    # logiquement jamais.
                    
                    raise Exception
                
                elif taille_de_code > derniere_taille_de_code:
                    code = (dernier_code + 1) << (taille_de_code - derniere_taille_de_code)
                
                else:
                    code = dernier_code + 1
            
            # On transforme notre code en chaîne contenant des chiffres
            # binaires : on utilise "format" pour créer la chaîne (qui
            # sera la plus courte possible), puis "zfill" qui va nous
            # ajouter le bon nombre de zéros à gauche
            
            chaine_de_code = format(code, 'b').zfill(taille_de_code)
            
            self.chaine_de_code_vers_valeur[chaine_de_code] = valeur
            
            dernier_code = code
            derniere_taille_de_code = taille_de_code
        
        """
            Maintenant qu'on a constitué tous nos codes dans 
            "chaine_de_code_vers_valeur", il va falloir créer la structure
            représentant notre arbre.
            
            Notre arbre sera fait de dictionnaires Python imbriqués entre eux,
            avec pour chacun au plus deux clefs, "0" et "1", et dont la valeur
            correspond soit à notre valeur Huffman (un entier), soit à un
            autre dictionnaire s'il faut continuer dans l'arbre.
        """
        
        # L'anotation de type "Union" signifie « une de ces valeurs » (comme
        # toutes les annotations de type, elle a uniquement une valeur
        # d'information)
        
        self.mon_arbre : Dict[int, Union[dict, int]] = {}
        
        for chaine_de_code, valeur in self.chaine_de_code_vers_valeur.items():
            
            # Pour chaque code, on va descendre dans l'arbre, en créant au
            # passage les dictionaires imbiqués successifs pour chaque bit
            # (sauf le dernier), et assigner le dernier bit à la valeur
            # associée au code
            
            noeud_actuel_dans_l_arbre = self.mon_arbre
            
            # Descendre dans l'arbre...
            
            for chiffre_bit in chaine_de_code[:-1]:
                bit = int(chiffre_bit) # Conversion de chaîne en entier
                
                if bit not in noeud_actuel_dans_l_arbre:
                    
                    noeud_actuel_dans_l_arbre[bit] = {} # Création de dictionnaire
                
                noeud_actuel_dans_l_arbre = noeud_actuel_dans_l_arbre[bit]
            
            # ...on est au bout
            
            dernier_bit = int(chaine_de_code[-1])
            
            noeud_actuel_dans_l_arbre[dernier_bit] = valeur
        
        # On a créé notre arbre, c'est tout bon !
    
    """
        Maintenant : voici la fonction qu'on va appeler tout le temps pour
        faire le décodage : elle va lire le code à partir du lecteur de bits.
        
        Elle commence en haut de l'arbre, puis elle avance d'un bit à la fois,
        jusqu'à avoir trouvé la valeur.
    """
    
    def lire_prochaine_valeur(self) -> int:
        
        noeud_actuel_dans_l_arbre = self.mon_arbre
        
        while True:
            prochain_bit : int = self.lecteur_de_bits.lire_bits(1)
            
            noeud_actuel_dans_l_arbre = noeud_actuel_dans_l_arbre[prochain_bit]
            
            # On est tombé sur un entier ? Chouette, on a trouvé notre valeur !
            
            if type(noeud_actuel_dans_l_arbre) == int:
                
                return noeud_actuel_dans_l_arbre


"""
    Notre décodeur pour l'encodage MTF (Move-to-Front)
"""

class DecodeurMTF:
    
    def decode(self, entree : List[int], dictionnaire_initial : List[int]) -> List[int]:
        
        dictionnaire_mtf : List[int] = list(dictionnaire_initial)
        
        sortie : List[int] = []
        
        for index_dictionnaire in entree:
            
            entree_dictionnaire = dictionnaire_mtf.pop(index_dictionnaire)
            dictionnaire_mtf.insert(0, entree_dictionnaire)
            
            sortie.append(entree_dictionnaire)
        
        return sortie


"""
    Notre CRC32 compatible BZ2
"""

def mon_crc32_msb(entree : bytes) -> int:
    
    def retourner_bits(bits : int, nombre_bits : int):
    
        return int(''.join(reversed(format(bits, 'b').zfill(nombre_bits))), 2)
    
    
    return retourner_bits(crc32(
        bytes(retourner_bits(octet, 8) for octet in entree)
    ) & 0xffffffff, 32)


"""
    Notre décodeur pour le format BZ2
"""

class DecodeurBZ2:
    
    def decompress(self, entree : bytes):
    
        lecteur_octets = BytesIO(entree)
        
        fichier_decompresse = b''

        # Lecture de chaque flux au sein du fichier

        while True:
            
            flux_decompresse = b''
        
            header_magic = lecteur_octets.read(2)
            version = lecteur_octets.read(1)
            compression_level : bytes = lecteur_octets.read(1)
            
            if not header_magic:
                
                break
            
            if header_magic != b'BZ' or version != b'h' or not (b'0' <= compression_level <= b'9'):
                
                raise ValueError('En-tête de flux BZ2 invalide')
            
            compression_level : int = int(compression_level)
            
            # Lecture de chaque bloc
            
            lecteur_de_bits = LecteurDeBitsMSB(lecteur_octets.read())
            
            while True:
                next_magic = lecteur_de_bits.lire_bits(48)
                
                BLOCK_MAGIC = 0x314159265359
                FOOTER_MAGIC = 0x177245385090
                
                if next_magic == BLOCK_MAGIC:
                    block_crc = lecteur_de_bits.lire_bits(32)
                    randomized = lecteur_de_bits.lire_bits(1)
                    orig_ptr = lecteur_de_bits.lire_bits(24)
                    
                    assert randomized == 0
                    
                    # Lire les symboles MTF utilisés au sein du bloc
                    
                    map_l1 = lecteur_de_bits.lire_bits(16) # Par exemple : 0b00110011
                    map_l1_bitfield : str = format(map_l1, 'b').zfill(16) # Par exemple : '00110011'
                    
                    
                    symboles_mtf_presents : List[str] = []
                    
                    for tranche_de_16_octets, tranche_presente in enumerate(map(int, map_l1_bitfield)):
                        
                        if tranche_presente:
                            map_l2 = lecteur_de_bits.lire_bits(16)
                            map_l2_bitfield : str = format(map_l2, 'b').zfill(16)
                            
                            for symbole_au_sein_de_la_tranche, symbole_present in enumerate(map(int, map_l2_bitfield)):
                                
                                if symbole_present:
                                    
                                    symboles_mtf_presents.append(tranche_de_16_octets * 16 + symbole_au_sein_de_la_tranche)
                                            
                    """
                        Décodage des informations relatives à l'arbre Huffman
                    """
                    
                    """
                        Décodage des correspondances entre sélecteur (groupe de
                        50 codes) et arbre
                    """
                    
                    num_trees = lecteur_de_bits.lire_bits(3)
                    num_selecteurs = lecteur_de_bits.lire_bits(15)
                    
                    decodeur_tree_index = DecodeurHuffman(
                        valeur_vers_taille_de_code = {tree_index: tree_index + 1 for tree_index in range(6)},
                        lecteur_de_bits = lecteur_de_bits)
                    
                    selecteur_vers_tree_index_avec_mtf = [decodeur_tree_index.lire_prochaine_valeur() for i in range(num_selecteurs)]
                    
                    selecteur_vers_tree_index = DecodeurMTF().decode(
                        entree = selecteur_vers_tree_index_avec_mtf,
                        dictionnaire_initial = [0, 1, 2, 3, 4, 5]
                    )
                    
                    
                    """
                        Décodage des tailles de code de chaque arbre
                    """
                    
                    # Les symboles Huffman spéciaux :
                    
                    self.RUNA = -2 # Bit "0" pour coder une taille de répétation RLE
                    self.RUNB = -1 # Bit "1" pour coder une taille de répétition RLE
                    self.EOB = 99999999 # End of block - marqueur de fin de bloc
                    
                    symboles_huffman_presents = [self.RUNA, self.RUNB]
                    symboles_huffman_presents += [index for index in range(1, len(symboles_mtf_presents))]
                    symboles_huffman_presents += [self.EOB]
                    
                    # Les symboles Huffman qui codent le delta encoding :
                    
                    DELTA_ENCODING_FIN_DE_TAILLE = 0b0
                    DELTA_ENCODING_INCREMENTER_TAILLE = 0b10
                    DELTA_ENCODING_DECREMENTER_TAILLE = 0b11
                    
                    decodeur_delta_encoding = DecodeurHuffman(
                        valeur_vers_taille_de_code = {
                            DELTA_ENCODING_FIN_DE_TAILLE: 1,
                            DELTA_ENCODING_INCREMENTER_TAILLE: 2,
                            DELTA_ENCODING_DECREMENTER_TAILLE: 2
                        },
                        lecteur_de_bits = lecteur_de_bits)
                    
                    tree_index_vers_arbre : List[DecodeurHuffman] = []
                    
                    for num_tree in range(num_trees):
                        
                        bit_len = lecteur_de_bits.lire_bits(5)
                        
                        tailles_symboles : List[int] = []
                        
                        for num_sym in range(len(symboles_huffman_presents)):
                            
                            taille_symbole = tailles_symboles[-1] if tailles_symboles else bit_len
                            
                            while True:
                                prochaine_instruction_delta = decodeur_delta_encoding.lire_prochaine_valeur()
                                
                                if prochaine_instruction_delta == DELTA_ENCODING_FIN_DE_TAILLE:
                                    break
                                elif prochaine_instruction_delta == DELTA_ENCODING_INCREMENTER_TAILLE:
                                    taille_symbole += 1
                                elif prochaine_instruction_delta == DELTA_ENCODING_DECREMENTER_TAILLE:
                                    taille_symbole -= 1
                            
                            tailles_symboles.append(taille_symbole)
                    
                        # En Python, la fonction "zip" permet de créer une liste de tuples à partir de
                        # plusieurs itérables. Ici, elle nous permet de créer un dictionnaire à partir
                        # de d'un côté les clefs, de l'autre côté les valeurs.
                    
                        tree_index_vers_arbre.append(DecodeurHuffman(
                            valeur_vers_taille_de_code = dict(zip(symboles_huffman_presents, tailles_symboles)),
                            lecteur_de_bits = lecteur_de_bits))
                        
                    """
                        On a décodé notre en-tête de bloc, on peut se mettre à décompresser
                        le contenu du flux !
                    """
                    
                    self.selecteur_vers_tree_index : List[int] = selecteur_vers_tree_index
                    
                    self.tree_index_vers_arbre : List[DecodeurHuffman] = tree_index_vers_arbre
                    self.symboles_mtf_presents : List[int] = symboles_mtf_presents
                    self.orig_ptr : int = orig_ptr
                    
                    self.decompresser_HUFF()
                    self.decompresser_RLE2()
                    self.decompresser_MTF()
                    self.decompresser_BWT()
                    self.decompresser_RLE1()
                    
                    
                    flux_decompresse += self.sortie_decompressee
                        
                    
                    assert mon_crc32_msb(self.sortie_decompressee) == block_crc
                
                elif next_magic == FOOTER_MAGIC:
                    stream_crc = lecteur_de_bits.lire_bits(32)
                    
                    
                    assert mon_crc32_msb(flux_decompresse) == stream_crc
                    break
                
                else:
                    raise ValueError('En-tête de bloc ou en-pied de flux BZ2 invalide')
                
            
            fichier_decompresse += flux_decompresse
            
            lecteur_de_bits.aligner_bits_sur_octet()
            entree = lecteur_de_bits.octets.read()
        
        return fichier_decompresse


    def decompresser_HUFF(self):
        
        symboles_lus : List[int] = []
        
        while True:
            selecteur = len(symboles_lus) // 50
            
            symbole = self.tree_index_vers_arbre[self.selecteur_vers_tree_index[selecteur]].lire_prochaine_valeur()
            
            symboles_lus.append(symbole)
            
            if symbole == self.EOB:
                break
        
        self.symboles_lus = symboles_lus
        

    def decompresser_RLE2(self):
        
        symboles_lus_apres_rle : List[int] = []
        
        symboles_runa_et_runb_a_traiter : List[int] = []
        
        for symbole_lu in self.symboles_lus:
            
            if symbole_lu in (self.RUNA, self.RUNB):
                symboles_runa_et_runb_a_traiter.append(symbole_lu)
            
            else:
                
                if symboles_runa_et_runb_a_traiter:
                    
                    nombre_de_zeros_a_repeter = 0b1
                    
                    for runa_ou_runb in reversed(symboles_runa_et_runb_a_traiter): # TODO reversed est-il bien nécessaire ici ?
                        
                        nombre_de_zeros_a_repeter <<= 1
                        nombre_de_zeros_a_repeter |= 1 if runa_ou_runb == self.RUNB else 0
                    
                    nombre_de_zeros_a_repeter -= 1
                    
                    symboles_lus_apres_rle += [0] * nombre_de_zeros_a_repeter
                    
                    
                    symboles_runa_et_runb_a_traiter = []
                
                if symbole_lu != self.EOB:
                    symboles_lus_apres_rle.append(symbole_lu)
                
                        
        
        self.symboles_lus : List[int] = symboles_lus_apres_rle
        

    def decompresser_MTF(self):
        
        self.symboles_lus : List[int] = DecodeurMTF().decode(
            entree = self.symboles_lus,
            dictionnaire_initial = self.symboles_mtf_presents
        )

    def decompresser_BWT(self):
        
        avant_tri = bytes(self.symboles_lus)
        apres_tri = sorted(avant_tri)
        
        
        # Nos paires Burrows-Wheeler à remettre dans l'ordre :
        
        class PaireBurrowsWheeler:
            
            char_1_valeur : str = None # Cette chaîne fera toujours un caractère
            char_1_position_origine : int = None # Cette position sera connue dès le début
            
            char_2_valeur : str = None # Cette chaîne fera toujours un caractère 
            char_2_position_origine : Optional[int] = None # Cette position sera déterminée après tri
        
        paires_a_trier = []
        
        for char_1_position_origine, (char_1_valeur, char_2_valeur) in enumerate(zip(avant_tri, apres_tri)):
            
            paire = PaireBurrowsWheeler()
            
            paire.char_1_valeur = bytes([char_1_valeur])
            paire.char_1_position_origine = char_1_position_origine
            
            paire.char_2_valeur = bytes([char_2_valeur])
            
            paires_a_trier.append(paire)
        
        
        # Reconstituer les positions d'origine des deuxièmes caractères de chaque paire
        
        paires_triees = sorted(paires_a_trier, key = lambda paire: (paire.char_1_valeur, paire.char_1_position_origine))
        
        for position_tri, paire in enumerate(paires_triees):
            paires_a_trier[position_tri].char_2_position_origine = paire.char_1_position_origine
        
        
        # Reconstituer la chaîne, maintenant qu'on a associé les paires entre elles
        
        donnees_sortie = b''
        
        premiere_paire = paires_triees[self.orig_ptr]
        paire_suivante = premiere_paire
        
        while not donnees_sortie or premiere_paire != paire_suivante:
            
            donnees_sortie += paire_suivante.char_1_valeur
            
            paire_suivante = paires_a_trier[paire_suivante.char_2_position_origine]
        

        self.sortie_decompressee : bytes = donnees_sortie
        
        
        """
            C'est tout bon, on a fait notre BWT !
        """
        

    def decompresser_RLE1(self):
        
        # Par concision, on va se permettre d'effectuer cette passe de RLE
        # avec des expressions régulières :
        # https://docs.python.org/3/library/re.html
        
        def effectuer_repetition(match):
            return match.group(1) * (ord(match.group(2)) + 4)
        
        self.sortie_decompressee : bytes = sub(rb'(.)\1\1\1(.)', effectuer_repetition, self.sortie_decompressee, flags = DOTALL)
        
        
        
print(DecodeurBZ2().decompress(bytes.fromhex('425a68393141592653593a1fb6180000285f940010408600008a0804002e21be00402008003000ac6193547a8d32681a7a20a0006819320488809313434c876f22a986198370606d342d7738f4e832122e26e1116dd24651715bc577384c4616e99a08ac9ca7b2e86b51648248e28ae6143b968594c582c2a6b43619a06461457633aeaa545f44ce49c372aebe6849774b04d957b4e2177301f92f45564f8bb9229c28481d0fdb0c00')).decode('utf8'))