huffman-archiver/internal/huffman/tree.go

package huffman

import (
	"huffman/internal/io"
)

type TreeNode struct {
	Value     byte
	Size      int16
	Frequency uint64
	Left      *TreeNode
	Right     *TreeNode
	Parent    *TreeNode
}

func (n *TreeNode) IsTerminate() bool {
	return n.Left == nil && n.Right == nil
}

func (n *TreeNode) CalcSize() {
	n.Size = 1

	if n.Left != nil {
		n.Left.CalcSize()
		n.Size += n.Left.Size
	}

	if n.Right != nil {
		n.Right.CalcSize()
		n.Size += n.Right.Size
	}
}

func DecodeTree(reader *io.BufferedReadSeeker) *TreeNode {
	var stack []*TreeNode

	for {
		node := &TreeNode{}

		if len(stack) > 0 {
			node.Parent = stack[len(stack)-1]

			if node.Parent.Left == nil {
				node.Parent.Left = node
			} else {
				node.Parent.Right = node
			}
		}

		stack = append(stack, node)
		if reader.ReadBit() == 0x0 {
			node.Value = reader.ReadByte()

			for len(stack) > 1 && (stack[len(stack)-1].Right != nil || stack[len(stack)-1].IsTerminate()) {
				stack = stack[:len(stack)-1]
			}

			if len(stack) == 1 && (stack[len(stack)-1].Right != nil || stack[len(stack)-1].IsTerminate()) {
				break
			}
		}
	}

	stack[0].CalcSize()
	return stack[0]
}

func encodeTreeRecursive(writer *io.BufferedWriter, root *TreeNode) {
	if root.IsTerminate() {
		writer.WriteBit(0x0)
		writer.WriteByte(root.Value)
		return
	}

	writer.WriteBit(0x1)

	encodeTreeRecursive(writer, root.Left)
	encodeTreeRecursive(writer, root.Right)
}

func EncodeTree(writer *io.BufferedWriter, root *TreeNode) {
	encodeTreeRecursive(writer, root)
}