Merge pull request #23 from Microsoft/wim_opts

Optimize WIM parsing

wim/lzx/lzx.go

@@ -6,7 +6,6 @@
 package lzx
 
 import (
-	"bufio"
 	"bytes"
 	"encoding/binary"
 	"errors"
@@ -17,6 +16,10 @@ const (
 	maincodecount = 496
 	maincodesplit = 256
 	lencodecount  = 249
+	lenshift      = 9
+	codemask      = 0x1ff
+	tablebits     = 9
+	tablesize     = 1 << tablebits
 
 	maxBlockSize = 32768
 	windowSize   = 32768
@@ -58,7 +61,7 @@ type Reader interface {
 }
 
 type decompressor struct {
-	r            Reader
+	r            io.Reader
 	err          error
 	unaligned    bool
 	nbits        byte
@@ -69,28 +72,59 @@ type decompressor struct {
 	mainlens     [maincodecount]byte
 	lenlens      [lencodecount]byte
 	window       [windowSize]byte
+	b            []byte
+	bv           int
+	bo           int
+}
+
+//go:noinline
+func (f *decompressor) fail(err error) {
+	if f.err == nil {
+		f.err = err
+	}
+	f.bo = 0
+	f.bv = 0
+}
+
+func (f *decompressor) ensureAtLeast(n int) error {
+	if f.bv-f.bo >= n {
+		return nil
+	}
+
+	if f.err != nil {
+		return f.err
+	}
+
+	if f.bv != f.bo {
+		copy(f.b[:f.bv-f.bo], f.b[f.bo:f.bv])
+	}
+	n, err := io.ReadAtLeast(f.r, f.b[f.bv-f.bo:], n)
+	if err != nil {
+		if err == io.EOF {
+			err = io.ErrUnexpectedEOF
+		} else {
+			f.fail(err)
+		}
+		return err
+	}
+	f.bv = f.bv - f.bo + n
+	f.bo = 0
+	return nil
 }
 
 // feed retrieves another 16-bit word from the stream and consumes
 // it into f.c. It returns false if there are no more bytes available.
 // Otherwise, on error, it sets f.err.
 func (f *decompressor) feed() bool {
-	if f.err != nil {
+	err := f.ensureAtLeast(2)
-		return true
-	}
-	var b0, b1 byte
-	b0, err := f.r.ReadByte()
-	if err == nil {
-		b1, err = f.r.ReadByte()
-	}
 	if err != nil {
-		if err == io.EOF {
+		if err == io.ErrUnexpectedEOF {
 			return false
 		}
-		f.err = err
 	}
-	f.c |= (uint32(b1)<<8 | uint32(b0)) << (16 - f.nbits)
+	f.c |= (uint32(f.b[f.bo+1])<<8 | uint32(f.b[f.bo])) << (16 - f.nbits)
 	f.nbits += 16
+	f.bo += 2
 	return true
 }
 
@@ -99,7 +133,7 @@ func (f *decompressor) feed() bool {
 func (f *decompressor) getBits(n byte) uint16 {
 	if f.nbits < n {
 		if !f.feed() {
-			f.err = io.ErrUnexpectedEOF
+			f.fail(io.ErrUnexpectedEOF)
 		}
 	}
 	c := uint16(f.c >> (32 - n))
@@ -109,9 +143,9 @@ func (f *decompressor) getBits(n byte) uint16 {
 }
 
 type huffman struct {
-	lens    []byte
+	extra   [][]uint16
-	table   []uint16
 	maxbits byte
+	table   [tablesize]uint16
 }
 
 // buildTable builds a huffman decoding table from a slice of code lengths,
@@ -148,54 +182,78 @@ func buildTable(codelens []byte) *huffman {
 
 	// Build a table for code lookup. For code sizes < max,
 	// put all possible suffixes for the code into the table, too.
-	// Typically a huffman implementation will only do this up to
+	// For max > tablebits, split long codes into additional tables
-	// a small code length maximum, then fall back to a different
+	// of suffixes of max-tablebits length.
-	// mechanism; this would probably improve performance.
+	h := &huffman{maxbits: max}
-	table := make([]uint16, 1<<max)
+	if max > tablebits {
-	for i, cl := range codelens {
+		core := first[tablebits+1] / 2 // Number of codes that fit without extra tables
-		if cl != 0 {
+		nextra := 1<<tablebits - core  // Number of extra entries
-			code := first[cl]
+		h.extra = make([][]uint16, nextra)
-			extendedCode := code << (max - cl)
+		for code := core; code < 1<<tablebits; code++ {
-			for j := uint(0); j < 1<<(max-cl); j++ {
+			h.table[code] = uint16(code - core)
-				table[extendedCode+j] = uint16(i)
+			h.extra[code-core] = make([]uint16, 1<<(max-tablebits))
-			}
-			first[cl]++
 		}
 	}
 
-	return &huffman{
+	for i, cl := range codelens {
-		lens:    codelens,
+		if cl != 0 {
-		table:   table,
+			code := first[cl]
-		maxbits: max,
+			first[cl]++
+			v := uint16(cl)<<lenshift | uint16(i)
+			if cl <= tablebits {
+				extendedCode := code << (tablebits - cl)
+				for j := uint(0); j < 1<<(tablebits-cl); j++ {
+					h.table[extendedCode+j] = v
 				}
+			} else {
+				prefix := code >> (cl - tablebits)
+				suffix := code & (1<<(cl-tablebits) - 1)
+				extendedCode := suffix << (max - cl)
+				for j := uint(0); j < 1<<(max-cl); j++ {
+					h.extra[h.table[prefix]][extendedCode+j] = v
+				}
+			}
+		}
+	}
+
+	return h
 }
 
 // getCode retrieves the next code using the provided
 // huffman tree. It sets f.err on error.
 func (f *decompressor) getCode(h *huffman) uint16 {
-	if h.maxbits == 0 {
+	if h.maxbits > 0 {
-		// This is an empty tree. It should not be used.
-		f.err = errCorrupt
-		return 0
-	}
 		if f.nbits < maxTreePathLen {
 			f.feed()
 		}
-	// For codes with length < h.maxbits, it doesn't matter
+
+		// For codes with length < tablebits, it doesn't matter
 		// what the remainder of the bits used for table lookup
 		// are, since entries with all possible suffixes were
 		// added to the table.
-	c := h.table[f.c>>(32-h.maxbits)]
+		c := h.table[f.c>>(32-tablebits)]
-	n := h.lens[c]
+		if c >= 1<<lenshift {
-	if f.nbits < n {
+			// The code is already in c.
-		f.err = io.ErrUnexpectedEOF
+		} else {
-		return 0
+			c = h.extra[c][f.c<<tablebits>>(32-(h.maxbits-tablebits))]
 		}
+
+		n := byte(c >> lenshift)
+		if f.nbits >= n {
 			// Only consume the length of the code, not the maximum
 			// code length.
 			f.c <<= n
 			f.nbits -= n
-	return c
+			return c & codemask
+		}
+
+		f.fail(io.ErrUnexpectedEOF)
+		return 0
+	}
+
+	// This is an empty tree. It should not be used.
+	f.fail(errCorrupt)
+	return 0
 }
 
 // mod17 computes the value mod 17.
@@ -279,13 +337,11 @@ func (f *decompressor) readBlockHeader() (byte, uint16, error) {
 	// If the previous block was an unaligned uncompressed block, restore
 	// 2-byte alignment.
 	if f.unaligned {
-		_, err := f.r.ReadByte()
+		err := f.ensureAtLeast(1)
 		if err != nil {
-			if err == io.EOF {
-				err = io.ErrUnexpectedEOF
-			}
 			return 0, 0, err
 		}
+		f.bo++
 		f.unaligned = false
 	}
 
@@ -321,19 +377,17 @@ func (f *decompressor) readBlockHeader() (byte, uint16, error) {
 		}
 
 		f.getBits(n)
-		if f.err != nil {
-			return 0, 0, f.err
-		}
 
 		// Read the LRU values for the next block.
-		var lru [12]byte
+		err := f.ensureAtLeast(12)
-		_, err := io.ReadFull(f.r, lru[:])
 		if err != nil {
 			return 0, 0, err
 		}
-		f.lru[0] = uint16(binary.LittleEndian.Uint32(lru[0:4]))
+
-		f.lru[1] = uint16(binary.LittleEndian.Uint32(lru[4:8]))
+		f.lru[0] = uint16(binary.LittleEndian.Uint32(f.b[f.bo : f.bo+4]))
-		f.lru[2] = uint16(binary.LittleEndian.Uint32(lru[8:12]))
+		f.lru[1] = uint16(binary.LittleEndian.Uint32(f.b[f.bo+4 : f.bo+8]))
+		f.lru[2] = uint16(binary.LittleEndian.Uint32(f.b[f.bo+8 : f.bo+12]))
+		f.bo += 12
 
 	default:
 		return 0, 0, errCorrupt
@@ -393,10 +447,11 @@ func (f *decompressor) readTrees(readAligned bool) (main *huffman, length *huffm
 // readCompressedBlock decodes a compressed block, writing into the window
 // starting at start and ending at end, and using the provided huffman trees.
 func (f *decompressor) readCompressedBlock(start, end uint16, hmain, hlength, haligned *huffman) (int, error) {
-	for i := start; i < end; {
+	i := start
+	for i < end {
 		main := f.getCode(hmain)
 		if f.err != nil {
-			return int(i - start), f.err
+			break
 		}
 		if main < 256 {
 			// Literal byte.
@@ -407,16 +462,13 @@ func (f *decompressor) readCompressedBlock(start, end uint16, hmain, hlength, ha
 
 		// This is a match backward in the window. Determine
 		// the offset and dlength.
-		lenheader := (main - 256) % 8
+		matchlen := (main - 256) % 8
 		slot := (main - 256) / 8
 
 		// The length is either the low bits of the code,
 		// or if this is 7, is encoded with the length tree.
-		var matchlen uint16
+		if matchlen == 7 {
-		if lenheader == 7 {
+			matchlen += f.getCode(hlength)
-			matchlen = f.getCode(hlength) + 7
-		} else {
-			matchlen = lenheader
 		}
 		matchlen += 2
 
@@ -452,16 +504,17 @@ func (f *decompressor) readCompressedBlock(start, end uint16, hmain, hlength, ha
 			f.lru[0] = matchoffset
 		}
 
-		if matchoffset > i || matchlen > end-i {
+		if matchoffset <= i && matchlen <= end-i {
-			return int(i - start), errCorrupt
+			copyend := i + matchlen
+			for ; i < copyend; i++ {
+				f.window[i] = f.window[i-matchoffset]
 			}
-
+		} else {
-		for j := uint16(0); j < matchlen; j++ {
+			f.fail(errCorrupt)
-			f.window[i+j] = f.window[i+j-matchoffset]
+			break
 		}
-		i += matchlen
 	}
-	return int(end - start), nil
+	return int(i - start), f.err
 }
 
 // readBlock decodes the current block and returns the number of uncompressed bytes.
@@ -476,7 +529,18 @@ func (f *decompressor) readBlock(start uint16) (int, error) {
 			// Remember to realign the byte stream at the next block.
 			f.unaligned = true
 		}
-		return io.ReadFull(f.r, f.window[start:start+size])
+		copied := 0
+		if f.bo < f.bv {
+			copied = int(size)
+			s := int(start)
+			if copied > f.bv-f.bo {
+				copied = f.bv - f.bo
+			}
+			copy(f.window[s:s+copied], f.b[f.bo:f.bo+copied])
+			f.bo += copied
+		}
+		n, err := io.ReadFull(f.r, f.window[start+uint16(copied):start+size])
+		return copied + n, err
 	}
 
 	hmain, hlength, haligned, err := f.readTrees(blockType == alignedOffsetBlock)
@@ -543,11 +607,8 @@ func NewReader(r io.Reader, uncompressedSize int) (io.ReadCloser, error) {
 	f := &decompressor{
 		lru:          [3]uint16{1, 1, 1},
 		uncompressed: uncompressedSize,
-	}
+		b:            make([]byte, 4096),
-	if br, ok := r.(Reader); ok {
+		r:            r,
-		f.r = br
-	} else {
-		f.r = bufio.NewReader(r)
 	}
 	return f, nil
 }

wim/wim.go

@@ -5,7 +5,6 @@
 package wim
 
 import (
-	"bufio"
 	"bytes"
 	"crypto/sha1"
 	"encoding/binary"
@@ -15,6 +14,7 @@ import (
 	"io"
 	"io/ioutil"
 	"strconv"
+	"sync"
 	"time"
 	"unicode/utf16"
 )
@@ -164,7 +164,6 @@ type securityblockDisk struct {
 const securityblockDiskSize = 8
 
 type direntry struct {
-	Length           int64
 	Attributes       uint32
 	SecurityID       uint32
 	SubdirOffset     int64
@@ -180,16 +179,15 @@ type direntry struct {
 	FileNameLength   uint16
 }
 
-const direntrySize = 102
+var direntrySize = int64(binary.Size(direntry{}) + 8) // includes an 8-byte length prefix
 
 type streamentry struct {
-	Length     int64
 	Unused     int64
 	Hash       SHA1Hash
 	NameLength int16
 }
 
-const streamentrySize = 38
+var streamentrySize = int64(binary.Size(streamentry{}) + 8) // includes an 8-byte length prefix
 
 // Filetime represents a Windows time.
 type Filetime struct {
@@ -299,6 +297,9 @@ type Image struct {
 	offset     resourceDescriptor
 	sds        [][]byte
 	rootOffset int64
+	r          io.ReadCloser
+	curOffset  int64
+	m          sync.Mutex
 
 	ImageInfo
 }
@@ -398,6 +399,14 @@ func NewReader(f io.ReaderAt) (*Reader, error) {
 	return r, nil
 }
 
+// Close releases resources associated with the Reader.
+func (r *Reader) Close() error {
+	for _, img := range r.Image {
+		img.reset()
+	}
+	return nil
+}
+
 func (r *Reader) resourceReader(hdr *resourceDescriptor) (io.ReadCloser, error) {
 	return r.resourceReaderWithOffset(hdr, 0)
 }
@@ -559,22 +568,23 @@ func (r *Reader) readSecurityDescriptors(rsrc io.Reader) (sds [][]byte, n int64,
 
 // Open parses the image and returns the root directory.
 func (img *Image) Open() (*File, error) {
+	if img.sds == nil {
 		rsrc, err := img.wim.resourceReaderWithOffset(&img.offset, img.rootOffset)
 		if err != nil {
 			return nil, err
 		}
-	defer rsrc.Close()
-
-	if img.sds == nil {
 		sds, n, err := img.wim.readSecurityDescriptors(rsrc)
 		if err != nil {
+			rsrc.Close()
 			return nil, err
 		}
 		img.sds = sds
+		img.r = rsrc
 		img.rootOffset = n
+		img.curOffset = n
 	}
 
-	f, err := img.readdir(rsrc)
+	f, err := img.readdir(img.rootOffset)
 	if err != nil {
 		return nil, err
 	}
@@ -584,16 +594,50 @@ func (img *Image) Open() (*File, error) {
 	return f[0], err
 }
 
-func (img *Image) readdir(rsrc io.Reader) ([]*File, error) {
+func (img *Image) reset() {
-	r := bufio.NewReader(rsrc)
+	if img.r != nil {
+		img.r.Close()
+		img.r = nil
+	}
+	img.curOffset = -1
+}
+
+func (img *Image) readdir(offset int64) ([]*File, error) {
+	img.m.Lock()
+	defer img.m.Unlock()
+
+	if offset < img.curOffset || offset > img.curOffset+chunkSize {
+		// Reset to seek backward or to seek forward very far.
+		img.reset()
+	}
+	if img.r == nil {
+		rsrc, err := img.wim.resourceReaderWithOffset(&img.offset, offset)
+		if err != nil {
+			return nil, err
+		}
+		img.r = rsrc
+		img.curOffset = offset
+	}
+	if offset > img.curOffset {
+		_, err := io.CopyN(ioutil.Discard, img.r, offset-img.curOffset)
+		if err != nil {
+			img.reset()
+			if err == io.EOF {
+				err = io.ErrUnexpectedEOF
+			}
+			return nil, err
+		}
+	}
 
 	var entries []*File
 	for {
-		e, err := img.readNextEntry(r)
+		e, n, err := img.readNextEntry(img.r)
+		img.curOffset += n
 		if err == io.EOF {
 			break
 		}
 		if err != nil {
+			img.reset()
 			return nil, err
 		}
 		entries = append(entries, e)
@@ -601,38 +645,39 @@ func (img *Image) readdir(rsrc io.Reader) ([]*File, error) {
 	return entries, nil
 }
 
-func (img *Image) readNextEntry(r *bufio.Reader) (*File, error) {
+func (img *Image) readNextEntry(r io.Reader) (*File, int64, error) {
-	lengthBuf, err := r.Peek(8)
+	var length int64
+	err := binary.Read(r, binary.LittleEndian, &length)
 	if err != nil {
-		return nil, &ParseError{Oper: "directory length check", Err: err}
+		return nil, 0, &ParseError{Oper: "directory length check", Err: err}
 	}
 
-	left := int(binary.LittleEndian.Uint64(lengthBuf))
+	if length == 0 {
-	if left == 0 {
+		return nil, 8, io.EOF
-		return nil, io.EOF
 	}
 
+	left := length
 	if left < direntrySize {
-		return nil, &ParseError{Oper: "directory entry", Err: errors.New("size too short")}
+		return nil, 0, &ParseError{Oper: "directory entry", Err: errors.New("size too short")}
 	}
 
 	var dentry direntry
 	err = binary.Read(r, binary.LittleEndian, &dentry)
 	if err != nil {
-		return nil, &ParseError{Oper: "directory entry", Err: err}
+		return nil, 0, &ParseError{Oper: "directory entry", Err: err}
 	}
 
 	left -= direntrySize
 
-	namesLen := int(dentry.FileNameLength + 2 + dentry.ShortNameLength)
+	namesLen := int64(dentry.FileNameLength + 2 + dentry.ShortNameLength)
 	if left < namesLen {
-		return nil, &ParseError{Oper: "directory entry", Err: errors.New("size too short for names")}
+		return nil, 0, &ParseError{Oper: "directory entry", Err: errors.New("size too short for names")}
 	}
 
 	names := make([]uint16, namesLen/2)
 	err = binary.Read(r, binary.LittleEndian, names)
 	if err != nil {
-		return nil, &ParseError{Oper: "file name", Err: err}
+		return nil, 0, &ParseError{Oper: "file name", Err: err}
 	}
 
 	left -= namesLen
@@ -652,7 +697,7 @@ func (img *Image) readNextEntry(r *bufio.Reader) (*File, error) {
 		var ok bool
 		offset, ok = img.wim.fileData[dentry.Hash]
 		if !ok {
-			return nil, &ParseError{Oper: "directory entry", Path: name, Err: fmt.Errorf("could not find file data matching hash %#v", dentry)}
+			return nil, 0, &ParseError{Oper: "directory entry", Path: name, Err: fmt.Errorf("could not find file data matching hash %#v", dentry)}
 		}
 	}
 
@@ -686,26 +731,30 @@ func (img *Image) readNextEntry(r *bufio.Reader) (*File, error) {
 	}
 
 	if isDir && f.subdirOffset == 0 {
-		return nil, &ParseError{Oper: "directory entry", Path: name, Err: errors.New("no subdirectory data for directory")}
+		return nil, 0, &ParseError{Oper: "directory entry", Path: name, Err: errors.New("no subdirectory data for directory")}
 	} else if !isDir && f.subdirOffset != 0 {
-		return nil, &ParseError{Oper: "directory entry", Path: name, Err: errors.New("unexpected subdirectory data for non-directory")}
+		return nil, 0, &ParseError{Oper: "directory entry", Path: name, Err: errors.New("unexpected subdirectory data for non-directory")}
 	}
 
 	if dentry.SecurityID != 0xffffffff {
 		f.SecurityDescriptor = img.sds[dentry.SecurityID]
 	}
 
-	_, err = r.Discard(left)
+	_, err = io.CopyN(ioutil.Discard, r, left)
 	if err != nil {
-		return nil, err
+		if err == io.EOF {
+			err = io.ErrUnexpectedEOF
+		}
+		return nil, 0, err
 	}
 
 	if dentry.StreamCount > 0 {
 		var streams []*Stream
 		for i := uint16(0); i < dentry.StreamCount; i++ {
-			s, err := img.readNextStream(r)
+			s, n, err := img.readNextStream(r)
+			length += n
 			if err != nil {
-				return nil, err
+				return nil, 0, err
 			}
 			// The first unnamed stream should be treated as the file stream.
 			if i == 0 && s.Name == "" {
@@ -720,42 +769,46 @@ func (img *Image) readNextEntry(r *bufio.Reader) (*File, error) {
 	}
 
 	if dentry.Attributes&FILE_ATTRIBUTE_REPARSE_POINT != 0 && f.Size == 0 {
-		return nil, &ParseError{Oper: "directory entry", Path: name, Err: errors.New("reparse point is missing reparse stream")}
+		return nil, 0, &ParseError{Oper: "directory entry", Path: name, Err: errors.New("reparse point is missing reparse stream")}
 	}
 
-	return f, nil
+	return f, length, nil
 }
 
-func (img *Image) readNextStream(r *bufio.Reader) (*Stream, error) {
+func (img *Image) readNextStream(r io.Reader) (*Stream, int64, error) {
-	lengthBuf, err := r.Peek(8)
+	var length int64
+	err := binary.Read(r, binary.LittleEndian, &length)
 	if err != nil {
-		return nil, &ParseError{Oper: "stream length check", Err: err}
+		if err == io.EOF {
+			err = io.ErrUnexpectedEOF
+		}
+		return nil, 0, &ParseError{Oper: "stream length check", Err: err}
 	}
 
-	left := int(binary.LittleEndian.Uint64(lengthBuf))
+	left := length
 	if left < streamentrySize {
-		return nil, &ParseError{Oper: "stream entry", Err: errors.New("size too short")}
+		return nil, 0, &ParseError{Oper: "stream entry", Err: errors.New("size too short")}
 	}
 
 	var sentry streamentry
 	err = binary.Read(r, binary.LittleEndian, &sentry)
 	if err != nil {
-		return nil, &ParseError{Oper: "stream entry", Err: err}
+		return nil, 0, &ParseError{Oper: "stream entry", Err: err}
 	}
 
 	left -= streamentrySize
 
-	if left < int(sentry.NameLength) {
+	if left < int64(sentry.NameLength) {
-		return nil, &ParseError{Oper: "stream entry", Err: errors.New("size too short for name")}
+		return nil, 0, &ParseError{Oper: "stream entry", Err: errors.New("size too short for name")}
 	}
 
 	names := make([]uint16, sentry.NameLength/2)
 	err = binary.Read(r, binary.LittleEndian, names)
 	if err != nil {
-		return nil, &ParseError{Oper: "file name", Err: err}
+		return nil, 0, &ParseError{Oper: "file name", Err: err}
 	}
 
-	left -= int(sentry.NameLength)
+	left -= int64(sentry.NameLength)
 	name := string(utf16.Decode(names))
 
 	var offset resourceDescriptor
@@ -763,7 +816,7 @@ func (img *Image) readNextStream(r *bufio.Reader) (*Stream, error) {
 		var ok bool
 		offset, ok = img.wim.fileData[sentry.Hash]
 		if !ok {
-			return nil, &ParseError{Oper: "stream entry", Path: name, Err: fmt.Errorf("could not find file data matching hash %v", sentry.Hash)}
+			return nil, 0, &ParseError{Oper: "stream entry", Path: name, Err: fmt.Errorf("could not find file data matching hash %v", sentry.Hash)}
 		}
 	}
 
@@ -777,12 +830,15 @@ func (img *Image) readNextStream(r *bufio.Reader) (*Stream, error) {
 		offset: offset,
 	}
 
-	_, err = r.Discard(left)
+	_, err = io.CopyN(ioutil.Discard, r, left)
 	if err != nil {
-		return nil, err
+		if err == io.EOF {
+			err = io.ErrUnexpectedEOF
+		}
+		return nil, 0, err
 	}
 
-	return s, nil
+	return s, length, nil
 }
 
 // Open returns an io.ReadCloser that can be used to read the stream's contents.
@@ -800,12 +856,7 @@ func (f *File) Readdir() ([]*File, error) {
 	if !f.IsDir() {
 		return nil, errors.New("not a directory")
 	}
-	rsrc, err := f.img.wim.resourceReaderWithOffset(&f.img.offset, f.subdirOffset)
+	return f.img.readdir(f.subdirOffset)
-	if err != nil {
-		return nil, err
-	}
-	defer rsrc.Close()
-	return f.img.readdir(rsrc)
 }
 
 // IsDir returns whether the given file is a directory. It returns false when it