huffman-archiver/internal/huffman/encode.go

package huffman

import (
	. "huffman/internal/io"
	"io"
	"math"
)

func encodeNode(root, prev *TreeNode) (res []byte) {
	if root.Parent != nil {
		res = encodeNode(root.Parent, root)
	}

	if prev == nil {
		return
	}

	if root.Left == prev {
		res = append(res, 0x0)
	} else {
		res = append(res, 0x1)
	}

	return
}

func initEncodeMap(root *TreeNode) map[byte][]byte {
	m := map[byte][]byte{}
	q := []*TreeNode{root}

	for len(q) > 0 {
		cur := q[0]
		q = q[1:]

		if cur.Left != nil {
			q = append(q, cur.Left)
		}

		if cur.Right != nil {
			q = append(q, cur.Right)
		}

		if cur.IsTerminate() {
			m[cur.Value] = encodeNode(cur, nil)
		}
	}

	return m
}

func calculateContentLen(root *TreeNode, len uint64) uint64 {
	if root.IsTerminate() {
		return len * root.Frequency
	}

	return calculateContentLen(root.Left, len+1) + calculateContentLen(root.Right, len+1)
}

func Encode(reader io.ReadSeeker, writer io.Writer) {
	bufReader := NewBufferedReader(reader)
	bufWriter := NewBufferedWriter(writer)

	tree, size := CreateTree(bufReader)
	treeMap := initEncodeMap(tree)
	lg := int(math.Ceil(math.Log2(float64(size)))) + 7

	EncodeTree(bufWriter, tree)

	if tree.Size == 1 {
		bufWriter.WriteByte(byte(lg))
		bufWriter.WriteBitsFromNumber(size, lg)
	} else {
		treeBitLen := uint64(5*tree.Size + 4)
		contentBitLen := calculateContentLen(tree, 0)

		additionalBits := (8 - (contentBitLen+treeBitLen+3)%8) % 8

		bufWriter.WriteBitsFromNumber(additionalBits, 3)
	}

	bufReader.SeekFromStart(0)

	for !bufReader.IsEOF() {
		node := treeMap[bufReader.ReadByte()]
		bufWriter.WriteBits(node)
	}

	bufWriter.Flush()
}