SWCompression/Sources/Common/CodingTree/DecodingTree.swift

// Copyright (c) 2026 Timofey Solomko
// Licensed under MIT License
//
// See LICENSE for license information

import BitByteData

struct DecodingTree {

    private let bitReader: BitReader

    private let tree: [Int]
    private let leafCount: Int

    init(_ huffmanCodes: HuffmanCodes, _ bitReader: BitReader) {
        self.bitReader = bitReader

        // Maximum amount of leaves in a binary tree with the height equal to `huffmanCodes.maxBits`.
        self.leafCount = (1 << (huffmanCodes.maxBits + 1)) - 1
        var tree = Array(repeating: -1, count: leafCount)

        for code in huffmanCodes.codes {
            // Put code in its place in the tree.
            var treeCode = code.code
            var index = 0
            for _ in 0..<code.bits {
                let bit = treeCode & 1
                index = bit == 0 ? 2 * index + 1 : 2 * index + 2
                treeCode >>= 1
            }
            tree[index] = code.symbol
        }
        self.tree = tree
    }

    func findNextSymbol() -> Int {
        var bitsLeft = bitReader.bitsLeft
        var index = 0
        while bitsLeft > 0 {
            let bit = bitReader.bit()
            index = bit == 0 ? 2 * index + 1 : 2 * index + 2
            bitsLeft -= 1
            guard index < self.leafCount
                else { return -1 }
            if self.tree[index] > -1 {
                return self.tree[index]
            }
        }
        return -1
    }

}