fse_decoder.go raw

   1  // Copyright 2019+ Klaus Post. All rights reserved.
   2  // License information can be found in the LICENSE file.
   3  // Based on work by Yann Collet, released under BSD License.
   4  
   5  package zstd
   6  
   7  import (
   8  	"encoding/binary"
   9  	"errors"
  10  	"fmt"
  11  	"io"
  12  )
  13  
  14  const (
  15  	tablelogAbsoluteMax = 9
  16  )
  17  
  18  const (
  19  	/*!MEMORY_USAGE :
  20  	 *  Memory usage formula : N->2^N Bytes (examples : 10 -> 1KB; 12 -> 4KB ; 16 -> 64KB; 20 -> 1MB; etc.)
  21  	 *  Increasing memory usage improves compression ratio
  22  	 *  Reduced memory usage can improve speed, due to cache effect
  23  	 *  Recommended max value is 14, for 16KB, which nicely fits into Intel x86 L1 cache */
  24  	maxMemoryUsage = tablelogAbsoluteMax + 2
  25  
  26  	maxTableLog    = maxMemoryUsage - 2
  27  	maxTablesize   = 1 << maxTableLog
  28  	maxTableMask   = (1 << maxTableLog) - 1
  29  	minTablelog    = 5
  30  	maxSymbolValue = 255
  31  )
  32  
  33  // fseDecoder provides temporary storage for compression and decompression.
  34  type fseDecoder struct {
  35  	dt             [maxTablesize]decSymbol // Decompression table.
  36  	symbolLen      uint16                  // Length of active part of the symbol table.
  37  	actualTableLog uint8                   // Selected tablelog.
  38  	maxBits        uint8                   // Maximum number of additional bits
  39  
  40  	// used for table creation to avoid allocations.
  41  	stateTable [256]uint16
  42  	norm       [maxSymbolValue + 1]int16
  43  	preDefined bool
  44  }
  45  
  46  // tableStep returns the next table index.
  47  func tableStep(tableSize uint32) uint32 {
  48  	return (tableSize >> 1) + (tableSize >> 3) + 3
  49  }
  50  
  51  // readNCount will read the symbol distribution so decoding tables can be constructed.
  52  func (s *fseDecoder) readNCount(b *byteReader, maxSymbol uint16) error {
  53  	var (
  54  		charnum   uint16
  55  		previous0 bool
  56  	)
  57  	if b.remain() < 4 {
  58  		return errors.New("input too small")
  59  	}
  60  	bitStream := b.Uint32NC()
  61  	nbBits := uint((bitStream & 0xF) + minTablelog) // extract tableLog
  62  	if nbBits > tablelogAbsoluteMax {
  63  		println("Invalid tablelog:", nbBits)
  64  		return errors.New("tableLog too large")
  65  	}
  66  	bitStream >>= 4
  67  	bitCount := uint(4)
  68  
  69  	s.actualTableLog = uint8(nbBits)
  70  	remaining := int32((1 << nbBits) + 1)
  71  	threshold := int32(1 << nbBits)
  72  	gotTotal := int32(0)
  73  	nbBits++
  74  
  75  	for remaining > 1 && charnum <= maxSymbol {
  76  		if previous0 {
  77  			//println("prev0")
  78  			n0 := charnum
  79  			for (bitStream & 0xFFFF) == 0xFFFF {
  80  				//println("24 x 0")
  81  				n0 += 24
  82  				if r := b.remain(); r > 5 {
  83  					b.advance(2)
  84  					// The check above should make sure we can read 32 bits
  85  					bitStream = b.Uint32NC() >> bitCount
  86  				} else {
  87  					// end of bit stream
  88  					bitStream >>= 16
  89  					bitCount += 16
  90  				}
  91  			}
  92  			//printf("bitstream: %d, 0b%b", bitStream&3, bitStream)
  93  			for (bitStream & 3) == 3 {
  94  				n0 += 3
  95  				bitStream >>= 2
  96  				bitCount += 2
  97  			}
  98  			n0 += uint16(bitStream & 3)
  99  			bitCount += 2
 100  
 101  			if n0 > maxSymbolValue {
 102  				return errors.New("maxSymbolValue too small")
 103  			}
 104  			//println("inserting ", n0-charnum, "zeroes from idx", charnum, "ending before", n0)
 105  			for charnum < n0 {
 106  				s.norm[uint8(charnum)] = 0
 107  				charnum++
 108  			}
 109  
 110  			if r := b.remain(); r >= 7 || r-int(bitCount>>3) >= 4 {
 111  				b.advance(bitCount >> 3)
 112  				bitCount &= 7
 113  				// The check above should make sure we can read 32 bits
 114  				bitStream = b.Uint32NC() >> bitCount
 115  			} else {
 116  				bitStream >>= 2
 117  			}
 118  		}
 119  
 120  		max := (2*threshold - 1) - remaining
 121  		var count int32
 122  
 123  		if int32(bitStream)&(threshold-1) < max {
 124  			count = int32(bitStream) & (threshold - 1)
 125  			if debugAsserts && nbBits < 1 {
 126  				panic("nbBits underflow")
 127  			}
 128  			bitCount += nbBits - 1
 129  		} else {
 130  			count = int32(bitStream) & (2*threshold - 1)
 131  			if count >= threshold {
 132  				count -= max
 133  			}
 134  			bitCount += nbBits
 135  		}
 136  
 137  		// extra accuracy
 138  		count--
 139  		if count < 0 {
 140  			// -1 means +1
 141  			remaining += count
 142  			gotTotal -= count
 143  		} else {
 144  			remaining -= count
 145  			gotTotal += count
 146  		}
 147  		s.norm[charnum&0xff] = int16(count)
 148  		charnum++
 149  		previous0 = count == 0
 150  		for remaining < threshold {
 151  			nbBits--
 152  			threshold >>= 1
 153  		}
 154  
 155  		if r := b.remain(); r >= 7 || r-int(bitCount>>3) >= 4 {
 156  			b.advance(bitCount >> 3)
 157  			bitCount &= 7
 158  			// The check above should make sure we can read 32 bits
 159  			bitStream = b.Uint32NC() >> (bitCount & 31)
 160  		} else {
 161  			bitCount -= (uint)(8 * (len(b.b) - 4 - b.off))
 162  			b.off = len(b.b) - 4
 163  			bitStream = b.Uint32() >> (bitCount & 31)
 164  		}
 165  	}
 166  	s.symbolLen = charnum
 167  	if s.symbolLen <= 1 {
 168  		return fmt.Errorf("symbolLen (%d) too small", s.symbolLen)
 169  	}
 170  	if s.symbolLen > maxSymbolValue+1 {
 171  		return fmt.Errorf("symbolLen (%d) too big", s.symbolLen)
 172  	}
 173  	if remaining != 1 {
 174  		return fmt.Errorf("corruption detected (remaining %d != 1)", remaining)
 175  	}
 176  	if bitCount > 32 {
 177  		return fmt.Errorf("corruption detected (bitCount %d > 32)", bitCount)
 178  	}
 179  	if gotTotal != 1<<s.actualTableLog {
 180  		return fmt.Errorf("corruption detected (total %d != %d)", gotTotal, 1<<s.actualTableLog)
 181  	}
 182  	b.advance((bitCount + 7) >> 3)
 183  	return s.buildDtable()
 184  }
 185  
 186  func (s *fseDecoder) mustReadFrom(r io.Reader) {
 187  	fatalErr := func(err error) {
 188  		if err != nil {
 189  			panic(err)
 190  		}
 191  	}
 192  	// 	dt             [maxTablesize]decSymbol // Decompression table.
 193  	//	symbolLen      uint16                  // Length of active part of the symbol table.
 194  	//	actualTableLog uint8                   // Selected tablelog.
 195  	//	maxBits        uint8                   // Maximum number of additional bits
 196  	//	// used for table creation to avoid allocations.
 197  	//	stateTable [256]uint16
 198  	//	norm       [maxSymbolValue + 1]int16
 199  	//	preDefined bool
 200  	fatalErr(binary.Read(r, binary.LittleEndian, &s.dt))
 201  	fatalErr(binary.Read(r, binary.LittleEndian, &s.symbolLen))
 202  	fatalErr(binary.Read(r, binary.LittleEndian, &s.actualTableLog))
 203  	fatalErr(binary.Read(r, binary.LittleEndian, &s.maxBits))
 204  	fatalErr(binary.Read(r, binary.LittleEndian, &s.stateTable))
 205  	fatalErr(binary.Read(r, binary.LittleEndian, &s.norm))
 206  	fatalErr(binary.Read(r, binary.LittleEndian, &s.preDefined))
 207  }
 208  
 209  // decSymbol contains information about a state entry,
 210  // Including the state offset base, the output symbol and
 211  // the number of bits to read for the low part of the destination state.
 212  // Using a composite uint64 is faster than a struct with separate members.
 213  type decSymbol uint64
 214  
 215  func newDecSymbol(nbits, addBits uint8, newState uint16, baseline uint32) decSymbol {
 216  	return decSymbol(nbits) | (decSymbol(addBits) << 8) | (decSymbol(newState) << 16) | (decSymbol(baseline) << 32)
 217  }
 218  
 219  func (d decSymbol) nbBits() uint8 {
 220  	return uint8(d)
 221  }
 222  
 223  func (d decSymbol) addBits() uint8 {
 224  	return uint8(d >> 8)
 225  }
 226  
 227  func (d decSymbol) newState() uint16 {
 228  	return uint16(d >> 16)
 229  }
 230  
 231  func (d decSymbol) baselineInt() int {
 232  	return int(d >> 32)
 233  }
 234  
 235  func (d *decSymbol) setNBits(nBits uint8) {
 236  	const mask = 0xffffffffffffff00
 237  	*d = (*d & mask) | decSymbol(nBits)
 238  }
 239  
 240  func (d *decSymbol) setAddBits(addBits uint8) {
 241  	const mask = 0xffffffffffff00ff
 242  	*d = (*d & mask) | (decSymbol(addBits) << 8)
 243  }
 244  
 245  func (d *decSymbol) setNewState(state uint16) {
 246  	const mask = 0xffffffff0000ffff
 247  	*d = (*d & mask) | decSymbol(state)<<16
 248  }
 249  
 250  func (d *decSymbol) setExt(addBits uint8, baseline uint32) {
 251  	const mask = 0xffff00ff
 252  	*d = (*d & mask) | (decSymbol(addBits) << 8) | (decSymbol(baseline) << 32)
 253  }
 254  
 255  // decSymbolValue returns the transformed decSymbol for the given symbol.
 256  func decSymbolValue(symb uint8, t []baseOffset) (decSymbol, error) {
 257  	if int(symb) >= len(t) {
 258  		return 0, fmt.Errorf("rle symbol %d >= max %d", symb, len(t))
 259  	}
 260  	lu := t[symb]
 261  	return newDecSymbol(0, lu.addBits, 0, lu.baseLine), nil
 262  }
 263  
 264  // setRLE will set the decoder til RLE mode.
 265  func (s *fseDecoder) setRLE(symbol decSymbol) {
 266  	s.actualTableLog = 0
 267  	s.maxBits = symbol.addBits()
 268  	s.dt[0] = symbol
 269  }
 270  
 271  // transform will transform the decoder table into a table usable for
 272  // decoding without having to apply the transformation while decoding.
 273  // The state will contain the base value and the number of bits to read.
 274  func (s *fseDecoder) transform(t []baseOffset) error {
 275  	tableSize := uint16(1 << s.actualTableLog)
 276  	s.maxBits = 0
 277  	for i, v := range s.dt[:tableSize] {
 278  		add := v.addBits()
 279  		if int(add) >= len(t) {
 280  			return fmt.Errorf("invalid decoding table entry %d, symbol %d >= max (%d)", i, v.addBits(), len(t))
 281  		}
 282  		lu := t[add]
 283  		if lu.addBits > s.maxBits {
 284  			s.maxBits = lu.addBits
 285  		}
 286  		v.setExt(lu.addBits, lu.baseLine)
 287  		s.dt[i] = v
 288  	}
 289  	return nil
 290  }
 291  
 292  type fseState struct {
 293  	dt    []decSymbol
 294  	state decSymbol
 295  }
 296  
 297  // Initialize and decodeAsync first state and symbol.
 298  func (s *fseState) init(br *bitReader, tableLog uint8, dt []decSymbol) {
 299  	s.dt = dt
 300  	br.fill()
 301  	s.state = dt[br.getBits(tableLog)]
 302  }
 303  
 304  // final returns the current state symbol without decoding the next.
 305  func (s decSymbol) final() (int, uint8) {
 306  	return s.baselineInt(), s.addBits()
 307  }
 308