Source file src/pkg/compress/bzip2/huffman.go

     1	// Copyright 2011 The Go Authors. All rights reserved.
     2	// Use of this source code is governed by a BSD-style
     3	// license that can be found in the LICENSE file.
     4	
     5	package bzip2
     6	
     7	import "sort"
     8	
     9	// A huffmanTree is a binary tree which is navigated, bit-by-bit to reach a
    10	// symbol.
    11	type huffmanTree struct {
    12		// nodes contains all the non-leaf nodes in the tree. nodes[0] is the
    13		// root of the tree and nextNode contains the index of the next element
    14		// of nodes to use when the tree is being constructed.
    15		nodes    []huffmanNode
    16		nextNode int
    17	}
    18	
    19	// A huffmanNode is a node in the tree. left and right contain indexes into the
    20	// nodes slice of the tree. If left or right is invalidNodeValue then the child
    21	// is a left node and its value is in leftValue/rightValue.
    22	//
    23	// The symbols are uint16s because bzip2 encodes not only MTF indexes in the
    24	// tree, but also two magic values for run-length encoding and an EOF symbol.
    25	// Thus there are more than 256 possible symbols.
    26	type huffmanNode struct {
    27		left, right           uint16
    28		leftValue, rightValue uint16
    29	}
    30	
    31	// invalidNodeValue is an invalid index which marks a leaf node in the tree.
    32	const invalidNodeValue = 0xffff
    33	
    34	// Decode reads bits from the given bitReader and navigates the tree until a
    35	// symbol is found.
    36	func (t huffmanTree) Decode(br *bitReader) (v uint16) {
    37		nodeIndex := uint16(0) // node 0 is the root of the tree.
    38	
    39		for {
    40			node := &t.nodes[nodeIndex]
    41			bit := br.ReadBit()
    42			// bzip2 encodes left as a true bit.
    43			if bit {
    44				// left
    45				if node.left == invalidNodeValue {
    46					return node.leftValue
    47				}
    48				nodeIndex = node.left
    49			} else {
    50				// right
    51				if node.right == invalidNodeValue {
    52					return node.rightValue
    53				}
    54				nodeIndex = node.right
    55			}
    56		}
    57	
    58		panic("unreachable")
    59	}
    60	
    61	// newHuffmanTree builds a Huffman tree from a slice containing the code
    62	// lengths of each symbol. The maximum code length is 32 bits.
    63	func newHuffmanTree(lengths []uint8) (huffmanTree, error) {
    64		// There are many possible trees that assign the same code length to
    65		// each symbol (consider reflecting a tree down the middle, for
    66		// example). Since the code length assignments determine the
    67		// efficiency of the tree, each of these trees is equally good. In
    68		// order to minimize the amount of information needed to build a tree
    69		// bzip2 uses a canonical tree so that it can be reconstructed given
    70		// only the code length assignments.
    71	
    72		if len(lengths) < 2 {
    73			panic("newHuffmanTree: too few symbols")
    74		}
    75	
    76		var t huffmanTree
    77	
    78		// First we sort the code length assignments by ascending code length,
    79		// using the symbol value to break ties.
    80		pairs := huffmanSymbolLengthPairs(make([]huffmanSymbolLengthPair, len(lengths)))
    81		for i, length := range lengths {
    82			pairs[i].value = uint16(i)
    83			pairs[i].length = length
    84		}
    85	
    86		sort.Sort(pairs)
    87	
    88		// Now we assign codes to the symbols, starting with the longest code.
    89		// We keep the codes packed into a uint32, at the most-significant end.
    90		// So branches are taken from the MSB downwards. This makes it easy to
    91		// sort them later.
    92		code := uint32(0)
    93		length := uint8(32)
    94	
    95		codes := huffmanCodes(make([]huffmanCode, len(lengths)))
    96		for i := len(pairs) - 1; i >= 0; i-- {
    97			if length > pairs[i].length {
    98				// If the code length decreases we shift in order to
    99				// zero any bits beyond the end of the code.
   100				length >>= 32 - pairs[i].length
   101				length <<= 32 - pairs[i].length
   102				length = pairs[i].length
   103			}
   104			codes[i].code = code
   105			codes[i].codeLen = length
   106			codes[i].value = pairs[i].value
   107			// We need to 'increment' the code, which means treating |code|
   108			// like a |length| bit number.
   109			code += 1 << (32 - length)
   110		}
   111	
   112		// Now we can sort by the code so that the left half of each branch are
   113		// grouped together, recursively.
   114		sort.Sort(codes)
   115	
   116		t.nodes = make([]huffmanNode, len(codes))
   117		_, err := buildHuffmanNode(&t, codes, 0)
   118		return t, err
   119	}
   120	
   121	// huffmanSymbolLengthPair contains a symbol and its code length.
   122	type huffmanSymbolLengthPair struct {
   123		value  uint16
   124		length uint8
   125	}
   126	
   127	// huffmanSymbolLengthPair is used to provide an interface for sorting.
   128	type huffmanSymbolLengthPairs []huffmanSymbolLengthPair
   129	
   130	func (h huffmanSymbolLengthPairs) Len() int {
   131		return len(h)
   132	}
   133	
   134	func (h huffmanSymbolLengthPairs) Less(i, j int) bool {
   135		if h[i].length < h[j].length {
   136			return true
   137		}
   138		if h[i].length > h[j].length {
   139			return false
   140		}
   141		if h[i].value < h[j].value {
   142			return true
   143		}
   144		return false
   145	}
   146	
   147	func (h huffmanSymbolLengthPairs) Swap(i, j int) {
   148		h[i], h[j] = h[j], h[i]
   149	}
   150	
   151	// huffmanCode contains a symbol, its code and code length.
   152	type huffmanCode struct {
   153		code    uint32
   154		codeLen uint8
   155		value   uint16
   156	}
   157	
   158	// huffmanCodes is used to provide an interface for sorting.
   159	type huffmanCodes []huffmanCode
   160	
   161	func (n huffmanCodes) Len() int {
   162		return len(n)
   163	}
   164	
   165	func (n huffmanCodes) Less(i, j int) bool {
   166		return n[i].code < n[j].code
   167	}
   168	
   169	func (n huffmanCodes) Swap(i, j int) {
   170		n[i], n[j] = n[j], n[i]
   171	}
   172	
   173	// buildHuffmanNode takes a slice of sorted huffmanCodes and builds a node in
   174	// the Huffman tree at the given level. It returns the index of the newly
   175	// constructed node.
   176	func buildHuffmanNode(t *huffmanTree, codes []huffmanCode, level uint32) (nodeIndex uint16, err error) {
   177		test := uint32(1) << (31 - level)
   178	
   179		// We have to search the list of codes to find the divide between the left and right sides.
   180		firstRightIndex := len(codes)
   181		for i, code := range codes {
   182			if code.code&test != 0 {
   183				firstRightIndex = i
   184				break
   185			}
   186		}
   187	
   188		left := codes[:firstRightIndex]
   189		right := codes[firstRightIndex:]
   190	
   191		if len(left) == 0 || len(right) == 0 {
   192			return 0, StructuralError("superfluous level in Huffman tree")
   193		}
   194	
   195		nodeIndex = uint16(t.nextNode)
   196		node := &t.nodes[t.nextNode]
   197		t.nextNode++
   198	
   199		if len(left) == 1 {
   200			// leaf node
   201			node.left = invalidNodeValue
   202			node.leftValue = left[0].value
   203		} else {
   204			node.left, err = buildHuffmanNode(t, left, level+1)
   205		}
   206	
   207		if err != nil {
   208			return
   209		}
   210	
   211		if len(right) == 1 {
   212			// leaf node
   213			node.right = invalidNodeValue
   214			node.rightValue = right[0].value
   215		} else {
   216			node.right, err = buildHuffmanNode(t, right, level+1)
   217		}
   218	
   219		return
   220	}